背景
本文基于 Spark 3.5
关于ShuffleLocalRead的作用简单的来说，就是会按照一定的规则，从一个 map Task 中连续读取多个 reduce数据 的任务，（正常的情况下是读取所有map Task中特定的一个reduce数据任务），具体可以参考Spark AQE中的CoalesceShufflePartitions和OptimizeLocalShuffleReader
分析
直接上OptimizeShuffleWithLocalRead代码：
override def apply(plan: SparkPlan): SparkPlan = {if (!conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)) {return plan}plan match {case s: SparkPlan if canUseLocalShuffleRead(s) =>createLocalRead(s)case s: SparkPlan =>createProbeSideLocalRead(s)}}...def canUseLocalShuffleRead(plan: SparkPlan): Boolean = plan match {case s: ShuffleQueryStageExec =>s.mapStats.isDefined && isSupported(s.shuffle)case AQEShuffleReadExec(s: ShuffleQueryStageExec, _) =>s.mapStats.isDefined && isSupported(s.shuffle) &&s.shuffle.shuffleOrigin == ENSURE_REQUIREMENTScase _ => false}...private def createLocalRead(plan: SparkPlan): AQEShuffleReadExec = {plan match {case c @ AQEShuffleReadExec(s: ShuffleQueryStageExec, _) =>AQEShuffleReadExec(s, getPartitionSpecs(s, Some(c.partitionSpecs.length)))case s: ShuffleQueryStageExec =>AQEShuffleReadExec(s, getPartitionSpecs(s, None))}
}
这里有两种情况会引入LocalshuffleRead：
第一种是引入了REBALANCE hint的场景。这种情况下，在Spark的内部表示 ShuffleOrigin 为 REBALANCE_PARTITIONS_BY_NONE，这种情况下 是hint为REBALANCE而不是REBALANCE(c)或者REBALANCE(num)的情况；
第二种是SMJ 转变为 BHJ的场景。
至于为啥会存在AQEShuffleReadExec(s: ShuffleQueryStageExec, _)这种情况是因为CoalesceShufflePartitions 这个规则会进行分区的合并等
所以在代码中会有两个case：
-
SparkPlan if canUseLocalShuffleRead(s)
如果满足是REBALANCE hint的情况或者是Spark内部加的（为了满足Shuffle上下算子的数据分布要求）就强加上AQEShuffleReadExec -
createProbeSideLocalRead
这里是进行SMJ 转 BHJ的BuildBroadcast的另一边进行ShuffleLocalRead的情况，这种情况下，因为已经进行broadcast了，所以参与BuildBroadcast的另一边也可以进行shufflelocalRead的
针对于第一种情况 强制加上 AQEShuffleReadExec , 这种情况下在ensureRequirements规则下，有可能会增加额外的Shuffle操作，这种情况就是负优化了，所以在进行了reOptimize操作后，会进行一个判断是否有增益:
val afterReOptimize = reOptimize(logicalPlan)if (afterReOptimize.isDefined) {val (newPhysicalPlan, newLogicalPlan) = afterReOptimize.getval origCost = costEvaluator.evaluateCost(currentPhysicalPlan)val newCost = costEvaluator.evaluateCost(newPhysicalPlan)if (newCost < origCost ||(newCost == origCost && currentPhysicalPlan != newPhysicalPlan)) {logOnLevel("Plan changed:\n" +sideBySide(currentPhysicalPlan.treeString, newPhysicalPlan.treeString).mkString("\n"))cleanUpTempTags(newPhysicalPlan)currentPhysicalPlan = newPhysicalPlancurrentLogicalPlan = newLogicalPlanstagesToReplace = Seq.empty[QueryStageExec]}
这里的条件默认是根据shuffle的个数来计算的，如果优化后的shuffle数有增加，则会回退到之前的物理计划中去，当然用户也可以配置spark.sql.adaptive.customCostEvaluatorClass来实现自己的是否有增益的逻辑。
针对第二种情况，这种情况一般来说都是有正向的提升效果的，但是也会经过第一种情况的逻辑判断。
