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源码

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

核心类：AdaptiveSparkPlanExec

Spark经典的doExecute()

获取更新后的plan数

getFinalPhysicalPlan?是核心方法

看着行代码

var result = createQueryStages(currentPhysicalPlan)

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

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

Exchange 节点：Exchange 节点被替换为?QueryStageExec ，这个类有两个实现类

ShuffleQueryStageExec??用于shuffle 数据，BroadcastQueryStageExec?用于广播数据

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

然后也会替换plan

replaceWithQueryStagesInLogicalPlan 做了

以上是里面的方法调用和处理事情

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