如何在Spark Scala中动态实现带Case When的列聚合？

需求说明 #

按col_A分组，依次判断sum(amt_M)、sum(amt_N)等列的求和结果是否不为0.0，取第一个符合条件的求和值并保留一位小数，命名为pct，且需支持后续列数动态增加。

参考Spark SQL实现 #

Select  col_A,
        round(case when sum(amt_M)   <> 0.0 then sum(amt_M) 
                   when sum(amt_N)   <> 0.0 then sum(amt_N)
                   when sum(amt_P)   <> 0.0 then sum(amt_P) 
              end,1) as pct 
from table_T1
group by col_A

问题分析 #

你尝试的Scala代码仅生成了多个单独的sum判断列，没有实现SQL中case when的顺序优先级判断逻辑，无法将结果合并为单一的pct列。

正确Spark-Scala实现 #

步骤说明 #

1. 先对指定列分组计算求和值，生成各列的sum结果
2. 使用coalesce结合when函数，按预设顺序依次判断每个sum是否非零，取第一个满足条件的值
3. 对最终结果做四舍五入保留一位小数，命名为pct

完整代码 #

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

// 读取源表
val df = spark.table("database.table_T1")

// 定义需要判断的列名（支持后续动态添加）
val amtCols = Seq("amt_M", "amt_N", "amt_P")

// 第一步：分组计算各列的求和值，生成sum_前缀的列
val sumExprs = amtCols.map(colName => sum(col(colName)).cast(DoubleType).alias(s"sum_$colName"))
val groupedDf = df.groupBy(col("col_A")).agg(sumExprs.head, sumExprs.tail:_*)

// 第二步：按顺序构造case when逻辑，取第一个非零的sum值
val pctExpr = amtCols.foldLeft(lit(null).cast(DoubleType)) { (acc, colName) =>
  coalesce(acc, when(col(s"sum_$colName") =!= 0.0, col(s"sum_$colName")))
}

// 第三步：保留一位小数，生成最终结果
val resultDf = groupedDf.withColumn("pct", round(pctExpr, 1)).select("col_A", "pct")

// 查看结果
resultDf.show()

结果验证 #

针对输入数据，执行上述代码后将得到预期输出：

+--------+---+
|   col_A|pct|
+--------+---+
|5C-SVS-1|1.0|
+--------+---+

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