Блог

Предыстория

У меня есть следующий pyspark dataframe

  ----- ----- 
|key_1|key_2|
 ----- ----- 
|True |True |
|False|True |
|False|False|
 ----- ----- 

  

Моя цель

И я хотел бы создать еще один столбец на основе значения в столбцах key_1 и key_2, как показано ниже.

  ----- ----- -------- 
|key_1|key_2|combined|
 ----- ----- -------- 
|True |True |  (1,2) |
|False|True |  (2,)  |
|False|False|  ()    |
 ----- ----- -------- 

  

Моя попытка и вопрос

Я попытался перебрать столбцы

 combined = []
for row in df.rdd.collections():
    tmp = []
    if row[0] is True:
        tmp.append(1)
    if row[1] is True:
        tmp.append(2)
    combined.append(tuple(tmp))
  

Тогда у меня есть список combined = [(1,2),(2,),()]

Как я могу вставить этот combined список, чтобы создать новый столбец в том же pyspark dataframe?

пример данных

 sdf1 = sc.parallelize([[True, True],[False, True],[False, False]]).toDF(["key_1", "key_2"])
sdf1.show()
 ----- ----- 
|key_1|key_2|
 ----- ----- 
| true| true|
|false| true|
|false|false|
 ----- ----- 
  

udf

 schema = StructType([
    StructField("value_1", IntegerType()),
    StructField("value_2", IntegerType())
])

def get_tuple(key_1, key_2):
    if key_1:
        value_1 = 1
    else:
        value_1 = None
    if key_2:
        value_2 = 2
    else:
        value_2 = None
    return (value_1, value_2)

get_tuple_udf = udf(get_tuple, schema)
  

Результат

 sdf1.withColumn("value", get_tuple_udf(sf.col("key_1"), sf.col("key_2"))).show()
 ----- ----- ------ 
|key_1|key_2| value|
 ----- ----- ------ 
| true| true|[1, 2]|
|false| true| [, 2]|
|false|false|   [,]|
 ----- ----- ------ 
  

Более простое решение:

 import pyspark.sql.functions as F

df = sc.parallelize([[True, True],
                       [False, True],
                       [False, False]]
                     ).toDF(["key_1", "key_2"])

df = df.withColumn("value", F.array(F.when(F.col("key_1"), 1), F.when(F.col("key_2"), 2)))

>>> df.show()
 ----- ----- ------ 
|key_1|key_2| value|
 ----- ----- ------ 
| true| true|[1, 2]|
|false| true| [, 2]|
|false|false|   [,]|
 ----- ----- ------