python - Spark DataFrame 聚合和分组多个列，同时保留顺序

Question

我有以下数据

id | value1 | value2 
-----------------------
1         A       red
1         B       red
1         C      blue
2         A      blue
2         B      blue
2         C     green

我需要的结果是:

id |                       values
---------------------------------
 1      [[A,red],[B,red][C,blue]]
 2   [[A,blue],[B,blue][C,green]]

到目前为止，我的方法是对单独的数组中的 value1 和 value2 进行分组和聚合，然后将它们合并在一起，如 Combine PySpark DataFrame ArrayType fields into single ArrayType field 中所述。

df.groupBy(["id"]).agg(*[F.collect_list("value1"), F.collect_list("value2")])

但是，由于 collect_list() 中不保证顺序(请参阅 here )，我如何确保 value1 和 value2 都与正确的值匹配？

这可能会导致两个列表的顺序不同，并且后续合并会匹配错误的值？

Answer 1

正如@Raphael所评论的，您可以首先将value1和value2列组合成一个struct类型列，然后收集列表:

import pyspark.sql.functions as F

(df.withColumn('values', F.struct(df.value1, df.value2))
   .groupBy('id')
   .agg(F.collect_list('values').alias('values'))).show()

+---+--------------------+
| id|              values|
+---+--------------------+
|  1|[[A,red], [B,red]...|
|  2|[[A,blue], [B,blu...|
+---+--------------------+