如何将多维度聚合SQL转换为Spark Scala API实现？

尝试将以下SQL转换为Spark Scala API：

select beg_dt, col1, col2, count(distinct col3)
from tbl
group by 1 ,2,3
union all 
select beg_dt, col1, "xyx" as col2 ,count(distinct col3)
from tbl
group by 1,2
union all
select beg_dt, "abc" as col1, "xyx" as col2 , count(distinct col3)
from tbl
group by 1

本质是按不同维度层级进行聚合。尝试通过如下维度列表循环处理，但无法在每次迭代中正确添加静态列：

val dimCols: List[List[String]] =  List(
  List("col1", "col2" ), 
  List("col1", "'xyx' as col2"), 
  List("'abc' as col1","'xyx' as col2"  )
) 

val df = for (dimCol <- dimCols) yield { val x = myDF.groupBy( ( $"beg_dt" +: dimCol ).map(col): _* ).
                                                         agg(
                                                             countDistinct($"col4").as("count"),
                            
                               )
                        x       
              }

请问如何简洁地完成这个实现？

解决方案 #

你的问题出在直接用字符串（带别名的静态值表达式）调用col()方法——col()只能解析现有列名，无法处理'xyx' as col2这类静态值+别名的语法。可以通过定义清晰的聚合层级配置，用Spark原生API构造列来实现：

import org.apache.spark.sql.functions.{col, lit, countDistinct}

// 定义每个聚合层级的规则：(分组列列表, 需要设置静态值的列映射)
val aggLevels = List(
  // 层级1：按beg_dt、col1、col2分组，无静态列
  (List("beg_dt", "col1", "col2"), Map.empty[String, String]),
  // 层级2：按beg_dt、col1分组，col2固定为"xyx"
  (List("beg_dt", "col1"), Map("col2" -> "xyx")),
  // 层级3：按beg_dt分组，col1固定为"abc"，col2固定为"xyx"
  (List("beg_dt"), Map("col1" -> "abc", "col2" -> "xyx"))
)

// 遍历每个层级生成聚合DF，最终合并为一个结果
val finalDF = aggLevels.map { case (groupCols, staticCols) =>
  // 第一步：按指定列分组聚合
  val aggDF = myDF.groupBy(groupCols.map(col): _*)
    .agg(countDistinct(col("col3")).as("count"))
  
  // 第二步：添加静态列（分组列里没有的列，用lit生成固定值）
  staticCols.foldLeft(aggDF) { case (df, (colName, value)) =>
    df.withColumn(colName, lit(value))
  }
// 用unionByName确保列名匹配，不受列顺序影响
}.reduce(_ unionByName _)

// 可选：调整列顺序与原SQL一致
finalDF.select("beg_dt", "col1", "col2", "count")

关键说明 #

1. 用aggLevels清晰定义每个层级的分组规则和静态列，比纯字符串列表更易维护和扩展
2. 先完成分组聚合，再通过foldLeft动态添加静态列，避免了字符串解析的问题
3. 使用unionByName替代普通union，无需担心不同层级DF的列顺序差异

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