小知识点源码解析-STAGE由最后一个RDD确定并行度的源码解析

以ShuffleMapStage 为例进行解析。

1      假设

RDD A – ShuffleDependency – RDD B

即 RDD B依赖 RDD A，并且依赖关系为宽依赖—— 依赖，针对的是两个RDD之间的关系，RDD可以有多个父依赖RDD，但针对每个父依赖RDD都会有对应的具体依赖。

2      关键源码及其解析

关键源码为DAGScheduler类中构建ShuffleMapStage 的代码—— 该方法中同时包含了Stage复用的概念。

  /**

   * Create a shuffle map Stage for the given RDD.  The stage will also be associated with the

   * provided firstJobId.  If a stage for the shuffleId existed previously so that the shuffleId is

   * present in the MapOutputTracker, then the number and location of available outputs are

   * recovered from the MapOutputTracker

   */

  private def newOrUsedShuffleStage(

      shuffleDep: ShuffleDependency[_, _, _],

      firstJobId: Int): ShuffleMapStage = {

 

// 根据假设： RDD B 依赖 RDD A，此时，shuffleDep表示两者直接的依赖关系，

// 对应的shuffleDep.rdd 为 RDD A， 即RDD A 是当前Stage的最后一个RDD

val rdd = shuffleDep.rdd

// 传入的任务数numTasks，对应的是Stage 最后一个RDD（即RDD A）

// 的分区个数。

// 而在后续由 ShuffleMapStage构建 ShuffleMapTask 时numTasks会

// 控制Task的个数，因此该Stage 的并行度是由其最后一个RDD的分区控制。

// 对应也体现了，是被宽依赖的 RDD A 需要根据 RDD B 所需的数据结构进行重组

// 重组的数据存入磁盘中（当前实现）

val numTasks = rdd.partitions.length

    val stage = newShuffleMapStage(rdd, numTasks, shuffleDep, firstJobId, rdd.creationSite)

    if (mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {

      val serLocs = mapOutputTracker.getSerializedMapOutputStatuses(shuffleDep.shuffleId)

      val locs = MapOutputTracker.deserializeMapStatuses(serLocs)

      (0 until locs.length).foreach { i =>

        if (locs(i) ne null) {

          // locs(i) will be null if missing

          stage.addOutputLoc(i, locs(i))

        }

      }

    } else {

      // Kind of ugly: need to register RDDs with the cache and map output tracker here

      // since we can't do it in the RDD constructor because # of partitions is unknown

      logInfo("Registering RDD " + rdd.id + " (" + rdd.getCreationSite + ")")

      mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)

    }

    stage

  }

 

ShuffleMapTask构建的对应代码在DAGScheduler类的submitMissingTasks方法中，具体如下：

    val tasks: Seq[Task[_]] = try {

      stage match {

        case stage: ShuffleMapStage =>

          partitionsToCompute.map { id =>

            val locs = taskIdToLocations(id)

            val part = stage.rdd.partitions(id)

            new ShuffleMapTask(stage.id, stage.latestInfo.attemptId,

              taskBinary, part, locs, stage.internalAccumulators)

          }

 

3      扩展

Stage 是由最后一个RDD的分区数控制并行度，因此在进行重分区时，需要注意最后的分区个数如果过小，则并行度降低，进而影响效率。此时可以采用宽依赖，切分Stage，保证前面的并行度不会降低。

比如，使用coalesce将大分区数合并成小分区数时，由于没有Shuffle（默认），此时最终执行的并行度为合并后的小分区个数 —— 可以考虑默认Shuffle为true的 repartition方法，认为加上宽依赖，保证前面计算的高并行度。