PySpark -拆分所有列中的数组并合并为行

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

spark = SparkSession.builder.appName("Array Split and Merge").getOrCreate()

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

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

df_exploded = df.select("id", explode("array_col").alias("array_value"))
df_exploded.show()

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

df_merged = df_exploded.groupBy("id").agg(collect_list("array_value").alias("merged_array"))
df_merged.show()

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