上篇文章 spark 原始碼分析之十八 -- Spark儲存體系剖析 重點剖析了 Spark的儲存體系。從本篇文章開始，剖析Spark作業的排程和計算體系。
在說DAG之前，先簡單說一下RDD。
對RDD的整體概括
文件說明如下：
RDD全稱Resilient Distributed Dataset，即分散式彈性資料集。它是Spark的基本抽象，代表不可變的可分割槽的可平行計算的資料集。
RDD的特點：
1. 包含了一系列的分割槽
2. 在每一個split上執行函式計算
3. 依賴於其他的RDD
4. 對於key-value對的有partitioner
5. 每一個計算有優先計算位置
更多內容可以去看Spark的論文：http://people.csail.mit.edu/matei/papers/2012/nsdi_spark.pdf
RDD的操作
RDD支援兩種型別的操作：
- transformation：它從已存在的資料集中建立一個新的資料集。它是懶執行的，即生成RDD的所有操作都是懶執行的，也就是說不會馬上計算出結果，它們只會記住它們依賴的基礎資料集（檔案、MQ等等），等到一個action需要結果返回到driver端的時候，才會執行transform的計算。這種設計使得RDD計算更加高效。
- action：它在資料集上執行計算之後給driver端返回一個值。
注意：reduce 是一個action，而 reduceByKey 則是一個transform，因為它返回的是一個分散式資料集，並沒有把資料返回給driver節點。
Action函式
官方提供了RDD的action函式，如下：
注意：這只是常見的函式，並沒有列舉所有的action函式。
Action函式的特點
那麼action函式有哪些特點呢？
根據上面介紹的，即action會返回一個值給driver節點。即它們的函式返回值是一個具體的非RDD型別的值或Unit，而不是RDD型別的值。
Transformation函式
官方提供了Transform 函式，如下：
Transformation函式的特點
上文提到，transformation接收一個存在的資料集，並將計算結果作為新的RDD返回。也是就說，它的返回結果是RDD。
總結
其實，理解了action和transformation的特點，看函式的定義就知道是action還是transformation。
RDD的依賴關係
官方文件裡，聊完RDD的操作，緊接著就聊了一下shuffle，我們按照這樣的順序來做一下說明。
Shuffle
官方給出的shuffle的解釋如下：
注意：shuffle是特定操作才會發生的事情，這跟action和transformation劃分沒有關係。
官方給出了一些常見的例子。
Operations which can cause a shuffle include repartition operations like repartition and coalesce, ByKey operations (except for counting) like groupByKey and reduceByKey, and join operations like cogroup and join.
RDD的四種依賴關係
那麼shuffle跟什麼有關係呢？
shuffle跟依賴有關係。在 spark 原始碼分析之一 -- RDD的四種依賴關係 中，說到 RDD 分為寬依賴和窄依賴，其中窄依賴有三種，一對一依賴、Range依賴、Prune 依賴。寬依賴只有一種，那就是 shuffle 依賴。
即RDD跟父RDD的依賴關係是寬依賴，那麼就是父RDD在生成新的子RDD的過程中是存在shuffle過程的。
如圖：
這張圖也說明了一個結論，並不是所有的join都是寬依賴。
依賴關係在原始碼中的體現
我們通常說的 RDD，在Spark中具體表現為一個抽象類，所有的RDD子類繼承自該RDD，全稱為 org.apache.spark.rdd.RDD，如下：
它有兩個引數，一個引數是SparkContext，另一個是deps，即Dependency集合，Dependency是所有依賴的公共父類，即deps儲存了父類的依賴關係。
其中，窄依賴的父類是 NarrowDependency， 它的構造方法裡是由父RDD這個引數的，寬依賴 ShuffleDependency ，它的構造方法裡也是有父RDD這個引數的。
RDD 依賴關係的不確定性
getDependencies 方法
獲取抽象的方法是 getDependencies 方法，如下：
這只是定義在RDD抽象父類中的預設方法，不同的子類會有不同的實現。
它在如下類中又重新實現了這個方法，如下：
是否是shuffle依賴，跟分割槽的數量也有一定的關係，具體可以看下面的幾個RDD的依賴的實現：
CoGroupedRDD
SubtractedRDD
DAG在Spark作業中的重要性
如下圖，一個application的執行過程被劃分為四個階段：
階段一：我們編寫driver程式，定義RDD的action和transformation操作。這些依賴關係形成操作的DAG。
階段二：根據形成的DAG，DAGScheduler將其劃分為不同的stage。
階段三：每一個stage中有一個TaskSet，DAGScheduler將TaskSet交給TaskScheduler去執行，TaskScheduler將任務執行完畢之後結果返回給DAGSCheduler。
階段四：TaskScheduler將任務分發到每一個Worker節點去執行，並將結果返回給TaskScheduler。
本篇文章的定位就是階段一和階段二。後面會介紹階段三和階段四。
注：圖片不知出處。
DAG的建立
我們先來分析一個top N案例。
一個真實的TopN案例
需求：一個大檔案裡有很多的重複整數，現在求出重複次數最多的前10個數。
程式碼如下（為了多幾個stage，特意加了幾個repartition）：
scala> val sourceRdd = sc.textFile("/tmp/hive/hive/result",10).repartition(5)
sourceRdd: org.apache.spark.rdd.RDD[String] = MapPartitionsRDD[5] at repartition at <console>:27
scala> val allTopNs = sourceRdd.flatMap(line => line.split(" ")).map(word => (word, 1)).reduceByKey(_+_).repartition(10).sortByKey(ascending = true, 100).map(tup => (tup._2, tup._1)).mapPartitions(
| iter => {
| iter.toList.sortBy(tup => tup._1).takeRight(100).iterator
| }
| ).collect()
// 結果略
scala> val finalTopN = scala.collection.SortedMap.empty[Int, String].++(allTopNs)
//結果略
scala> finalTopN.takeRight(10).foreach(tup => {println(tup._2 + " occurs times : " + tup._1)})
53 occurs times : 1070
147 occurs times : 1072
567 occurs times : 1073
931 occurs times : 1075
267 occurs times : 1077
768 occurs times : 1080
612 occurs times : 1081
877 occurs times : 1082
459 occurs times : 1084
514 occurs times : 1087
下面看一下生成的DAG和Stage
任務概覽
Description描述的就是每一個job的最後一個方法。
stage 0 到 3的DAG圖：
stage 4 到 8的DAG圖:
每一個stage的Description描述的是stage的最後一個方法。
總結
可以看出，RDD的依賴關係是有driver端對RDD的操作形成的。
一個Stage中DAG的是根據RDD的依賴來構建的。
我們來看一下原始碼。
Stage
定義
Stage是一組並行任務，它們都計算需要作為Spark作業的一部分執行的相同功能，其中所有任務具有相同的shuffle依賴。由排程程式執行的每個DAG任務在發生shuffle的邊界處被分成多個階段，然後DAGScheduler以拓撲順序執行這些階段。每個Stage都可以是一個shuffle map階段，在這種情況下，其任務的結果是為其他階段或結果階段輸入的，在這種情況下，其任務在RDD上執行函式直接計算Spark action（例如count()，save()等）。對於shuffle map階段，我們還跟蹤每個輸出分割槽所在的節點。每個stage還有一個firstJobId，用於識別首次提交stage的作業。使用FIFO排程時，這允許首先計算先前作業的階段，或者在失敗時更快地恢復。最後，由於故障恢復，可以在多次嘗試中重新執行單個stage。在這種情況下，Stage物件將跟蹤多個StageInfo物件以傳遞給listener 或Web UI。最近的一個將通過latestInfo訪問。
構造方法
Stage是一個抽象類，構造方法如下：
引數介紹如下：
id – Unique stage ID
rdd – RDD that this stage runs on: for a shuffle map stage, it's the RDD we run map tasks on, while for a result stage, it's the target RDD that we ran an action on
numTasks – Total number of tasks in stage; result stages in particular may not need to compute all partitions, e.g. for first(), lookup(), and take().
parents – List of stages that this stage depends on (through shuffle dependencies).
firstJobId – ID of the first job this stage was part of, for FIFO scheduling.
callSite – Location in the user program associated with this stage: either where the target RDD was created, for a shuffle map stage, or where the action for a result stage was called
callSite其實記錄的就是stage使用者程式碼的位置。
成員變數
成員方法
其實相對來說比較簡單。
Stage的子類
它有兩個子類，如下：
ResultStage
類說明：
ResultStages apply a function on some partitions of an RDD to compute the result of an action.
The ResultStage object captures the function to execute, func, which will be applied to each partition, and the set of partition IDs, partitions.
Some stages may not run on all partitions of the RDD, for actions like first() and lookup().
ResultStage在RDD的某些分割槽上應用函式來計算action操作的結果。 對於諸如first（）和lookup（）之類的操作，某些stage可能無法在RDD的所有分割槽上執行。
簡言之，ResultStage是應用action操作在action上進而得出計算結果。
原始碼如下：
ShuffleMapStage
類說明
ShuffleMapStages are intermediate stages in the execution DAG that produce data for a shuffle.
They occur right before each shuffle operation, and might contain multiple pipelined operations before that (e.g. map and filter).
When executed, they save map output files that can later be fetched by reduce tasks.
The shuffleDep field describes the shuffle each stage is part of, and variables like outputLocs and numAvailableOutputs track how many map outputs are ready.
ShuffleMapStages can also be submitted independently as jobs with DAGScheduler.submitMapStage.
For such stages, the ActiveJobs that submitted them are tracked in mapStageJobs.
Note that there can be multiple ActiveJobs trying to compute the same shuffle map stage.
ShuffleMapStage 是中間的stage，為shuffle生產資料。它們在shuffle之前出現。當執行完畢之後，結果資料被儲存，以便reduce 任務可以獲取到。
構造方法
shuffleDep記錄了每一個stage所屬的shuffle。
Stage的劃分
在上面我們提到，每一個RDD都有對父RDD的依賴關係，這樣的依賴關係形成了一個有向無環圖。即DAG。
當一個使用者在一個RDD上執行一個action時，排程會檢查RDD的血緣關係（即依賴關係）來建立一個stage中的DAG圖來執行。
如下圖：
在說stage劃分之前先，剖析一下跟DAGScheduler相關的類。
EventLoop
類說明
An event loop to receive events from the caller and process all events in the event thread. It will start an exclusive event thread to process all events.
Note: The event queue will grow indefinitely. So subclasses should make sure onReceive can handle events in time to avoid the potential OOM.
它定義了非同步訊息處理機制框架。
訊息佇列
其內部有一個阻塞雙端佇列，用於存放訊息：
post到訊息佇列
外部執行緒呼叫 post 方法將事件post到堵塞佇列中：
消費執行緒
有一個訊息的消費執行緒：
onReceive 方法是一個抽象方法，由子類來實現。
下面來看其實現類 -- DAGSchedulerEventProcessLoop。
其接收的是DAGSchedulerEvent型別的事件。DAGSchedulerEvent 是一個sealed trait，其實現如下：
它的每一個子類事件，在doOnReceive 方法中都有體現，如下：
DAGScheduler
這個類的定義已經超過2k行了。所以也不打算全部介紹，本篇文章只介紹跟stage任務的生成相關的屬性和方法。
類說明
The high-level scheduling layer that implements stage-oriented scheduling. It computes a DAG of stages for each job, keeps track of which RDDs and stage outputs are materialized, and finds a minimal schedule to run the job. It then submits stages as TaskSets to an underlying TaskScheduler implementation that runs them on the cluster. A TaskSet contains fully independent tasks that can run right away based on the data that's already on the cluster (e.g. map output files from previous stages), though it may fail if this data becomes unavailable.
Spark stages are created by breaking the RDD graph at shuffle boundaries. RDD operations with "narrow" dependencies, like map() and filter(), are pipelined together into one set of tasks in each stage, but operations with shuffle dependencies require multiple stages (one to write a set of map output files, and another to read those files after a barrier). In the end, every stage will have only shuffle dependencies on other stages, and may compute multiple operations inside it. The actual pipelining of these operations happens in the RDD.compute() functions of various RDDs
In addition to coming up with a DAG of stages, the DAGScheduler also determines the preferred locations to run each task on, based on the current cache status, and passes these to the low-level TaskScheduler. Furthermore, it handles failures due to shuffle output files being lost, in which case old stages may need to be resubmitted. Failures within a stage that are not caused by shuffle file loss are handled by the TaskScheduler, which will retry each task a small number of times before cancelling the whole stage. When looking through this code, there are several key concepts:
- Jobs (represented by ActiveJob) are the top-level work items submitted to the scheduler. For example, when the user calls an action, like count(), a job will be submitted through submitJob. Each Job may require the execution of multiple stages to build intermediate data.
- Stages (Stage) are sets of tasks that compute intermediate results in jobs, where each task computes the same function on partitions of the same RDD. Stages are separated at shuffle boundaries, which introduce a barrier (where we must wait for the previous stage to finish to fetch outputs). There are two types of stages: ResultStage, for the final stage that executes an action, and ShuffleMapStage, which writes map output files for a shuffle. Stages are often shared across multiple jobs, if these jobs reuse the same RDDs.
- Tasks are individual units of work, each sent to one machine.
- Cache tracking: the DAGScheduler figures out which RDDs are cached to avoid recomputing them and likewise remembers which shuffle map stages have already produced output files to avoid redoing the map side of a shuffle.
- Preferred locations: the DAGScheduler also computes where to run each task in a stage based on the preferred locations of its underlying RDDs, or the location of cached or shuffle data.
- Cleanup: all data structures are cleared when the running jobs that depend on them finish, to prevent memory leaks in a long-running application.
To recover from failures, the same stage might need to run multiple times, which are called "attempts". If the TaskScheduler reports that a task failed because a map output file from a previous stage was lost, the DAGScheduler resubmits that lost stage. This is detected through a CompletionEvent with FetchFailed, or an ExecutorLost event. The DAGScheduler will wait a small amount of time to see whether other nodes or tasks fail, then resubmit TaskSets for any lost stage(s) that compute the missing tasks. As part of this process, we might also have to create Stage objects for old (finished) stages where we previously cleaned up the Stage object. Since tasks from the old attempt of a stage could still be running, care must be taken to map any events received in the correct Stage object.
Here's a checklist to use when making or reviewing changes to this class:
- All data structures should be cleared when the jobs involving them end to avoid indefinite accumulation of state in long-running programs.
- When adding a new data structure, update DAGSchedulerSuite.assertDataStructuresEmpty to include the new structure. This will help to catch memory leaks.
下面直接來看stage的劃分
從原始碼看Stage的劃分
從action函式到DAGScheduler
以collect函式為例。
collect 函式定義如下：
其呼叫了SparkContext的 runJob 方法，又呼叫了幾次其過載方法最終呼叫的runJob 方法如下：
其內部呼叫了DAGScheduler的runJob 方法
DAGScheduler對stage的劃分
DAGScheduler的runJob 方法如下：
思路，提交方法後返回一個JobWaiter 物件，等待任務執行完成，然後根據任務執行狀態去執行對應的成功或失敗的方法。
submitJob 如下：
最終任務被封裝進了JobSubmitted 事件訊息體中，最終該事件訊息被放入了eventProcessLoop 物件中，eventProcessLoop定義如下：
即事件被放入到了上面我們提到的 DAGSchedulerEventProcessLoop 非同步訊息處理模型中。
DAGSchedulerEventProcessLoop 的 doOnReceive 中，發現了 JobSubmitted 事件對應的分支為：
即會執行DAGScheduler的handleJobSubmitted方法，如下：
這個方法裡面有兩步：
- 建立ResultStage
- 提交Stage
本篇文章，我們只分析第一步，第二步在下篇文章分析。
createResultStage 方法如下：
getOrCreateParentStage 方法建立或獲取該RDD的Shuffle依賴關係，然後根據shuffle依賴進而劃分stage，原始碼如下：
獲取其所有父類的shuffle依賴，getShuffleDependency 方法如下，類似於樹的深度遍歷。
getOrCreateShuffleMapStage方法根據shuffle依賴建立ShuffleMapStage，如下，思路，先檢視當前stage是否已經記錄在shuffleIdToMapStage變數中，若存在，表示已經建立過了，否則需要根據依賴的RDD去找其RDD的shuffle依賴，然後再建立shuffleMapStage。
shuffleIdToMapStage定義如下：
這個map中只包含正在執行的job的stage資訊。
其中shuffle 依賴的唯一id 是：shuffleId，這個id 是 SpackContext 生成的全域性shuffleId。
getMissingAncestorShuffleDependencies 方法如下，思路：深度遍歷依賴關係，把所有未執行的shuffle依賴都找到。
到此，所有尋找shuffle依賴關係的的邏輯都已經剖析完畢，下面看建立MapShuffleStage的方法，
思路：生成ShuffleMapStage，並更新 stageIdToStage變數，更新shuffleIdToMapStage變數，如果 MapOutputTrackerMaster 中沒有註冊過該shuffle，需要註冊，最後返回ShuffleMapStage物件。
updateJobIdStageIdMaps方法如下，思路該ResultStage依賴的所有ShuffleMapStage的jobId設定為指定的jobId，即跟ResultStage一致的jobId：
至此，stage的劃分邏輯剖析完畢。
總結
本篇文章對照官方文件，說明了RDD的主要操作，action和transformation，進一步引出了RDD的依賴關係，最後剖析了DAGScheduler根據shuffle依賴劃分stage的邏輯。
注：文章中圖片來源於 Spark 論文，論文地址：http://people.csail.mit.edu/matei/papers/2012/nsdi_spark.pdf