Python PySpark SQL Functions collect_set方法用法及代码示例

例子

考虑以下PySpark DataFrame：

data = [("Alex", "A"), ("Alex", "B"), ("Bob", "A"), ("Cathy", "C"), ("Dave", None)]
df = spark.createDataFrame(data, ["name", "group"])
df.show()



+-----+-----+
| name|group|
+-----+-----+
| Alex|    A|
| Alex|    B|
|  Bob|    A|
|Cathy|    C|
| Dave| null|
+-----+-----+

获取 PySpark 中的一组列值

要获取 group 列中唯一的一组值：

import pyspark.sql.functions as F
df.select(F.collect_set("group")).show()



+------------------+
|collect_set(group)|
+------------------+
|         [C, B, A]|
+------------------+

同样，您也可以将 Column 对象传递给 collect_set(~)：

import pyspark.sql.functions as F
df.select(F.collect_set(df.group)).show()



+------------------+
|collect_set(group)|
+------------------+
|         [C, B, A]|
+------------------+

请注意null 值没有出现在结果集中。

将集合作为标准列表

要将集合作为标准列表：

list_rows = df.select(F.collect_set(df.group)).collect()
list_rows[0][0]



['C', 'B', 'A']

这里，PySpark DataFrame 的 collect() 方法返回Row 对象的列表。由于 collect_set(~) 的性质，该列表保证长度为一。 Row 对象包含该列表，因此我们需要包含另一个 [0] 。

获取PySpark中每个组的一组列值

方法 collect_set(~) 通常在聚合上下文中使用。考虑与之前相同的 PySpark DataFrame：

df.show()



+-----+-----+
| name|group|
+-----+-----+
| Alex|    A|
| Alex|    B|
|  Bob|    A|
|Cathy|    C|
| Dave| null|
+-----+-----+

要将 group 列展平为每个 name 的单个集合：

import pyspark.sql.functions as F
df.groupby("name").agg(F.collect_set("group")).show()



+-----+------------------+
| name|collect_set(group)|
+-----+------------------+
| Alex|            [B, A]|
|  Bob|               [A]|
|Cathy|               [C]|
+-----+------------------+