pyspark 数据框中的元组数据类型答案

【问题标题】：Tuple-like datatype in pyspark dataframepyspark 数据框中的元组数据类型
【发布时间】：2020-11-06 14:51:40
【问题描述】：

背景

我有以下 pyspark 数据框

+-----+-----+
|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 数据框中创建新列？

【问题讨论】：

标签： python pyspark

【解决方案1】：

样本数据

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|   [,]|
+-----+-----+------+

【讨论】：

【解决方案2】：

一个更简单的解决方案：

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|   [,]|
+-----+-----+------+

【讨论】：