Pyspark：将两个数组类型列转换为 Array Type(Struct()) 类型的单个列

Question

这里是示例数据：

id sub_id       score      
A  [1, 4]   [0.9, 0.2]
B  [5, 7]   [0.1, 0.5]

我希望生成的列如下所示：

id sub_id       score        result
A  [1, 4]   [0.9, 0.2]   [Struct{id = 1, score = 0.9}, Struct{id = 4 , score = 0.2}]
B  [5, 7]   [0.1, 0.5]   [Struct{id = 5, score = 0.1}, Struct{id = 7 , score = 0.5}]

我知道如何做到这一点的唯一方法是：

1. 分解两列
2. 创建一个包含两个分解列的结构
3. 按id 分组以创建result 列。

我想知道是否有更有效的方法来做到这一点。

Answer 1

arrays_zip 函数，压缩两个数组列并创建一个数组结构。

from pyspark.sql import functions as F

data = [("A", [1, 4], [0.9, 0.2],),
        ("B", [5, 7], [0.1, 0.5],), ]

df = spark.createDataFrame(data, ("id", "sub_id", "score", ))

result = df.withColumn("result", F.arrays_zip(F.col("sub_id").alias("id"), F.col("score")))

result.printSchema()

result.show()

输出

root
 |-- id: string (nullable = true)
 |-- sub_id: array (nullable = true)
 |    |-- element: long (containsNull = true)
 |-- score: array (nullable = true)
 |    |-- element: double (containsNull = true)
 |-- result: array (nullable = true)
 |    |-- element: struct (containsNull = false)
 |    |    |-- id: long (nullable = true)
 |    |    |-- score: double (nullable = true)



+---+------+----------+--------------------+
| id|sub_id|     score|              result|
+---+------+----------+--------------------+
|  A|[1, 4]|[0.9, 0.2]|[{1, 0.9}, {4, 0.2}]|
|  B|[5, 7]|[0.1, 0.5]|[{5, 0.1}, {7, 0.5}]|
+---+------+----------+--------------------+