优化Spark SQL重复groupBy/count调用效率并实现结果透视

问题分析 #

你的原方案最大的性能瓶颈在于对每一列单独执行groupBy再union——1000列就要生成1000个小DataFrame，再做999次union操作，这会触发大量的Shuffle和Job调度，完全是资源浪费。尤其是40万行的大表，重复的Shuffle操作直接把时间拉到了2小时以上。

高效解决方案思路 #

咱们换个思路，用宽表转长表的方式一次性整合所有列信息，全程只需要一次Shuffle操作，性能会提升几个数量级：

1. 用stack函数将所有列转换为window（原列名）和severity（列值）的长格式DataFrame
2. 按window和severity分组统计计数
3. 对severity进行透视，得到目标格式

实现代码（Scala） #

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

// 第一步：生成stack表达式，一次性把所有列转成长格式
val numCols = df.columns.length
val stackExpr = s"stack($numCols, ${df.columns.flatMap(c => s"'$c', $c").mkString(", ")}) as (window, severity)"

// 第二步：宽转长 + 一次分组计数
val longDF = df.selectExpr(stackExpr).groupBy("window", "severity").count()

// 第三步：透视得到最终结果，指定0-3的顺序保证列顺序符合预期
val results = longDF.groupBy("window").pivot("severity", Seq(0,1,2,3)).agg(first("count"))

代码细节解释 #

1. stack表达式生成：stack(n, 列1名称, 列1值, 列2名称, 列2值...)是Spark内置的宽转长函数，我们通过动态拼接字符串，把所有列的名称和对应值传入，一次性完成所有列的格式转换，避免了循环处理每一列。
2. 分组计数：这一步只需要一次groupBy操作，统计每个列每个值的出现次数，全程仅触发一次Shuffle。
3. 指定透视顺序：pivot的第二个参数传入Seq(0,1,2,3)，可以强制结果列的顺序，避免默认排序带来的不确定性。

性能对比 #

• 原方案：1000列需要1000次groupBy + 999次union，大量重复Shuffle和Job调度，资源利用率极低
• 新方案：仅1次宽转长 + 1次groupBy + 1次透视，Shuffle次数仅1次，在1000列40万行的场景下，预计耗时会从2小时压缩到几分钟甚至更短。

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