Spark中基于ADTs实现Map转CaseWhen的优化问题

背景回顾 #

Spark的CaseWhen接收Seq[(Expression, Expression)]参数，第一个Expression为条件，第二个为条件满足时的取值：

CaseWhen(
    branches: Seq[(Expression, Expression)],
    elseValue: Option[Expression] = None): ...

我希望通过Map对象实现Spark的CaseWhen逻辑，Map支持两种结构：

• 单层结构：

val spec = Map(
    ($"column_one" === 1) -> lit(2),
    ($"column_one" === 2 && $"column_two" === 1) -> lit(1),
    ($"column_one" === 3) -> lit(4),
)

• 嵌套与单层混合结构：

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)
)

当前实现代码 #

基于之前的思路，我实现了以下代码：

sealed trait ConditionValue  {
  def enumerate(reduceFunc: (Column, Column) => Column): Seq[(Expression, Expression)]
}
object ConditionValue {

  object implicits{

    implicit def test(condition: Column, value: Column): ConditionValue = {
        print("test")
        SingleLevelCaseWhen(Map(condition -> value))
        
    }
    
    implicit def testTuple(conditionValue: (Column, Column)): ConditionValue = {
        print("testTuple")
        SingleLevelCaseWhen(Map(conditionValue))
        
    }
        
    implicit def testNested(spec: Map[Column, ConditionValue]): ConditionValue = {
        print("testNested")
        NestedCaseWhen(spec)
        
    }
        
    implicit def testMap(spec: Map[Column, Column]): ConditionValue = {
        print("testMap")
        SingleLevelCaseWhen(spec)
        
    }
    
    implicit def expressionToColumn(expr: Expression): Column = new Column(expr)

    implicit def columnToExpression(col: Column): Expression = col.expr
  }
    
  import implicits._

  final case class SingleLevelCaseWhen(specificationMap: Map[Column, Column]) extends ConditionValue{
    override def enumerate(reduceFunc: (Column, Column) => Column): Seq[(Expression, Expression)] =
      specificationMap.map(x => (x._1.expr, x._2.expr)).toSeq
  }
  
  final case class NestedCaseWhen(specificationMap: Map[Column, ConditionValue]) extends ConditionValue{
    override def enumerate(reduceFunc: (Column, Column) => Column): Seq[(Expression, Expression)] =
      specificationMap.mapValues(_.enumerate(reduceFunc)).map{
        case (outerCondition, innerExpressions) => innerExpressions.map{
          case (innerCondition, innerValue) =>
            val conditions: Expression = reduceFunc(outerCondition, innerCondition)
            conditions -> innerValue

        }
      }.reduce(_ ++ _)
  }

}

验证情况 #

• 单层Map转换正常：

import ConditionValue.implicits._
val spec = Map(
    ($"column_one" === 1) -> lit(2),
    ($"column_one" === 2 && $"column_two" === 1) -> lit(1),
    ($"column_one" === 3) -> lit(4)
)
val d: ConditionValue= spec
>> d: ConditionValue= SingleLevelCaseWhen(Map((column_one = 1) -> 2, ((column_one = 2) AND (column_two = 1)) -> 1, (column_one = 3) -> 4))

• 纯嵌套Map转换正常（需显式指定类型）：

val spec = Map[Column, ConditionValue](
    ($"column_one" === 1) -> Map(
        ($"column_two" === 2) -> lit(54),
        ($"column_two" === 5) -> lit(524)
    ),
    ($"column_one" === 2) -> Map[Column, ConditionValue](
        ($"column_two" === 7) -> Map(
            ($"whatever_column" === "whatever") -> lit(12),
            ($"whatever_column" === "whatever_two") -> lit(13)
        )
    )
)
val d: ConditionValue= spec
>>d: ConditionValue= NestedCaseWhen(Map((column_one = 1) -> SingleLevelCaseWhen(Map((column_two = 2) -> 54, (column_two = 5) -> 524)), (column_one = 2) -> NestedCaseWhen(Map((column_two = 7) -> SingleLevelCaseWhen(Map((whatever_column = whatever) -> 12, (whatever_column = whatever_two) -> 13))))))

现存问题 #

• 混合Map（同时包含嵌套与单层结构）无法正常转换
• 处理嵌套Map时，必须显式指定Map类型Map[Column, ConditionValue]

临时解决方案 #

我新增了如下隐式转换，暂时解决了混合Map的转换问题，但不确定这是否是最优方案：

implicit def testVal(value: Column): ConditionValue = {
        testMap(Map(lit(true) -> value))
    }

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