spark sql 2.3 源码解读 - whole stage codegen (8)

最新推荐文章于 2025-10-14 19:12:00 发布

原创

最新推荐文章于 2025-10-14 19:12:00 发布 · 1.2k 阅读

本文介绍了Spark SQL 2.3中的关键特性——Whole Stage Codegen，该特性显著提高了性能。文章讨论了codegen的背景、作用，以及如何通过生成代码段来提升执行效率。内容包括WholeStageCodegenId的计数逻辑，对支持codegen的物理计划的处理，以及如何通过doProduce和doConsume函数将计划串联起来生成执行代码。最后，文章概述了代码生成和执行的过程，强调理解源码有助于掌握这一技术。

whole stage codegen 是spark 2.0 以后引入的新特性，所以在最后单独把这一块拿出来讲一下。

相关背景可以看spark官方的jira：https://issues.apache.org/jira/browse/SPARK-12795a

whole stage codegen对性能有很大的提升。

如下图所示，将一棵树翻译成了一段代码执行，性能肯定会大幅提升。

whole-stage-code-generation-model

codegen的更多原理以及测试结果：

https://databricks.com/blog/2016/05/23/apache-spark-as-a-compiler-joining-a-billion-rows-per-second-on-a-laptop.html

https://databricks.com/blog/2015/04/13/deep-dive-into-spark-sqls-catalyst-optimizer.html

Whole stage codegen是默认开启的：

`val WHOLESTAGE_CODEGEN_ENABLED = buildConf("spark.sql.codegen.wholeStage")`
    .internal()
    .doc("When true, the whole stage (of multiple operators) will be compiled into single java" +
      " method.")
    .booleanConf
    `.createWithDefault(true)`

其入口逻辑在preparations里：

protected def preparations: Seq[Rule[SparkPlan]] = Seq(
  python.ExtractPythonUDFs,
  PlanSubqueries(sparkSession),
  EnsureRequirements(sparkSession.sessionState.conf),
  CollapseCodegenStages(sparkSession.sessionState.conf),
  ReuseExchange(sparkSession.sessionState.conf),
  ReuseSubquery(sparkSession.sessionState.conf))

其中的CollapseCodegenStages是codegen优化的入口。

他的apply方法，如果开启了whole stage codegen，则执行相关的逻辑：

def apply(plan: SparkPlan): SparkPlan = {
  if (conf.wholeStageEnabled) {
    WholeStageCodegenId.resetPerQuery()
    insertWholeStageCodegen(plan)
  } else {
    plan
  }
}

WholeStageCodegenId就是一个递增的计数器，用来计数，resetPerQuery重置为1：

object WholeStageCodegenId {
  private val codegenStageCounter = ThreadLocal.withInitial(new Supplier[Integer] {
    override def get() = 1  // TODO: change to Scala lambda syntax when upgraded to Scala 2.12+
   })

  def resetPerQuery(): Unit = codegenStageCounter.set(1)

  def getNextStageId(): Int = {
    val counter = codegenStageCounter
    val id = counter.get()
    counter.set(id + 1)
    id
  }
}

还记得前面的physical plan 每一个stage前面带的数字1，2，... 5么，这个就是WholeStageCodegenId，用来将codegen生成的class和operator关联；前面的*号代表这个stage进行了codegen。可以看到Exchange是没有codegen的，因为它没有计算，只是一个shuffle过程。

*(5) Project [B#6]
+- *(5) SortMergeJoin [B#6], [B#14], Inner
   :- *(2) Sort [B#6 ASC NULLS FIRST], false, 0

最低0.47元/天解锁文章