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

最新推荐文章于 2022-09-20 12:29:26 发布

sddyljsx

最新推荐文章于 2022-09-20 12:29:26 发布

阅读量3.9k

点赞数

分类专栏： Spark Spark Sql 2.3源码解读文章标签： spark sql big data

本文链接：https://blog.csdn.net/sddyljsx/article/details/81784911

版权

本文介绍了Spark SQL 2.3中的Whole Stage Codegen特性，它显著提升了性能。文章详细阐述了codegen的入口、优化逻辑、Physical Plan中的应用，以及如何通过实现CodegenSupport接口使计划支持codegen。通过SortExec的例子展示了codegen的工作流程，解释了doProduce和doConsume的角色，并简述了WholeStageCodegenExec的执行过程。通过对源码的深入解读，帮助读者理解codegen的实现原理。

摘要由CSDN通过智能技术生成

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
   :  +- Exchange(coordinator id: 1121577170) hashpartitioning(B#6, 200), coordinator[target post-shuffle partition size: 67108864]
   :     +- *(1) Project [B#6]
   :        +- *(1) Filter isnotnull(B#6)
   :           +- *(1) FileScan json [B#6] Batched: false, Format: JSON, Location: InMemoryFileIndex[file:examples/src/main/resources/test.json], PartitionFilters: [], PushedFilters: [IsNotNull(B)], ReadSchema