PySpark按指定列分组并将多列按序转换为列表的实现

问题场景 #

你现在有这样的Spark DataFrame：

需要按a分组，把每行的b和c按原始行顺序拼接成逗号分隔的字符串，最终输出如下：

解决方案 #

要严格保证顺序，核心是先保留原始行的顺序标记，再分组拼接。下面是完整的PySpark代码实现：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, array, collect_list, flatten, concat_ws, monotonically_increasing_id

# 初始化SparkSession
spark = SparkSession.builder.appName("OrderPreserveConcat").getOrCreate()

# 创建示例DataFrame
data = [(5, 2, 1), (5, 4, 3), (2, 4, 2), (2, 3, 7)]
df = spark.createDataFrame(data, ["a", "b", "c"])

# 步骤1：添加行号，确保原始顺序可追溯
df_with_row_num = df.withColumn("row_num", monotonically_increasing_id())

# 步骤2：将每行的b和c组合成数组
df_with_array = df_with_row_num.withColumn("bc_array", array(col("b"), col("c")))

# 步骤3：按a分组，按行号排序后收集数组，扁平化并拼接成字符串
result_df = df_with_array.groupBy("a") \
    .agg(collect_list("bc_array").orderBy("row_num").alias("collected_arrays")) \
    .withColumn("b_list", concat_ws(",", flatten(col("collected_arrays")))) \
    .select("a", "b_list")

# 展示结果
result_df.show(truncate=False)

代码解释 #

• 添加行号：使用monotonically_increasing_id()生成唯一递增的行号，确保后续分组时能按原始行的顺序排序。如果你的DataFrame有天然的排序键（比如时间戳），也可以用那个键来替代行号。
• 组合数组：array(col("b"), col("c"))把每行的b和c打包成一个二元数组，保证每行内部b在前、c在后的顺序。
• 分组拼接：
  • groupBy("a")按目标字段分组；
  • collect_list("bc_array").orderBy("row_num")按行号排序后收集所有行的数组，保证原始行的顺序不会乱；
  • flatten()把嵌套的数组（比如[[2,1],[4,3]]）转成一维数组[2,1,4,3]；
  • concat_ws(",", ...)把数组元素用逗号拼接成字符串，得到最终的b_list。

输出验证 #

运行上述代码后，输出结果和你期望的完全一致：

+---+--------+
|a  |b_list  |
+---+--------+
|5  |2,1,4,3 |
|2  |4,2,3,7 |
+---+--------+

内容的提问来源于stack exchange，提问作者YOLO