如何实现支持嵌套Map的Spark CaseWhen外观函数？

需求背景 #

已实现支持单层Map[Column, Column]的caseWhenFacade函数，但需要扩展为支持嵌套Map结构：键为条件Column，值可以是结果Column，也可以是嵌套的子条件Map[Column, Any]。

现有单层实现代码 #

def withExpr(expr: Expression): Column = new Column(expr)

def caseWhenFacade(
       outputColName: String,
       conditionsValues: Map[Column, Column],
       defaultValue: Option[Column] = None
)(df: DataFrame): DataFrame = {
    require(conditionsValues.nonEmpty, "Cannot call caseWhenFacade with an Empty Map")

    val conditionalMap =  conditionsValues.map(x => (x._1.expr, x._2.expr)).toSeq

    defaultValue match {
      case Some(value) => df.withColumn(
        outputColName,
        withExpr {CaseWhen(
            conditionalMap,
            value.expr
            )}
      )
      case None => df.withColumn(
        outputColName,
        withExpr {CaseWhen(
          conditionalMap
          )}
      )
  }
}

嵌套条件映射示例 #

需要适配的嵌套结构如下：

val spec: Map[Column, Any] = Map(
    ($"column_one" === 1) -> Map(
        ($"column_two" === 2) -> lit(54),
        ($"column_two" === 5) -> lit(524)
    ),
    ($"column_one" === 2) -> Map(
        ($"column_two" === 7) -> Map(
            ($"whatever_column" === "whatever") -> lit(12),
            ($"whatever_column" === "whatever_two") -> lit(13)
        ),
        ($"column_two" === 8) -> lit(524)
    ),
    ($"column_one" === 3) -> lit(4)
)

解决方案：递归处理嵌套结构 #

核心思路是递归解析嵌套Map：当值是子Map时，将外层条件与内层条件组合，生成嵌套的CaseWhen表达式；当值是Column时直接作为结果。

完整实现代码 #

import org.apache.spark.sql.{Column, DataFrame}
import org.apache.spark.sql.catalyst.expressions.{CaseWhen, Expression}
import org.apache.spark.sql.functions.lit

def withExpr(expr: Expression): Column = new Column(expr)

// 递归解析嵌套条件，生成CaseWhen对应的(条件Expression, 结果Expression)序列
private def parseNestedConditions(conditions: Map[Column, Any]): Seq[(Expression, Expression)] = {
  conditions.map {
    case (condCol, resultCol: Column) =>
      (condCol.expr, resultCol.expr)
    case (condCol, nestedMap: Map[_, _]) =>
      // 递归处理嵌套Map，将外层条件与内层条件组合
      val nestedCases = parseNestedConditions(nestedMap.asInstanceOf[Map[Column, Any]])
      val nestedCaseWhen = CaseWhen(nestedCases)
      (condCol.expr, nestedCaseWhen)
    case (_, invalidValue) =>
      throw new IllegalArgumentException(s"Unsupported value type: ${invalidValue.getClass.getName}. Must be Column or Map[Column, Any]")
  }.toSeq
}

def caseWhenFacade(
       outputColName: String,
       conditionsValues: Map[Column, Any],
       defaultValue: Option[Column] = None
)(df: DataFrame): DataFrame = {
    require(conditionsValues.nonEmpty, "Cannot call caseWhenFacade with an Empty Map")

    val conditionalSeq = parseNestedConditions(conditionsValues)
    val caseWhenExpr = defaultValue match {
      case Some(value) => CaseWhen(conditionalSeq, value.expr)
      case None => CaseWhen(conditionalSeq)
    }

    df.withColumn(outputColName, withExpr(caseWhenExpr))
}

代码说明 #

1. 递归解析函数parseNestedConditions：
   • 遍历输入Map，判断每个值的类型：
     • 如果是Column，直接转换为(条件Expression, 结果Expression)对；
     • 如果是嵌套Map，递归调用自身解析内层条件，然后将外层条件与内层的CaseWhen表达式组合；
     • 其他类型直接抛出非法参数异常。
2. 修改后的caseWhenFacade：
   • 参数类型改为Map[Column, Any]，适配嵌套场景；
   • 调用递归函数生成完整的条件序列，再根据是否有默认值创建对应的CaseWhen表达式；
   • 最后将表达式添加到DataFrame中。

调用示例 #

直接使用给定的嵌套spec即可：

df.transform(caseWhenFacade("column_three", spec))

该调用会生成等价于多层嵌套CASE WHEN的SQL逻辑，例如：

CASE
  WHEN column_one = 1 THEN CASE
                             WHEN column_two = 2 THEN 54
                             WHEN column_two = 5 THEN 524
                           END
  WHEN column_one = 2 THEN CASE
                             WHEN column_two = 7 THEN CASE
                                                        WHEN whatever_column = 'whatever' THEN 12
                                                        WHEN whatever_column = 'whatever_two' THEN 13
                                                      END
                             WHEN column_two = 8 THEN 524
                           END
  WHEN column_one = 3 THEN 4
END AS column_three

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