Spark Scala：无需explode实现col1、col2、col3分组聚合

问题背景 #

给定DataFrame结构：

col1: String
col2: String
col3: Array[String]
col4: Long

原方案通过explode(col3)后按col1、col2、col3分组聚合col4，但explode导致Executor内存溢出，且无法调整集群资源配置，需要实现等价逻辑但避免explode的提前展开。

可行解决方案 #

核心思路 #

先按col1、col2做粗粒度分组，在组内处理col3数组与col4的关联聚合，最后再展开结果——这种方式避免了原始数据集因explode产生的量级膨胀，将内存压力集中在分组后的小数据集上。

代码实现（Scala，Spark 3.0+）

import org.apache.spark.sql.functions._

// 1. 将每行的col3数组元素与对应col4值绑定为结构体列表
val mappedDF = df.withColumn(
  "elem_col4",
  transform(col("col3"), elem => struct(elem as "col3", col("col4") as "col4"))
)

// 2. 按col1、col2分组，收集所有结构体并按col3聚合sum(col4)
val aggregatedDF = mappedDF.groupBy("col1", "col2")
  .agg(collect_flat("elem_col4").alias("all_elem_col4"))
  .select(
    col("col1"),
    col("col2"),
    explode(
      map_from_entries(
        group_by_key(col("all_elem_col4"), sum("col4"))
      )
    ).alias("col3", "col4")
  )

低版本Spark兼容方案（自定义UDF）

如果使用Spark 3.0以下版本，可通过自定义UDF实现组内聚合：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.Row

val sumByElem = udf((elems: Seq[Row]) => {
  elems.groupBy(_.getAs[String]("col3"))
       .mapValues(_.map(_.getAs[Long]("col4")).sum)
       .toSeq
})

val aggregatedDF = mappedDF.groupBy("col1", "col2")
  .agg(collect_flat("elem_col4").alias("all_elem_col4"))
  .select(
    col("col1"),
    col("col2"),
    explode(sumByElem(col("all_elem_col4"))).alias("col3", "col4")
  )

关于你的拆分方案的问题 #

你提到的拆分df1和df2并用collect_set的思路不可行：collect_set会对col3元素去重，导致同一col1、col2下重复出现的数组元素对应的col4值无法累加，无法得到正确的聚合结果。

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