源碼系列1.DAG之createResultStage-《每日五分鐘搞定大數據》

原創，轉載請標明出處： https://www.jianshu.com/p/c39596da86bb

本文主要關於stage的劃分，createResultStage方法裏包含了整個的劃分流程，代碼包含大量的遞歸，第一次看看起來是比較噁心的，這裏做一個整理記錄和分享。

這篇文章的一切都是從這句代碼開始：

finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)

一、總覽

圖畫得有點醜。。。。。見諒。。。。

DAG_STAGE.png

解釋一下：
1. 傳入rdd8 獲得首個寬依賴3(getShuffleDependencies)，傳入寬依賴3調用getMissingAncestorShuffleDependencies獲得3以前的所有的寬依賴關係(包括寬依賴1和2...)，由於存放的結構是先進後出的stack，所以從寬依賴1開始遍歷寬依賴3以前的每一個寬依賴關係，調用createShuffleMapStage。

2. createShuffleMapStage
     1)第一次調用createShuffleMapStage時傳進來的是寬依賴1，沒有父stage，調用getOrCreateParentStages返回空list。那麼從1開始根據“寬依賴1”，“父stage(ShuffleMapStage(1)沒有父stage)” 來 new ShuffleMapStage 獲得ShuffleMapStage(1)添加進 shuffleIdToMapStage ;
     2)第二次調用createShuffleMapStage同樣調用getOrCreateParentStages()傳入rdd4，獲得寬依賴1;getOrCreateShuffleMapStage傳入寬依賴1，shuffleIdToMapStage.get(shuffleDep.shuffleId)直接獲得ShuffleMapStage(1)。此時拿到了ShuffleMapStage(1)，根據“寬依賴2”，“父stage:ShuffleMapStage(1)” 來 new ShuffleMapStage(2) 同樣添加進 shuffleIdToMapStage。
   ............
   3. 遍歷完寬依賴3以前的所有寬依賴後，再次調用createShuffleMapStage把寬依賴3傳進去，同樣很容易根據rdd6獲得寬依賴2再獲得ShuffleMapStage(2)，返回ShuffleMapStage(3)。(此時ResultStage之前的所有stage都已經劃分好並添加進shuffleIdToMapStage)
   4. 回到createResultStage的getOrCreateParentStages方法，獲得了ShuffleMapStage(3)，創建ResultStage：new ResultStage(id, rdd, func, partitions, parents, jobId, callSite)。

總結：根據finalRdd容易得到寬依賴3，再根據寬依賴1得到所有寬依賴關係，放進stack遍歷。拿出第一個寬依賴1(沒有父stage)創建ShuffleMapStage1，添加進shuffleIdToMapStage維護寬依賴和stage的關係。拿出第二個寬依賴2推出寬依賴1後，直接從shuffleIdToMapStage拿到ShuffleMapStage1，創建ShuffleMapStage2添加進shuffleIdToMapStage。

注意：
1.關鍵的map
stageIdToStage ：自增id對應stage
shuffleIdToMapStage ：shuffleId來自shuffleDep.shuffleId 維護寬依賴和stage的關係
2.dependency和rdd的關係
rdd.dependencies:通過rdd獲得依賴關係
dependency.rdd:通過依賴關係獲得rdd
3.創建ResultStage和ShuffleMapStage

1. ResultStage：每個job只且只有一個，爲action算子前的一個stage。(主要構成條件：父stage)
   new ResultStage(id, rdd, func, partitions, parents, jobId, callSite)
2. ShuffleMapStage：每個job有0+個。 (主要構成條件：stage(可能沒有)和shuffleDep寬依賴(一定有))
   new ShuffleMapStage(id, rdd, numTasks, parents, jobId, rdd.creationSite, shuffleDep)
   4.各個類的依賴關係
   1)createResultStage(傳入finalRDD獲得ResultStage) ->2
   2)getOrCreateParentStages(傳入rdd獲得父stage) ->3->4
    3)getShuffleDependencies(傳入rdd獲得寬依賴)
    4)getOrCreateShuffleMapStage(傳入寬依賴獲得ShuffleMapStage) ->5->6
     5)getMissingAncestorShuffleDependencies(傳入一個rdd獲得所有寬依賴) ->3
     6)createShuffleMapStage(傳入寬依賴獲得ShuffleMapStage) ->2

下面來看源碼

1.createResultStage(line 354)

獲得ResultStage

private def createResultStage(rdd: RDD[_],func: (TaskContext, Iterator[_]) => _,
    partitions: Array[Int],jobId: Int,callSite: CallSite): ResultStage = {
  val parents = getOrCreateParentStages(rdd, jobId)//獲得父stage,若沒有shuffle則返回空List
  val id = nextStageId.getAndIncrement() // 到這裏時已經分好stage，這個id是ResultStage的id
   // 主要根據父stage獲得ResultStage
  val stage = new ResultStage(id, rdd, func, partitions, parents, jobId, callSite)
  stageIdToStage(id) = stage// 把ResultStage和它的ID加入stageIdToStage
  updateJobIdStageIdMaps(jobId, stage) // 更新jobIds和jobIdToStageIds
  stage// 返回這個ResultStage
}

2.getOrCreateParentStages(line 372)

獲得父stage

private def getOrCreateParentStages(rdd: RDD[_],
    firstJobId: Int): List[Stage] = {
  // 遍歷RDD的依賴關係，找到第一個ShuffleDependency(可能多個，也可能沒有)然後放入HashSet並返回
  // 如果獲取不到ShuffleDependency，邏輯在此終止返回空list
  getShuffleDependencies(rdd).map { shuffleDep =>
    // 裏面會創建當前ShuffleDependency的所有父ShuffleMapStage
    getOrCreateShuffleMapStage(shuffleDep, firstJobId)
  }.toList
}

3.getShuffleDependencies(line 414)

獲得單個rdd的所有寬依賴關係(一般只有一個寬依賴；像CoGroupedRDD有多個)

private[scheduler] def getShuffleDependencies(rdd: RDD[_])
                : HashSet[ShuffleDependency[_, _, _]] = {
  // 用來存放ShuffleDependency的HashSet
  val parents = new HashSet[ShuffleDependency[_, _, _]]
  val visited = new HashSet[RDD[_]] // 遍歷過的RDD
  val waitingForVisit = new Stack[RDD[_]]//後進先出的數據結構
  waitingForVisit.push(rdd)
  while (waitingForVisit.nonEmpty) {
    val toVisit = waitingForVisit.pop()// 取出頂部的第一個元素 RDD
    if (!visited(toVisit)) {// 若這個rdd沒有被訪問過
      visited += toVisit// 標記已訪問
//rdd.dependencies獲得當前這個rdd的所有依賴,這個依賴可能爲一個或多個
//(一般只有一個寬依賴；像CoGroupedRDD有多個)，返回的是Seq[Dependency[_]]序列類型
      toVisit.dependencies.foreach {
        case shuffleDep: ShuffleDependency[_, _, _] =>
          parents += shuffleDep
        case dependency =>// 若不是寬依賴就把這個RDD的父RDD push進waitingForVisit
          waitingForVisit.push(dependency.rdd)
      }
    }
  }
  parents//獲得傳進來的finalRDD的第一個ShuffleDependency並不是所有
}

4.getOrCreateShuffleMapStage(line 289)

根據寬依賴關係獲得stage

//第一次來到這裏傳進來的是左到右首個shuffleDep，沒有父stage，
//shuffleIdToMapStage也沒有記錄，會調getMissingAncestorShuffleDependencies並創建stage
//之後進來可直接根據shuffleIdToMapStage提取到ShuffleMapStage
private def getOrCreateShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _],
            firstJobId: Int): ShuffleMapStage = {
  //  通過shuffleDep獲得之前添加的ShuffleMapStage
  shuffleIdToMapStage.get(shuffleDep.shuffleId) match {
    case Some(stage) =>
      stage
    case None =>
//獲得所有父ShuffleDependencies遍歷，先進後出，從首個寬依賴開始遍歷直到構建所有父ShuffleMapStage
      getMissingAncestorShuffleDependencies(shuffleDep.rdd).foreach { dep =>
        if (!shuffleIdToMapStage.contains(dep.shuffleId)) {
          createShuffleMapStage(dep, firstJobId)  //裏面會遞歸調用，創建所有的ShuffleMapStage
        }
      }
      // 這個shuffleDep爲finallRDD找到的首個shuffleDep，創建這個shuffleDep的ShuffleMapStage
      createShuffleMapStage(shuffleDep, firstJobId)
  }
}

5.getMissingAncestorShuffleDependencies(line 379)

獲得所有rdd的所有寬依賴

private def getMissingAncestorShuffleDependencies(rdd: RDD[_])
                    : Stack[ShuffleDependency[_, _, _]] = {

  val ancestors = new Stack[ShuffleDependency[_, _, _]]
  val visited = new HashSet[RDD[_]]
  val waitingForVisit = new Stack[RDD[_]]
  waitingForVisit.push(rdd)
  while (waitingForVisit.nonEmpty) {
    val toVisit = waitingForVisit.pop()
    if (!visited(toVisit)) {
      visited += toVisit
      //getShuffleDependencies獲得當前rdd的shuffleDependencies
      getShuffleDependencies(toVisit).foreach { shuffleDep =>
        if (!shuffleIdToMapStage.contains(shuffleDep.shuffleId)) {
        // 若shuffleIdToMapStage不包含ShuffleDep的shuffleId，把這個ShuffleDep存進ancestors
          ancestors.push(shuffleDep)
          //把ShuffleDep的父RDD再傳進waitingForVisit，準備再遍歷
          waitingForVisit.push(shuffleDep.rdd)
        }
      }
    }
  }
  ancestors// 包含(除finalRDD的第一個ShuffleDep以外)所有的ShuffleDep(可能爲空)
}

6.createShuffleMapStage(line 319)

獲得ShuffleMapStage 添加進stageIdToStage

def createShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _],
            jobId: Int): ShuffleMapStage = {

  val rdd = shuffleDep.rdd//父RDD
  val numTasks = rdd.partitions.length// 分區數
   //獲得父stage，第一次到這裏時返回空list，沒有父stage
  val parents = getOrCreateParentStages(rdd, jobId)
  //當沒有parents時當前的stage爲第一個stage，第一次調用到這裏，StageId編號1
  val id = nextStageId.getAndIncrement()
  val stage =
        new ShuffleMapStage(id, rdd, numTasks, parents, jobId, rdd.creationSite, shuffleDep)

  stageIdToStage(id) = stage// 自增id和stage
  shuffleIdToMapStage(shuffleDep.shuffleId) = stage //  維護寬依賴和stage的關係
  updateJobIdStageIdMaps(jobId, stage)// 更新jobIds和jobIdToStageIds
  // 把shuffle信息註冊到Driver上的MapOutputTrackerMaster的shuffleStatuses
  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 + ")")
    // 把Shuffle信息註冊到自己Driver的MapOutputTrackerMaster上，
    // 生成的是shuffleId和ShuffleStatus的映射關係
    // 在後面提交Job的時候還會根據它來的驗證map stage是否已經準備好
    mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)
  }
  stage // 返回生成的ShuffleMapStage
}

公衆號：大叔據 。

評論不能及時回覆可直接加公衆號提問或交流，知無不答，謝謝 。

歡迎關注大叔