Spark 3.x 的 WSCG 机制源码解析

前言

本文隶属于专栏《大数据技术体系》，该专栏为笔者原创，引用请注明来源，不足和错误之处请在评论区帮忙指出，谢谢！

本专栏目录结构和参考文献请见大数据技术体系

---

火山迭代模型

火山迭代模型来自论文《Volcano – An Extensible and Parallel Query Evaluation System》，它是 Goetz Graefe 的 Volcano Project 的执行器框架，其概念已被广泛接受和使用，也就是我们最为熟悉的 Volcano iterator （火山迭代）的执行框架。

该模型中的每个操作都由 3 种方法组成：

• open() - 初始化状态
• next() - 产生了一个输出
• close() - 清理状态

知道这一点，我们可以简单地翻译由filter和select操作组成的 Apache Spark 代码，例如：

scan iterator -- next() --> select iterator -- next() --> filter iterator -- next() --> output iterator

可以看到，这种迭代器模型通过每次调用 next() 方法来获取数据。

火山迭代模型虽然简单却很强大，非常灵活而具有扩展性，比如单个算子的执行逻辑完全不需要考虑其上下游是什么，也不需要考虑自身是否是并行在执行，这些逻辑都被放到了外部，而自身的策略也是注入式的，可以由外层灵活修改，整个迭代树只负责整体处理流程。

---

问题

Spark 1.x 版本的查询计划执行过程就是完全遵循着火山迭代模型，但是在使用过程中发现了很多问题：

1. 虚函数调用

在火山迭代模型中，处理一次数据最少需要调用一次next()函数。

这些函数的调用是由编译器通过虚函数调度实现的。

虽然虚函数调度是现代计算机体系结构中的重点优化部分，它仍然需要消耗很多CPU指令而且相当慢，特别是调度数达十亿次。

关于 Java 中的虚函数调用请参考我的博客——虚方法调用在Java虚拟机中的实现方式？

---

2. 内存相比于 CPU 寄存器的开销

在火山迭代模型中，每次算子将数据传递给另一个算子时，都需要将算子放入内存（函数调用堆栈）。

相比之下，在 WSCG 版本中，编译器（在这种情况下是 JVM JIT）实际上将中间数据放置在 CPU 寄存器中。

同样，CPU 访问内存中数据所需的周期数比寄存器大几个数量。

为了理解这一部分，我们下面给出了 CPU 访问不同存储数据的延迟：

ns 代表纳秒，即一秒的十亿分之一秒（1 / 1 00 000 000）。
ps代表一皮秒，是一秒的万亿分之一（1 / 1 000 000 000 000）。

火山迭代模型的数据（main memory）离寄存器（register）非常远，其访问时间会显著增加。

---

3. 循环展开和 SIMD

运行简单的循环时，现代编译器和 CPU 的效率高得令人难以置信。

编译器会自动展开简单的循环，甚至在每个 CPU 指令中产生 SIMD 指令来处理多组数据。

CPU 的特性，如管道（pipelining）、预取（prefetching）以及指令重排序（instruction reordering）使得运行简单的循环非常高效。

然而，这些编译器和 CPU 对复杂函数调用的优化极少，而这些函数正是火山迭代模型依赖的。

关于循环展开可以参考我的博客——在什么情况下循环代码会被优化？JVM 针对循环代码有哪些优化？
关于 SIMD 可以参考我的博客——即时编译器的向量化优化是什么？SIMD 到底是什么？

---

WSCG

Spark 2.x 版本使用了 WSCG（Whole Stage Code Generation，即 全阶段代码生成） 技术，它可以：

1. 消除虚函数调度。
2. 将中间数据从内存移动到 CPU 寄存器。
3. 利用现代 CPU 功能循环展开和使用 SIMD。通过向量化技术，引擎将加快对复杂操作代码生成运行的速度。对于许多数据处理的核心算子，新引擎的运行速度要提升一个数量级。

---

实践

我们可以通过 Spark 配置参数：spark.sql.codegen.wholeStage 来控制是否使用WSCG

package com.shockang.study.spark.sql.codegen

import org.apache.spark.sql.SparkSession

/**
 *
 * @author Shockang
 */
object CodeGenExample {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("CodeGenExample")
      .master("local[*]")
      // .config("spark.sql.codegen.wholeStage", "false")
      .config("spark.sql.codegen.wholeStage", "true")
      .getOrCreate()

    spark.sparkContext.setLogLevel("DEBUG")

    import spark.implicits._

    val df = Seq(("A", 1, 1), ("B", 2, 1), ("C", 3, 1), ("D", 4, 1), ("E", 5, 1))
      .toDF("letter", "nr", "a_flag")

    df.filter("letter != 'A'")
      .map(row => row.getAs[String]("letter")).count()

    spark.stop()
  }
}

---

源码下载

spark-examples 代码已开源，本项目致力于提供最具实践性的 Apache Spark 代码开发学习指南。

点击链接前往 github 下载源码：spark-examples

---

对比

无 WSCG：

有 WSCG

虽然执行非常相似，因为框架调用每个物理节点的doConsume，但实际上它不同。

区别来自生成的代码，该代码将所有操作连接在一个代码单元中（为了清晰起见，方法省略了正文）：

尽管存在这些差异，但物理计划几乎保持不变（* 表示 WSCG 标记）：

无 WSCG：

== Physical Plan ==
SerializeFromObject [staticinvoke(class org.apache.spark.unsafe.types.UTF8String, StringType, fromString, input[0, java.lang.String, true], true, false, true) AS value#19]
+- MapElements com.shockang.study.spark.sql.codegen.CodeGenExample$$$Lambda$1401/355366659@46a388cc, obj#18: java.lang.String
   +- DeserializeToObject createexternalrow(letter#10.toString, nr#11, a_flag#12, StructField(letter,StringType,true), StructField(nr,IntegerType,false), StructField(a_flag,IntegerType,false)), obj#17: org.apache.spark.sql.Row
      +- LocalTableScan [letter#10, nr#11, a_flag#12]

有 WSCG：

== Physical Plan ==
*(1) SerializeFromObject [staticinvoke(class org.apache.spark.unsafe.types.UTF8String, StringType, fromString, input[0, java.lang.String, true], true, false, true) AS value#19]
+- *(1) MapElements com.shockang.study.spark.sql.codegen.CodeGenExample$$$Lambda$1401/355366659@46a388cc, obj#18: java.lang.String
   +- *(1) DeserializeToObject createexternalrow(letter#10.toString, nr#11, a_flag#12, StructField(letter,StringType,true), StructField(nr,IntegerType,false), StructField(a_flag,IntegerType,false)), obj#17: org.apache.spark.sql.Row
      +- *(1) LocalTableScan [letter#10, nr#11, a_flag#12]

---

生成代码

相信很多同学都很好奇通过 WSCG 最终生成的代码长什么样子呢？

我们通过控制台的 DEBUG 日志可以得到下面的代码：

通过 DataFrame.explain("codegen") 也可以得到我们想要的代码。

 public Object generate(Object[] references) {
   return new GeneratedIteratorForCodegenStage2(references);
 }

 // codegenStageId=2
 final class GeneratedIteratorForCodegenStage2 extends org.apache.spark.sql.execution.BufferedRowIterator {
   private Object[] references;
   private scala.collection.Iterator[] inputs;
   private boolean hashAgg_initAgg_0;
   private boolean hashAgg_bufIsNull_0;
   private long hashAgg_bufValue_0;
   private scala.collection.Iterator inputadapter_input_0;
   private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] hashAgg_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[1];

   public GeneratedIteratorForCodegenStage2(Object[] references) {
     this.references = references;
   }

   public void init(int index, scala.collection.Iterator[] inputs) {
     partitionIndex = index;
     this.inputs = inputs;

     inputadapter_input_0 = inputs[0];
     hashAgg_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);

   }

   private void hashAgg_doAggregateWithoutKey_0() throws java.io.IOException {
     // initialize aggregation buffer
     hashAgg_bufIsNull_0 = false;
     hashAgg_bufValue_0 = 0L;

     while ( inputadapter_input_0.hasNext()) {
       InternalRow inputadapter_row_0 = (InternalRow) inputadapter_input_0.next();

       long inputadapter_value_0 = inputadapter_row_0.getLong(0);

       hashAgg_doConsume_0(inputadapter_row_0, inputadapter_value_0);
       // shouldStop check is eliminated
     }

   }

   private void hashAgg_doConsume_0(InternalRow inputadapter_row_0, long hashAgg_expr_0_0) throws java.io.IOException {
     // do aggregate
     // common sub-expressions

     // evaluate aggregate functions and update aggregation buffers

     long hashAgg_value_3 = -1L;

     hashAgg_value_3 = hashAgg_bufValue_0 + hashAgg_expr_0_0;

     hashAgg_bufIsNull_0 = false;
     hashAgg_bufValue_0 = hashAgg_value_3;

   }

   protected void processNext() throws java.io.IOException {
     while (!hashAgg_initAgg_0) {
       hashAgg_initAgg_0 = true;

       long hashAgg_beforeAgg_0 = System.nanoTime();
       hashAgg_doAggregateWithoutKey_0();
       ((org.apache.spark.sql.execution.metric.SQLMetric) references[1] /* aggTime */).add((System.nanoTime() - hashAgg_beforeAgg_0) / 1000000);

       // output the result

       ((org.apache.spark.sql.execution.metric.SQLMetric) references[0] /* numOutputRows */).add(1);
       hashAgg_mutableStateArray_0[0].reset();

       hashAgg_mutableStateArray_0[0].zeroOutNullBytes();

       hashAgg_mutableStateArray_0[0].write(0, hashAgg_bufValue_0);
       append((hashAgg_mutableStateArray_0[0].getRow()));
     }
   }

 }

---

源码解析

WSCG 的源码集中在WholeStageCodegenExec.scala这个文件中。

---

类图

---

CodegenSupport

CodegenSupport特质表示的是那些支持 CodeGen 的物理算子。这样的算子共有 10 个：

关于 CodeGen 可以参考我的博客——批处理中如何提高CPU利用率？

1. HashAggregateExec：基于哈希的聚合算子，当数据超过内存大小时，也可以回退到排序。
2. SortAggregateExec：基于排序的聚合算子
3. BroadcastHashJoinExec：执行两个子关系的 INNER HASH JOIN。构造此算子的输出 RDD 时，将异步启动 Spark 作业，以计算广播关系的值。然后将这些数据放入 Spark 广播变量中。流式关系不会被 Shuffle。
4. ShuffledHashJoinExec：首先使用 JOIN KEY Shuffle 数据，执行两个子关系的HASH JOIN。
5. SortMergeJoinExec：执行两个子关系的SORT MERGE JOIN。
6. BroadcastNestedLoopJoinExec：以广播的方式执行两个子关系的嵌套循环 JOIN。
7. RDDScanExec：用来扫描来自 InternalRow 类型 RDD 数据的物理计划节点。
8. DataSourceScanExec：用来扫描数据源的物理计划节点。
9. InMemoryTableScanExec：用来处理内存数据库关系InMemoryRelation中表的扫描操作。
10. WholeStageCodegenExec：将支持 codegen 的计划子树编译成单个 Java 方法。

其中，WholeStageCodegenExec起到了一个汇总的作用，它会将支持 codegen 的计划子树，也就是其他 9 种CodegenSupport的实现类，编译成一个 Java 方法。

我们重点看下它是如何实现的：

---

WholeStageCodegenExec

类结构

从上面的思维导图中可以看出：

WholeStageCodegenExec 这个类总共有 17 个方法/属性，其中 15 个是重写过的，这 15 个 override 方法按照来自哪个父类/接口以及异常的划分办法，划分成 5 个类型：

1. TreeNode

TreeNode 是 Spark SQL 内部所有树形结构节点的父类。

• nodeName 表示的是节点的名称。
• generateTreeString 是用来以树形结构的方法打印计划树的，我们使用TreeNode.toString方法底层就是调用的它，所以很明显，当我们想要在控制台或者日志中打印出计划树的结构时，就会用到。比如常见的Dataset.explain底层就是它。
• otherCopyArgs 表示的是那些应该直接复制不需要任何转换的，待添加到构造函数的参数。这些参数会由 makeCopy方法自动附加到转换后的参数中。

makeCopy 同样来自TreeNode，表示在转换后创建此类型树节点的副本。

• withNewChildInternal 表示我们内部在生成一个新的子节点的过程中可以顺便会对当前子节点做点事情。

2. QueryPlan

Spark SQL 查询计划树的抽象，包括常见的逻辑计划（LogicalPlan）/物理计划（SparkPlan）都是继承自这个抽象类。

这个类定义了查询计划节点的一些基本属性，以及一些新的转换 API 来转换计划节点的表达式。

• output 方法定义了查询计划的输出。

3. SparkPlan

SparkPlan是 Spark SQL 中物理算子的基类。

其中物理算子的命名约定是以“Exec”后缀结尾，例如ProjectExec。

• outputPartitioning指定如何在集群中的不同节点之间对数据进行分区。注意：如果在应用EnsureRequirements之前调用此方法，则可能会失败，因为PartitioningCollection要求其所有分区具有相同数量的分区。

关于EnsureRequirements请参考我的博客——Spark SQL 工作流程源码解析（五）planning 阶段（基于 Spark 3.3.0）

• outputOrdering 指定数据在每个分区中的排序方式。
• supportsColumnar 表示如果计划的此阶段支持列式执行，则返回true。计划还可以支持基于行的执行（参考supportsRowBased）。Spark 将决定在查询规划期间调用哪个执行。

关于列式存储/行式存储请参考我的博客——列式存储和行式存储有什么区别？

• metrics 表示的是当前SparkPlan中包含的所有度量信息。
• doExecuteColumnar表示如果supportsColumnar返回true，则以RDD[ColumnarBatch]的形式生成查询结果。默认情况下，创建ColumnarBatch的Executor负责在不再需要时关闭它。这使得输入格式能够在需要时重用批处理。

ColumnarBatch类将多个ColumnVector包装为一个基于行的表。它提供了该批处理的行视图，以便 Spark 可以逐行访问数据。它的实例将在整个数据加载过程中重用。数据源可以使用自定义逻辑扩展这个类。

ColumnVector在 Spark 中表示内存中的列数据的接口。

• doExecute会将查询结果生成为RDD[InternalRow]，这也是物理算子中最核心的方法，表示的是物理计划该怎样落地。

4. CodegenSupport

• doConsume 用来生成 Java 源代码以处理子 SparkPlan 中的数据行（Row）。只能通过 CodegenSupport.consume方法来调用。

consume方法使用当前 SparkPlan 中生成的列或行。

• needStopCheck表示：当消费输入行（Row）时，这个算子的子级是否应生成一个停止检查。在一个WholeStageCodegen循环中可以用来抑制shouldStop()。 如果一个算子启动一个新的 pipeline，这应该是错误的，这意味着它会消费子级生成的所有行，但是不通过调用append()方法将数据行输出到缓冲区，因此子级不需要在生成行的循环中使用shouldStop()。
• limitNotReachedChecks表示如果下游的Limit算子没有收到足够的记录并达到限值，则会求值为true的一系列检查。如果当前节点是数据生成节点，它可以利用此信息停止生成数据并提前完成数据流。常见的数据生成节点有Range和Scan等叶节点，以及Sort和Aggregate等阻塞节点。这些检查应放入数据生成循环的循环条件中。

5. UnsupportedOperationException

这里表示的是 WholeStageCodegenExec 不支持的操作，此处不再具体分析。

---

此外，还有 2 个非 override 方法：

doCodeGen

这个方法会为当前算子的子树生成代码。

这是实际进行代码生成操作的方法。

generatedClassName

这个方法是在上面调用 doCodeGen 方法的时候，生成类名。

这里面会涉及到一个 Spark 配置参数：spark.sql.codegen.useIdInClassName，默认值为 true。

当配置参数为 true 的时候，生成的类名为：s"GeneratedIteratorForCodegenStage$codegenStageId"格式，其中codegenStageId表示的是 WSCG 中的 Stage ID。

当配置参数为 false 的时候，类名会直接写死成GeneratedIterator。

---

针对上面类结构中涉及到的 2 个核心方法，我们将具体地进行源码分析：

doExecute

override def doExecute(): RDD[InternalRow] = {
    // 先进行代码生成
    val (ctx, cleanedSource) = doCodeGen()
    // 尝试使用 Janino 编译，如果失败的话进行回退
    val (_, compiledCodeStats) = try {
      CodeGenerator.compile(cleanedSource)
    } catch {
      case NonFatal(_) if !Utils.isTesting && conf.codegenFallback =>
        logWarning(s"Whole-stage codegen disabled for plan (id=$codegenStageId):\n $treeString")
        return child.execute()
    }

    // 检查编译代码是否有一个代码很长的函数
    if (compiledCodeStats.maxMethodCodeSize > conf.hugeMethodLimit) {
      // 发现了太长的生成代码，JIT 优化器可能不会工作：
      // 字节码尺寸超过上限，当前计划的全阶段代码生成（WSCG）被禁用。
      // 为了避免这种情况，建议设置配置参数：spark.sql.codegen.hugeMethodLimit 来提高上限。
      logInfo(s"Found too long generated codes and JIT optimization might not work: " +
        s"the bytecode size (${compiledCodeStats.maxMethodCodeSize}) is above the limit " +
        s"${conf.hugeMethodLimit}, and the whole-stage codegen was disabled " +
        s"for this plan (id=$codegenStageId). To avoid this, you can raise the limit " +
        s"`${SQLConf.WHOLESTAGE_HUGE_METHOD_LIMIT.key}`:\n$treeString")
      return child.execute()
    }

    // 保存可用于传递到生成类中的对象数组。
    val references = ctx.references.toArray

    // pipeline 的时间度量信息
    val durationMs = longMetric("pipelineTime")

    // 尽管 rdds 是一个`RDD[InternalRow]`，但它实际上可能是一个类型擦除的`RDD[ColumnarBatch]`。
    //  这允许代码生成阶段的输入是列式结构的，但是输出必须是行。
    val rdds = child.asInstanceOf[CodegenSupport].inputRDDs()
    // 最多支持 2 个输入的 RDD
    assert(rdds.size <= 2, "Up to two input RDDs can be supported")
    if (rdds.length == 1) {
      rdds.head.mapPartitionsWithIndex { (index, iter) =>
        val (clazz, _) = CodeGenerator.compile(cleanedSource)
        val buffer = clazz.generate(references).asInstanceOf[BufferedRowIterator]
        buffer.init(index, Array(iter))
        new Iterator[InternalRow] {
          override def hasNext: Boolean = {
            val v = buffer.hasNext
            if (!v) durationMs += buffer.durationMs()
            v
          }
          override def next: InternalRow = buffer.next()
        }
      }
    } else {
      // 当前，我们最多支持 2 个输入的 RDD
      rdds.head.zipPartitions(rdds(1)) { (leftIter, rightIter) =>
        Iterator((leftIter, rightIter))
        // 一个小的修改程序来获取正确的分区索引
      }.mapPartitionsWithIndex { (index, zippedIter) =>
        val (leftIter, rightIter) = zippedIter.next()
        val (clazz, _) = CodeGenerator.compile(cleanedSource)
        val buffer = clazz.generate(references).asInstanceOf[BufferedRowIterator]
        buffer.init(index, Array(leftIter, rightIter))
        new Iterator[InternalRow] {
          override def hasNext: Boolean = {
            val v = buffer.hasNext
            if (!v) durationMs += buffer.durationMs()
            v
          }
          override def next: InternalRow = buffer.next()
        }
      }
    }
  }

可以看到，doExecute方法的核心步骤就 2 步:

1. 进行代码生成，即调用doCodeGen。
2. 使用 Janino 将 Java 源码编译成 Java 类。

---

doCodeGen

  def doCodeGen(): (CodegenContext, CodeAndComment) = {
    // 代码生成的开始时间
    val startTime = System.nanoTime()
    // 代码生成的上下文
    val ctx = new CodegenContext
    // 返回 Java 源码，来处理输入 RDD 的数据行
    // 此处会触发 SparkPlan 的执行查询逻辑
    val code = child.asInstanceOf[CodegenSupport].produce(ctx, this)

    // 主要的 next 函数
    ctx.addNewFunction("processNext",
      s"""
        protected void processNext() throws java.io.IOException {
          ${code.trim}
        }
       """, inlineToOuterClass = true)

    // 生成类名
    val className = generatedClassName()

    // 通过字符串插值的方式来定义好源码。
    // 可以看到包含 generate 方法以及一个 BufferedRowIterator 的子类。
    // 子类中定义了类变量、构造器和 init 方法。
    // 此外，emitExtraCode 方法用来添加额外的内部类
    // declareAddedFunctions 方法会声明所有函数代码。如果添加的函数太多，要将它们拆分为嵌套的子类，以避免遇到 Java 编译器常量池限制。
    val source = s"""
      public Object generate(Object[] references) {
        return new $className(references);
      }

      ${ctx.registerComment(
        s"""Codegened pipeline for stage (id=$codegenStageId) |${this.treeString.trim}""".stripMargin,
         "wsc_codegenPipeline")}
      ${ctx.registerComment(s"codegenStageId=$codegenStageId", "wsc_codegenStageId", true)}
      final class $className extends ${classOf[BufferedRowIterator].getName} {

        private Object[] references;
        private scala.collection.Iterator[] inputs;
        ${ctx.declareMutableStates()}

        public $className(Object[] references) {
          this.references = references;
        }

        public void init(int index, scala.collection.Iterator[] inputs) {
          partitionIndex = index;
          this.inputs = inputs;
          ${ctx.initMutableStates()}
          ${ctx.initPartition()}
        }

        ${ctx.emitExtraCode()}

        ${ctx.declareAddedFunctions()}
      }
      """.trim

    // 这就是我们最终要返回的源代码
    val cleanedSource = CodeFormatter.stripOverlappingComments(
      new CodeAndComment(CodeFormatter.stripExtraNewLines(source), ctx.getPlaceHolderToComments()))

    val duration = System.nanoTime() - startTime
    // 此处通过原子变量的方法来保存代码生成的时间
    WholeStageCodegenExec.increaseCodeGenTime(duration)

    logDebug(s"\n${CodeFormatter.format(cleanedSource)}")
    (ctx, cleanedSource)
  }

代码生成的逻辑也比较清晰，就是利用 Scala 的字符串插值方式，先定义好模板，然后根据代码生成上下文CodegenContext中的数据来填充，最终形成 Java 源代码。

建议对比上面实践环节最终生成的源代码，相信你能对代码生成的逻辑有更清晰的认识。