Spark3 AQE之自動合并Shuffle partition原始碼解讀

Branch：spark-3.0
有不對的地方歡迎各位大佬批評指正！
相關引數：
spark.sql.adaptive.enabled AQE是否開啟
spark.sql.adaptive.coalescePartitions.enabled 磁區合并是否開啟
spark.sql.adaptive.coalescePartitions.minPartitionNum 合并后最小的磁區數，下文我們簡稱為minPartitionNum
spark.sql.adaptive.advisoryPartitionSizeInBytes 開發者建議的磁區大小，下文我們簡稱為advisoryPartitionSizeInBytes

一、代碼入口

類：AdaptiveSparkPlanExec.scala
入口：queryStageOptimizerRules —> CoalesceShufflePartitions —> coalescePartitions

  @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq(
    ReuseAdaptiveSubquery(conf, context.subqueryCache),
    CoalesceShufflePartitions(context.session),
    // The following two rules need to make use of 'CustomShuffleReaderExec.partitionSpecs'
    // added by `CoalesceShufflePartitions`. So they must be executed after it.
    OptimizeSkewedJoin(conf),
    OptimizeLocalShuffleReader(conf)
  )

二、流程

1、判斷

判斷AQE是否開啟、所有葉子結點是否都是查詢階段（如果不是的話合并磁區會破壞所有子節點具有相同數量的輸出磁區假設）

if (!conf.coalesceShufflePartitionsEnabled) {
      return plan
    }
if (!plan.collectLeaves().forall(_.isInstanceOf[QueryStageExec])
        || plan.find(_.isInstanceOf[CustomShuffleReaderExec]).isDefined) {
      // If not all leaf nodes are query stages, it's not safe to reduce the number of
      // shuffle partitions, because we may break the assumption that all children of a spark plan
      // have same number of output partitions. 
      return plan
}

2、收集Stage

將一棵樹所有節點的ShuffleStage收集起來，為接下來磁區合并使用

    def collectShuffleStages(plan: SparkPlan): Seq[ShuffleQueryStageExec] = plan match {
      case stage: ShuffleQueryStageExec => Seq(stage)
      case _ => plan.children.flatMap(collectShuffleStages)
    }

    val shuffleStages = collectShuffleStages(plan)

3、執行

首先會判斷這些Shuffle是否能夠進行磁區合并，如果不能的話會直接將plan回傳

if (!shuffleStages.forall(_.shuffle.canChangeNumPartitions)) {
      plan

若Shuffle的磁區已被確定好，則此Stage也會跳過

val validMetrics = shuffleStages.flatMap(_.mapStats)

多個Task進行Shuffle時，每個Task需要具有相同的磁區數才能進行合并

val distinctNumPreShufflePartitions =
        validMetrics.map(stats => stats.bytesByPartitionId.length).distinct
if (validMetrics.nonEmpty && distinctNumPreShufflePartitions.length == 1){
........
}

取最小的磁區數，如果未定義則取Spark默認的并行度（為了避免磁區合并后的性能退化）

val minPartitionNum = conf.getConf(SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM)
          .getOrElse(session.sparkContext.defaultParallelism)

進入真正執行的方法
引數：validMetrics.toArray（shuffleID和每個磁區的大小）
advisoryTargetSize（開發者定義的磁區大小默認值64M）
minNumPartitions（最小磁區數）

val partitionSpecs = ShufflePartitionsUtil.coalescePartitions(
          validMetrics.toArray,
          advisoryTargetSize = conf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES),
          minNumPartitions = minPartitionNum)

4、coalescePartitions真正執行磁區合并的方法

advisoryTargetSize數值的重新設定
如果inputSize=1000M 10磁區 而設定advisoryTargetSize為200M 則通過一下計算會排除200M這個設定
advisoryTargetSize=maxTargetSize

// 所有磁區的總大小
val totalPostShuffleInputSize = mapOutputStatistics.map(_.bytesByPartitionId.sum).sum
// 磁區大小的最大值
val maxTargetSize = math.max(
      math.ceil(totalPostShuffleInputSize / minNumPartitions.toDouble).toLong, 16)
// 確定真正的磁區大小(避免上述例子的情況出現)
val targetSize = math.min(maxTargetSize, advisoryTargetSize)

磁區合并

    while (i < numPartitions) {// 我們從所有shuffle中計算第i個shuffle磁區的總大小 對于每個task中的每個相鄰磁區合并，直到不大于targetSize
      // 從所有洗牌中計算第i次洗牌磁區的總大小，
      var totalSizeOfCurrentPartition = 0L
      var j = 0
      while (j < mapOutputStatistics.length) {// 對每個shuffle中的partition進行合并
        totalSizeOfCurrentPartition += mapOutputStatistics(j).bytesByPartitionId(i)
        j += 1
      }

      // 如果包含' totalSizeOfCurrentPartition '將超過目標大小，則啟動一個新的合并磁區，
      if (i > latestSplitPoint && coalescedSize + totalSizeOfCurrentPartition > targetSize) {
        partitionSpecs += CoalescedPartitionSpec(latestSplitPoint, i)
        latestSplitPoint = i
        // 重置postShuffleInputSize
        coalescedSize = totalSizeOfCurrentPartition
      } else {
        coalescedSize += totalSizeOfCurrentPartition
      }
      i += 1
    }

最后 將合并的磁區回傳

    partitionSpecs += CoalescedPartitionSpec(latestSplitPoint, numPartitions)

    partitionSpecs

官方例子（說實話這例子我還沒驗證成功）

 * For example, we have two shuffles with the following partition size statistics:
   *  - shuffle 1 (5 partitions): [100 MiB, 20 MiB, 100 MiB, 10MiB, 30 MiB]
   *  - shuffle 2 (5 partitions): [10 MiB,  10 MiB, 70 MiB,  5 MiB, 5 MiB]
   * Assuming the target size is 128 MiB, we will have 4 coalesced partitions, which are:
   *  - coalesced partition 0: shuffle partition 0 (size 110 MiB)
   *  - coalesced partition 1: shuffle partition 1 (size 30 MiB)
   *  - coalesced partition 2: shuffle partition 2 (size 170 MiB)
   *  - coalesced partition 3: shuffle partition 3 and 4 (size 50 MiB)
   *
   *  @return A sequence of [[CoalescedPartitionSpec]]s. For example, if partitions [0, 1, 2, 3, 4]
   *          split at indices [0, 2, 3], the returned partition specs will be:
   *          CoalescedPartitionSpec(0, 2), CoalescedPartitionSpec(2, 3) and
   *          CoalescedPartitionSpec(3, 5).

三、結束

結合官方給出的對應單側看原始碼能夠更快的理解
磁區合并對應的測驗類ShufflePartitionsUtilSuite

引流一波，強烈推薦極客時間吳磊老師講的spark課程，非常非常好！看完后受益匪淺，而且從每節課下面的評論都能學到很多東西！！！
強烈推薦！！！
做與Spark相關作業的小伙伴可無腦入手！