PySpark -按数组分组列

from pyspark.sql import SparkSession
from pyspark.sql.functions import explode

spark = SparkSession.builder.getOrCreate()

data = [(1, [1, 2, 3]), (2, [4, 5, 6]), (3, [7, 8, 9])]
df = spark.createDataFrame(data, ["id", "values"])
df.show()

+---+---------+
|id |values   |
+---+---------+
|1  |[1, 2, 3]|
|2  |[4, 5, 6]|
|3  |[7, 8, 9]|
+---+---------+

df_exploded = df.select("id", explode("values").alias("value"))
df_exploded.show()

+---+-----+
|id |value|
+---+-----+
|1  |1    |
|1  |2    |
|1  |3    |
|2  |4    |
|2  |5    |
|2  |6    |
|3  |7    |
|3  |8    |
|3  |9    |
+---+-----+

df_grouped = df_exploded.groupBy("id").agg({"value": "collect_list"})
df_grouped.show()

+---+-----------+
|id |collect_list(value)|
+---+-----------+
|1  |[1, 2, 3]  |
|3  |[7, 8, 9]  |
|2  |[4, 5, 6]  |
+---+-----------+