Spark Codegen浅析

作者：周克勇，花名一锤，阿里巴巴计算平台事业部EMR团队技术专家，大数据领域技术爱好者，对Spark有浓厚兴趣和一定的了解，目前主要专注于EMR产品中开源计算引擎的优化工作。

背景介绍

SparkSQL的优越性能背后有两大技术支柱：Optimizer和Runtime。前者致力于寻找最优的执行计划，后者则致力于把既定的执行计划尽可能快地执行出来。Runtime的多种优化可概括为两个层面：
1. 全局优化。从提升全局资源利用率、消除数据倾斜、降低IO等角度做优化，包括自适应执行(Adaptive Execution), Shuffle Removal等。
2. 局部优化。优化具体的Task的执行效率，主要依赖Codegen技术，具体包括Expression级别和WholeStage级别的Codegen。
本文介绍Spark Codegen的技术原理。

Case Study

本节通过两个具体case介绍Codegen的做法。

Expression级别

考虑下面的表达式计算：x + (1 + 2)，用scala代码表达如下:

Add(Attribute(x), Add(Literal(1), Literal(2)))

语法树如下:

递归求值这棵语法树的常规代码如下:

tree.transformUp {
       case Attribute(idx) => Literal(row.getValue(idx))  case Add(Literal(c1),Literal(c2)) => Literal(c1+c2)  case Literal(c) => Literal(c)}

执行上述代码需要做很多类型匹配、虚函数调用、对象创建等额外逻辑，这些overhead远超对表达式求值本身。
为了消除这些overhead，Spark Codegen直接拼成求值表达式的java代码并进行即时编译。具体分为三个步骤:
1. 代码生成。根据语法树生成java代码，封装在wrapper类中:

... // class wrapperrow.getValue(idx) + (1 + 2)... // class wrapper

2. 即时编译。使用Janino框架把生成代码编译成class文件。
3. 加载执行。最后加载并执行。
优化前后性能有数量级的提升。

WholeStage级别

考虑如下的sql语句:

select count(*) from store_saleswhere ss_item_sk=1000;

生成的物理执行计划如下:

执行该计划的常规做法是使用火山模型(vocano model)，每个Operator都继承了Iterator接口，其next()方法首先驱动上游执行拿到输入，然后执行自己的逻辑。代码示例如下:

class Agg extends Iterator[Row] {
       def doAgg() {
         while (child.hasNext()) {
           val row = child.next();      // do aggregation      ...    }  }  def next(): Row {
         if (!doneAgg) {
           doAgg();    }    return aggIter.next();  }}

class Filter extends Iterator[Row] {
       def next(): Row {
         var current = child.next()    while (current != null && !predicate(current)) {
           current = child.next()    }    return current;  }}

从上述代码可知，火山模型会有大量类型转换和虚函数调用。虚函数调用会导致CPU分支预测失败，从而导致严重的性能回退。
为了消除这些overhead，Spark WholestageCodegen会为该物理计划生成类型确定的java代码，然后类似Expression的做法即时编译和加载执行。本例生成的java代码示例如下(非真实代码，真实代码片段见后文):

var count = 0for (ss_item_sk in store_sales) {
       if (ss_item_sk == 1000) {
         count += 1  }}

优化前后性能