当前位置：首页 > article >正文

spark的stage划分的原理

article 2024/11/18 18:29:09

在 Apache Spark 中，stage 是执行作业时的重要执行单元。一个 Spark 作业会被划分为若干个 stage，每个 stage 由一组可以并行执行的任务组成。这种划分主要依赖于 RDD 中的操作类型（窄依赖和宽依赖）。下面我们来讨论 Spark stage 的创建和划分的原理以及代码实现的核心逻辑。

Spark Stage 划分的原理

RDD 依赖（窄依赖和宽依赖）：
- Spark 中，RDD 可以有两种依赖关系：
  - 窄依赖（narrow dependency）：父 RDD 的每个分区至多被子 RDD 的一个分区使用，典型的操作如 map、filter 等。
  - 宽依赖（wide dependency）：父 RDD 的每个分区可能被多个子 RDD 的分区使用，典型的操作如 reduceByKey、groupByKey 等，这类操作会触发 shuffle。
- 窄依赖的 RDD 操作可以被划分到同一个 stage 中，而宽依赖的 RDD 操作会触发 shuffle，导致 stage 划分。
DAG（有向无环图）：
Spark 的作业会构建一个 RDD 的依赖图（DAG）。这个 DAG 中每个 RDD 的窄依赖操作会被合并成一个 stage，宽依赖操作会划分出不同的 stage，并在两个 stage 之间插入 shuffle。
Stage 划分规则：
- 每当遇到一个宽依赖（如 reduceByKey、join、groupByKey 等），Spark 会创建一个新的 stage，并将之前的 RDD 操作划分到一个 stage 中，形成一个有序的 stage 执行链。
- stage 划分的核心任务是：将窄依赖操作尽可能合并到一起，直到遇到需要 shuffle 的宽依赖操作。

Spark `Stage` 划分的核心代码逻辑

Spark 的 DAG 划分及 stage 划分主要在 DAGScheduler 中实现。DAGScheduler 是 Spark 作业调度的核心组件，负责将逻辑作业（job）划分为多个 stage，并调度这些 stage 执行。

以下是 Spark 3.x 版本中有关 stage 划分的核心逻辑及其简化代码片段。

1. DAGScheduler 类

DAGScheduler 类位于 org.apache.spark.scheduler 包下，它负责管理 RDD 依赖关系并创建 stage。DAGScheduler 会根据 RDD 的依赖图和操作类型，生成任务的 DAG 并划分 stage。

class DAGScheduler(
    // 参数略...
) extends Logging {

  // stage 列表
  private val stages = new HashMap[StageId, Stage]()

  // 提交 Job 时触发的函数
  def submitJob(
      rdd: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Seq[Int],
      callSite: CallSite,
      allowLocal: Boolean,
      resultHandler: (Int, _) => Unit,
      properties: Properties = null): JobWaiter[_] = {

    // 根据 RDD 和依赖关系生成最终的 ResultStage
    val finalStage = createFinalStage(rdd, partitions, callSite)

    // 提交该 stage 执行
    submitStage(finalStage)
  }

  // 创建 ResultStage 和后续的 Stage
  private def createFinalStage(
      rdd: RDD[_],
      partitions: Seq[Int],
      callSite: CallSite): ResultStage = {
    // 创建该作业的最终的 stage，并递归创建所有依赖的 stage
    val finalStage = newStage(rdd, partitions)
    finalStage
  }

  // 递归生成各个 Stage，核心逻辑
  private def newStage(rdd: RDD[_], partitions: Seq[Int]): Stage = {
    // 检查缓存，避免重复生成 Stage
    stages.getOrElseUpdate(rdd.id, {
      val shuffleDeps = getShuffleDependencies(rdd)

      // 如果存在宽依赖，则要划分为不同的 stage
      if (shuffleDeps.nonEmpty) {
        val parentStages = shuffleDeps.map { dep =>
          newStage(dep.rdd, dep.rdd.partitions.indices)
        }
        val newStage = new ShuffleMapStage(rdd, parentStages)
        stages(newStage.id) = newStage
        newStage
      } else {
        // 如果只有窄依赖，当前操作在同一个 stage 内
        val parentStages = getNarrowDependencies(rdd).map { dep =>
          newStage(dep.rdd, dep.rdd.partitions.indices)
        }
        val newStage = new ResultStage(rdd, parentStages)
        stages(newStage.id) = newStage
        newStage
      }
    })
  }

  // 获取 RDD 的 shuffle 依赖（宽依赖）
  private def getShuffleDependencies(rdd: RDD[_]): List[ShuffleDependency[_, _, _]] = {
    rdd.dependencies.collect {
      case shuffleDep: ShuffleDependency[_, _, _] => shuffleDep
    }
  }

  // 获取 RDD 的窄依赖
  private def getNarrowDependencies(rdd: RDD[_]): List[Dependency[_]] = {
    rdd.dependencies.collect {
      case narrowDep: NarrowDependency[_] => narrowDep
    }
  }
}

2. Stage 划分的基本过程

RDD 依赖遍历：通过 newStage 函数递归遍历 RDD 的依赖关系，将遇到的每一个 shuffle 依赖（宽依赖）创建一个新的 ShuffleMapStage，而 ResultStage 则用于最终计算结果。
宽依赖处理：当遇到宽依赖（ShuffleDependency），说明需要进行 shuffle，因此要创建一个新的 stage。
窄依赖处理：当只有窄依赖时，RDD 可以继续合并在当前的 stage 中。

3. ShuffleMapStage 和 ResultStage

ShuffleMapStage 和 ResultStage 是 Spark 中两种类型的 Stage：

ShuffleMapStage：处理宽依赖（shuffle），该 stage 会产生 shuffle 文件供下游 stage 使用。
ResultStage：最终计算 Action（如 collect、saveAsTextFile 等）结果的 stage，是 DAG 中的最后一个 stage。