前言在 Spark SQL 物理计划已经生成之后、真正执行之前还会通过一系列 rule如 AQE、EnsureRequirements、WSCG 等对整棵物理计划进行全局整合或优化。普通 Join、Group By 等算子所需的隐式 Shuffle通常是在EnsureRequirements中补充的。例如一个聚合 SQL 生成的初始物理计划是HashAggregateExec(final) - HashAggregateExec(partial) - FileSourceScanExec经过EnsureRequirements后HashAggregateExec(final) - ShuffleExchangeExec(HashPartitioning(user_id, 200)) - HashAggregateExec(partial) - FileSourceScanExecSpark SQL 的 Shuffle 分区数通常是在执行EnsureRequirements并插入ShuffleExchangeExec时确定的。接下来本文基于 Spark 3.1.2 源码视角详细剖析 Spark SQL 发生 Shuffle 时的分区数决定机制。EnsureRequirements介绍org.apache.spark.sql.execution.exchange.EnsureRequirementsobjectEnsureRequirementsextendsRule[SparkPlan]{privatedefensureDistributionAndOrdering(operator:SparkPlan):SparkPlan{// 1、获取当前算子要求 child 需要满足的数据分布和排序要求valrequiredChildDistributions:Seq[Distribution]operator.requiredChildDistributionvalrequiredChildOrderings:Seq[Seq[SortOrder]]operator.requiredChildOrderingvarchildren:Seq[SparkPlan]operator.children assert(requiredChildDistributions.lengthchildren.length)assert(requiredChildOrderings.lengthchildren.length)// 2、判断当前算子 child 的数据分布是否满足当前算子的输入要求并根据不同结果作处理// Ensure that the operators children satisfy their output distribution requirements.childrenchildren.zip(requiredChildDistributions).map{case(child,distribution)ifchild.outputPartitioning.satisfies(distribution)childcase(child,BroadcastDistribution(mode))BroadcastExchangeExec(mode,child)case(child,distribution)valnumPartitionsdistribution.requiredNumPartitions.getOrElse(conf.numShufflePartitions)ShuffleExchangeExec(distribution.createPartitioning(numPartitions),child)}//....}这段代码的作用是对当前物理算子operator的每个 child 进行检查判断 child 的输出分布是否满足父算子的输入要求。如果不满足就在已有 child 上方插入BroadcastExchangeExec或ShuffleExchangeExec。operator.childrenvarchildren:Seq[SparkPlan]operator.childrenoperator.children表示当前物理算子的输入子计划。单 child 场景HashAggregateExec - FileSourceScanExec对HashAggregateExec来说children Seq(FileSourceScanExec)多 child 场景SortMergeJoinExec :- leftPlan - rightPlan对SortMergeJoinExec来说children Seq(leftPlan, rightPlan)叶子节点没有 child例如FileSourceScanExec其自身的children Seq.empty。而ProjectExec、FilterExec这类一元算子通常会有一个 child。operator.requiredChildDistributionvalrequiredChildDistributions:Seq[Distribution]operator.requiredChildDistributionrequiredChildDistribution定义了当前物理算子要求所有 child 需要满足的数据分布Distribution。requiredChildDistribution在SparkPlan中定义所有物理算子都会包含该方法。org.apache.spark.sql.execution.SparkPlandefrequiredChildDistribution:Seq[Distribution]Seq.fill(children.size)(UnspecifiedDistribution)默认含义是有几个 child就返回几个UnspecifiedDistribution对 child 的数据分区不关心。如果算子没有重写requiredChildDistribution通常表示它对 child 数据分区没有特殊要求比如ProjectExec、FilterExec这类一元算子。而常见的会发生 Shuffle 的算子如HashAggregateExec、SortMergeJoinExec会要求 child 满足一定的数据分布因此会重写requiredChildDistribution。常见 Shuffle 算子要求 child 满足的数据分布Distribution不同算子对 child 的数据分布要求是不一样的下面列举几种常见 Shuffle 算子所要求的数据分布。① HashAggregateExecHashAggregateExec的requiredChildDistribution方法在BaseAggregateExec中实现。org/apache/spark/sql/execution/aggregate/BaseAggregateExecoverridedefrequiredChildDistribution:List[Distribution]{requiredChildDistributionExpressionsmatch{caseSome(exprs)ifexprs.isEmptyAllTuples::NilcaseSome(exprs)ClusteredDistribution(exprs)::NilcaseNoneUnspecifiedDistribution::Nil}}有指定表达式且为空这种场景常见于全局聚合如SELECT COUNT(*) FROM t要求所有数据进入同一个分区Distribution数据分布要求为AllTuples。指定了表达式且非空这种场景常见于普通分组聚合如SELECT key, COUNT(*) FROM t GROUP BY key要求 child 节点输出的数据按照这些表达式聚集或分区数据分布要求为ClusteredDistribution。没有指定任何分布要求对 child 节点的数据分布没有要求。② SortMergeJoinExec 或 ShuffledHashJoinExec对于 Shuffle Join例如SortMergeJoinExec或ShuffledHashJoinExec它们有两个 child并且对左右 child 都有ClusteredDistribution数据分布要求按照 Join key 做相同的分布。org/apache/spark/sql/execution/joins/ShuffledJoinoverridedefrequiredChildDistribution:Seq[Distribution]{ClusteredDistribution(leftKeys)::ClusteredDistribution(rightKeys)::Nil}上面频繁出现Distribution这个概念它到底是什么下面一起看看。Distributionorg/apache/spark/sql/catalyst/plans/physical/DistributionsealedtraitDistribution{// 该 Distribution 要求的分区数量。如果是 None表示任意分区数量都可以满足这个 Distribution。defrequiredNumPartitions:Option[Int]// 为该 Distribution 创建一个默认的 Partitioning。这个 Partitioning 既能满足当前 Distribution 的要求又能匹配给定的分区数量defcreatePartitioning(numPartitions:Int):Partitioning}Distribution表示“数据应该怎么分布”当一个查询在多台机器上并行执行时具有相同表达式值的元组应该如何分布。前面涉及了不同算子所要求的Distribution下面将介绍常见的Distribution类型UnspecifiedDistribution不对 child 输入数据的分区方式提出要求如普通投影、过滤。caseobjectUnspecifiedDistributionextendsDistribution{overridedefrequiredNumPartitions:Option[Int]NoneoverridedefcreatePartitioning(numPartitions:Int):Partitioning{thrownewIllegalStateException(UnspecifiedDistribution does not have default partitioning.)}}AllTuples只有一个分区的分布表示整个数据集必须放在一个分区里通常用于需要全局处理的操作。caseobjectAllTuplesextendsDistribution{overridedefrequiredNumPartitions:Option[Int]Some(1)overridedefcreatePartitioning(numPartitions:Int):Partitioning{assert(numPartitions1,The default partitioning of AllTuples can only have 1 partition.)SinglePartition}}ClusteredDistribution要求相同 key 的数据落在同一分区通常用于 Group By、窗口函数。caseclassClusteredDistribution(clustering:Seq[Expression],requiredNumPartitions:Option[Int]None)extendsDistribution{require(clustering!Nil,// 参与数据分布的表达式不能为空The clustering expressions of a ClusteredDistribution should not be Nil.)overridedefcreatePartitioning(numPartitions:Int):Partitioning{assert(requiredNumPartitions.isEmpty||requiredNumPartitions.getnumPartitions,sThis ClusteredDistribution requires${requiredNumPartitions.get}partitions, but sthe actual number of partitions is$numPartitions.)HashPartitioning(clustering,numPartitions)}}HashClusteredDistribution不仅要求相同 key 的数据在一起还要求按 hash 规则进入指定分区常用于 join 算子。相比ClusteredDistribution它提供了更强的保障。caseclassHashClusteredDistribution(expressions:Seq[Expression],requiredNumPartitions:Option[Int]None)extendsDistribution{require(expressions!Nil,// 参与数据分布的表达式不能为空The expressions of a HashClusteredDistribution should not be Nil.)overridedefcreatePartitioning(numPartitions:Int):Partitioning{assert(requiredNumPartitions.isEmpty||requiredNumPartitions.getnumPartitions,sThis HashClusteredDistribution requires${requiredNumPartitions.get}partitions, but sthe actual number of partitions is$numPartitions.)HashPartitioning(expressions,numPartitions)}}BroadcastDistribution是一种广播分布要求数据会被广播到每个执行节点上。caseclassBroadcastDistribution(mode:BroadcastMode)extendsDistribution{overridedefrequiredNumPartitions:Option[Int]Some(1)// 广播数据在被广播前通常会先收集成一个整体然后作为一个广播变量发到各个 executor。所以从 Spark 物理分区语义上看默认分区数要求是 1。overridedefcreatePartitioning(numPartitions:Int):Partitioning{assert(numPartitions1,The default partitioning of BroadcastDistribution can only have 1 partition.)BroadcastPartitioning(mode)}}Partitioningorg/apache/spark/sql/catalyst/plans/physical/partitioningtraitPartitioning{// 该 SparkPlan 输出 RDD 的分区数目valnumPartitions:Int// 当前 child 的 Partitioning 是否能够满足下游要求的 Distribution不满足时返回 false。// 要满足需要两个条件① 分区数 numPartitions 要相等// ② satisfies0 方法返回 true包括两种情况1. 对子节点的分布没有要求2. 全局处理。finaldefsatisfies(required:Distribution):Boolean{required.requiredNumPartitions.forall(_numPartitions)satisfies0(required)}// 1、如果 requiredChildDistribution 为 UnspecifiedDistribution则说明对子节点的分布没有要求返回 true// 2、如果 requiredChildDistribution 为 AllTuples则只要 numPartitions 1返回 true// 3、其他情况返回 false。// 具体的 Partitioning 类型会对 satisfies0 进行重写。protecteddefsatisfies0(required:Distribution):Booleanrequiredmatch{caseUnspecifiedDistributiontruecaseAllTuplesnumPartitions1case_false}}Partitioning定义了一个物理算子输出数据的分区方式。常见Partitioning类型有UnknownPartitioning未知分区方式SinglePartition单分区BroadcastPartitioning广播分区HashPartitioning基于哈希的分区方式RangePartitioning基于范围的分区方式通过确定分区键是否在某个范围内来选择分区Distribution 和 Partitioning 的关系Distribution是“要求”Partitioning是“实现”。比如某个物理算子要求 child 中相同 key 的数据必须在同一个分区这是一种Distribution要求。Spark 为了满足它可能会创建HashPartitioning(keys, numPartitions)这就是具体的Partitioning实现也就是通过hash(key)的方式把数据分到指定数量的分区中。简单总结Distribution 算子需要什么样的数据分布 Partitioning 当前数据实际是什么样的分区方式 Shuffle 当现有 Partitioning 不能满足 Distribution 时Spark 插入 Exchange 来重新分区children.zip(requiredChildDistributions)childrenchildren.zip(requiredChildDistributions).map{case(child,distribution)ifchild.outputPartitioning.satisfies(distribution)childcase(child,BroadcastDistribution(mode))BroadcastExchangeExec(mode,child)case(child,distribution)valnumPartitionsdistribution.requiredNumPartitions.getOrElse(conf.numShufflePartitions)ShuffleExchangeExec(distribution.createPartitioning(numPartitions),child)}这一步会把当前算子的每个 child 和它要求的Distribution配对。例如SortMergeJoinExec算子children: leftPlan rightPlan requiredChildDistributions: ClusteredDistribution(leftJoinKeys) ClusteredDistribution(rightJoinKeys)配对后变成(leftPlan, ClusteredDistribution(leftJoinKeys)) (rightPlan, ClusteredDistribution(rightJoinKeys))后续map会逐个检查每个 child 是否满足对应的Distribution。第一分支child 已经满足分布要求case(child,distribution)ifchild.outputPartitioning.satisfies(distribution)child如果 child 的输出分区方式已经满足父算子的分布要求就直接返回原 child。例如ProjectExec、FilterExec等物理算子并不要求 child 具备特定的输入分布。第二分支需要 BroadcastExchangeExeccase(child,BroadcastDistribution(mode))BroadcastExchangeExec(mode,child)如果父算子要求 child 满足BroadcastDistribution说明该 child 不需要普通 Shuffle而是需要被广播。典型场景是 Broadcast Hash JoinSELECT/* BROADCAST(s) */*FROMbig_table bJOINsmall_table sONb.ids.id;对于小表s父算子会要求BroadcastDistribution(HashedRelationBroadcastMode(keys))第三分支需要 ShuffleExchangeExeccase(child,distribution)valnumPartitionsdistribution.requiredNumPartitions.getOrElse(conf.numShufflePartitions)ShuffleExchangeExec(distribution.createPartitioning(numPartitions),child)如果 child 输出分区不满足当前算子的Distribution且当前算子要求的不是 broadcast就进入普通 Shuffle 分支。1 确定 Shuffle 分区数valnumPartitionsdistribution.requiredNumPartitions.getOrElse(conf.numShufflePartitions)如果distribution.requiredNumPartitions有值则优先使用它。例如AllTuples.requiredNumPartitions Some(1)表示强制单分区。如果没有强制分区数则使用conf.numShufflePartitions。org/apache/spark/sql/internal/SQLConfdefdefaultNumShufflePartitions:IntgetConf(SHUFFLE_PARTITIONS)defnumShufflePartitions:Int{if(adaptiveExecutionEnabledcoalesceShufflePartitionsEnabled){getConf(COALESCE_PARTITIONS_INITIAL_PARTITION_NUM).getOrElse(defaultNumShufflePartitions)}else{defaultNumShufflePartitions}}因此未开启 AQE 时shuffle 分区数等于spark.sql.shuffle.partitions。开启 AQE 且开启 Shuffle 分区合并时初始分区数可能来自spark.sql.adaptive.coalescePartitions.initialPartitionNum如果没有设置initialPartitionNum仍然回退到spark.sql.shuffle.partitions。相关参数①spark.sql.shuffle.partitions该参数表示在进行 join 或聚合操作且需要 shuffle 数据时默认使用的分区数量。默认值是 200。valSHUFFLE_PARTITIONSbuildConf(spark.sql.shuffle.partitions).doc(The default number of partitions to use when shuffling data for joins or aggregations.).version(1.1.0).intConf.checkValue(_0,The value of spark.sql.shuffle.partitions must be positive).createWithDefault(200)②spark.sql.adaptive.coalescePartitions.initialPartitionNum该参数表示在进行分区合并之前shuffle 分区的初始数量。如果没有设置该值则它等于spark.sql.shuffle.partitions。这个配置只有在spark.sql.adaptive.enabled和spark.sql.adaptive.coalescePartitions.enabled都为true时才会生效。valCOALESCE_PARTITIONS_INITIAL_PARTITION_NUMbuildConf(spark.sql.adaptive.coalescePartitions.initialPartitionNum).doc(The initial number of shuffle partitions before coalescing. If not set, it equals to s${SHUFFLE_PARTITIONS.key}. This configuration only has an effect when s${ADAPTIVE_EXECUTION_ENABLED.key} and ${COALESCE_PARTITIONS_ENABLED.key} are both true.).version(3.0.0).intConf.checkValue(_0,The initial number of partitions must be positive.).createOptionalspark.sql.adaptive.enabled表示是否开启 Spark SQL 的 AQE即 Adaptive Query Execution自适应查询执行默认关闭。spark.sql.adaptive.coalescePartitions.enabled表示是否开启 AQE 下的 shuffle 分区合并功能默认开启。2 创建具体的分区distribution.createPartitioning(numPartitions)创建满足当前算子要求的Distribution分布。常见映射关系ClusteredDistribution(keys) - HashPartitioning(keys, numPartitions) OrderedDistribution(ordering) - RangePartitioning(ordering, numPartitions) AllTuples - SinglePartition3 创建 ShuffleExchangeExecShuffleExchangeExec(partitioning,child)在当前物理算子的子节点child外面包一层ShuffleExchangeExec强制对child的输出数据做一次 shuffle。在生成新的物理计划树时当前物理算子的 child 会被替换为这个新的ShuffleExchangeExec从而在当前算子和原计划之间增加一层 shuffle。例如一个聚合 SQL 生成的初始物理计划是HashAggregateExec(final) - HashAggregateExec(partial) - FileSourceScanExec经过EnsureRequirements后HashAggregateExec(final) - ShuffleExchangeExec(HashPartitioning(user_id, 200)) - HashAggregateExec(partial) - FileSourceScanExecShuffleExchangeExec 将分区数带入 Spark CoreShuffleExchangeExec持有最终的输出分区方式org/apache/spark/sql/execution/exchange/ShuffleExchangeExeccaseclassShuffleExchangeExec(overridevaloutputPartitioning:Partitioning,child:SparkPlan,shuffleOrigin:ShuffleOriginENSURE_REQUIREMENTS)extendsShuffleExchangeLike其中outputPartitioning.numPartitions就是 SQL 物理计划层面已经确定的 Shuffle 分区数。执行时会构造 Shuffle 依赖valshuffleDependencyShuffleExchangeExec.prepareShuffleDependency(inputRDD,child.output,outputPartitioning,serializer,writeMetrics)进入 Spark Core 后分区数会体现在ShuffleDependency的partitioner上。core/src/main/scala/org/apache/spark/Dependency.scalaclassShuffleDependency[K:ClassTag,V:ClassTag,C:ClassTag](_rdd:RDD[_:Product2[K,V]],valpartitioner:Partitioner,...)其中partitioner.numPartitions决定了 Shuffle reduce 端的分区数也通常对应 reduce task 的数量。核心结论Spark SQL Shuffle 分区数决定机制主流程可以概括为父算子声明 requiredChildDistribution - EnsureRequirements 检查 child.outputPartitioning 是否满足要求 - 满足则复用 child - 要求 BroadcastDistribution 则插入 BroadcastExchangeExec - 否则插入 ShuffleExchangeExec - Shuffle 分区数优先取 distribution.requiredNumPartitions - 没有强制要求时取 conf.numShufflePartitions - conf.numShufflePartitions 通常来自 spark.sql.shuffle.partitions默认 200Spark SQL 的 Shuffle 分区数不是由单一位置决定的而是由逻辑计划、物理计划、分布需求、用户显式设置和 AQE 共同决定。分区数的优先级可以按下面理解显式重分区优先例如df.repartition(n)、repartitionByRange(n, ...)分区数直接使用用户指定的n。算子强制分区数优先例如AllTuples.requiredNumPartitions Some(1)会生成单分区 Shuffle。普通 Join、Group By、Distinct、Window 等场景通常使用conf.numShufflePartitions也就是spark.sql.shuffle.partitions。开启 AQE 后spark.sql.shuffle.partitions更多表示初始 Shuffle 分区数最终 Shuffle Read task 数可能被合并或因倾斜拆分而变化。参考Spark SQL 源码研读系列 06Executed Plan