Spark/Scala：如何按cytid将多列值合并到单个列中

需求说明 #

按cytid分组，将每组内所有行的percentile和mean列值按顺序合并为逗号分隔的字符串，存入ipm列（注：输出结果未包含itemcode数值，推测需求描述中提及的itemcode为笔误，实际仅需合并percentile和mean）。

输入数据集 #

+-------+--------+----------+----+----------+ 
|cytid  |itemcode|percentile|mean|DTXN      |
+-------+--------+----------+----+----------+
|800    |7       |x4        |y#  |2/1/2017  |
|700    |13      |y#        |oz  |3/1/2017  |
|800    |24      |oz        |oz  |8/1/2017  |
|900    |747     |oz        |IH  |27/9/2017 |
|700    |798     |IH        |yc  |19/10/2017| 
|800    |841     |yc        |g5  |8/11/2017 |
|800    |867     |g5        |IH  |19/11/2017| 
|900    |911     |IH        |x4  |2/12/2017 |
|800    |945     |x4        |ko  |13/12/2017| 
|800    |984     |ko        |cd  |27/12/2017| 
+-------+--------+----------+----+----------+

期望输出数据集 #

+-------+--------------------------------------+
|cytid  |ipm                                   |
+-------+--------------------------------------+
|800    |x4,y#,oz,oz,yc,g5,g5,IH,x4,ko,ko,cd   |
|700    |y#,oz,IH,yc                           |
|900    |oz,IH,IH,x4                           |
+-------+--------------------------------------+

Scala实现代码 #

import org.apache.spark.sql.functions.{array, collect_list, concat_ws, flatten}

// 假设原始DataFrame已加载完成，命名为df
val resultDF = df
  // 每行生成包含percentile和mean值的数组
  .withColumn("temp_array", array($"percentile", $"mean"))
  // 按cytid分组，收集所有行的数组并扁平化
  .groupBy("cytid")
  .agg(flatten(collect_list($"temp_array")).alias("ipm_array"))
  // 将扁平化后的数组转为逗号分隔的字符串
  .withColumn("ipm", concat_ws(",", $"ipm_array"))
  // 保留需要的列
  .select("cytid", "ipm")

// 打印结果
resultDF.show(false)

代码说明 #

1. 生成临时数组列：使用array函数将每行的percentile和mean值打包成一个数组，便于后续批量收集。
2. 分组收集并扁平化：按cytid分组后，collect_list会把每组内的所有临时数组收集为嵌套数组，再通过flatten将嵌套数组转为一维数组，保证值的顺序与原始数据一致。
3. 转为字符串：使用concat_ws将一维数组的元素用逗号连接，生成最终的ipm列。

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