Spark AQE源码探索

介绍

AQE全称是Adaptive Query Execution，官网介绍如下

Performance Tuning - Spark 3.5.0 Documentation

AQE做了什么

AQE 是 Spark SQL 的一种动态优化机制，在运行时，每当 Shuffle Map 阶段执行完毕，AQE 都会结合这个阶段的统计信息，基于既定的规则动态地调整、修正尚未执行的逻辑计划和物理计划，来完成对原始查询语句的运行时优化

特性

1. 自动分区合并：在 Shuffle 过后，Reduce Task 数据分布参差不齐，AQE 将自动合并过小的数据分区。
2. Join 策略调整：如果某张表在过滤之后，尺寸小于广播变量阈值，这张表参与的数据关联就会从 Shuffle Sort Merge Join 降级（Demote）为执行效率更高的 Broadcast Hash Join。
3. 自动倾斜处理：结合配置项，AQE 自动拆分 Reduce 阶段过大的数据分区，降低单个 Reduce Task 的工作负载。

AQE开启前后对比：

在非AQE的情况下，Spark会在规划阶段确定了物理执行计划后，根据每个算子的定义生成RDD对应的DAG。然后 Spark DAGScheduler通过shuffle来划分RDD Graph并创建stage，然后提交Stage以供执行

AQE：

首先会将逻辑树拆分为多个QueryStages, 在执行时先将它的子 QueryStages 被提交（先提交mapStage），收集它们的MapOutputStatistics对象。根据收集到的 shuffle 数据统计信息，将当前 QueryStage 的执行计划优化为更好的执行计划。然后转换为DAG图再执行Stage

源码

策略源码

AQE做了这么多事情，肯定有对应的执行方法（AdaptiveSparkPlanExec类），介绍几个核心的

规则方法	描述
AdaptiveExecutionRule	所有自适应执行规则的基类。所有的规则都需要继承这个类，并实现applyInternal方法来对物理计划进行优化
CustomShuffleReaderExec	用于处理倾斜的join和shuffle分区合并
CoalesceShufflePartitions	动态合并shuffle分区
OptimizeSkewedJoin	优化倾斜的join
OptimizeLocalShuffleReader	优化本地shuffle读取的策略

AQE执行源码

首先AQE 只作用在 exchange 阶段，即需要发生数据交换的阶段，spark AQE 优化都是发生在 shuffle map 之后。

核心类：AdaptiveSparkPlanExec，AQE的主要执行类，它会首先执行初始物理计划，并且在执行过程中收集有关数据统计信息，然后基于这些信息对物理计划进行优化

Spark经典的doExecute()

获取更新后的plan数

getFinalPhysicalPlan 是核心方法

看这行代码

var result = createQueryStages(currentPhysicalPlan)

将原始的物理计划通过 createQueryStages 方法进行替换==========================

对参数中传递的 sparkplan 的所有节点进行自下而上的递归。根据节点类型的不同进行针对处理

Exchange 节点：Exchange 节点被替换为 QueryStageExec ，QueryStageExec是每个查询阶段的抽象基类

最后调用CreateStageResult，参数中包含 QueryStageExec，同时 allChildStagesMaterialized 参数表示当前 plan 的所有子节点的输出是否已经具像化

========================== createQueryStages 返回 CreateStageResult ，会在 AQE 的循环中判断是否可以结束。这部分逻辑又回到 getFinalPhysicalPlan 中

然后也会替换plan，replaceWithQueryStagesInLogicalPlan和reOptimize都很关键，一个一个看

replaceWithQueryStagesInLogicalPlan 做了

再返回上诉代码，再看看reoptimize干了什么

这个里面其实又有策略执行了

返回主程序，然后接着往下走

然后这个方法会执行策略

大致上是这样的流程，当然还有很多细节，代码还是很庞大的，只能给大家看看整体流程和关键的方法，其余有深入可以关注我的公众号一起探讨

总体来看还是很多固定策略的，可以根据上一stage的write的统计信息去动态修改下游的执行计划

闲聊

spark的树，在aqe节点有两个num partition，但是会出现两者不一致且有时会相差很大，即AQE不被采用，源码可以看到实际执行和打印日志确实是独立的，但是没搞明白这个设计的特殊含义