Scala Spark中groupBy与pivot性能对比及多聚合pivot实现问题

测试数据集 #

首先给出测试用的DataFrame生成代码：

val data = Seq((1, "k1", "measureA", 2), (1, "k1", "measureA", 4), (1, "k1", "measureB", 5), (1, "k1", "measureB", 7), (1, "k1", "measureC", 7), (1, "k1", "measureC", 1), (2, "k1", "measureB", 8), (2, "k1", "measureC", 9), (2, "k2", "measureA", 5), (2, "k2", "measureC", 5), (2, "k2", "measureC", 8))
val df = data.toDF("ts","key","measure_type","value")

原始数据格式如下：

+---+---+------------+-----+
| ts|key|measure_type|value|
+---+---+------------+-----+
|  1| k1|    measureA|    2|
|  1| k1|    measureA|    4|
|  1| k1|    measureB|    5|
|  1| k1|    measureB|    7|
|  1| k1|    measureC|    7|
|  1| k1|    measureC|    1|
|  2| k1|    measureB|    8|
|  2| k1|    measureC|    9|
|  2| k2|    measureA|    5|
|  2| k2|    measureC|    5|
|  2| k2|    measureC|    8|
+---+---+------------+-----+

两种实现方案 #

方案1：groupBy + 条件聚合 #

代码：

val ddf = df.groupBy("ts", "key").agg(
 sum(when(col("measure_type") === "measureA",col("value"))).as("measureA"),
 avg(when(col("measure_type") === "measureB",col("value"))).as("measureB"),
 max(when(col("measure_type") === "measureC",col("value"))).as("measureC"))

输出结果：

+---+---+--------+--------+--------+
| ts|key|measureA|measureB|measureC|
+---+---+--------+--------+--------+
|  1| k1|       6|     6.0|       7|
|  2| k1|    null|     8.0|       9|
|  2| k2|       5|    null|       8|
+---+---+--------+--------+--------+

方案2：pivot + 条件聚合 #

代码：

val listA = Seq("measureA")
val listB = Seq("measureB")
val listC = Seq("measureC")

val ddf = df.groupBy("ts", "key").pivot(col("measure_type"), Seq("measureA", "measureB", "measureC")).agg(
 sum(when(col("measure_type").isInCollection(listA),col("value"))).as("measureA"),
 avg(when(col("measure_type").isInCollection(listB),col("value"))).as("measureB"),
 max(when(col("measure_type").isInCollection(listC),col("value"))).as("measureC"))

输出会生成3个pivot值 × 3个聚合函数的9列冗余结果，大部分列全为null，过滤后才能得到和方案1一致的结果。

问题解答 #

1. 性能对比 #

方案1性能明显更优，原因如下：

• 方案1仅执行一次分组聚合，只计算3个必要的聚合逻辑，无冗余计算，shuffle数据量最小
• 方案2的pivot操作本身会增加列展开的额外开销，同时需要计算9个聚合表达式，其中6个完全是无用的null值计算，后续还要额外做列过滤，整体计算量、shuffle量都远高于方案1

2. 方案2的写法优化 #

你当前的写法存在逻辑冗余：pivot操作已经自动按measure_type的取值做了分组，每个pivot值对应的聚合逻辑不需要再额外加when判断measure_type的取值。
如果必须使用pivot实现，可以直接丢弃冗余列、重命名得到目标结果，无需额外判断全null列：

// 仅保留pivot值和聚合别名一致的列，并重命名为目标列名
val targetCols = Seq("measureA", "measureB", "measureC").map(m => col(s"${m}_${m}").as(m))
val finalDf = ddf.select("ts", "key", targetCols:_*)

原生pivot搭配多聚合函数的写法必然会生成笛卡尔积列，无法在生成阶段完全避免冗余列，只能生成后按规则筛选保留需要的列。

最优实现建议 #

当前场景下直接使用方案1的groupBy+条件聚合是最优解，不需要引入pivot操作的额外开销。

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