前言
本文隶属于专栏《大数据技术体系》，该专栏为笔者原创，引用请注明来源，不足和错误之处请在评论区帮忙指出，谢谢！
本专栏目录结构和参考文献请见大数据技术体系
目录
Spark SQL 工作流程源码解析（一）总览（基于 Spark 3.3.0）
Spark SQL 工作流程源码解析（二）parsing 阶段（基于 Spark 3.3.0）
Spark SQL 工作流程源码解析（三）analysis 阶段（基于 Spark 3.3.0）
Spark SQL 工作流程源码解析（四）optimization 阶段（基于 Spark 3.3.0）
Spark SQL 工作流程源码解析（五）planning 阶段（基于 Spark 3.3.0）
关联
我们要做什么？
在前面的几讲中，我们详细描述了 Spark SQL 中一条 SQL 语句是如何一步步变成一个已优化过的逻辑计划的。
本讲，我们将致力于搞懂物理计划如何而来。
物理计划阶段是 Spark SQL 整个查询处理流程的最后一步。
不同于逻辑计划（LogicalPlan）的平台无关性，物理计划（SparkPlan）是与底层运行的平台密切相关的。
在这个阶段中，Spark SQL 会对优化后的逻辑算子树进行进一步处理，得到一个可以顺畅执行的物理算子树。
我们有了什么？
我们有了下面这个已经优化后的逻辑计划树：
入口
在上一讲中，我们已经找到了planning阶段的入口，也就是QueryExecution.executedPlan 方法，我们从这个方法出发。
/**
* executedPlan 不应该用于初始化一切 SparkPlan，它应该只能用来执行
*/
lazy val executedPlan: SparkPlan = withCteMap {
// 此处我们需要物化 optimizedPlan，在追踪 planning 阶段之前，要确定优化的时间不会计算到 planning 阶段里面
assertOptimized()
executePhase(QueryPlanningTracker.PLANNING) {
// 克隆计划来避免在不同的阶段比如：analyzing/optimizing/planning 之间共享计划实例
QueryExecution.prepareForExecution(preparations, sparkPlan.clone())
}
}
下面这行代码，实际上将整个planning阶段分成了 3 个部分：
- preparations
- sparkPlan.clone()
- prepareForExecution
QueryExecution.prepareForExecution(preparations, sparkPlan.clone())
我们将分别介绍这 3 个部分分别表示什么意思以及背后的实现逻辑~
1. preparations
protected def preparations: Seq[Rule[SparkPlan]] = {
QueryExecution.preparations(sparkSession,
Option(InsertAdaptiveSparkPlan(AdaptiveExecutionContext(sparkSession, this))), false)
}
从preparations 方法的返回值就可以看出，这个方法准备的是物理计划的执行规则序列，这些规则序列包含 2 个部分：
- AQE 相关的执行规则
- 其他执行规则
AQE
AQE是 Spark SQL 的一种动态优化机制，在运行时，当 Shuffle Map 阶段执行完成的时候，AQE 结合这个阶段的统计信息，基于既定的规则动态地来调整和修正还没有执行的逻辑/物理计划，从而实现对原始查询语句运行时优化的目的。
关于 AQE 的更多细节请参考我的博客——Spark SQL 的 AQE 机制
AQE 相关的执行规则到底是什么？
上面的preparations方法涉及到 AQE 相关的 2 个类，即：
- AdaptiveExecutionContext
- InsertAdaptiveSparkPlan
我们来一一分析~
AdaptiveExecutionContext
case class AdaptiveExecutionContext(session: SparkSession, qe: QueryExecution) {
/**
* 整个查询中共享的子查询重用 Map
*/
val subqueryCache: TrieMap[SparkPlan, BaseSubqueryExec] =
new TrieMap[SparkPlan, BaseSubqueryExec]()
/**
* 整个查询中共享的 Exchange 重用 Map，包括子查询
*/
val stageCache: TrieMap[SparkPlan, QueryStageExec] =
new TrieMap[SparkPlan, QueryStageExec]()
}
AdaptiveExecutionContext 是在主查询和所有子查询之间共享的执行上下文，它由 2 个 Map 构成：一个用来缓存子查询，另一个用来缓存 Exchange。
Exchange指的是那些可以在多线程/进程之间交换数据的算子，它是 Spark 能够实现并行处理的关键所在。
关于
Exchange的详情请参考我的博客——Spark 3.x 的 Exchange 体系源码解析
InsertAdaptiveSparkPlan
case class InsertAdaptiveSparkPlan(
adaptiveExecutionContext: AdaptiveExecutionContext) extends Rule[SparkPlan]
从继承就可以看出来，这是个物理计划的执行规则，它底层使用AdaptiveSparkPlanExec来包装查询计划，而AdaptiveSparkPlanExec用来执行查询计划，并在执行期间根据运行时数据统计信息重新优化计划。
注意，这个规则是有状态的，因此不应该在查询执行期间重用。
下面我们分析一下它是如何实现的：
apply
相信看了前面几讲的同学一定对apply方法不陌生，Spark 中应用规则默认就会调用apply方法。
override def apply(plan: SparkPlan): SparkPlan = applyInternal(plan, false)
这个规则的apply方法会调用内部的applyInternal
private def applyInternal(plan: SparkPlan, isSubquery: Boolean): SparkPlan = plan match {
// 如果没有开启 AQE 的话，直接返回
case _ if !conf.adaptiveExecutionEnabled => plan
case _: ExecutedCommandExec => plan
case _: CommandResultExec => 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 {
// 递归地规划子查询，并将其传入共享阶段缓存，以供`Exchange`重用。
// 如果在任何子查询中不支持 AQE，请返回`non-AQE`模式。
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 {
logDebug(s"${SQLConf.ADAPTIVE_EXECUTION_ENABLED.key} is enabled " +
s"but is not supported for query: $plan.")
plan
}
case _ => plan
}
这个方法的执行流程如下所示：
我们按照上面的流程图来分别讲解每一个流程节点的含义：
是否开启 AQE
我们通过 Spark 的配置参数：spark.sql.adaptive.enabled 来判断是否开启 AQE，这个参数默认为 true。
ExecutedCommandExec
case class ExecutedCommandExec(cmd: RunnableCommand) extends LeafExecNode
ExecutedCommandExec是用来执行RunnableCommand的run方法并保存结果以防止多次执行的物理算子。
cmd 表示这个算子将要运行的 RunnableCommand。
RunnableCommand是因为要用到它的副作用（side-effects）从而执行的逻辑命令。
在执行期间，RunnableCommand被包装在ExecutedCommand中。
CommandResultExec
case class CommandResultExec(
output: Seq[Attribute],
@transient commandPhysicalPlan: SparkPlan,
@transient rows: Seq[InternalRow]) extends LeafExecNode with InputRDDCodegen
CommandResultExec是用于保存来自命令的数据的物理计划节点。
commandPhysicalPlan仅用于显示计划树以进行解释。
rows可能无法序列化，理想情况下，我们不应该将rows发送给Executor，因此将它们标记为transient
DataWritingCommandExec
case class DataWritingCommandExec(cmd: DataWritingCommand, child: SparkPlan)
extends UnaryExecNode
DataWritingCommandExec是用来执行DataWritingCommand的run方法并保存结果以防止多次执行的物理算子。
cmd表示这个算子将要运行的数据写入命令。
child表示由DataWritingCommand运行的物理计划子项。
DataWritingCommand 是一种特殊的命令，用来写出数据并更新度量（metrics）。
V2CommandExec
V2CommandExec执行run并保存结果以防止多次执行的物理算子。
任何不需要触发 spark 作业的 V2 命令都应该扩展这个类。
这里可以参考我的博客——Spark DataSource API v2 版本有哪些改进？v1 版本和 v2 版本有什么区别？
应该用 AQE 吗？
这里指代的即InsertAdaptiveSparkPlan.shouldApplyAQE方法，我们看一下它的源码：
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
}
}
上面的配置
ADAPTIVE_EXECUTION_FORCE_APPLY来自 Spark 配置参数spark.sql.adaptive.forceApply，表示：当查询没有Exchange或子查询，AQE 会被跳过。通过将此配置设置为true（以及spark.sql.adaptive.enabled设置为true），Spark 将强制地为所有支持的查询应用 AQE。
通过分析源码，很容易得出这样的结论：
AQE 仅在查询具有Exchange或子查询时有用。
如果满足以下条件之一则方法会返回true：
- 配置
ADAPTIVE_EXECUTION_FORCE_APPLY为true。 - 输入查询来自子查询。当这种情况发生时，这意味着我们已经决定对主查询应用 AQE，我们必须继续这样做。
- 该查询包含
Exchange。 - 查询可能需要添加
Exchange。在这里运行EnsureRequirements太过分了，所以我们只是检查一下SparkPlan.requiredChildDistribution并查看
查询是不是需要稍后添加Exchange。 - 该查询包含子查询。
能用 AQE 吗？
其实和上面的区别在于：
应该表示是被动的：外部环境允不允许我对这个物理计划使用 AQE ？能表示的是主动的：这个物理计划本身支不支持对它使用 AQE ？
这里指代的即InsertAdaptiveSparkPlan.supportAdaptive方法，我们看一下它的源码：
private def supportAdaptive(plan: SparkPlan): Boolean = {
sanityCheck(plan) &&
!plan.logicalLink.exists(_.isStreaming) &&
plan.children.forall(supportAdaptive)
}
其中sanityCheck的源码如下：
private def sanityCheck(plan: SparkPlan): Boolean =
plan.logicalLink.isDefined
我们再看下SparkPlan.logicalLink 表示什么意思：
def logicalLink: Option[LogicalPlan] =
getTagValue(SparkPlan.LOGICAL_PLAN_TAG)
.orElse(getTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG))
其中
SparkPlan.LOGICAL_PLAN_TAG用来记录当前的SparkPlan是从哪个LogicalPlan转换而来的，而如果找不到对应的逻辑计划，我们就会给它一个默认值SparkPlan.LOGICAL_PLAN_INHERITED_TAG，表示的是继承自祖先的逻辑计划。
同时，我们在转换逻辑计划到物理计划的时候，可以选择继承或者非继承，继承就会用到
SparkPlan.LOGICAL_PLAN_INHERITED_TAG，非继承就会用到SparkPlan.LOGICAL_PLAN_TAG。
从上面的源码中可以看出，只有当前的物理计划满足下面3 个条件时，才能支持自适应查询执行（AQE）。
- 当前的物理计划需要关联一个逻辑计划，如果连继承自祖先的逻辑计划都没有定义的时候，就不能使用 AQE 了。
- 当前的物理计划关联的逻辑计划不能来自流式数据源。
- 当前的物理计划的所有子节点都必须满足上面 1、2 两个条件。
预处理子查询规则
从val preprocessingRules = Seq(planSubqueriesRule)可以看出来，预处理规则只有一个，即自适应子查询规则——PlanAdaptiveSubqueries，这个规则先要构建一个子查询 Map：
private def buildSubqueryMap(plan: SparkPlan): Map[Long, BaseSubqueryExec] = {
val subqueryMap = mutable.HashMap.empty[Long, BaseSubqueryExec]
if (!plan.containsAnyPattern(SCALAR_SUBQUERY, IN_SUBQUERY, DYNAMIC_PRUNING_SUBQUERY)) {
return subqueryMap.toMap
}
// PLAN_EXPRESSION
plan.foreach(_.expressions.filter(_.containsPattern(PLAN_EXPRESSION)).foreach(_.foreach {
// 只返回一行和一列的子查询。这将在`planning`期间转换为物理标量子查询。
case expressions.ScalarSubquery(p, _, exprId, _)
if !subqueryMap.contains(exprId.id) =>
val executedPlan = compileSubquery(p)
verifyAdaptivePlan(executedPlan, p)
val subquery = SubqueryExec.createForScalarSubquery(
s"subquery#${exprId.id}", executedPlan)
subqueryMap.put(exprId.id, subquery)
// 如果在`query`的结果集中返回`values`，则计算结果为`true`。
case expressions.InSubquery(_, ListQuery(query, _, exprId, _, _))
if !subqueryMap.contains(exprId.id) =>
val executedPlan = compileSubquery(query)
verifyAdaptivePlan(executedPlan, query)
val subquery = SubqueryExec(s"subquery#${exprId.id}", executedPlan)
subqueryMap.put(exprId.id, subquery)
// `DynamicPruningSubquery`表达式仅在`JOIN`操作中使用，通过另一侧的筛选器来修剪`JOIN`的一侧。它是在可以应用分区裁剪的情况下插入的。
case expressions.DynamicPruningSubquery(value, buildPlan,
buildKeys, broadcastKeyIndex, onlyInBroadcast, exprId)
if !subqueryMap.contains(exprId.id) =>
val executedPlan = compileSubquery(buildPlan)
verifyAdaptivePlan(executedPlan, buildPlan)
val name = s"dynamicpruning#${exprId.id}"
val subquery = SubqueryAdaptiveBroadcastExec(
name, broadcastKeyIndex, onlyInBroadcast,
buildPlan, buildKeys, executedPlan)
subqueryMap.put(exprId.id, subquery)
case _ =>
}))
subqueryMap.toMap
}
buildSubqueryMap为所有子查询返回 表达式 ID 到 执行计划的映射。
对于每个子查询，通过应用此规则为每个子查询生成自适应执行计划，或者尽可能从另一个具有相同语义的子查询重用执行计划。
然后，我们会将子查询 Map 传给自适应子查询规则——PlanAdaptiveSubqueries
case class PlanAdaptiveSubqueries(
subqueryMap: Map[Long, BaseSubqueryExec]) extends Rule[SparkPlan] {
def apply(plan: SparkPlan): SparkPlan = {
plan.transformAllExpressionsWithPruning(
_.containsAnyPattern(SCALAR_SUBQUERY, IN_SUBQUERY, DYNAMIC_PRUNING_SUBQUERY)) {
case expressions.ScalarSubquery(_, _, exprId, _) =>
execution.ScalarSubquery(subqueryMap(exprId.id), exprId)
case expressions.InSubquery(values, ListQuery(_, _, exprId, _, _)) =>
val expr = if (values.length == 1) {
values.head
} else {
CreateNamedStruct(
values.zipWithIndex.flatMap { case (v, index) =>
Seq(Literal(s"col_$index"), v)
}
)
}
InSubqueryExec(expr, subqueryMap(exprId.id), exprId, shouldBroadcast = true)
case expressions.DynamicPruningSubquery(value, _, _, _, _, exprId) =>
DynamicPruningExpression(InSubqueryExec(value, subqueryMap(exprId.id), exprId))
}
}
}
这个规则会递归地检查物理计划树中是否存在ScalarSubquery、InSubquery和 DynamicPruningSubquery，并进行针对性处理。
ScalarSubquery即只返回一行一列的子查询。这将在planning阶段期间转换为物理标量子查询。
InSubquery表示如果指定的值存在于查询的结果集中，则计算结果为true。
DynamicRuningSubQuery仅在Join操作中使用，通过使用Join另一侧的过滤器来裁剪Join的一侧。
构建好预处理规则后，我们通过AdaptiveSparkPlanExec.applyPhysicalRules方法来执行。
def applyPhysicalRules(
plan: SparkPlan,
rules: Seq[Rule[SparkPlan]],
loggerAndBatchName: Option[(PlanChangeLogger[SparkPlan], String)] = None): SparkPlan = {
if (loggerAndBatchName.isEmpty) {
rules.foldLeft(plan) { case (sp, rule) => rule.apply(sp) }
} else {
val (logger, batchName) = loggerAndBatchName.get
val newPlan = rules.foldLeft(plan) { case (sp, rule) =>
val result = rule.apply(sp)
logger.logRule(rule.ruleName, sp, result)
result
}
logger.logBatch(batchName, plan, newPlan)
newPlan
}
}
这个方法会应用 SparkPlan 上的一系列物理算子规则。
这个方法中，我们看到了一个熟悉的身影——PlanChangeLogger，忘记的小伙伴建议回看一下前两讲。
这就说明，我们通过配置 Spark 参数：spark.sql.planChangeLog.level，就可以实现在控制台中看到物理计划的变更过程。
我们会在InsertAdaptiveSparkPlan.applyInternal方法的最后将执行结果传给 AQE 的物理执行算子AdaptiveSparkPlanExec：
AdaptiveSparkPlanExec
case class AdaptiveSparkPlanExec(
inputPlan: SparkPlan,
@transient context: AdaptiveExecutionContext,
@transient preprocessingRules: Seq[Rule[SparkPlan]],
@transient isSubquery: Boolean,
@transient override val supportsColumnar: Boolean = false)
extends LeafExecNode
AdaptiveSparkPlanExec是用于自适应执行查询计划的根节点。它会把查询计划拆分为独立的阶段（Stage），并根据它们的依赖关系按顺序执行它们。查询阶段在最后实现其输出。当一个阶段完成时，物化输出的数据统计信息将用于优化查询的其余部分。
为了创建查询阶段，我们需要自下而上遍历地查询树。当我们点击一个Exchange节点时，如果这个Exchange节点的所有子查询阶段都被物化，我们将为这个Exchange节点创建一个新的查询阶段。新阶段一旦创建，就会异步实现。
当一个查询阶段完成物化时，其余的查询将根据所有物化阶段提供的最新统计信息进行重新优化和规划。然后，我们再次遍历查询计划，如果可能的话，创建更多的阶段。在所有阶段都实现后，我们执行计划的其余部分。
既然是一个物理计划执行算子，那么它的核心就在于执行，即方法——doExecute，我们看下这块的源码：
可以看看我的这篇博客——Spark 3.x 的 WSCG 机制源码解析
override def doExecute(): RDD[InternalRow] = {
withFinalPlanUpdate(_.execute())
}
这个方法调用了withFinalPlanUpdate。
private def withFinalPlanUpdate[T](fun: SparkPlan => T): T = {
val plan = getFinalPhysicalPlan()
val result = fun(plan)
finalPlanUpdate
result
}
这个方法先调用getFinalPhysicalPlan获取最终的物理计划，然后执行它，在最终结果返回前，调用finalPlanUpdate方法。
getFinalPhysicalPlan
private def getFinalPhysicalPlan(): SparkPlan = lock.synchronized {
if (isFinalPlan) return currentPhysicalPlan
// 如果在`withActive`范围外执行此自适应计划，
// 例如：`plan.queryExecution.rdd`，我们需要在这里设置活跃的会话，
// 因为在执行中间可以创建新的计划节点
context.session.withActive {
val executionId = getExecutionId
// 此处使用 inputPlan 的 logicalLink ，为了避免在`initialPlan`期间一些顶层的物理节点被移除
var currentLogicalPlan = inputPlan.logicalLink.get
var result = createQueryStages(currentPhysicalPlan)
val events = new LinkedBlockingQueue[StageMaterializationEvent]()
val errors = new mutable.ArrayBuffer[Throwable]()
var stagesToReplace = Seq.empty[QueryStageExec]
while (!result.allChildStagesMaterialized) {
currentPhysicalPlan = result.newPlan
if (result.newStages.nonEmpty) {
stagesToReplace = result.newStages ++ stagesToReplace
executionId.foreach(onUpdatePlan(_, result.newStages.map(_.plan)))
// SPARK-33933: 我们应该首先提交广播阶段的任务，以避免等待用于计划任务并导致广播超时。
// 这个部分修复只保证`BroadcastQueryStage`的开始物化优先于其他，但因为广播收集作业的提交在另一个线程中运行时，问题并没有完全解决。
val reorderedNewStages = result.newStages
.sortWith {
case (_: BroadcastQueryStageExec, _: BroadcastQueryStageExec) => false
case (_: BroadcastQueryStageExec, _) => true
case _ => false
}
// 开始所有新阶段的物化，如果任何阶段急切地失败，则 fail fast
reorderedNewStages.foreach { stage =>
try {
stage.materialize().onComplete { res =>
if (res.isSuccess) {
events.offer(StageSuccess(stage, res.get))
} else {
events.offer(StageFailure(stage, res.failed.get))
}
}(AdaptiveSparkPlanExec.executionContext)
} catch {
case e: Throwable =>
cleanUpAndThrowException(Seq(e), Some(stage.id))
}
}
}
// 等待下一个完成的阶段，这表明新的统计数据可用，并且可能可以创建新的阶段。
// 可能还有其他阶段在大致相同的时间结束，所以我们也处理这些阶段，以减少重新规划。
val nextMsg = events.take()
val rem = new util.ArrayList[StageMaterializationEvent]()
events.drainTo(rem)
(Seq(nextMsg) ++ rem.asScala).foreach {
case StageSuccess(stage, res) =>
stage.resultOption.set(Some(res))
case StageFailure(stage, ex) =>
errors.append(ex)
}
// 如果出现错误，我们将取消所有正在运行的阶段并抛出异常。
if (errors.nonEmpty) {
cleanUpAndThrowException(errors.toSeq, None)
}
// 尝试重新优化和重新规划。
// 如果新计划的成本小于或者等于当前的计划就采用新计划；
// 否则，就保持当前的物理计划和逻辑计划，因为物理计划的`logicalLink`指向它起源的逻辑计划。
// 同时，我们保留了自上次计划更新以来创建的查询阶段的列表，这代表了当前的逻辑计划和物理计划之间的“语义鸿沟”。
// 每次重新规划之前，我们都会用逻辑查询阶段来替换当前逻辑计划的对应节点，这会使得在语义上和当前的物理计划保持同步。
// 一旦一个新的计划被采纳，逻辑计划和物理计划都更新后，我们可以清除查询阶段列表，因为此时两个计划在语义和物理上又再次同步了。
val logicalPlan = replaceWithQueryStagesInLogicalPlan(currentLogicalPlan, stagesToReplace)
val (newPhysicalPlan, newLogicalPlan) = reOptimize(logicalPlan)
val origCost = costEvaluator.evaluateCost(currentPhysicalPlan)
val newCost = costEvaluator.evaluateCost(newPhysicalPlan)
if (newCost < origCost ||
(newCost == origCost && currentPhysicalPlan != newPhysicalPlan)) {
logOnLevel(s"Plan changed from $currentPhysicalPlan to $newPhysicalPlan")
cleanUpTempTags(newPhysicalPlan)
currentPhysicalPlan = newPhysicalPlan
currentLogicalPlan = newLogicalPlan
stagesToReplace = Seq.empty[QueryStageExec]
}
// 现在一些阶段已经结束了，我们可以创建新的阶段。
result = createQueryStages(currentPhysicalPlan)
}
// 当没有未完成的阶段时运行 final plan
currentPhysicalPlan = applyPhysicalRules(
optimizeQueryStage(result.newPlan, isFinalStage = true),
postStageCreationRules(supportsColumnar),
Some((planChangeLogger, "AQE Post Stage Creation")))
isFinalPlan = true
executionId.foreach(onUpdatePlan(_, Seq(currentPhysicalPlan)))
currentPhysicalPlan
}
}
这个方法的流程如下：
finalPlanUpdate
@transient private lazy val finalPlanUpdate: Unit = {
if (!isSubquery && currentPhysicalPlan.find(_.subqueries.nonEmpty).isDefined) {
getExecutionId.foreach(onUpdatePlan(_, Seq.empty))
}
logOnLevel(s"Final plan: $currentPhysicalPlan")
}
这里使用lazy val来避免被调用超过一次。
不属于主查询的任何查询阶段的子查询，都会在getFinalPhysicalPlan中最后一次更新 UI，所以我们需要在这里再次更新 UI 以确保这些子查询的新生成节点被更新。
onUpdatePlan 方法在getFinalPhysicalPlan中也出现过，我们再看下它的源码：
onUpdatePlan
private def onUpdatePlan(executionId: Long, newSubPlans: Seq[SparkPlan]): Unit = {
if (isSubquery) {
// 在执行子查询时，我们不能更新 UI 中的查询计划，因为 UI 还不支持部分更新。
// 但是，子查询可能已被优化到不同的计划中，我们必须让 UI 知道新计划节点的 SQL 度量信息，以便以后可以跟踪有效的累加器更新并正确地显示 SQL 度量。
val newMetrics = newSubPlans.flatMap { p =>
p.flatMap(_.metrics.values.map(m => SQLPlanMetric(m.name.get, m.id, m.metricType)))
}
context.session.sparkContext.listenerBus.post(SparkListenerSQLAdaptiveSQLMetricUpdates(
executionId.toLong, newMetrics))
} else {
val planDescriptionMode = ExplainMode.fromString(conf.uiExplainMode)
context.session.sparkContext.listenerBus.post(SparkListenerSQLAdaptiveExecutionUpdate(
executionId,
context.qe.explainString(planDescriptionMode),
SparkPlanInfo.fromSparkPlan(context.qe.executedPlan)))
}
}
这个方法用来通知监听器们物理计划已经改变了。
底层是通过将事件添加到sparkContext.listenerBus的事件队列中，同时SQLAppStatusListener在监听这个事件，最终呈现在 Spark UI 中的SQL页面中。
分析完了 AQE 相关的执行规则，我们再来看看其他的执行规则，这些规则都定义在QueryExecution.preparations方法里面：
QueryExecution.preparations
private[execution] def preparations(
sparkSession: SparkSession,
adaptiveExecutionRule: Option[InsertAdaptiveSparkPlan] = None,
subquery: Boolean): Seq[Rule[SparkPlan]] = {
// `AdaptiveSparkPlanExec` 是叶子节点。 如果被插入了，所有下面的规则都会变成 `no-op`（无操作）
// 因为原始的计划会隐藏在 `AdaptiveSparkPlanExec`后面。
adaptiveExecutionRule.toSeq ++
Seq(
CoalesceBucketsInJoin,
PlanDynamicPruningFilters(sparkSession),
PlanSubqueries(sparkSession),
RemoveRedundantProjects,
EnsureRequirements(),
// `ReplaceHashWithSortAgg` 需要在 `EnsureRequirements` 后面添加，这是为了确保检查每个节点的排序顺序是否有效。
ReplaceHashWithSortAgg,
// `RemoveRedundantSorts` 需要在 `EnsureRequirements` 后面添加， 这是为了确保实例化`PartitioningCollection`时有相同的分区数。
RemoveRedundantSorts,
DisableUnnecessaryBucketedScan,
ApplyColumnarRulesAndInsertTransitions(
sparkSession.sessionState.columnarRules, outputsColumnar = false),
CollapseCodegenStages()) ++
(if (subquery) {
Nil
} else {
Seq(ReuseExchangeAndSubquery)
})
}
为了执行规划好的 SparkPlan，需要做一些准备，这个方法会构建一系列规则。
这些规则会确保子查询得到规划，使用正确的数据分区和排序，插入 WSCG，并尝试通过重用Exchange和子查询来减少所做的工作。
可以看到，除了 AQE 相关的规则外，剩下的执行规则有：
- CoalesceBucketsInJoin
- PlanDynamicPruningFilters
- PlanSubqueries
- RemoveRedundantProjects
- EnsureRequirements
- ReplaceHashWithSortAgg
- RemoveRedundantSorts
- DisableUnnecessaryBucketedScan
- ApplyColumnarRulesAndInsertTransitions
- CollapseCodegenStages
- ReuseExchangeAndSubquery
我们来一一进行解析~
1. CoalesceBucketsInJoin
object CoalesceBucketsInJoin extends Rule[SparkPlan]
如果满足以下条件，此规则将合并SortMergeJoin和ShuffledHashJoin的一侧：
- 两张桶表进行
JOIN。 JOIN KEY在各自的边上与输出分区表达式匹配。- 较大的桶数可被较小的桶数整除。
COALESCE_BUCKETS_IN_JOIN_ENABLED设置为true。- 桶数的比率小于设定值
COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO。
COALESCE_BUCKETS_IN_JOIN_ENABLED即 Spark 配置参数spark.sql.bucketing.coalesceBucketsInJoin.enabled，它表示：如果为true，且如果两个有不同的桶数的桶表进行JOIN，则桶数较多的一侧会合并成和另一侧相同的桶数。桶的较大数量可被桶的较小数量整除。桶合并会应用于sort-merge join和shuffle hash join。
注意：合并桶表可以避免
JOIN不必要的 Shuffle，但它也会降低并行性，并可能在shuffle hash join的时候导致 OOM。
COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO 即 Spark 配置参数spark.sql.bucketing.coalesceBucketsInJoin.maxBucketRatio，表示桶表合并的两桶的数量之比应小于或等于要应用的桶合并值。
2. PlanDynamicPruningFilters
case class PlanDynamicPruningFilters(sparkSession: SparkSession)
extends Rule[SparkPlan] with PredicateHelper
该规则旨在重写动态裁剪的谓词，以便重用广播的结果。
对于那些没有规划成broadcast hash join的join，我们使用子查询复制来保留回退机制。
3. PlanSubqueries
case class PlanSubqueries(sparkSession: SparkSession) extends Rule[SparkPlan]
这个规则用来规划给定 SparkPlan 中存在的子查询，包括ScalarSubquery（只返回一行一列的标量子查询）和InSubquery（如果指定的值存在于查询的结果集中，则计算结果为 true）。
4. RemoveRedundantProjects
object RemoveRedundantProjects extends Rule[SparkPlan]
从 spark plan 中删除冗余的 ProjectExec节点。
ProjectExec节点在以下情况下是多余的
- 它与子节点的输出具有相同的输出属性和顺序，并且属性的顺序是必需的。
- 当不需要属性输出顺序时，它的输出属性与其子级的输出属性相同。
该规则必须是物理规则，因为在逻辑优化过程中，project节点对修剪数据非常有用。
在物理规划期间，可以删除冗余的project节点，以简化查询计划。
5. EnsureRequirements
case class EnsureRequirements(
optimizeOutRepartition: Boolean = true,
requiredDistribution: Option[Distribution] = None)
extends Rule[SparkPlan]
通过在需要时插入ShuffleExchangeExec算子，确保输入数据的Partitioning满足每个算子的Distribution要求。还要确保满足输入分区排序要求。
关于
ShuffleExchangeExec请查看我的博客————Spark 3.x 的 Exchange 体系源码解析
参数：
- optimizeOutRepartition：一个标志，指示此规则是否应优化用户指定的
repartition shuffle。这在很大程度上是正确的，但在 AQE 中可能是错误的，因为 AQE 优化可能会更改计划输出分区，并且需要在计划中保留用户指定的repartition shuffle。 - requiredDistribution：我们应该确保根节点所需的
distribution。在 AQE 中使用该值，以防我们更改final stage的输出分区。
6. ReplaceHashWithSortAgg
object ReplaceHashWithSortAgg extends Rule[SparkPlan]
在以下情况下，将物理计划（SparkPlan）中将基于哈希的聚合替换为排序聚合：
- 计划是
HashAggregateExec或者ObjectHashAggregateExec，并且子计划中存在部分聚合并且哈希的情况，这时候还要求部分聚合的子级满足相应SortAggregateExec的排序顺序。 - 计划是
HashAggregateExec或ObjectHashAggregateExec，子计划满足相应SortAggregateExec的排序顺序。
示例
JOIN后聚合：
HashAggregate(t1.i, SUM, final)
| SortAggregate(t1.i, SUM, complete)
HashAggregate(t1.i, SUM, partial) => |
| SortMergeJoin(t1.i = t2.j)
SortMergeJoin(t1.i = t2.j)
- 排序后聚合：
HashAggregate(t1.i, SUM, partial) SortAggregate(t1.i, SUM, partial)
| => |
Sort(t1.i) Sort(t1.i)
当其子级满足相应的排序聚合的排序顺序时，可以替换基于哈希的聚合。
排序聚合速度更快，因为它没有哈希聚合的哈希开销。
7. RemoveRedundantSorts
object RemoveRedundantSorts extends Rule[SparkPlan] {
def apply(plan: SparkPlan): SparkPlan = {
if (!conf.getConf(SQLConf.REMOVE_REDUNDANT_SORTS_ENABLED)) {
plan
} else {
removeSorts(plan)
}
}
private def removeSorts(plan: SparkPlan): SparkPlan = plan transform {
case s @ SortExec(orders, _, child, _)
if SortOrder.orderingSatisfies(child.outputOrdering, orders) &&
child.outputPartitioning.satisfies(s.requiredChildDistribution.head) =>
child
}
}
这个规则会从物理计划（SparkPlan）中删除多余的SortExec节点。
当排序节点的子节点同时满足它的排序顺序和它要求的子节点分布的时候，排序节点是冗余的。
注意：这个规则与优化规则
EliminateSorts的不同之处在于，此规则还检查子项是否满足所需的分布，以便在其子项已经满足所需的排序顺序时，不仅可以安全地删除本地排序，还可以安全地删除全局排序。
优化规则
EliminateSorts的作用是：如果排序操作不影响最终的输出顺序，就删除它们。详情请见上一讲 —— Spark SQL 工作流程源码解析（四）optimization 阶段（基于 Spark 3.3.0）
上面的REMOVE_REDUNDANT_SORTS_ENABLED即 Spark 配置参数spark.sql.execution.removeRedundantSorts，表示是否移除冗余的物理排序节点，默认是 true。
方法 SortOrder.orderingSatisfies 表示如果一个SortOrder序列满足另一个SortOrder序列，则返回true。
SortOrder序列 A 满足SortOrder序列 B 当且仅当 B 是 A 或 A 前缀的等价物。
以下是排序 A 满足 排序 B 的示例：
- 排序 A 是 [x, y] 并且排序 B 是 [x]
- 排序 A 是 [x(sameOrderExpressions=x1)] 并且排序 B 是 [x1]
- 排序 A 是 [x(sameOrderExpressions=x1), y] 并且排序 B 是 [x1]
SortOrder
case class SortOrder(
child: Expression,
direction: SortDirection,
nullOrdering: NullOrdering,
sameOrderExpressions: Seq[Expression])
extends Expression with Unevaluable
SortOrder是用于对元组排序的表达式。这个类主要是扩展了表达式，以便表达式上的转换可以下降到它的子级。其中sameOrderExpressions是一组与子表达式具有相同排序顺序的表达式。它源自算子中的等价关系，例如内部sort merge join的左/右键。
8. DisableUnnecessaryBucketedScan
object DisableUnnecessaryBucketedScan extends Rule[SparkPlan]
这个规则会根据实际的物理查询计划禁用不必要的桶表扫描。
注意：此规则被设计成在
EnsureRequirements之后应用，此时所有ShuffleExchangeExec和SortExec都已正确添加到计划中。
当BUCKETING_ENABLED和AUTO_BUCKETED_SCAN_ENABLED设置为 true 时，遍历查询计划来检查哪里不需要扫描桶表，并在以下情况下禁用桶表扫描：
- 从根节点到桶表扫描的子计划，不包含
hasInterestingPartition算子。 - 从最近的下游
hasInterestingPartition算子到桶表扫描的子计划，只包含isAllowedUnaryExecNode算子和至少一个Exchange。
示例
- 没有
hasInterestingPartition算子：
Project
|
Filter
|
Scan(t1: i, j)
(bucketed on column j, DISABLE bucketed scan)
- JOIN ：
SortMergeJoin(t1.i = t2.j)
/ \
Sort(i) Sort(j)
/ \
Shuffle(i) Scan(t2: i, j)
/ (bucketed on column j, enable bucketed scan)
Scan(t1: i, j)
(bucketed on column j, DISABLE bucketed scan)
- 聚合：
HashAggregate(i, ..., Final)
|
Shuffle(i)
|
HashAggregate(i, ..., Partial)
|
Filter
|
Scan(t1: i, j)
(bucketed on column j, DISABLE bucketed scan)
hasInterestingPartition 的灵感来自论文《Access Path Selection in a Relational Database Management System》中的interesting order。
9. ApplyColumnarRulesAndInsertTransitions
case class ApplyColumnarRulesAndInsertTransitions(
columnarRules: Seq[ColumnarRule],
outputsColumnar: Boolean)
extends Rule[SparkPlan]
这个规则会应用任何用户定义的ColumnarRules，并找到正确的位置来插入来自列式数据的转换或者转换成列式数据。
参数：
- columnarRules：自定义列规则
- outputsColumnar：生成的计划是否应输出列式格式。
10. CollapseCodegenStages
case class CollapseCodegenStages(
codegenStageCounter: AtomicInteger = new AtomicInteger(0))
extends Rule[SparkPlan]
找到支持codegen的链式计划，将它们合并为WholeStageCodegen（全阶段代码生成）。
关于
WholeStageCodegen的详细内容请参考我的博客——Spark 3.x 的 WSCG 机制源码解析
codegenStageCounter为查询计划中的codegen阶段生成 ID。它不影响相等性，也不参与对WholeStageCodegenExec的模式匹配的解构。 此 ID 用于帮助区分codegen阶段。
它是物理计划解释输出的一部分。
例如：
== Physical Plan ==
*(5) SortMergeJoin [x#3L], [y#9L], Inner
:- *(2) Sort [x#3L ASC NULLS FIRST], false, 0
: +- Exchange hashpartitioning(x#3L, 200)
: +- *(1) Project [(id#0L % 2) AS x#3L]
: +- *(1) Filter isnotnull((id#0L % 2))
: +- *(1) Range (0, 5, step=1, splits=8)
+- *(4) Sort [y#9L ASC NULLS FIRST], false, 0
+- Exchange hashpartitioning(y#9L, 200)
+- *(3) Project [(id#6L % 2) AS y#9L]
+- *(3) Filter isnotnull((id#6L % 2))
+- *(3) Range (0, 5, step=1, splits=8)
ID 表明并非所有相邻的codegen计划算子都属于同一个codegen阶段。
codegen阶段 ID 还可以作为后缀选择性地包含在生成的类的名称中，以便更容易将生成的类与物理算子关联起来。
这是由 Spark SQL 配置参数：spark.sql.codegen.useIdInClassName控制的。
ID 也包含在各种日志消息中。
在查询中，计划中的codegen阶段从 1 开始按“插入顺序”计数。
WholeStageCodegenExec算子会以深度优先后序遍历的方式被插入到计划中。
参考CollapseCodegenStages.insertWholeStageCodegen来获取插入顺序的定义。
0 保留为一个特殊 ID 值，用来指示创建了一个临时的WholeStageCodegenExec对象，例如，当现有的WholeStageCodegenExec无法生成/编译代码时，可以进行特殊的回退处理。
11. ReuseExchangeAndSubquery
case object ReuseExchangeAndSubquery extends Rule[SparkPlan] {
def apply(plan: SparkPlan): SparkPlan = {
if (conf.exchangeReuseEnabled || conf.subqueryReuseEnabled) {
val exchanges = mutable.Map.empty[SparkPlan, Exchange]
val subqueries = mutable.Map.empty[SparkPlan, BaseSubqueryExec]
def reuse(plan: SparkPlan): SparkPlan = {
plan.transformUpWithPruning(_.containsAnyPattern(EXCHANGE, PLAN_EXPRESSION)) {
case exchange: Exchange if conf.exchangeReuseEnabled =>
val cachedExchange = exchanges.getOrElseUpdate(exchange.canonicalized, exchange)
if (cachedExchange.ne(exchange)) {
ReusedExchangeExec(exchange.output, cachedExchange)
} else {
cachedExchange
}
case other =>
other.transformExpressionsUpWithPruning(_.containsPattern(PLAN_EXPRESSION)) {
case sub: ExecSubqueryExpression =>
val subquery = reuse(sub.plan).asInstanceOf[BaseSubqueryExec]
val newSubquery = if (conf.subqueryReuseEnabled) {
val cachedSubquery = subqueries.getOrElseUpdate(subquery.canonicalized, subquery)
if (cachedSubquery.ne(subquery)) {
ReusedSubqueryExec(cachedSubquery)
} else {
cachedSubquery
}
} else {
subquery
}
sub.withNewPlan(newSubquery)
}
}
}
reuse(plan)
} else {
plan
}
}
}
这个规则会在包括子查询在内的整个物理计划（SparkPlan）中找出重复的exchange和子查询，然后对于所有的引用我们使用相同的exchange或子查询。
关于
exchange请查看我的博客————Spark 3.x 的 Exchange 体系源码解析
注意，SparkPlan是一个相互递归的数据结构：
SparkPlan -> Expr -> Subquery -> SparkPlan -> Expr -> Subquery -> ...
因此，在这个规则中，我们以自下而上的方式一次递归地重写exchange和子查询。
上面的源码中还涉及了 2 个 Spark 配置参数：
conf.exchangeReuseEnabled即spark.sql.exchange.reuse，表示为 true 的时候，planner 会找出重复的exchange然后重用，默认为 true。conf.subqueryReuseEnabled即spark.sql.execution.reuseSubquery，表示为 true 的时候，planner 会找出重复的子查询然后重用，默认为 true。
2. sparkPlan.clone()
scala 中的clone方法底层是调用 Java 的Object.clone()方法实现的（前提是先实现 cloneable 接口）。
这里克隆的目的是为了避免在不同的阶段（analyzing/optimizing/planning）之间共享计划实例。
故我们分析的重点在sparkPlan上，我们看下QueryExecution.sparkPlan的源码：
lazy val sparkPlan: SparkPlan = withCteMap {
// 此处我们需要物化 optimizedPlan，在追踪 planning 阶段之前，要确定优化的时间不会计算到 planning 阶段里面
assertOptimized()
executePhase(QueryPlanningTracker.PLANNING) {
// 克隆计划来避免在不同的阶段比如：analyzing/optimizing/planning 之间共享计划实例
QueryExecution.createSparkPlan(sparkSession, planner, optimizedPlan.clone())
}
}
这里会通过createSparkPlan来将逻辑计划转换成物理计划。
/**
* 将逻辑计划转换为 Spark Plan。
* 请注意，返回的物理计划仍需准备好执行。
*/
def createSparkPlan(
sparkSession: SparkSession,
planner: SparkPlanner,
plan: LogicalPlan): SparkPlan = {
// 这里目前使用的 next()，会选择计划器返回的第一个计划
// 但是未来我们会实现选择最佳的计划。
planner.plan(ReturnAnswer(plan)).next()
}
可以看到，我们首先使用了ReturnAnswer来包装我们逻辑计划。
ReturnAnswer
当我们在 take()或collect()操作时，在调用查询计划器之前，会在逻辑计划的顶部插入这个特殊节点。
规则可以在这个节点上进行模式匹配，这样的话可以应用那些只有在逻辑查询计划顶部才能生效的转换。
case class ReturnAnswer(child: LogicalPlan) extends UnaryNode {
// 最大数据行数
override def maxRows: Option[Long] = child.maxRows
// 查询计划的输出
override def output: Seq[Attribute] = child.output
// 生成一个新的子节点的过程中顺便会对当前子节点做的事情
override protected def withNewChildInternal(newChild: LogicalPlan): ReturnAnswer =
copy(child = newChild)
}
包装我们的逻辑计划后，我们的逻辑计划会发生改变：
变化
逻辑计划变化前：
GlobalLimit 21
+- LocalLimit 21
+- Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11]
+- Relation [addr#8,age#9L,name#10,sex#11] json
逻辑计划变化后：
ReturnAnswer
+- GlobalLimit 21
+- LocalLimit 21
+- Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11]
+- Relation [addr#8,age#9L,name#10,sex#11] json
可以看到，我们的逻辑计划放到了 ReturnAnswer 这个节点下面，其他内容没有发生过改变。
使用ReturnAnswer包装我们逻辑计划后，我们会调用SparkStrategies.plan方法。
plan
override def plan(plan: LogicalPlan): Iterator[SparkPlan] = {
super.plan(plan).map { p =>
val logicalPlan = plan match {
case ReturnAnswer(rootPlan) => rootPlan
case _ => plan
}
p.setLogicalLink(logicalPlan)
p
}
}
这里会先调用其父类QueryPlanner的plan方法，然后对于生成的物理计划，先脱去 ReturnAnswer这层外壳，然后setLogicalLink，我们先看看后者的源码：
private def setLogicalLink(logicalPlan: LogicalPlan, inherited: Boolean = false): Unit = {
// 在从另一个逻辑节点转换而来的子树的根上停止。
if (inherited && getTagValue(SparkPlan.LOGICAL_PLAN_TAG).isDefined) {
return
}
val tag = if (inherited) {
SparkPlan.LOGICAL_PLAN_INHERITED_TAG
} else {
SparkPlan.LOGICAL_PLAN_TAG
}
setTagValue(tag, logicalPlan)
children.foreach(_.setLogicalLink(logicalPlan, true))
}
在上面的能用 AQE 吗？解析中，我们实际上已经对logicalLink有所了解，本质上就是用来记录逻辑计划到物理计划的转换的。
接着，我们再来看看QueryPlanner.plan 方法的源码，这块内容比较复杂，我们着重来讲解：
def plan(plan: LogicalPlan): Iterator[PhysicalPlan] = {
// 收集物理计划的候选者
val candidates = strategies.iterator.flatMap(_(plan))
// 候选者必须包含被标记为[[planLater]]的占位符
// 因此试着用子计划来替代它们
val plans = candidates.flatMap { candidate =>
val placeholders = collectPlaceholders(candidate)
if (placeholders.isEmpty) {
// 按原样接受候选者，因为它不包含占位符。
Iterator(candidate)
} else {
// 把逻辑计划标记为 [[planLater]] ，同时替换占位符
placeholders.iterator.foldLeft(Iterator(candidate)) {
case (candidatesWithPlaceholders, (placeholder, logicalPlan)) =>
// 继续规划那些有占位符的逻辑计划
val childPlans = this.plan(logicalPlan)
candidatesWithPlaceholders.flatMap { candidateWithPlaceholders =>
childPlans.map { childPlan =>
// 使用子计划来替代占位符
candidateWithPlaceholders.transformUp {
case p if p.eq(placeholder) => childPlan
}
}
}
}
}
}
// 修剪物理计划在当前 3.3.0 版本中还未实现
val pruned = prunePlans(plans)
assert(pruned.hasNext, s"No plan for $plan")
pruned
}
通过上面的注释也可以看出，我们先找出那些物理计划的候选者，然后通过递归的方式逐步地先占位再使用子计划替换直至占位符为空，最后再修剪一下后返回物理计划的列表。
上面源码的核心在于这一行：
val candidates = strategies.iterator.flatMap(_(plan))
部分同学可能看不懂这里的语法，实际上问题的答案在GenericStrategy的源码中：
策略序列的迭代器扁平化后对其中每一次都调用它的apply方法，这个方法会将逻辑计划转换成物理计划的序列。
/**
* 给定一个 LogicalPlan，返回可用于执行的物理计划列表。
* 如果此策略不适用于给定的逻辑操作，则应返回空列表
*/
abstract class GenericStrategy[PhysicalPlan <: TreeNode[PhysicalPlan]] extends Logging {
/**
* 返回执行计划的物理计划的占位符。
* QueryPlanner将使用其他可用的执行策略自动填写此占位符。
*/
protected def planLater(plan: LogicalPlan): PhysicalPlan
def apply(plan: LogicalPlan): Seq[PhysicalPlan]
}
strategies
QueryPlanner.plan 方法的重点在于获取策略序列，即QueryPlanner.strategies方法，这个方法在它的子类SparkPlanner中实现。
override def strategies: Seq[Strategy] =
experimentalMethods.extraStrategies ++
extraPlanningStrategies ++ (
LogicalQueryStageStrategy ::
PythonEvals ::
new DataSourceV2Strategy(session) ::
FileSourceStrategy ::
DataSourceStrategy ::
SpecialLimits ::
Aggregation ::
Window ::
JoinSelection ::
InMemoryScans ::
SparkScripts ::
WithCTEStrategy ::
BasicOperators :: Nil)
在详细解析上面每一项策略之前，我们先来看看 Strategy 体系的类图：
Strategy 体系类图
上面的类图中有 2 点值得注意：
GenericStrategy和QueryPlanner中的PhysicalPlan只是一个泛型占位符，表示物理计划，对应的实际类型为：SparkPlanSparkPlanner和ExperimentalMethods中的Strategy实际上就是SparkStrategy，这在 SQL 模块的包对象中有说明：type Strategy = SparkStrategy。
相信上面的类图能够帮助大家理解 Spark Strategy 体系的整体代码架构，这对于我们接下来的分析会大有助益。
好了，看完了类图，我们来一一分析：
experimentalMethods.extraStrategies
experimentalMethods.extraStrategies 表示允许在运行时将额外的策略注入查询计划器。
extraPlanningStrategies
extraPlanningStrategies这个是用来进行方法重写的，以便向计划器添加额外的计划策略。
这些策略在上面的ExperimentalMethods中定义的策略之后，在常规策略之前进行尝试。
LogicalQueryStageStrategy
object LogicalQueryStageStrategy extends Strategy with PredicateHelper
这是包含LogicalQueryStage节点的计划策略：
- 将
LogicalQueryStage转换为正在执行或已完成执行的相应物理计划。 - 转换已计划并作为
BroadcastQueryStageExec执行了一个子关系的join。为了防止较大的子级关系在较小的关系之前结束，我们会阻止将广播阶段（Stage）转换成 shuffle 阶段。请注意，在常规的join策略之前，需要应用此规则。
LogicalQueryStage是QueryStageExec的逻辑计划包装，或者包含QueryStageExec的物理计划片段，QueryStageExec的所有祖先节点都链接到同一逻辑节点。例如，可以将一个逻辑聚合转换为FinalAgg - Shuffle - PartialAgg，其中 Shuffle 将被包装为QueryStageExec，因此LogicalQueryStage将FinalAgg - QueryStageExec作为其物理计划。
上面的链接也就是前面
能用 AQE 吗？中提到的SparkPlan.logicalLink。
QueryStageExec查询阶段（Stage）是查询计划的独立子图。在继续执行查询计划的其他算子之前，查询阶段将实现其输出。物化输出的数据统计信息可用于优化后续的查询阶段。 有两种类型的查询阶段：1. Shuffle 查询阶段：这个阶段将其输出具体化为 Shuffle 文件，Spark 启动另一个作业来执行进一步的算子。2.广播查询阶段：这个阶段将其输出具体化为 Driver JVM 中的一个数组。Spark 在执行进一步的算子之前先广播数组。
PythonEvals
object PythonEvals extends Strategy
将EvalPython逻辑算子转换为物理算子的策略。
object PythonEvals extends Strategy {
override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
// 使用 Apache Arrow 来求值一个`PythonUDF`的逻辑计划
case ArrowEvalPython(udfs, output, child, evalType) =>
ArrowEvalPythonExec(udfs, output, planLater(child), evalType) :: Nil
// 求值`PythonUDF`的逻辑计划
case BatchEvalPython(udfs, output, child) =>
BatchEvalPythonExec(udfs, output, planLater(child)) :: Nil
case _ =>
Nil
}
}
DataSourceV2Strategy
class DataSourceV2Strategy(session: SparkSession) extends Strategy with PredicateHelper
DataSource V2 的策略。
关于
DataSource V2请参考我的博客——Spark DataSource API v2 版本有哪些改进？v1 版本和 v2 版本有什么区别？
这个策略里面每一种类型的操作和与之对应的物理执行算子如下所示：
FileSourceStrategy
object FileSourceStrategy extends Strategy with PredicateHelper with Logging
这种策略可以用来规划对可能由用户指定的列进行分区或分桶的文件集合的扫描。
在高层次的计划分为几个阶段：
- 按需要求值的时间拆分过滤器。
- 根据现有的任何投影（project）修剪所请求数据的 Schema。如今，这种修剪只在顶层列上进行，但格式也应该支持对嵌套列进行修剪。
- 通过将过滤器和 schema 传递到文件格式中来构造读取器函数。
- 使用分区修剪谓词，枚举应该读取的文件列表。
- 将文件拆分为任务，并构造一个
FileScanRDD。 - 添加扫描后必须求值的任何投影（project）或过滤器。
使用以下算法将文件分配到任务中：
- 如果表是分桶的，则按
bucket id将文件分组到正确数量的分区中。 - 如果表没有分桶或者分桶功能是关闭的：
-
- 如果任何文件大于阈值，那么根据阈值将其拆分为多个部分。
-
- 通过减小文件大小对文件进行排序。
-
- 使用以下算法将已排序的文件分配给桶：如果当前分区在添加下一个文件时低于阈值，添加它。如果没有，打开一个新的桶并添加它。继续下一个文件。
DataSourceStrategy
object DataSourceStrategy
extends Strategy with Logging with CastSupport with PredicateHelper with SQLConfHelper
计划对使用 sources API 定义的数据源进行扫描的策略。
SpecialLimits
规划 Limit 算子的特殊场景。
object SpecialLimits extends Strategy {
override def apply(plan: LogicalPlan): Seq[SparkPlan]
Aggregation
用于为基于AggregateFunction2接口的表达式规划聚合算子。
object Aggregation extends Strategy
Window
窗口算子
object Window extends Strategy
JoinSelection
object JoinSelection extends Strategy
with PredicateHelper
with JoinSelectionHelper
根据join strategy hints、equi-join keys的可用性和joining relations的大小，选择合适的JOIN物理计划。
建议配合我的这篇博客一起来分析：Apache Spark 怎么选择 JOIN 策略？
以下是现有的JOIN策略、特点和局限性。
Broadcast hash join（BHJ）
仅支持等值 join ，而join key不需要可排序。
支持除full outer join以外的所有 join 类型。
当广播端较小时，BHJ 通常比其他 join 算法执行得更快。
然而，广播表是一种网络密集型操作，在某些情况下可能会导致 OOM 或性能不佳，尤其是在构建/广播端较大的情况下。
Shuffle hash join
仅支持等值 join，而 join key不需要可排序。
支持所有 join 类型。
从表中构建哈希映射是一项内存密集型操作，当构建端很大时，它可能会导致 OOM。
Shuffle sort merge join（SMJ）
仅支持等值 join，join key必须是可排序的。
支持所有 join 类型。
Broadcast nested loop join（BNLJ）
支持等值 join 和不等值 join。
支持所有 join 类型，但实现优化为：
- 在
right outer join中广播左侧； - 在
left outer、left semi、left anti或existence join中广播右侧； - 在一个类
inner join中广播任意一侧。
对于其他情况，我们需要多次扫描数据，这可能会非常缓慢。
Shuffle-and-replicate nested loop join（又称 cartesian product join，CPJ）
支持等值 join 和不等值 join。
只支持类 inner join。
InMemoryScans
用来在内存数据库关系InMemoryRelation中查找表等元数据信息。
object InMemoryScans extends Strategy {
def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case PhysicalOperation(projectList, filters, mem: InMemoryRelation) =>
pruneFilterProject(
projectList,
filters,
identity[Seq[Expression]], // 所有的 filter 算子仍然需要被求值
InMemoryTableScanExec(_, filters, mem)) :: Nil
case _ => Nil
}
}
其中，PhysicalOperation表示在另一个关系算子上匹配任意数量的Project或Filter操作的模式。收集所有Filter算子，将它们的条件分解，并与顶层的Project算子一起返回。
SparkScripts
用来处理一些特定脚本的策略，这里的脚本指的是应该被执行的命令。
object SparkScripts extends Strategy {
def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case logical.ScriptTransformation(script, output, child, ioschema) =>
SparkScriptTransformationExec(
script,
output,
planLater(child),
ScriptTransformationIOSchema(ioschema)
) :: Nil
case _ => Nil
}
}
ScriptTransformation表示通过分叉并运行指定的脚本来转换输入。参数：script–应该执行的命令。 output–脚本生成的属性。 ioschema–脚本执行中应用的输入和输出模式。
case class ScriptTransformation(
script: String,
output: Seq[Attribute],
child: LogicalPlan,
ioschema: ScriptInputOutputSchema) extends UnaryNode
SparkScriptTransformationExec表示通过分叉并运行指定的脚本来转换输入的物理执行算子。参数和上面的类似：script–应该执行的命令。 output–脚本生成的属性。 child–输出被转换的逻辑计划。 ioschema–定义如何处理输入/输出数据的类集。
case class SparkScriptTransformationExec(
script: String,
output: Seq[Attribute],
child: SparkPlan,
ioschema: ScriptTransformationIOSchema)
extends BaseScriptTransformationExec
WithCTEStrategy
规划 LEFT 没有内联的 CTE 关系的策略。
关于 CTE 请参考我的博客——一篇文章了解 SQL 中的 CTE
object WithCTEStrategy extends Strategy {
override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case WithCTE(plan, cteDefs) =>
val cteMap = QueryExecution.cteMap
cteDefs.foreach { cteDef =>
cteMap.put(cteDef.id, cteDef)
}
planLater(plan) :: Nil
case r: CTERelationRef =>
val ctePlan = QueryExecution.cteMap(r.cteId).child
val projectList = r.output.zip(ctePlan.output).map { case (tgtAttr, srcAttr) =>
Alias(srcAttr, tgtAttr.name)(exprId = tgtAttr.exprId)
}
val newPlan = Project(projectList, ctePlan)
// 将 CTE 引用规划为 repartition shuffle，以便同一个 CTE 定义的所有引用在运行时共享 Exchange 重用。
exchange.ShuffleExchangeExec(
RoundRobinPartitioning(conf.numShufflePartitions),
planLater(newPlan),
REPARTITION_BY_COL) :: Nil
case _ => Nil
}
}
WithCTE表示UnresolvedWith的解析版本，CTE 引用通过唯一 ID 而不是关系别名链接到 CTE 定义。 参数： plan–查询计划。 cteDefs–CTE定义。
case class WithCTE(plan: LogicalPlan, cteDefs: Seq[CTERelationDef]) extends LogicalPlan
UnresolvedWith用于保存命名公共表表达式（CTE）和查询计划的容器。在分析过程中，将删除此算子，并将关系替换为子关系。
CTERelationRef表示 CTE 引用的关系。 参数：cteId – 相应 CTE 定义的 ID。 _resolved – 此引用是否已解决。 output – 此 CTE 引用的输出属性，可以不同于属性重复消除后其相应 CTE 定义的输出。 statsOpt – 根据相应的 CTE 定义推断出的可选统计数据。
case class CTERelationRef(
cteId: Long,
_resolved: Boolean,
override val output: Seq[Attribute],
statsOpt: Option[Statistics] = None) extends LeafNode with MultiInstanceRelation
BasicOperators
只剩下最后一块“难啃的骨头”了——BasicOperators，这代表的是一些基本算子的处理策略（前面的都是特例），可以看到，基本上就是利用 Scala 的模式匹配机制，也是偏函数，碰到某种类型的逻辑计划，就将其转换成相应的物理计划，当我们想要新增一个逻辑计划时，只需要新增一个 case 分支即可。
object BasicOperators extends Strategy {
def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case d: DataWritingCommand => DataWritingCommandExec(d, planLater(d.query)) :: Nil
case r: RunnableCommand => ExecutedCommandExec(r) :: Nil
case MemoryPlan(sink, output) =>
val encoder = RowEncoder(StructType.fromAttributes(output))
val toRow = encoder.createSerializer()
LocalTableScanExec(output, sink.allData.map(r => toRow(r).copy())) :: Nil
case logical.Distinct(child) =>
throw new IllegalStateException(
"logical distinct operator should have been replaced by aggregate in the optimizer")
case logical.Intersect(left, right, false) =>
throw new IllegalStateException(
"logical intersect operator should have been replaced by semi-join in the optimizer")
case logical.Intersect(left, right, true) =>
throw new IllegalStateException(
"logical intersect operator should have been replaced by union, aggregate" +
" and generate operators in the optimizer")
case logical.Except(left, right, false) =>
throw new IllegalStateException(
"logical except operator should have been replaced by anti-join in the optimizer")
case logical.Except(left, right, true) =>
throw new IllegalStateException(
"logical except (all) operator should have been replaced by union, aggregate" +
" and generate operators in the optimizer")
case logical.ResolvedHint(child, hints) =>
throw new IllegalStateException(
"ResolvedHint operator should have been replaced by join hint in the optimizer")
case Deduplicate(_, child) if !child.isStreaming =>
throw new IllegalStateException(
"Deduplicate operator for non streaming data source should have been replaced " +
"by aggregate in the optimizer")
case logical.DeserializeToObject(deserializer, objAttr, child) =>
execution.DeserializeToObjectExec(deserializer, objAttr, planLater(child)) :: Nil
case logical.SerializeFromObject(serializer, child) =>
execution.SerializeFromObjectExec(serializer, planLater(child)) :: Nil
case logical.MapPartitions(f, objAttr, child) =>
execution.MapPartitionsExec(f, objAttr, planLater(child)) :: Nil
case logical.MapPartitionsInR(f, p, b, is, os, objAttr, child) =>
execution.MapPartitionsExec(
execution.r.MapPartitionsRWrapper(f, p, b, is, os), objAttr, planLater(child)) :: Nil
case logical.FlatMapGroupsInR(f, p, b, is, os, key, value, grouping, data, objAttr, child) =>
execution.FlatMapGroupsInRExec(f, p, b, is, os, key, value, grouping,
data, objAttr, planLater(child)) :: Nil
case logical.FlatMapGroupsInRWithArrow(f, p, b, is, ot, key, grouping, child) =>
execution.FlatMapGroupsInRWithArrowExec(
f, p, b, is, ot, key, grouping, planLater(child)) :: Nil
case logical.MapPartitionsInRWithArrow(f, p, b, is, ot, child) =>
execution.MapPartitionsInRWithArrowExec(
f, p, b, is, ot, planLater(child)) :: Nil
case logical.FlatMapGroupsInPandas(grouping, func, output, child) =>
execution.python.FlatMapGroupsInPandasExec(grouping, func, output, planLater(child)) :: Nil
case f @ logical.FlatMapCoGroupsInPandas(_, _, func, output, left, right) =>
execution.python.FlatMapCoGroupsInPandasExec(
f.leftAttributes, f.rightAttributes,
func, output, planLater(left), planLater(right)) :: Nil
case logical.MapInPandas(func, output, child) =>
execution.python.MapInPandasExec(func, output, planLater(child)) :: Nil
case logical.PythonMapInArrow(func, output, child) =>
execution.python.PythonMapInArrowExec(func, output, planLater(child)) :: Nil
case logical.AttachDistributedSequence(attr, child) =>
execution.python.AttachDistributedSequenceExec(attr, planLater(child)) :: Nil
case logical.MapElements(f, _, _, objAttr, child) =>
execution.MapElementsExec(f, objAttr, planLater(child)) :: Nil
case logical.AppendColumns(f, _, _, in, out, child) =>
execution.AppendColumnsExec(f, in, out, planLater(child)) :: Nil
case logical.AppendColumnsWithObject(f, childSer, newSer, child) =>
execution.AppendColumnsWithObjectExec(f, childSer, newSer, planLater(child)) :: Nil
case logical.MapGroups(f, key, value, grouping, data, objAttr, child) =>
execution.MapGroupsExec(f, key, value, grouping, data, objAttr, planLater(child)) :: Nil
case logical.FlatMapGroupsWithState(
f, keyDeserializer, valueDeserializer, grouping, data, output, stateEncoder, outputMode,
isFlatMapGroupsWithState, timeout, hasInitialState, initialStateGroupAttrs,
initialStateDataAttrs, initialStateDeserializer, initialState, child) =>
FlatMapGroupsWithStateExec.generateSparkPlanForBatchQueries(
f, keyDeserializer, valueDeserializer, initialStateDeserializer, grouping,
initialStateGroupAttrs, data, initialStateDataAttrs, output, timeout,
hasInitialState, planLater(initialState), planLater(child)
) :: Nil
case logical.CoGroup(f, key, lObj, rObj, lGroup, rGroup, lAttr, rAttr, oAttr, left, right) =>
execution.CoGroupExec(
f, key, lObj, rObj, lGroup, rGroup, lAttr, rAttr, oAttr,
planLater(left), planLater(right)) :: Nil
case r @ logical.Repartition(numPartitions, shuffle, child) =>
if (shuffle) {
ShuffleExchangeExec(r.partitioning, planLater(child), REPARTITION_BY_NUM) :: Nil
} else {
execution.CoalesceExec(numPartitions, planLater(child)) :: Nil
}
case logical.Sort(sortExprs, global, child) =>
execution.SortExec(sortExprs, global, planLater(child)) :: Nil
case logical.Project(projectList, child) =>
execution.ProjectExec(projectList, planLater(child)) :: Nil
case logical.Filter(condition, child) =>
execution.FilterExec(condition, planLater(child)) :: Nil
case f: logical.TypedFilter =>
execution.FilterExec(f.typedCondition(f.deserializer), planLater(f.child)) :: Nil
case e @ logical.Expand(_, _, child) =>
execution.ExpandExec(e.projections, e.output, planLater(child)) :: Nil
case logical.Sample(lb, ub, withReplacement, seed, child) =>
execution.SampleExec(lb, ub, withReplacement, seed, planLater(child)) :: Nil
case logical.LocalRelation(output, data, _) =>
LocalTableScanExec(output, data) :: Nil
case CommandResult(output, _, plan, data) => CommandResultExec(output, plan, data) :: Nil
case logical.LocalLimit(IntegerLiteral(limit), child) =>
execution.LocalLimitExec(limit, planLater(child)) :: Nil
case logical.GlobalLimit(IntegerLiteral(limit), child) =>
execution.GlobalLimitExec(limit, planLater(child)) :: Nil
case union: logical.Union =>
execution.UnionExec(union.children.map(planLater)) :: Nil
case g @ logical.Generate(generator, _, outer, _, _, child) =>
execution.GenerateExec(
generator, g.requiredChildOutput, outer,
g.qualifiedGeneratorOutput, planLater(child)) :: Nil
case _: logical.OneRowRelation =>
execution.RDDScanExec(Nil, singleRowRdd, "OneRowRelation") :: Nil
case r: logical.Range =>
execution.RangeExec(r) :: Nil
case r: logical.RepartitionByExpression =>
val shuffleOrigin = if (r.partitionExpressions.isEmpty && r.optNumPartitions.isEmpty) {
REBALANCE_PARTITIONS_BY_NONE
} else if (r.optNumPartitions.isEmpty) {
REPARTITION_BY_COL
} else {
REPARTITION_BY_NUM
}
exchange.ShuffleExchangeExec(r.partitioning, planLater(r.child), shuffleOrigin) :: Nil
case r: logical.RebalancePartitions =>
val shuffleOrigin = if (r.partitionExpressions.isEmpty) {
REBALANCE_PARTITIONS_BY_NONE
} else {
REBALANCE_PARTITIONS_BY_COL
}
exchange.ShuffleExchangeExec(r.partitioning, planLater(r.child), shuffleOrigin) :: Nil
case ExternalRDD(outputObjAttr, rdd) => ExternalRDDScanExec(outputObjAttr, rdd) :: Nil
case r: LogicalRDD =>
RDDScanExec(r.output, r.rdd, "ExistingRDD", r.outputPartitioning, r.outputOrdering) :: Nil
case _: UpdateTable =>
throw QueryExecutionErrors.ddlUnsupportedTemporarilyError("UPDATE TABLE")
case _: MergeIntoTable =>
throw QueryExecutionErrors.ddlUnsupportedTemporarilyError("MERGE INTO TABLE")
case logical.CollectMetrics(name, metrics, child) =>
execution.CollectMetricsExec(name, metrics, planLater(child)) :: Nil
case _ => Nil
}
}
上面的源码看起来又长又难以理解，所以，我们按照顺序将上面的转换逻辑放到下面的表格中：
| 转换前的 LogicalPlan | LogicalPlan 的解释说明 | 转换后的 SparkPlan | SparkPlan 的解释说明 | 补充说明 |
|---|---|---|---|---|
| DataWritingCommand | 一种特殊的命令，用于写出数据并更新度量。 | DataWritingCommandExec | 执行DataWritingCommand的run方法并保存结果以防止多次执行的物理算子。 | |
| RunnableCommand | 因其副作用而执行的逻辑命令。在执行期间，RunnableCommands被包装在ExecutedCommand中。 | ExecutedCommandExec | 执行RunnableCommand的run方法并保存结果以防止多次执行的物理运算符。 | |
| MemoryPlan | 用来查询写入到MemorySink中的数据。 | LocalTableScanExec | 用于扫描本地集合中的数据的物理计划节点。 | MemorySink是将结果存储在内存中的Sink。当前的Sink主要用于单元测试，不提供持久性保障。 |
| Distinct | 返回删除输入行重复数据的新逻辑计划。 | - | - | 抛出异常：在 optimizer 中，逻辑 distinct 算子应该被聚合替换 |
| Intersect(left, right, false) | 返回left和right两个逻辑计划的交集，等同于 SQL 中的INTERSECT，其中 false 表示不带ALL，即删除重复行。 | - | - | 抛出异常：在 optimizer 中，逻辑 intersect 算子应该被 semi-join 替换 |
| Intersect(left, right, true) | 返回left和right两个逻辑计划的交集，等同于 SQL 中的INTERSECT ALL，其中 true 表示带ALL，即保留重复行，返回全部数据。 | - | - | 抛出异常：在 optimizer 中，逻辑 intersect 算子应该被 union，aggregate，generate 算子替换 |
| Except(left, right, false) | 返回left逻辑计划中排除right逻辑计划的部分，等同于 SQL 中的EXCEPT，其中 false 表示不带ALL，即删除重复行。 | - | - | 抛出异常： 在 optimizer 中，逻辑 except 算子应该被 anti-join 替换 |
| Except(left, right, true) | 返回left逻辑计划中排除right逻辑计划的部分，等同于 SQL 中的EXCEPT ALL，其中 true 表示带ALL，即保留重复行，返回全部数据。 | - | - | 抛出异常： 在 optimizer 中，逻辑 except(all) 算子应该被union，aggregate，generate 算子替换 |
| ResolvedHint | 已解析的Hint节点。analyzer应将所有UnsolvedHint转换为ResolvedHint。在optimization阶段开始之前，此节点将被删除。 | - | - | 抛出异常：在 optimizer 中，ResolvedHint 算子应该被 join hint替换 |
| Deduplicate(_, child) if !child.isStreaming | 删除重复数据的逻辑计划，且子节点数据不是来自流式数据源 | - | - | 抛出异常：在 optimizer 中，非流式数据源的 Deduplicate 算子应该被聚合替换 |
| DeserializeToObject | 从子节点获取输入行，并使用给定的反序列化器表达式将其转换为对象。 | DeserializeToObjectExec | 从子节点获取输入行，并使用给定的反序列化器表达式将其转换为对象。这个算子的输出是包含反序列化对象的单个字段safe row。 | safe row和unsafe row对应，即数据行存放在 JVM 内存，而不是存放在直接内存中。 |
| SerializeFromObject | 从子节点获取输入对象，并使用给定的序列化表达式将其转换为unsafe row。 | SerializeFromObjectExec | 从子节点获取输入对象，并使用给定的序列化表达式将其转换为unsafe row。其子节点的输出必须是包含输入对象的单个字段数据行。 | |
| MapPartitions | 将func函数应用于子节点的每个分区所产生的关系。 | MapPartitionsExec | 将给定函数应用于输入对象的迭代器。其子节点的输出必须是包含输入对象的单个字段数据行。 | |
| MapPartitionsInR | 将序列化的R函数func应用于子节点的每个分区所产生的关系。 | MapPartitionsExec(MapPartitionsRWrapper(f, p, b, is, os), objAttr, planLater(child)) | 将给定函数应用于输入对象的迭代器。其子节点的输出必须是包含输入对象的单个字段数据行。 | MapPartitionsRWrapper表示将给定的R函数应用于每个分区的函数包装器。 |
| FlatMapGroupsInR | R中的FlatMapGroups操作 | FlatMapGroupsInRExec | 将输入行分组，对每个组和包含组中所有元素的迭代器调用R函数。这个函数的结果会在输出前被flatten。 | FlatMapGroups会将给定函数应用于每组数据。对于每个唯一的组，函数将被传递分组键和包含组中所有元素的迭代器。这个函数可以返回包含任意类型元素的迭代器，这些元素将作为新数据集返回。 这个函数不支持部分聚合，因此需要 shuffle Dataset中的所有数据。如果应用程序打算对每个键执行聚合，最好使用reduce函数或org.apache.spark.sql.expressions#Aggregator。 在内部，如果任何给定的组太大而无法放入内存，则实现将溢出到磁盘。然而，用户必须注意避免为一个组具象化整个迭代器（例如，通过调用toList），除非给定集群内存受限是允许的。 |
| FlatMapGroupsInRWithArrow | 与FlatMapGroupsInR类似，但以Arrow格式序列化/反序列化输入和输出。这也和FlatMapGroupsInPandas有些相似。 | FlatMapGroupsInRWithArrowExec | 与FlatMapGroupsInRExec类似，但以Arrow格式序列化/反序列化输入和输出。这和FlatMapGroupsInPandasExec有些相似。 | Apache Arrow为平面和分层数据定义了一种独立于语言的列式内存格式，用于在 CPU 和 GPU 等现代硬件上进行高效的分析操作。Arrow内存格式还支持零拷贝读取（Zero Copy），以实现闪电般的数据访问，而无需序列化开销。 |
| MapPartitionsInRWithArrow | 与MapPartitionsInR类似，但以Arrow格式序列化/反序列化输入和输出。 这与org.apache.spark.sql.execution.python.ArrowEvalPython有些相似。 | MapPartitionsInRWithArrowExec | 与MapPartitionsExec和MapPartitionsRWrapper类似。但以Arrow格式序列化/反序列化输入和输出。这和ArrowEvalPythonExec有些相似。 | ArrowEvalPython是用来求值PythonUDF的物理计划。MapPartitionsRWrapper是将给定的R函数应用于每个分区的函数包装器。 |
| FlatMapGroupsInPandas | pandas中的FlatMapGroups操作 | FlatMapGroupsInPandasExec | FlatMapGroupsInPandas的物理节点。每个组中的数据行都作为Arrow记录批传递给Python worker。Python worker将记录批处理转换为pandas.DataFrame，调用用户自定义函数，并传递结果pandas.DataFrame作为Arrow记录批处理。最后，使用ColumnarBatch将每个记录批转换为Iterator[InternalRow]。 关于内存使用的注意事项：Python worker和Java executor都需要有足够的内存来容纳最大的组。Java 端的内存用于构造记录批（堆外内存）。Python 一侧的内存用于存放pandas.DataFrame。可以进一步将一个组拆分为多个记录批，以减少 Java 端的内存占用，这将留待以后来实现。 | |
| FlatMapCoGroupsInPandas | pandas的FlatMapCogroups操作 | FlatMapCoGroupsInPandasExec | FlatMapCoGroupsInPandas的物理节点。输入的DataFrame是第一次cogroup，cogroup中每一侧的数据行都通过Arrow传递给Python worker。因为cogroup的任一侧可能有不同的schema，我们把每一个分组都发送到他自己的Arrow流中。Python worker将记录批处理转换为pandas.DataFrame，调用用户自定义函数，并传递结果pandas.DataFrame作为Arrow记录批处理。最后，使用ColumnarBatch将每个记录批转换为Iterator[InternalRow]。 关于内存使用的注意事项：Python worker和Java executor都需要有足够的内存来容纳最大的cogroup。Java 端的内存用于构造记录批（堆外内存）。Python 一侧的内存用于存放pandas.DataFrame。可以进一步将一个组拆分为多个记录批，以减少 Java 端的内存占用，这将留待以后来实现。 | |
| MapInPandas | 函数Dataset.mapInPandas中用到。将标量迭代器PandasUDF 应用于每个分区。用户自定义函数定义了一个转换：iter(pandas.DataFrame) -> iter(pandas.DataFrame)。每个分区都是由DataFrame作为批组成的迭代器。 这个函数使用Apache Arrow作为 Java executor和Python worker之间的序列化格式。 | MapInPandasExec | 通过应用一个函数产生的关系，该函数接受pandas DataFrame的迭代器并输出pandas DataFrame的迭代器。 | |
| PythonMapInArrow | pyspark 的DataFrame.mapInArrow中使用。以Arrow格式将函数应用于每个分区。用户自定义函数定义了一个转换：iter(pyarrow.RecordBatch) -> iter(pyarrow.RecordBatch)。每个分区都是由pyarrow.RecordBatch作为批组成的迭代器。 | PythonMapInArrowExec | 通过应用一个函数生成的关系，该函数接受PyArrow记录批次的迭代器，并输出PyArrow记录批次的迭代器。 | |
| AttachDistributedSequence | 一种逻辑计划，它添加一个新的长列，其名称name会一个接一个地增加。这是 Spark 上pandas API中的distributed-sequence默认索引。 | AttachDistributedSequenceExec | 使用sequenceAttr添加一个新的长列的物理计划，该列将逐个增加。这是 Spark 上 pandas API中的distributed-sequence默认索引。 | |
| MapElements | 将func应用于child子节点的每个元素所产生的关系。 | MapElementsExec | 将给定函数应用于每个输入对象。其子对象的输出必须是包含输入对象的单个字段行。 这个算子是ProjectExec的一种安全版本，因为它的输出是自定义对象，所以我们需要使用safe row来包含它。 | safe row和unsafe row对应，即数据行存放在 JVM 内存，而不是存放在直接内存中。 |
| AppendColumns | 将func应用于child子节点的每个元素，并将结果列连接到输入行末尾所产生的关系。 | AppendColumnsExec | 将给定函数应用于每个输入行，并将编码后的结果追加到行的末尾。 | |
| AppendColumnsWithObject | AppendColumns的优化版本，可以直接在反序列化对象上执行。 | AppendColumnsWithObjectExec | AppendColumnsExec的优化版本，可直接在反序列化对象上执行。 | |
| MapGroups | 根据groupingAttributes的计算结果，将func应用于child子节点中的每个唯一组。func是用一个表示分组键的对象来调用的。分组键是一个迭代器，其中包含表示该键的所有行的对象。 | MapGroupsExec | 将输入行分组，并使用每个组和包含组中所有元素的迭代器调用函数。此函数的结果在输出前被flatten。 | |
| FlatMapGroupsWithState | 在使用状态数据时，基于对groupingAttributes的求值，将func应用于child子节点中的每个唯一组。func是用一个表示分组键的对象来调用的。分组键是一个迭代器，其中包含表示该键的所有行的对象。 | FlatMapGroupsWithStateExec | FlatMapGroupsWithState的执行物理算子。 | |
| CoGroup | 将func应用于每个分组键以及来自左、右子级的关联值所产生的关系。 | CoGroupExec | 将来自左、右子级的数据进行co-group，并使用每个组和 2 个迭代器调用该函数，其中包含来自左、右两侧的组中的所有元素。此函数的结果在输出前被flatten。 | |
| Repartition | 返回一个新的 RDD，该 RDD 正好包含numPartitions分区。与RepartitionByExpression不同，因为此方法由DataFrame直接调用，因为用户要求coalesce或者repartition。当输出的使用者需要对数据进行特定的排序或分发时，可以使用RepartitionByExpression。 | 如果Shuffle，则ShuffleExchangeExec；否则，CoalesceExec | 1. ShuffleExchangeExec表示执行一次Shuffle，将产生所需的分区。2. CoalesceExec表示返回一个新 RDD 的物理计划，该 RDD 正好有numPartitions个分区。与 RDD 上定义的coalesce类似，此操作会产生一个窄依赖，例如，如果从 1000 个分区到 100 个分区，则不会出现Shuffle，而是 100 个新分区中的每一个都会占用当前分区中的 10 个。如果请求更多的分区，它将保持当前的分区数。 但是，如果正在进行剧烈合并，例如 numPartitions=1，这可能会导致计算发生在希望的更少的节点上（例如，numPartitions=1的情况下是一个节点）。为了避免这种情况，可以看到ShuffleExchange。这将添加一个Shuffle步骤，但意味着当前的上游分区将并行执行（无论当前分区是什么）。 | |
| Sort | 排序操作 | SortExec | 执行（外部）排序操作的物理计划。 | |
| Project | 投影操作 | ProjectExec | Project的物理计划。 | |
| Filter | 过滤操作 | FilterExec | Filter的物理计划。 | |
| TypedFilter | 将func应用于child子节点的每个元素并按结果布尔值过滤它们所产生的关系。 这在逻辑上等于一个普通的Filter算子，其条件表达式是将输入行解码为对象，并将给定函数应用于解码的对象。然而，我们需要对TypedFilter进行封装，以使概念更清晰，并使编写优化器规则更容易。 | FilterExec | Filter的物理计划。 | typedCondition是一个条件表达式。 |
| Expand | 对每个输入行应用多个投影，因此我们将为一个输入行获得多个输出行。 | ExpandExec | 将所有GroupExpression应用于每个输入行，因此我们将为一个输入行获得多个输出行。 | |
| Sample | Dataset的取样操作。 | SampleExec | Dataset取样操作的物理计划。 | |
| LocalRelation | 从本地集合中扫描数据的逻辑计划节点。 | LocalTableScanExec | 用于扫描本地集合中的数据的物理计划节点。 | |
| CommandResult | 用于保存命令数据的逻辑计划节点。 | CommandResultExec | 用于保存来自命令的数据的物理计划节点。 | |
| LocalLimit | 分区局部的Limit | LocalLimitExec | 获取每个子分区的第一个Limit元素，但不收集或Shuffle它们。 | |
| GlobalLimit | 全局的Limit | GlobalLimitExec | 获取子节点的单个输出分区的第一个Limit元素。 | |
| Union | 用于联合多个计划的逻辑计划，没有去重。这是SQL中的UNION ALL。 | UnionExec | 合并两个计划的物理计划，没有去重。这是SQL中的UNION ALL。 | |
| Generate | 将生成器（Generator）应用于输入行流，并将每个行的输出组合成新的行流。此操作类似于函数式编程中的flatMap，它有一个重要的附加功能，即允许将输入行与其输出连接（join）起来。 | GenerateExec | 将生成器（Generator）应用于输入行流，并将每个行的输出组合成新的行流。此操作类似于函数式编程中的flatMap，它有一个重要的附加功能，即允许将输入行与其输出连接（join）起来。这个算子支持没有实现方法terminate()的生成器的全阶段代码生成（WSCG）。 | |
| OneRowRelation | 一个数据行的关系。这在没有FROM子句的SELECT ...子句中会用到。 | RDDScanExec | 扫描来自 RDD[InternalRow] 的数据的物理计划节点。 | |
| Range | 表示范围，类似 Scala 中的Range，类变量包括开始，结束和步长。 | RangeExec | Range 的物理计划节点（生成 64 位的数字范围）。 | |
| RepartitionByExpression | 此方法使用表达式（Expression）将数据重新分区为optNumPartitions，并在执行期间接收有关分区数的信息。当查询结果的使用者期望有特定的排序或者分布时使用。对类 RDD 的coalesce和repartition使用Repartition。如果没有给出optNumPartitions，默认情况下，它会将数据划分为SQLConf 中定义的numShufflePartitions，并可以由AQE合并。 | ShuffleExchangeExec | 执行一次Shuffle，将产生所需的分区。 | SQLConf 中的numShufflePartitions 在AQE开启并且spark.sql.adaptive.coalescePartitions.enabled为true情况下为配置参数：spark.sql.adaptive.coalescePartitions.initialPartitionNum，否则，为配置参数：spark.sql.shuffle.partitions。 |
| RebalancePartitions | 这个算子用于重新平衡给定子节点（child）的输出分区，以便每个分区都具有合理的大小（不太小也不太大）。它还尽力用partitionExpressions对子节点的输出进行分区。如果有数据倾斜，Spark 将切分倾斜的分区，使这些分区不会太大。当需要将子节点的结果写入表时，这个算子非常有用，以避免文件太小/太大。 请注意，这个算子仅在启用AQE时才有意义。 | ShuffleExchangeExec | 执行一次Shuffle，将产生所需的分区。 | |
| ExternalRDD | 扫描来自一个 RDD 的数据的逻辑计划节点。 | ExternalRDDScanExec | 扫描来自一个 RDD 的数据的物理计划执行节点。 | |
| LogicalRDD | 扫描来自 RDD[InternalRow] 的数据的逻辑计划节点 | RDDScanExec | 扫描来自一个 RDD 的数据的物理计划执行节点。 | |
| UpdateTable | UPDATE TABLE命令的逻辑计划。 | - | - | 抛出异常：目前不支持UPDATE TABLE。 |
| MergeIntoTable | MERGE INTO命令的逻辑计划 | - | - | 抛出异常：目前不支持MERGE INTO TABLE。 |
| CollectMetrics | 从数据集中收集任意（命名的）度量信息。一旦查询到达完成点（批处理查询完成或流式查询完成），Driver上就会发出一个事件，可以通过将监听器连接到spark session来观察该事件。这些度量信息被命名，因此我们可以在一个数据集中的多个位置收集它们。这个节点的行为类似于全局聚合。收集的所有度量必须是聚合函数或者文本值。 | CollectMetricsExec | 收集来自SparkPlan的任意（已命名）度量信息。 |
策略小结
通过前面的分析，我们知道：Spark 中负责将逻辑计划转换成物理计划的策略总共有 15 种类型。
- experimentalMethods.extraStrategies
- extraPlanningStrategies
- LogicalQueryStageStrategy
- PythonEvals
- new DataSourceV2Strategy(session)
- FileSourceStrategy
- DataSourceStrategy
- SpecialLimits
- Aggregation
- Window
- JoinSelection
- InMemoryScans
- SparkScripts
- WithCTEStrategy
- BasicOperators
其中，前 2 种类型策略是为了方便扩展。
第 3 针对的是包含LogicalQueryStage节点的逻辑计划。
第 4 针对的是 pyspark。
第 5 针对的是 DataSource V2。
第 6 针对的是可以通过指定列来分区/分桶的文件数据源。
第 7 针对的是 DataSource V1。
第 8 针对的是 Limit算子。
第 9 针对的是聚合算子。
第 10 针对的是窗口算子。
第 11 针对的是 Join 策略的选择。
第 12 针对的是内存中的数据库关系InMemoryRelation。
第 13 针对的是脚本命令。
第 14 针对的是 CTE。
最后一项是通用的逻辑计划算子，是用来处理前面策略中未曾覆盖的逻辑计划算子的策略。
我们将我们提到的策略画成思维导图，并将每一种策略的类型和与之对应的物理执行算子标识出来。
3. prepareForExecution
/**
* 根据需要，通过插入 shuffle 操作和内部行格式转换，准备一个计划好的 SparkPlan 以供执行。
*/
private[execution] def prepareForExecution(
preparations: Seq[Rule[SparkPlan]],
plan: SparkPlan): SparkPlan = {
// PlanChangeLogger 此处用来记录物理计划的变更
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
}
这部分内容比较简单，经过前面 2 个阶段我们得到了执行规则和物理计划之后，此处我们会将这些规则一一应用到物理计划上。
此处，我们又发现了一个熟悉的身影——PlanChangeLogger，我们通过配置 Spark 参数：spark.sql.planChangeLog.level，就可以实现在控制台中看到物理计划的变更过程。
planning 阶段小结
第一步preparations我们准备好物理计划的执行规则序列，这些规则会确保子查询得到规划，使用正确的数据分区和排序，插入 WSCG，并尝试通过重用Exchange和子查询来减少所做的工作。
第二步sparkPlan.clone()我们得到了物理计划的副本，这一步当中我们会应用不同的策略将逻辑计划转换成物理计划序列，在最终的序列中，当前版本的实现是只选择第一个物理计划。
第三步prepareForExecution我们所做的事情就是将第一步的规则一一应用到第二步中得到的物理计划中。
回到最初的例子
和前面 2 讲的结构类似，我们需要基于最初的例子来理解优化后的逻辑计划是怎样转换成物理计划的。
通过配置了spark.sql.planChangeLog.level，我们得到了下面控制台的日志：
22/04/26 20:00:25 INFO FileSourceStrategy: Pushed Filters:
22/04/26 20:00:25 INFO FileSourceStrategy: Post-Scan Filters:
22/04/26 20:00:25 INFO FileSourceStrategy: Output Data Schema: struct<addr: array<string>, age: bigint, name: string, sex: string ... 2 more fields>
22/04/26 20:00:25 INFO PlanChangeLogger:
=== Applying Rule org.apache.spark.sql.execution.CollapseCodegenStages ===
CollectLimit 21 CollectLimit 21
!+- Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11] +- *(1) Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11]
+- FileScan json [addr#8,age#9L,name#10,sex#11] Batched: false, DataFilters: [], Format: JSON, Location: InMemoryFileIndex(1 paths)[file:/Users/shockang/code/spark-examples/data/sql/user.json], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<addr:array<string>,age:bigint,name:string,sex:string> +- FileScan json [addr#8,age#9L,name#10,sex#11] Batched: false, DataFilters: [], Format: JSON, Location: InMemoryFileIndex(1 paths)[file:/Users/shockang/code/spark-examples/data/sql/user.json], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<addr:array<string>,age:bigint,name:string,sex:string>
22/04/26 20:00:25 INFO PlanChangeLogger:
=== Result of Batch Preparations ===
CollectLimit 21 CollectLimit 21
!+- Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11] +- *(1) Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11]
+- FileScan json [addr#8,age#9L,name#10,sex#11] Batched: false, DataFilters: [], Format: JSON, Location: InMemoryFileIndex(1 paths)[file:/Users/shockang/code/spark-examples/data/sql/user.json], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<addr:array<string>,age:bigint,name:string,sex:string> +- FileScan json [addr#8,age#9L,name#10,sex#11] Batched: false, DataFilters: [], Format: JSON, Location: InMemoryFileIndex(1 paths)[file:/Users/shockang/code/spark-examples/data/sql/user.json], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<addr:array<string>,age:bigint,name:string,sex:string>
上面的日志中涉及到 2 个类：
- FileSourceStrategy
- CollapseCodegenStages
我们来一一分析：
FileSourceStrategy
我们看下这个策略是如何实现的：
def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case ScanOperation(projects, filters,
l @ LogicalRelation(fsRelation: HadoopFsRelation, _, table, _)) =>
// 根据我们可以在哪里使用过滤器来避免读取不需要的数据，这个关系的过滤器可分为四类
// -仅分区键-用于修剪目录以读取
// -bucket key only（仅适用于bucket key）——可选地用于修剪要读取的文件
// -仅存储在数据中的键-可选地用于跳过文件中的数据组
// -扫描后需要再次评估的过滤器
val filterSet = ExpressionSet(filters)
// 如果查询`analyzer`不区分大小写，则属性名称可能与 schema 中的名称不同。
// 我们应该更改属性名称以匹配 schema 中的名称，这样就不必再担心大小写敏感了。
val normalizedFilters = DataSourceStrategy.normalizeExprs(
filters.filter(_.deterministic), l.output)
// 将给定 schema 解析为此查询计划中的具体属性（Attribute）引用。
// resolve 函数只应该在已解析（analyzed）的计划上调用，因为如果有未解析的属性会引发 AnalysisException。
val partitionColumns =
l.resolve(
fsRelation.partitionSchema, fsRelation.sparkSession.sessionState.analyzer.resolver)
val partitionSet = AttributeSet(partitionColumns)
// 此处的 partitionKeyFilters 应该和 PruneFileSourcePartitions中执行的相同
val partitionKeyFilters = DataSourceStrategy.getPushedDownFilters(partitionColumns,
normalizedFilters)
// 子查询表达式会被过滤掉，因为它们不能用于修剪桶或作为数据过滤器下推，它们会被执行
val normalizedFiltersWithoutSubqueries =
normalizedFilters.filterNot(SubqueryExpression.hasSubquery)
val bucketSpec: Option[BucketSpec] = fsRelation.bucketSpec
val bucketSet = if (shouldPruneBuckets(bucketSpec)) {
genBucketSet(normalizedFiltersWithoutSubqueries, bucketSpec.get)
} else {
None
}
val dataColumns =
l.resolve(fsRelation.dataSchema, fsRelation.sparkSession.sessionState.analyzer.resolver)
// 分区键在文件的统计信息中不可用。
// `dataColumns`可能有分区列，我们需要过滤掉它们。
val dataColumnsWithoutPartitionCols = dataColumns.filterNot(partitionColumns.contains)
val dataFilters = normalizedFiltersWithoutSubqueries.flatMap { f =>
if (f.references.intersect(partitionSet).nonEmpty) {
extractPredicatesWithinOutputSet(f, AttributeSet(dataColumnsWithoutPartitionCols))
} else {
Some(f)
}
}
val supportNestedPredicatePushdown =
DataSourceUtils.supportNestedPredicatePushdown(fsRelation)
val pushedFilters = dataFilters
.flatMap(DataSourceStrategy.translateFilter(_, supportNestedPredicatePushdown))
logInfo(s"Pushed Filters: ${pushedFilters.mkString(",")}")
// 分区键和属性的谓词需要在扫描后执行
val afterScanFilters = filterSet -- partitionKeyFilters.filter(_.references.nonEmpty)
logInfo(s"Post-Scan Filters: ${afterScanFilters.mkString(",")}")
val filterAttributes = AttributeSet(afterScanFilters)
val requiredExpressions: Seq[NamedExpression] = filterAttributes.toSeq ++ projects
val requiredAttributes = AttributeSet(requiredExpressions)
val readDataColumns =
dataColumns
.filter(requiredAttributes.contains)
.filterNot(partitionColumns.contains)
val outputSchema = readDataColumns.toStructType
logInfo(s"Output Data Schema: ${outputSchema.simpleString(5)}")
val metadataStructOpt = l.output.collectFirst {
case FileSourceMetadataAttribute(attr) => attr
}
val metadataColumns = metadataStructOpt.map { metadataStruct =>
metadataStruct.dataType.asInstanceOf[StructType].fields.map { field =>
FileSourceMetadataAttribute(field.name, field.dataType)
}.toSeq
}.getOrElse(Seq.empty)
// outputAttributes 的最后还应该包含元数据列
val outputAttributes = readDataColumns ++ partitionColumns ++ metadataColumns
val scan =
FileSourceScanExec(
fsRelation,
outputAttributes,
outputSchema,
partitionKeyFilters.toSeq,
bucketSet,
None,
dataFilters,
table.map(_.identifier))
// 额外的 Project 节点： 将平面的元数据列包装到一个元数据结构体中
val withMetadataProjections = metadataStructOpt.map { metadataStruct =>
val metadataAlias =
Alias(CreateStruct(metadataColumns), METADATA_NAME)(exprId = metadataStruct.exprId)
execution.ProjectExec(
scan.output.dropRight(metadataColumns.length) :+ metadataAlias, scan)
}.getOrElse(scan)
val afterScanFilter = afterScanFilters.toSeq.reduceOption(expressions.And)
val withFilter = afterScanFilter
.map(execution.FilterExec(_, withMetadataProjections))
.getOrElse(withMetadataProjections)
val withProjections = if (projects == withFilter.output) {
withFilter
} else {
execution.ProjectExec(projects, withFilter)
}
withProjections :: Nil
case _ => Nil
}
这个方法的执行流程图如下：
变化
转换前的逻辑计划：
ReturnAnswer
+- GlobalLimit 21
+- LocalLimit 21
+- Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11]
+- Relation [addr#8,age#9L,name#10,sex#11] json
转换后的物理计划：
CollectLimit 21
!+- Project [cast(addr#8 as string) AS addr#24, cast(age#9L as string) AS age#25, name#10, sex#11]
+- FileScan json [addr#8,age#9L,name#10,sex#11] Batched: false, DataFilters: [], Format: JSON, Location: InMemoryFileIndex(1 paths)[file:/Users/shockang/code/spark-examples/data/sql/user.json], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<addr:array<string>,age:bigint,name:string,sex:string>
CollapseCodegenStages
建议先阅读我的这篇博客——Spark 3.x 的 WSCG 机制源码解析
case class CollapseCodegenStages(
codegenStageCounter: AtomicInteger = new AtomicInteger(0))
extends Rule[SparkPlan] {
private def supportCodegen(e: Expression): Boolean = e match {
case e: LeafExpression => true
// CodegenFallback 要求输入为 InternalRow
case e: CodegenFallback => false
case _ => true
}
private def supportCodegen(plan: SparkPlan): Boolean = plan match {
case plan: CodegenSupport if plan.supportCodegen =>
val willFallback = plan.expressions.exists(_.find(e => !supportCodegen(e)).isDefined)
// 如果有太多列的话，生成的代码量会很大。
val hasTooManyOutputFields =
WholeStageCodegenExec.isTooManyFields(conf, plan.schema)
val hasTooManyInputFields =
plan.children.exists(p => WholeStageCodegenExec.isTooManyFields(conf, p.schema))
!willFallback && !hasTooManyOutputFields && !hasTooManyInputFields
case _ => false
}
/**
* 在那些不支持 codegen 的物理计划顶部上插入一个 InputAdapter。
*/
private def insertInputAdapter(plan: SparkPlan): SparkPlan = {
plan match {
case p if !supportCodegen(p) =>
// 递归地折叠它们
InputAdapter(insertWholeStageCodegen(p))
case j: SortMergeJoinExec =>
// SortMergeJoin 的子级应该单独地进行 codegen
j.withNewChildren(j.children.map(
child => InputAdapter(insertWholeStageCodegen(child))))
case j: ShuffledHashJoinExec =>
// ShuffledHashJoin 的子级应该单独地进行 codegen
j.withNewChildren(j.children.map(
child => InputAdapter(insertWholeStageCodegen(child))))
case p => p.withNewChildren(p.children.map(insertInputAdapter))
}
}
/**
* 在那些支持 codegen 的物理计划顶部上插入一个 WholeStageCodegen。
*/
private def insertWholeStageCodegen(plan: SparkPlan): SparkPlan = {
plan match {
// 对于将输出领域对象的算子，不要插入 WholeStageCodegen，因为它作为领域对象无法写入 unsafe row。
case plan if plan.output.length == 1 && plan.output.head.dataType.isInstanceOf[ObjectType] =>
plan.withNewChildren(plan.children.map(insertWholeStageCodegen))
case plan: LocalTableScanExec =>
// 不要因为要支持 Driver 本地的快速 collect/take 路径，从而将 LogicalTableScanExec 设置为 WholeStageCodegen 的根节点
plan
case plan: CommandResultExec =>
// 不要因为要支持 Driver 本地的快速 collect/take 路径，从而将 CommandResultExec 设置为 WholeStageCodegen 的根节点
plan
case plan: CodegenSupport if supportCodegen(plan) =>
// 全阶段代码生成的框架是基于行的。
// 如果一个计划支持列式执行，它就不能同时支持 WSCG。
assert(!plan.supportsColumnar)
WholeStageCodegenExec(insertInputAdapter(plan))(codegenStageCounter.incrementAndGet())
case other =>
other.withNewChildren(other.children.map(insertWholeStageCodegen))
}
}
def apply(plan: SparkPlan): SparkPlan = {
// 即 Spark 配置参数：spark.sql.codegen.wholeStage
if (conf.wholeStageEnabled) {
insertWholeStageCodegen(plan)
} else {
plan
}
}
}
代码生成结果：
public Object generate(Object[] references) {
return new GeneratedIteratorForCodegenStage1(references);
}
// codegenStageId=1
final class GeneratedIteratorForCodegenStage1 extends org.apache.spark.sql.execution.BufferedRowIterator {
private Object[] references;
private scala.collection.Iterator[] inputs;
private scala.collection.Iterator inputadapter_input_0;
private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] project_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[1];
public GeneratedIteratorForCodegenStage1(Object[] references) {
this.references = references;
}
public void init(int index, scala.collection.Iterator[] inputs) {
partitionIndex = index;
this.inputs = inputs;
inputadapter_input_0 = inputs[0];
project_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(4, 128);
}
private UTF8String project_elementToString_0(UTF8String element) {
UTF8String elementStr = null;
elementStr = UTF8String.fromString(String.valueOf(element));
return elementStr;
}
protected void processNext() throws java.io.IOException {
while ( inputadapter_input_0.hasNext()) {
InternalRow inputadapter_row_0 = (InternalRow) inputadapter_input_0.next();
// common sub-expressions
boolean inputadapter_isNull_0 = inputadapter_row_0.isNullAt(0);
ArrayData inputadapter_value_0 = inputadapter_isNull_0 ?
null : (inputadapter_row_0.getArray(0));
boolean project_isNull_0 = inputadapter_isNull_0;
UTF8String project_value_0 = null;
if (!inputadapter_isNull_0) {
org.apache.spark.unsafe.UTF8StringBuilder project_buffer_0 = new org.apache.spark.unsafe.UTF8StringBuilder();
project_buffer_0.append("[");
if (inputadapter_value_0.numElements() > 0) {
if (inputadapter_value_0.isNullAt(0)) {
project_buffer_0.append("null");
} else {
project_buffer_0.append(project_elementToString_0(inputadapter_value_0.getUTF8String(0)));
}
for (int project_loopIndex_0 = 1; project_loopIndex_0 < inputadapter_value_0.numElements(); project_loopIndex_0++) {
project_buffer_0.append(",");
if (inputadapter_value_0.isNullAt(project_loopIndex_0)) {
project_buffer_0.append(" null");
} else {
project_buffer_0.append(" ");
project_buffer_0.append(project_elementToString_0(inputadapter_value_0.getUTF8String(project_loopIndex_0)));
}
}
}
project_buffer_0.append("]");;
project_value_0 = project_buffer_0.build();
}
boolean inputadapter_isNull_1 = inputadapter_row_0.isNullAt(1);
long inputadapter_value_1 = inputadapter_isNull_1 ?
-1L : (inputadapter_row_0.getLong(1));
boolean project_isNull_2 = inputadapter_isNull_1;
UTF8String project_value_2 = null;
if (!inputadapter_isNull_1) {
project_value_2 = UTF8String.fromString(String.valueOf(inputadapter_value_1));
}
boolean inputadapter_isNull_2 = inputadapter_row_0.isNullAt(2);
UTF8String inputadapter_value_2 = inputadapter_isNull_2 ?
null : (inputadapter_row_0.getUTF8String(2));
boolean inputadapter_isNull_3 = inputadapter_row_0.isNullAt(3);
UTF8String inputadapter_value_3 = inputadapter_isNull_3 ?
null : (inputadapter_row_0.getUTF8String(3));
project_mutableStateArray_0[0].reset();
project_mutableStateArray_0[0].zeroOutNullBytes();
if (project_isNull_0) {
project_mutableStateArray_0[0].setNullAt(0);
} else {
project_mutableStateArray_0[0].write(0, project_value_0);
}
if (project_isNull_2) {
project_mutableStateArray_0[0].setNullAt(1);
} else {
project_mutableStateArray_0[0].write(1, project_value_2);
}
if (inputadapter_isNull_2) {
project_mutableStateArray_0[0].setNullAt(2);
} else {
project_mutableStateArray_0[0].write(2, inputadapter_value_2);
}
if (inputadapter_isNull_3) {
project_mutableStateArray_0[0].setNullAt(3);
} else {
project_mutableStateArray_0[0].write(3, inputadapter_value_3);
}
append((project_mutableStateArray_0[0].getRow()));
if (shouldStop()) return;
}
}
}
总结
本讲我们完成了 Spark SQL planning 阶段的源码解析。
从一开始的QueryExecution.executedPlan方法中，我们了解到，planning阶段实际上可以分成了 3 个部分：
- preparations
- sparkPlan.clone()
- prepareForExecution
第一部分我们准备好物理计划的执行规则序列，这些规则会确保子查询得到规划，使用正确的数据分区和排序，插入 WSCG，并尝试通过重用Exchange和子查询来减少所做的工作。
第二部分我们得到了物理计划的副本，这一步当中我们会应用不同的策略将逻辑计划转换成物理计划序列，在最终的序列中，当前版本的实现是只选择第一个物理计划。
第三部分我们所做的事情就是将第一步的规则一一应用到第二步中得到的物理计划中。
由于篇幅的限制，我们没法做到详解每一种策略和规则，故我们选择了我们 2 个典型的代表：
- FileSourceStrategy 策略
- CollapseCodegenStages 规则
相信搞懂了一个典型的策略加上一个典型的规则能够助我们触类旁通。
至此，Spark SQL 的 4 个阶段的源码解析我们已经全部搞定。
本文是第一个版本，后面有时间会慢慢修改完善。暂时先发出来(*￣︶￣)
扫一扫
944
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
