Spark SQL AQE工作原理源码剖析

前言

在把SparkPlan提交执行前，还需要完成若干的准备工作，对树形结构的物理计划进行全局的整合处理或者优化，通过一系列的规则序列，如AQE、EnsureRequirements等，保证SparkPlan的正确性和高效性。

整体执行逻辑如下：

org/apache/spark/sql/execution/QueryExecution

scala 复制代码

  lazy val executedPlan: SparkPlan = {
    // We need to materialize the optimizedPlan here, before tracking the planning phase, to ensure
    // that the optimization time is not counted as part of the planning phase.
    assertOptimized() // 生成优化物理执行计划
    executePhase(QueryPlanningTracker.PLANNING) {
      // clone the plan to avoid sharing the plan instance between different stages like analyzing,
      // optimizing and planning.
      QueryExecution.prepareForExecution(preparations, sparkPlan.clone()) // 1.生成优化规则；2.对物理计划树运用这些规则进行优化
    }
  }

// 1.生成规则  
protected def preparations: Seq[Rule[SparkPlan]] = {
    QueryExecution.preparations(sparkSession,
      Option(InsertAdaptiveSparkPlan(AdaptiveExecutionContext(sparkSession, this))))
  }

// 2.运用规则
private[execution] def prepareForExecution(
      preparations: Seq[Rule[SparkPlan]],
      plan: SparkPlan): SparkPlan = {
    val planChangeLogger = new PlanChangeLogger[SparkPlan]()
    val preparedPlan = preparations.foldLeft(plan) { case (sp, rule) =>
      val result = rule.apply(sp)
      planChangeLogger.logRule(rule.ruleName, sp, result)
      result
    }
    planChangeLogger.logBatch("Preparations", plan, preparedPlan)
    preparedPlan
  }

生成的 execution preparation rules 规则如下：

scala 复制代码

private[execution] def preparations(
    sparkSession: SparkSession,
    adaptiveExecutionRule: Option[InsertAdaptiveSparkPlan] = None): Seq[Rule[SparkPlan]] = {
  adaptiveExecutionRule.toSeq ++
  Seq(
    CoalesceBucketsInJoin,
    PlanDynamicPruningFilters(sparkSession),
    PlanSubqueries(sparkSession),
    RemoveRedundantProjects,
    EnsureRequirements,
    RemoveRedundantSorts,
    DisableUnnecessaryBucketedScan,
    ApplyColumnarRulesAndInsertTransitions(sparkSession.sessionState.columnarRules),
    CollapseCodegenStages(),
    ReuseExchange,
    ReuseSubquery
  )
}

这组规则的作用是把已经生成的物理计划进一步整理成可执行计划，其中adaptiveExecutionRule决定AQE 是否开启，AQE是否开启将影响最终的execution preparation rules ，本文接下来将基于Spark 3.1.2 源码深入分析Spark SQL AQE的执行路径与生效机制。

AQE简介

首先介绍一下AQE的产生背景。

在 Spark 3.0 之前，Spark SQL 的优化器核心是 Catalyst。Catalyst 是一个优秀的静态优化器：它在任务真正提交到集群运行之前，会根据表结构、过滤条件和基于成本的优化（CBO）生成物理执行计划。但这会存在一个问题，物理计划是在运行之前就已经制定好，无法在运行过程中根据实际的结果信息进行动态调整，难免会"闭门造车"。

在这种背景下，AQE出现了，是 Spark SQL 的一种动态优化机制。AQE 全称是 Adaptive Query Execution，即自适应查询执行。

自适应查询简单理解就是当 Shuffle Map 阶段执行完成的时候，AQE 结合这个阶段的统计信息，基于既定的规则动态地来调整和修正还没有执行的逻辑/物理计划，从而实现对原始查询语句运行时优化的目的。

AQE 主要提供了三大核心优化策略：① 动态合并 Shuffle 分区；② 动态优化倾斜 Join；③ 动态切换 Join 策略。

第一部分：Spark AQE 开启与不开启时的执行路径对比

1. AQE 未开启时的执行路径

当 AQE 没有开启时，插入的InsertAdaptiveSparkPlan规则内部不会对原计划做任何处理（等同于不生效），最终 preparation rules 就是普通规则序列：

text 复制代码

CoalesceBucketsInJoin
PlanDynamicPruningFilters
PlanSubqueries
RemoveRedundantProjects
EnsureRequirements
RemoveRedundantSorts
DisableUnnecessaryBucketedScan
ApplyColumnarRulesAndInsertTransitions
CollapseCodegenStages
ReuseExchange
ReuseSubquery

这种模式下，物理计划在执行前基本已经固定。Spark 会依赖优化阶段和 preparation 阶段已有的信息来决定：

是否插入 shuffle；
是否插入 sort；
是否使用 bucket join；
是否使用 broadcast join；
是否复用 exchange；
是否插入 codegen stage。

一旦计划进入执行阶段，普通情况下不会再根据运行时真实数据量改变 join 策略或 shuffle 分区数量。

2. AQE 开启时的 preparations 定义

开启 AQE 后，插入的InsertAdaptiveSparkPlan规则内部会对原计划做一些处理，最终 rules 顺序变成：

text 复制代码

InsertAdaptiveSparkPlan
CoalesceBucketsInJoin
PlanDynamicPruningFilters
PlanSubqueries
RemoveRedundantProjects
EnsureRequirements
RemoveRedundantSorts
DisableUnnecessaryBucketedScan
ApplyColumnarRulesAndInsertTransitions
CollapseCodegenStages
ReuseExchange
ReuseSubquery

也就是说，InsertAdaptiveSparkPlan 会被放在所有普通 preparation rules 之前。那InsertAdaptiveSparkPlan 规则将如何影响exeution sparkplan产生以及AQE配置如何开启呢？下面将详细介绍。

3. InsertAdaptiveSparkPlan 的作用

3.1 applyInternal 方法概况

applyInternal 是 InsertAdaptiveSparkPlan 的核心方法。它的职责不是执行 AQE 优化，而是判断当前物理计划是否应该进入 AQE，并在满足条件时把原始 SparkPlan 包装成 AdaptiveSparkPlanExec。

org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan

scala 复制代码

private def applyInternal(plan: SparkPlan, isSubquery: Boolean): SparkPlan = plan match {
  case _ if !conf.adaptiveExecutionEnabled => plan
  case _: ExecutedCommandExec => plan
  case c: DataWritingCommandExec => c.copy(child = apply(c.child))
  case c: V2CommandExec => c.withNewChildren(c.children.map(apply))
  case _ if shouldApplyAQE(plan, isSubquery) =>
    if (supportAdaptive(plan)) {
      try {
        val subqueryMap = buildSubqueryMap(plan)
        val planSubqueriesRule = PlanAdaptiveSubqueries(subqueryMap)
        val preprocessingRules = Seq(planSubqueriesRule)
        val newPlan = AdaptiveSparkPlanExec.applyPhysicalRules(plan, preprocessingRules)
        logDebug(s"Adaptive execution enabled for plan: $plan")
        AdaptiveSparkPlanExec(newPlan, adaptiveExecutionContext, preprocessingRules, isSubquery)
      } catch {
        case SubqueryAdaptiveNotSupportedException(subquery) =>
          logWarning(s"${SQLConf.ADAPTIVE_EXECUTION_ENABLED.key} is enabled " +
            s"but is not supported for sub-query: $subquery.")
          plan
      }
    } else {
      logWarning(s"${SQLConf.ADAPTIVE_EXECUTION_ENABLED.key} is enabled " +
        s"but is not supported for query: $plan.")
      plan
    }

  case _ => plan
}

整体判断链路可以概括为：

text 复制代码

applyInternal(plan, isSubquery)
  ↓
AQE 总开关是否开启？
  否：返回原计划
  是：
    是否是命令类计划？
      是：原样返回，或只递归处理查询 child
      否：
        当前计划是否值得应用 AQE？
          否：返回原计划
          是：
            当前计划是否支持 AQE？
              否：打印 warning，返回原计划
              是：
                处理 adaptive subquery
                应用预处理规则
                返回 AdaptiveSparkPlanExec

3.2 AQE 总开关判断

第一条分支：

scala 复制代码

case _ if !conf.adaptiveExecutionEnabled => plan

如果配置：spark.sql.adaptive.enabled = false，则 InsertAdaptiveSparkPlan 不做任何处理，直接返回原始 SparkPlan。因此如果要AQE生效，必须设置**spark.sql.adaptive.enabled = true**。

3.3 shouldApplyAQE：判断是否值得启用 AQE

主分支是：

scala 复制代码

case _ if shouldApplyAQE(plan, isSubquery) =>

shouldApplyAQE 判断的是：当前计划是否有必要进入 AQE。

源码如下：

scala 复制代码

private def shouldApplyAQE(plan: SparkPlan, isSubquery: Boolean): Boolean = {
  conf.getConf(SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY) || isSubquery || {
    plan.find {
      case _: Exchange => true
      case p if !p.requiredChildDistribution.forall(_ == UnspecifiedDistribution) => true
      case p => p.expressions.exists(_.find {
        case _: SubqueryExpression => true
        case _ => false
      }.isDefined)
    }.isDefined
  }
}

它满足以下任一条件就返回 true。

① 配置强制启用 AQE：conf.getConf(SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY)【对应参数spark.sql.adaptive.forceApply】

这种情况下，即使计划当前看起来没有明显的 shuffle 或 subquery，也会尝试套 AQE。

② 当前计划来自子查询：isSubquery

如果主查询已经决定进入 AQE，那么子查询也要继续沿用 AQE 处理路径。

③ 计划中已经存在 Exchange：

scala 复制代码

case _: Exchange => true

Exchange 是 AQE 的核心边界，包括 shuffle 和 broadcast。只要计划里已有 Exchange，AQE 就可以把它包装成 query stage，执行后收集运行时统计信息。

④ 计划后续可能需要插入 Exchange：

scala 复制代码

case p if !p.requiredChildDistribution.forall(_ == UnspecifiedDistribution) => true

有些算子当前还没有 Exchange，但它声明了对子节点分布的要求，例如 join、aggregate、window 等。后续 EnsureRequirements 很可能会为这些要求插入 ShuffleExchangeExec。

⑤ 计划中包含子查询表达式：

scala 复制代码

case p => p.expressions.exists(_.find {
  case _: SubqueryExpression => true
  case _ => false
}.isDefined)

如果表达式里有 SubqueryExpression，AQE 需要递归规划这些子查询。

如果以上条件都不满足，shouldApplyAQE 返回 false，applyInternal 最终走默认分支：

scala 复制代码

case _ => plan

也就是不套 AdaptiveSparkPlanExec。

3.4 supportAdaptive：判断是否安全支持 AQE

shouldApplyAQE 表示"有没有必要"，但还需要 supportAdaptive 判断"能不能安全使用"。

源码如下：

scala 复制代码

private def supportAdaptive(plan: SparkPlan): Boolean = {
  sanityCheck(plan) &&
    !plan.logicalLink.exists(_.isStreaming) &&
    !plan.expressions.exists(_.find(_.isInstanceOf[DynamicPruningSubquery]).isDefined) &&
    plan.children.forall(supportAdaptive)
}

它主要检查四类条件。

① 基础合法性检查：

scala 复制代码

sanityCheck(plan)

这通常用于判断当前物理计划是否满足 AQE 的基本约束。

② 不能是 streaming 查询：

scala 复制代码

!plan.logicalLink.exists(_.isStreaming)

AQE 主要面向 batch query。流式查询的执行模型不同，不适合使用这种 query stage 级别的运行时重优化。

③ 不能包含暂不支持的动态分区裁剪子查询：

scala 复制代码

!plan.expressions.exists(_.find(_.isInstanceOf[DynamicPruningSubquery]).isDefined)

源码注释中通常会提到动态分区裁剪还需要迁移到 adaptive execution 中，因此这里遇到 DynamicPruningSubquery 会认为不支持 AQE。

④ 所有子节点也必须支持 AQE：

scala 复制代码

plan.children.forall(supportAdaptive)

这是递归检查。只要某个 child 不支持 AQE，整个计划就不能直接包装成 AdaptiveSparkPlanExec。

3.5 创建 AdaptiveSparkPlanExec

最终返回：AdaptiveSparkPlanExec(newPlan, adaptiveExecutionContext, preprocessingRules, isSubquery)

这一步表示 AQE 正式接管当前物理计划。外层计划从：SparkPlan变成：AdaptiveSparkPlanExec。

AdaptiveSparkPlanExec 是一个叶子节点。一旦它被插入，原始计划就被隐藏在 AdaptiveSparkPlanExec 内部。

scala 复制代码

case class AdaptiveSparkPlanExec(
    inputPlan: SparkPlan,
    @transient context: AdaptiveExecutionContext,
    @transient preprocessingRules: Seq[Rule[SparkPlan]],
    @transient isSubquery: Boolean)
  extends LeafExecNode

此时后面的普通 preparation rules 虽然仍然会依次调用：

text 复制代码

CoalesceBucketsInJoin
PlanDynamicPruningFilters
PlanSubqueries
RemoveRedundantProjects
EnsureRequirements
RemoveRedundantSorts
DisableUnnecessaryBucketedScan
ApplyColumnarRulesAndInsertTransitions
CollapseCodegenStages
ReuseExchange
ReuseSubquery

但它们看到的已经不是原始完整计划树，而只是外层的：AdaptiveSparkPlanExec。而AdaptiveSparkPlanExec 又是一个叶子节点，所以大多数规则无法遍历到内部真实计划，即使被调度调用，但通常不会再改写原始物理计划，这就是源码里说的 no-op，表示"没有实际操作"或"没有产生效果"。

真正的 AQE 逻辑会在 AdaptiveSparkPlanExec 内部继续执行，后续将详细介绍。

4. AQE 开启后的实际执行路径

AQE 开启后，外层 preparation 路径可以概括为：

text 复制代码

原始 SparkPlan
  ↓
InsertAdaptiveSparkPlan
  ↓
AdaptiveSparkPlanExec
  ↓
后续普通 preparation rules 被调用，但多数 no-op
  ↓
执行 AdaptiveSparkPlanExec

5. 总结

spark.sql.adaptive.enabled参数控制AQE的开启（默认是关闭的）；

① AQE 未开启时，Spark 会让普通 preparation rules 直接作用于原始 SparkPlan，然后生成一个执行前基本固定的最终物理计划。

② AQE 开启时，InsertAdaptiveSparkPlan 规则会优先执行，它用一个AdaptiveSparkPlanExec实例来包装查询计划，后续普通 preparation rules 虽然仍会被调用，但由于原始计划已经被隐藏在 AdaptiveSparkPlanExec 内部，这些规则通常会变成 no-op。

那么当开启AQE时，目前看在外层preparation rules中的规则是没有生效的，是否会在其他地方生效呢？以及AQE 如何进行自适应优化的？ AQE优化有哪些策略？一切答案都藏在 AdaptiveSparkPlanExec 中。

第二部分：AdaptiveSparkPlanExec 内部执行路径详解

第一部分总结了 AQE 开启与不开启时的整体差异。本节进一步展开 AQE 开启后，AdaptiveSparkPlanExec 内部到底如何执行。

开启 AQE 后，外层 InsertAdaptiveSparkPlan 会把原始物理计划包装成：AdaptiveSparkPlanExec(inputPlan)。

整体路径可以概括为：

text 复制代码

1. QueryExecution.preparations
   ↓
2. InsertAdaptiveSparkPlan
   ↓
3. AdaptiveSparkPlanExec(inputPlan)
   ↓
4. AQE 内部 queryStagePreparationRules
   ↓
5. EnsureRequirements 插入 ShuffleExchangeExec / BroadcastExchangeExec
   ↓
6. getFinalPhysicalPlan()
   ↓
    7. createQueryStages(currentPhysicalPlan)
       ↓
        8. Exchange -> QueryStageExec
              ShuffleExchangeExec -> ShuffleQueryStageExec
              BroadcastExchangeExec -> BroadcastQueryStageExec
           ↓
    9. stage.materialize()
       ↓
    10. StageSuccess 后写入 resultOption
       ↓
    11. replaceWithQueryStagesInLogicalPlan 重新生成 logical plan
       ↓
    12. reOptimize 重新生成 physical plan
       ↓
    13. costEvaluator 比较新旧计划
       ↓
    14. 如果采用新计划，再次 createQueryStages
   ↓
15. 循环直到所有 stage 物化完成
   ↓
16. finalStageOptimizerRules
   ↓
17. 返回 finalPhysicalPlan

1. currentPhysicalPlan变量

AdaptiveSparkPlanExec类包含一个变量currentPhysicalPlan，这是AQE 内部维护一个当前物理计划。

scala 复制代码

@volatile private var currentPhysicalPlan = initialPlan

这个计划通常会经过 AQE 内部 preparation rules 处理：

scala 复制代码

@transient private val initialPlan = context.session.withActive {
    applyPhysicalRules(
      inputPlan, queryStagePreparationRules, Some((planChangeLogger, "AQE Preparations")))
  }

queryStagePreparationRules 中会包含 EnsureRequirements 这类规则，用来保证物理算子的分区和排序要求。

scala 复制代码

  private def queryStagePreparationRules: Seq[Rule[SparkPlan]] = Seq(
    RemoveRedundantProjects,
    EnsureRequirements,
    RemoveRedundantSorts,
    DisableUnnecessaryBucketedScan
  ) ++ context.session.sessionState.queryStagePrepRules

这意味着在AdaptiveSparkPlanExec 内部就会执行一些规则，把物理计划准备成满足执行要求的形态（如插入 ShuffleExchangeExec），而外层QueryExecution.preparations中的多数规则是 no-op，因此真正的计划准备是在AQE内部完成。

2. getFinalPhysicalPlan()方法

AdaptiveSparkPlanExec的执行入口是 doExecute()

scala 复制代码

override def doExecute(): RDD[InternalRow] = {
    val rdd = getFinalPhysicalPlan().execute()
    finalPlanUpdate
    rdd
}

核心逻辑是getFinalPhysicalPlan() ，它负责把当前 AQE 物理计划推进到最终物理计划。

org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec

scala 复制代码

private def getFinalPhysicalPlan(): SparkPlan = lock.synchronized {
  if (isFinalPlan) return currentPhysicalPlan

  context.session.withActive {

    var currentLogicalPlan = currentPhysicalPlan.logicalLink.get
    var result = createQueryStages(currentPhysicalPlan) // 把 Exchange 包装成 QueryStageExec

    while (!result.allChildStagesMaterialized) { // 只要当前计划中还有未物化完成的 child query stage，AQE 主循环就继续执行。
      ...
    }

    currentPhysicalPlan = applyPhysicalRules(
      result.newPlan,
      finalStageOptimizerRules,
      Some((planChangeLogger, "AQE Final Query Stage Optimization")))
    isFinalPlan = true
    currentPhysicalPlan
  }
}

2.1 第一次创建 query stage

scala 复制代码

var result = createQueryStages(currentPhysicalPlan)

这里会扫描当前物理计划，把已经存在的 Exchange 转换成 QueryStageExec。

text 复制代码

ShuffleExchangeExec -> ShuffleQueryStageExec
BroadcastExchangeExec -> BroadcastQueryStageExec

2.2 createQueryStages 的递归逻辑

createQueryStages 的源码注释说明它会自底向上遍历计划树：

scala 复制代码

private def createQueryStages(plan: SparkPlan): CreateStageResult = plan match {
  case e: Exchange =>
    ...
  case q: QueryStageExec =>
    ...
  case _ =>
    ...
}

2.2.1 为什么要 bottom-up

AQE 必须先处理子 stage，再处理父 stage，例如：

text 复制代码

Exchange
  +- Join
      :- Exchange
      +- Exchange

外层 Exchange 不能马上变成 query stage，因为它的 child 里还有两个内层 Exchange 没有完成。只有内层 stage 执行完成，AQE 拿到统计信息并完成可能的重优化后，父 stage 的输入才稳定。

所以 createQueryStages 是自底向上的。

2.2.2 遇到 Exchange

1、cache 未命中：递归处理 child

scala 复制代码

val result = createQueryStages(e.child)

如果没有可复用 stage，会先处理 Exchange 的 child，如果 child 里面还有更底层的 exchange，会先创建更底层的 query stage。

scala 复制代码

val newPlan = e.withNewChildren(Seq(result.newPlan)).asInstanceOf[Exchange]

递归处理 child 后，child 的执行计划可能已经改变（创建了 createQueryStages ），因此，用处理后的 child 重新替换旧的child构造一个新的Exchange 节点。

例如：

复制代码

Exchange_outer
      :- Exchange_left
      +- Exchange_right

第一步递归 child 后：

复制代码

Exchange_outer
      :- ShuffleQueryStageExec(left)
      +- ShuffleQueryStageExec(right)

这里的 Exchange_outer 还是原来的外层 exchange，只是它的 child 已经更新了。

2、child stage 全部物化完成后才创建当前 stage

scala 复制代码

if (result.allChildStagesMaterialized) {
  var newStage = newQueryStage(newPlan) // 只有 child query stage 都完成，当前 Exchange 才能被包装成 QueryStageExec。
  ...
   val isMaterialized = newStage.resultOption.get().isDefined // 当前stage是否物化（执行完成）
   CreateStageResult(
        newPlan = newStage,
        allChildStagesMaterialized = isMaterialized,
        newStages = if (isMaterialized) Seq.empty else Seq(newStage))
}

child stage 全部物化完成一般是最里层的exchange会发生。

注：对于最里层 Exchange 一开始是还没有物化的，但第一轮会被包装成 QueryStageExec，并加入 newStages，等待 AQE 主循环的stage.materialize()

3、child stage 未全部物化

scala 复制代码

CreateStageResult(newPlan = newPlan,allChildStagesMaterialized = false, newStages = result.newStages)

child stage 未全部物化一般是对于外层exchange，因为最里层的 Exchange 一开始是还没有物化的，因此会走到这个分支，结果如下：

复制代码

newPlan = 外层 Exchange，child 已经替换为包含内层 QueryStageExec 的 Join
allChildStagesMaterialized = false
newStages = 内层新创建的 stage

只有在AQE getFinalPhysicalPlan() 主循环中把内部 query stage 全部物化完成后*，*下一轮 createQueryStages 才会把外层 Exchange 包装成 QueryStageExec。【动态运行，是AQE自适应优化的关键】

2.2.3 普通节点递归处理 children

普通非 exchange 节点，例如 FilterExec、ProjectExec、SortExec、JoinExec，不会成为 query stage 边界，它们只递归处理 children，并把 children 的处理结果汇总回来：

scala 复制代码

if (plan.children.isEmpty) {// 1、叶子节点
    CreateStageResult(newPlan = plan, allChildStagesMaterialized = true, newStages = Seq.empty)
} else { // 2、非叶子节点
    val results = plan.children.map(createQueryStages) // 对所有 child 递归调用 createQueryStages
    CreateStageResult(
        newPlan = plan.withNewChildren(results.map(_.newPlan)),
        allChildStagesMaterialized = results.forall(_.allChildStagesMaterialized),
        newStages = results.flatMap(_.newStages))
}

① 叶子节点

text 复制代码

没有 children，
下面也没有需要等待物化的 child stage，所以 allChildStagesMaterialized = true
也不是 Exchange ，不会创建新的 query stage。

② 普通非叶子节点

复制代码

  自己不创建 query stage。
  递归调用 createQueryStages 处理所有 children。
  用处理后的 children 重建当前 plan。
  再把 children 返回的物化状态和新增 stage 汇总后向上传递。

因此，普通节点在 createQueryStages 中承担的是"递归透传和汇总"的角色：

text 复制代码

不切 stage，
不物化自己，
只负责把子树中的 QueryStageExec 创建结果和完成状态向上传递。

2.2.4 返回结果：`CreateStageResult`

CreateStageResult 是 createQueryStages(...) 递归处理物理计划树时返回的结果对象，定义如下：

scala 复制代码

CreateStageResult(
  newPlan: SparkPlan,
  allChildStagesMaterialized: Boolean,
  newStages: Seq[QueryStageExec])

1. newPlan

表示当前子树经过 createQueryStages 处理后的新物理计划。比如原来：

复制代码

SortMergeJoinExec
:- ShuffleExchangeExec
+- ShuffleExchangeExec

处理后可能变成：

复制代码

SortMergeJoinExec
:- ShuffleQueryStageExec
+- ShuffleQueryStageExec

这个替换后的子树就是 newPlan。

2.allChildStagesMaterialized

表示当前子树下面依赖的所有 child query stage 是否都已经物化完成，也就是当前节点往下看所有需要先执行的 QueryStageExec 是否都有结果了。

例如：

复制代码

SortMergeJoinExec
:- ShuffleQueryStageExec(left, 已完成)
+- ShuffleQueryStageExec(right, 未完成)

那么这个 SortMergeJoinExec 对应的：allChildStagesMaterialized = false；如果左右都完成，才是：allChildStagesMaterialized = true；

这样设计的原因是 AQE 必须保证父 stage 的输入已经稳定。如果子 stage 还没有执行完成，AQE 还拿不到 runtime statistics 无法进行优化。

3. newStages

表示本轮递归中新创建出来、需要提交执行的 QueryStageExec。

比如第一次扫描到ShuffleExchangeExec，并把它包装成ShuffleQueryStageExec，那么这个新的 ShuffleQueryStageExec 就会放进：newStages。

总结来说，CreateStageResult 可理解为：AQE 在扫描一棵 SparkPlan 子树后，返回"当前子树被改写成什么样了、子 stage 是否都完成了、本轮新创建了哪些 stage"。

2.3 newQueryStage()方法

newQueryStage 负责把一个具体的 Exchange 包装成 QueryStageExec。源码结构类似：

scala 复制代码

private def newQueryStage(e: Exchange): QueryStageExec = {
  val optimizedPlan = applyPhysicalRules(
    e.child,
    queryStageOptimizerRules,
    Some((planChangeLogger, "AQE Query Stage Optimization")))

  val queryStage = e match {
    case s: ShuffleExchangeLike =>
      ...
      ShuffleQueryStageExec(currentStageId, newShuffle)

    case b: BroadcastExchangeLike =>
      ...
      BroadcastQueryStageExec(currentStageId, newBroadcast)
  }

  currentStageId += 1
  setLogicalLinkForNewQueryStage(queryStage, e)
  queryStage
}

1、先优化 Exchange 的 child

scala 复制代码

val optimizedPlan = applyPhysicalRules(
  e.child,
  queryStageOptimizerRules,
  Some((planChangeLogger, "AQE Query Stage Optimization")))

创建 stage 前，AQE 会先对 Exchange 的 child 应用 query stage 内部优化规则。【动态合并 Shuffle 分区、动态优化倾斜 Join策略】

scala 复制代码

  @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq(
    ReuseAdaptiveSubquery(context.subqueryCache),
    CoalesceShufflePartitions(context.session), // 合并shuffle 分区
    // The following two rules need to make use of 'CustomShuffleReaderExec.partitionSpecs'
    // added by `CoalesceShufflePartitions`. So they must be executed after it.
    OptimizeSkewedJoin, // join倾斜优化
    OptimizeLocalShuffleReader
  )

2 ShuffleExchangeLike 变成 ShuffleQueryStageExec

如果当前 exchange 是 shuffle 类型：

scala 复制代码

case s: ShuffleExchangeLike =>

会先基于优化后的 child 重建 shuffle，然后创建：

scala 复制代码

ShuffleQueryStageExec(currentStageId, newShuffle)

3 BroadcastExchangeLike 变成 BroadcastQueryStageExec

如果当前 exchange 是 broadcast 类型，则创建：

scala 复制代码

BroadcastQueryStageExec(currentStageId, newBroadcast)

2.3.1 返回结果：`QueryStageExec`类

QueryStageExec 是 AQE 在物理计划中引入的运行时执行边界，能够把Exchange 边界包装成可以独立执行 、可以保存运行结果 、可以提供 runtime stats 的 stage。

它通常由 Exchange 包装而来：

text 复制代码

ShuffleExchangeExec   -> ShuffleQueryStageExec
BroadcastExchangeExec -> BroadcastQueryStageExec

它的核心作用包括：

独立执行某一段物理计划，例如 shuffle map stage 或 broadcast stage；
在 materialize() 完成后保存执行结果；
向 AQE 提供 runtime statistics，例如真实数据大小、shuffle partition 大小；
作为后续 reOptimize 的依据，让 AQE 可以基于真实统计信息重新选择 join 策略、合并 shuffle partition 或处理 skew join；
通过 stage cache 支持 exchange / query stage 复用，避免重复执行等价的 stage。

2.4 stage.materialize 与运行时统计信息

在 getFinalPhysicalPlan() 主循环中，如果本轮创建了新的 query stage，会启动它们的物化：

scala 复制代码

result.newStages.foreach { stage =>
  stage.materialize().onComplete { res =>
    if (res.isSuccess) {
      events.offer(StageSuccess(stage, res.get))
    } else {
      events.offer(StageFailure(stage, res.failed.get))
    }
  }(AdaptiveSparkPlanExec.executionContext)
}

stage.materialize() 会真正触发 stage 执行。

对于 ShuffleQueryStageExec：

text 复制代码

触发 shuffle map stage
  ↓
生成 shuffle 文件
  ↓
收集 MapOutputStatistics

对于 BroadcastQueryStageExec：

text 复制代码

执行 broadcast child plan
  ↓
构建 broadcast relation
  ↓
记录 broadcast 数据大小

stage 完成后，会把 stage 的执行结果写回 QueryStageExec

scala 复制代码

case StageSuccess(stage, res) =>
	stage.resultOption.set(Some(res))

之后 AQE 就可以从这个 stage 读取 runtime statistics，例如：① shuffle partition size ② broadcast data size ③row count / data size

这些统计信息后面会用于后续调优。

2.5 reOptimize：基于 runtime statistics 重新优化

stage 完成后，AQE 会执行：

scala 复制代码

val logicalPlan = replaceWithQueryStagesInLogicalPlan(currentLogicalPlan, stagesToReplace)
val (newPhysicalPlan, newLogicalPlan) = reOptimize(logicalPlan)

这里分两步。

2.5.1 replaceWithQueryStagesInLogicalPlan

当前 physical plan 中已经有一些完成的 query stage。为了重新优化时不丢失这些已经物化的边界，需要把 logical plan 中对应的部分替换成 logical query stage。而 LogicalQueryStage 背后关联的是已经物化完成的：QueryStageExec，这些 QueryStageExec 里已经有完成 stage的运行信息了。

2.5.2 reOptimize

reOptimize 会基于已经完成的 query stage 统计信息重新优化 logical plan，并重新生成 physical plan。【AQE自适应优化的关键】

scala 复制代码

  /**
   * Re-optimize and run physical planning on the current logical plan based on the latest stats.
   */
  private def reOptimize(logicalPlan: LogicalPlan): (SparkPlan, LogicalPlan) = {
    logicalPlan.invalidateStatsCache()
    val optimized = optimizer.execute(logicalPlan) // 生成 优化后的逻辑计划
    val sparkPlan = context.session.sessionState.planner.plan(ReturnAnswer(optimized)).next() // 生成 物理计划【在这里会优化 Join 策略】
    val newPlan = applyPhysicalRules(
      sparkPlan,
      preprocessingRules ++ queryStagePreparationRules,
      Some((planChangeLogger, "AQE Replanning")))
    (newPlan, optimized)
  }

其中，会再次运用queryStagePreparationRules规则，让reOptimize 重新生成的物理计划满足执行要求，例如插入 shuffle、sort、清理冗余节点等。

scala 复制代码

  private def queryStagePreparationRules: Seq[Rule[SparkPlan]] = Seq(
    RemoveRedundantProjects,
    EnsureRequirements,
    RemoveRedundantSorts,
    DisableUnnecessaryBucketedScan
  ) ++ context.session.sessionState.queryStagePrepRules

2.6 costEvaluator：是否采用新计划

AQE 不会无条件采用 reOptimize 后的新计划。它会比较新旧计划成本：

scala 复制代码

val origCost = costEvaluator.evaluateCost(currentPhysicalPlan)
val newCost = costEvaluator.evaluateCost(newPhysicalPlan)

采用新计划的条件是：

scala 复制代码

if (newCost < origCost || // 新计划成本更低 or 新旧成本相同，但计划结构发生变化
    (newCost == origCost && currentPhysicalPlan != newPhysicalPlan)) {
  currentPhysicalPlan = newPhysicalPlan
  currentLogicalPlan = newLogicalPlan
  stagesToReplace = Seq.empty[QueryStageExec]
}

2.7 再次执行createQueryStages(currentPhysicalPlan)

根据最新物理计划和最新 stage 完成状态，继续向上推进 query stage 创建，直到所有 child stage 都完成，即 result.allChildStagesMaterialized=true。

因为 AQE 是自底向上创建 stage 的，上一轮可能只创建并物化了底层 stage。等这些 stage 完成后，父节点的输入才稳定，才能继续创建更上层的 stage。

3. finalStageOptimizerRules：最终阶段优化

当循环条件结束：

scala 复制代码

while (!result.allChildStagesMaterialized)

说明当前计划中的 child query stage 都已经物化完成。

接下来 AQE 会执行最终阶段优化：

scala 复制代码

currentPhysicalPlan = applyPhysicalRules(
  result.newPlan,
  finalStageOptimizerRules,
  Some((planChangeLogger, "AQE Final Query Stage Optimization")))

这一步是在所有 query stage 都完成后，对最终计划做收尾处理。

4. 总结

开启 AQE 后，AdaptiveSparkPlanExec 是运行时自适应优化的控制器。

ShuffleQueryStageExec 和 BroadcastQueryStageExec 是 AQE 的执行边界，它们可以被异步物化，物化完成后会产生真实统计信息。AQE 再利用这些统计信息反复优化 reOptimize，直到生成最终物理计划。

第三部分：参考

Spark SQL源码研读系列06：Executed Plan

Spark SQL深入分析之自适应执行优化引擎（Adaptive Query Execution）的工作原理

为什么 Spark AQE 建议的 64MB 分区，会让你的 Driver 直接暴毙？

Spark SQL AQE工作原理源码剖析

前言

AQE简介

第一部分：Spark AQE 开启与不开启时的执行路径对比

1. AQE 未开启时的执行路径

2. AQE 开启时的 preparations 定义

3. InsertAdaptiveSparkPlan 的作用

3.1 applyInternal 方法概况

3.2 AQE 总开关判断

3.3 shouldApplyAQE：判断是否值得启用 AQE

3.4 supportAdaptive：判断是否安全支持 AQE

3.5 创建 AdaptiveSparkPlanExec

4. AQE 开启后的实际执行路径

5. 总结

第二部分：AdaptiveSparkPlanExec 内部执行路径详解

1. currentPhysicalPlan变量

2. getFinalPhysicalPlan()方法

2.1 第一次创建 query stage

2.2 createQueryStages 的递归逻辑

2.2.1 为什么要 bottom-up

2.2.2 遇到 Exchange

2.2.3 普通节点递归处理 children

2.2.4 返回结果：CreateStageResult

2.3 newQueryStage()方法

2.3.1 返回结果：QueryStageExec类

2.4 stage.materialize 与运行时统计信息

2.5 reOptimize：基于 runtime statistics 重新优化

2.5.1 replaceWithQueryStagesInLogicalPlan

2.5.2 reOptimize

2.6 costEvaluator：是否采用新计划

2.7 再次执行createQueryStages(currentPhysicalPlan)

3. finalStageOptimizerRules：最终阶段优化

4. 总结

第三部分：参考

2.2.4 返回结果：`CreateStageResult`

2.3.1 返回结果：`QueryStageExec`类