Pyspark - 按列分组并从一列整数数组中收集唯一的值集

Question

我需要一些帮助来解决以下问题。

我有一个包含两列的 pyspark 数据框：

+----------------------+------+
|             col_list | group|
+----------------------+------+
|[1, 2, 3, 4, 5, 6, 7] |group1|
|            [6, 7, 8] |group1|
|         [1, 2, 3, 4] |group2|
|             [10, 11] |group2|
+----------------------+------+

我想对名为group 的列进行分组，并将唯一值仅收集到来自col_list 列的一个列表中。

我试过这个：

df.groupby("group").agg(F.flatten(F.collect_set('col_list')))

它的回答是这样的：

+------+-------------------------------+
| group|flatten(collect_set(col_list)) |
+------+-------------------------------+
|group1|          [1,2,3,4,5,6,7,6,7,8]|
|group2|          [10, 11, 1, 2, 3, 4] |
+------+-------------------------------+

group1 展平列表有重复项，我需要一些帮助来仅返回唯一值，例如：

[1,2,3,4,5,6,7,8]

Answer 1

这应该可以解决问题 - 你需要先explode()，然后是collect_set()：

df.select("group", F.explode(F.col("col_list")).alias("col_list")).groupby("group").agg(F.collect_set('col_list'))