Spark 中的 DAG 逻辑视图

本部分记录DAG的生成机制，通过DAG Spark可以对计算流程进行优化，通过WC案例对DAG进行解析

DAG的生成机制

DAG（有向无环图），通过DAG，Spark可以对计算的流程进行优化，对于数据处理，可以将在单一节点上进行的操作进行合并，并且计算中间数据通过内存进行高效读写，对于数据处理，需要涉及Shuffle 操作的步骤划分Stage，从而使得计算资源的利用更加高效和合理，减少计算资源的等待过程，减少计算中间数据读写产生的时间浪费（基于内存的高效读写）。

Spark中DAG生成过程的重点是对Stage的划分，划分的依据是RDD的依赖关系，对于不同的依赖关系，调度器会进行不同的处理，对于窄依赖，RDD之间的数据不需要进行shuffle，多个数据处理可以在同一台机器的内存中完成，所以窄依赖在Spark中被划分为同一个Stage，对于宽依赖，由于Shuffle的存在，必须要等到父RDD的shuffle处理完成之后才能进行接下来的计算，所以此时会进行stage的划分。

在Spark中DAG生成的流程关键在于回溯，程序提交以后，调度器会将所有的Stage看成一个Stage，然后对此Stage进行从后往前的回溯，遇到Shuffle就进行断开，遇到窄依赖就划分到本Stage当中，等到所有的回溯完成之后，便生成一个DAG图。

Spark相关源码在DAGSchedule.scala 中，这里以Saprk3.0 版本为例

  /**
   * Get or create the list of parent stages for a given RDD.  The new Stages will be created with
   * the provided firstJobId.
   */
  private def getOrCreateParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
    getShuffleDependencies(rdd).map { shuffleDep =>
      getOrCreateShuffleMapStage(shuffleDep, firstJobId)
    }.toList
  }

getShuffleDependencies源码：

  /**
   * Returns shuffle dependencies that are immediate parents of the given RDD.
   *
   * This function will not return more distant ancestors.  For example, if C has a shuffle
   * dependency on B which has a shuffle dependency on A:
   *
   * A <-- B <-- C
   *
   * calling this function with rdd C will only return the B <-- C dependency.
   *
   * This function is scheduler-visible for the purpose of unit testing.
   */
  private[scheduler] def getShuffleDependencies(
      rdd: RDD[_]): HashSet[ShuffleDependency[_, _, _]] = {
    val parents = new HashSet[ShuffleDependency[_, _, _]]
    val visited = new HashSet[RDD[_]]
    val waitingForVisit = new ListBuffer[RDD[_]]
    waitingForVisit += rdd
    while (waitingForVisit.nonEmpty) {
      val toVisit = waitingForVisit.remove(0)
      if (!visited(toVisit)) {
        visited += toVisit
        toVisit.dependencies.foreach {
          case shuffleDep: ShuffleDependency[_, _, _] =>
            parents += shuffleDep
          case dependency =>
            waitingForVisit.prepend(dependency.rdd)
        }
      }
    }
    parents
  }

DAG逻辑视图解析

接下来通过一个简单的WordCount案例讲解DAG具体生成的流程和关系

  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setMaster("local[2]").setAppName(this.getClass.getSimpleName)
    val sc = new SparkContext(conf)
    val lineRDD: RDD[String] = sc.textFile("data/word.txt")
    // 由lineRDD 通过flatMap操作生成新的MapPartitionRDD
    val words: RDD[String] = lineRDD.flatMap(_.split(","))
    words.cache()
    //  通过flatMap操作生成新的MapPartitionRDD
    val pairs: RDD[(String, Int)] = words.map((_, 1))
    //  此步骤生成MapPartitionRDD 和ShuffleRDD
    val cnts: RDD[(String, Int)] = pairs.reduceByKey(_ + _)
    cnts.collect().foreach(println(_))
    Thread.sleep(1000000000)
    sc.stop()
  }

在程序开始运行之前，Spark的DAG调度器会将整个流程设置为一个Stage，此Stage包含3个操作，5个RDD，分别是MapPartitionRDD（读文件）、MapPartitionRDD（map操作）、MapPartitionRDD（flatMap）、MapPartitionRDD（reduceByKey 的local阶段操作，归并）和ShuffleRDD（reduceByKey的shuffle操作）。

回溯的流程如下：

1. 在shuffleRDD与MapPartitionRDD（reduceByKey）阶段中存在shuffle操作，整个RDD先从此切开形成两个stage
2. 继续向前面回溯，reduceByKey的local阶段和map操作之间不存在宽依赖，所以划分为同一个Stage
3. 继续回溯，发现前面的所有RDD都不存在shuffle，故前面的这些RDD都划分为同一个Stage
4. 回溯完成，形成DAG，该DAG由两个Stage构成

下面是两个Stage的详细信息：