java case when多分支判断_Apache Spark之单分支CaseWhen优化为IF

前几天在客户环境遇到一个Spark “CASE WHEN”语句的性能优化问题。

客户那边通过一个“时间范围筛选”控件来动态修改图表的数据。其很多指标的计算逻辑类似于：

CASE

WHEN `bizdate`

BETWEEN '2020-09-06' AND '2020-09-13'

THEN `sales_amount`

ELSE 0

END

CASE WHEN语句有些类似于编程语言中的Switch语句，当这里的 WHEN从句只有一个的时候，可以简化为IF语句(或者 IF-ELSE 语句)。

于是想：对于Spark(客户用的是2.4.x版本)， Spark会不会把这种只有一个WHEN分支的 CASE WHEN 语句优化为IF语句呢？ 于是试了一下性能，发现如果修改上面的SQL为：

IF(`bizdate`

BETWEEN '2020-09-06' AND '2020-09-13',

`sales_amount`,

0

)

那么执行速度将减少为原来的一半！ 原来Spark 2.4并没有做这个优化，突然感到有些兴奋。那是不是我的机会来了。

首先、这个应该是一个比较简单的优化，比如我是否可以通过增加一个Spark的优化器规则，来自动把一个分支的CASE WHEN转为IF，看着好像不难。

不过在真正动手前，先看看： Spark的最新版本是否已经有了这个修改？

于是先看看 CASE When 语句的实现，发现最新的发布版本(Spark 3.0.1)的代码是这样的：

override def doGenCode(ctx: CodegenContext,

ev: ExprCode): ExprCode = {

if (branches.length == 1) {

// If we have only single branch we can use If expression and its codeGen If(

branches(0)._1,

branches(0)._2,

elseValue

.getOrElse(Literal.create(null, branches(0)._2.dataType)))

.doGenCode(ctx, ev)

} else {

multiBranchesCodegen(ctx, ev)

}

}

发现，在Spark转化执行代码为 Java时(doGenCode)，其已经对于分支为1的情况，做了自动转化为 IF 语句的操作。

虽然我感觉更适合放在优化器中做，不过直接修改 CaseWhen 这个类的 doGenCode() 可能简单直接！ 从这个修改的PR的diff来看也确实如此(只改了几行代码)：

首先，发现其 “Fix Version/s: 3.0.0”，果然是3.0才优化的。

其描述问题时的重现步骤：

val df = spark.range(10000000000L).withColumn("x", rand)

val resultA = df.withColumn("r", when($"x" < 0.5, lit(1)).otherwise(lit(0))).agg(sum($"r"))

val resultB = df.withColumn("r", expr("if(x < 0.5, 1, 0)")).agg(sum($"r"))

resultA.collect() // takes 56s to finishresultB.collect() // takes 30s to finish

发现其在spark旧版本中 IF 比 CaseWhen 要快很多 (30秒 vs 56秒)

虽然没有为Spark贡献成，但是也了解到了Spark 3.0的一些细节优化已经可以解决现在的一些实际问题了，Spark 3.0.1 值得期待应用到产品中！