Scala Spark分组后展平collect_set生成的嵌套数组结构体列

问题场景 #

分组后使用collect_set操作，原本的array<struct<num: int, message: string>>类型列变成了array<array<struct<num: int, message: string>>>，需要将嵌套数组展平为单层数组，恢复分组前的格式。

解决方案 #

使用Spark的flatten函数可以直接将嵌套数组打平为单层数组，以下是具体实现代码：

import org.apache.spark.sql.functions.{collect_set, flatten, col}

// 假设分组字段为groupCol，原始DataFrame为df
val flattenedDF = df.groupBy("groupCol")
  .agg(collect_set("fullMessage").alias("nested_fullMessage"))
  // 展平嵌套数组
  .withColumn("fullMessage", flatten(col("nested_fullMessage")))
  // 移除中间生成的嵌套列
  .drop("nested_fullMessage")

额外去重处理 #

如果需要确保最终的fullMessage数组中所有struct元素完全唯一（避免不同子数组中存在重复的struct），可以在展平后再执行一次collect_set：

val flattenedUniqueDF = df.groupBy("groupCol")
  .agg(collect_set("fullMessage").alias("nested_fullMessage"))
  .withColumn("fullMessage", flatten(col("nested_fullMessage")))
  // 对展平后的数组再次去重
  .withColumn("fullMessage", collect_set(col("fullMessage")))
  .drop("nested_fullMessage")

flatten函数的作用是将array<array<T>>转换为array<T>，完美匹配本次嵌套数组的展平需求；二次collect_set则保证了最终数组内元素的唯一性。

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