从零开始手写 Spark 07｜划分 Stage 与 DAG

Shuffle 的物理边界已经清楚了：Map 端写出中间文件，Reduce 端再读取这些文件。

这个顺序带来一个必须解决的问题：

Reduce 端要读文件
文件必须先存在
所以写文件那一批计算，必须先完成

上一章为了先看清 Shuffle 文件，把写文件这件事放在了 ShuffledRDD.compute() 里面：Reduce 分区第一次被计算时，先推动上游 Map 分区写文件，然后自己再读文件。

从功能上看没有问题——Reduce 任务确实读到了数据。但调度职责发生了错位：Reduce 分区任务已经开始执行，才临时去准备上游数据。 现在要把"准备 Map 输出"这个职责从 RDD 层移出，从 ShuffledRDD.compute() 提升到调度层。 调度器先查看完整条血缘，发现中间有 Shuffle，就把作业切成两段：先主动触发写文件的那段，再提交读文件的那段。reduceByKey 仍然返回 ShuffledRDD，但它的执行职责变了：不再决定"什么时候写 Map 输出"，只记录自己依赖一条 Shuffle 边，并在 compute() 中读取已经写出的 Map 输出。

本章按这个顺序展开：先从中间文件推出 Stage 边界（7.1），再解释窄依赖为什么不切 Stage（7.2）；然后把 RDD、Dependency、Stage 这三个对象的结构介绍清楚（7.3）；接着看 DAGScheduler 怎么沿血缘把 RDD 图切成 Stage 图（7.4）；切完以后，看它怎么保证父 Stage 先写完文件、子 Stage 才开始读（7.5）；最后运行一次查看 Stage 树（7.6）。

7.1 从中间文件推出 Stage

这一章围绕一段 reduceByKey 示例展开：把一批 (单词, 1) 按 key 求和。

try (SparkContext sc = new SparkContext(2, true)) {
    RDD<KeyValuePair<String, Integer>> rdd = sc.parallelize(words, 3);

    MapPartitionsRDD<KeyValuePair<String, Integer>, KeyValuePair<String, Integer>> mapped =
            rdd.map(Function.identity());

    ShuffledRDD<String, Integer> shuffled = mapped.reduceByKey(
            Integer::sum, 2);

    System.out.println("reduceByKey 结果: " + shuffled.collect());
}

parallelize(words, 3) 把输入数据切成 3 个 Map 分区。中间那个 map(Function.identity()) 不改数据，只是让血缘里多一条窄依赖。reduceByKey(Integer::sum, 2) 把相同 key 的值相加，结果生成一个 ShuffledRDD。最后一行的 collect() 是 action，触发整条血缘实际执行。

输入数据被切成 3 个 Map 分区：

分区 0:  (hello,1), (world,1), (hello,1)
分区 1:  (spark,1), (world,1), (hello,1)
分区 2:  (java,1),  (spark,1), (hello,1)

reduceByKey(Integer::sum, 2) 会产生 2 个 Reduce 分区。Map 端写出的文件是：

map_0_reduce_0
map_0_reduce_1
map_1_reduce_0
map_1_reduce_1
map_2_reduce_0
map_2_reduce_1

3 个 Map 分区，2 个 Reduce 分区，所以一共 6 个文件。

Reduce 分区 0 只读所有 *_reduce_0 文件，Reduce 分区 1 只读所有 *_reduce_1 文件。

上一章让 Reduce 任务在 compute() 里临时补写 Map 文件——任务已经开始执行，才去准备上游数据。本章要把写文件这件事提前到 Reduce 分区任务提交之前。这样一来，执行顺序便分成两段：

第一段：把父 RDD 的所有 Map 分区算完，并写出 shuffle 文件
第二段：让下游 Reduce 分区读取这些文件，并合并相同 key

这两段之间的分界线，就是 Stage 边界。

更准确地说：

Stage 是沿 Shuffle 边界切出来的一段计算。

在 reduceByKey 这个例子里，前一段叫 ShuffleMapStage，后一段叫 ResultStage：

Stage 0: ShuffleMapStage，输出是 shuffle 中间文件
Stage 1: ResultStage，输出是用户要的最终结果

两种 Stage 的名字不同，但划分依据都是 Shuffle 文件的生成与消费：父 Stage 负责写，子 Stage 负责读。

7.2 窄依赖为什么不切

如果所有依赖都切 Stage，系统会失去流水线执行的机会。

比如这条链：

ListRDD -> map -> filter -> flatMap

这里没有跨分区重新分布。一个子分区只需要读自己对应的父分区，数据逐条流过整个算子链：

flowchart LR
    A[取一条数据] --> B[map] --> C[filter] --> D[flatMap] --> E[交给 action]
    E -.->|处理完后取下一条| A

整个过程不需要把整个父 RDD 算完，也不需要写入磁盘。数据像流水一样穿过整条窄依赖链。

所以，map、filter、flatMap 这类窄依赖应该留在同一个 Stage 里。

reduceByKey 的情况则不同。它的一个 Reduce 分区可能要读所有 Map 分区写给它的文件，不能只依赖一个父分区，必须等一批上游输出准备完成。

因此，窄依赖和宽依赖在调度上必须区别对待：

NarrowDependency：不切 Stage，继续往父 RDD 走。
ShuffleDependency：切 Stage，父 Stage 先把中间文件写出来。

引入 Dependency 时，血缘分成了这两类。现在，这个分类第一次成为实际的调度规则。

7.3 三个对象：RDD、Dependency、Stage

在讲调度器怎么切 Stage 之前，先把这一章涉及的三个对象介绍清楚。

• RDD 通过 dependencies() 报告自己依赖谁。
• Dependency 标明这条依赖是窄依赖还是宽依赖。
• Stage 是调度器沿 Shuffle 边界切出来的一段计算。

三者的关系如下：

flowchart LR
    R[RDD] -->|dependencies&#40&#41| D[Dependency]
    D -->|如果是 ShuffleDependency| S[Stage&#40 父 &#41]
    D -->|如果是窄依赖| R2[继续往父RDD走]

调度器沿 RDD 血缘回溯，读取 dependencies()，遇到 ShuffleDependency 就创建一个父 Stage。

7.3.1 RDD 的三个接口

先看所有 RDD 都有的接口：

public abstract class RDD<T> {
    public abstract List<Partition> partitions();

    public abstract Iterator<T> compute(Partition partition);

    public abstract List<Dependency<?>> dependencies();
}

partitions() 回答“这个 RDD 有几个分区”。

compute(partition) 回答“某个分区怎么算”。

dependencies() 回答“这个 RDD 依赖谁”。

调度器最关心的是第三个：这个 RDD 的依赖是什么类型？

对于普通的 map，依赖是窄依赖。一个子分区只读一个父分区，所以 MapPartitionsRDD 返回的是 OneToOneDependency：

private final List<Dependency<?>> dependencies;

public MapPartitionsRDD(
        RDD<T> parent,
        Function<Iterator<T>, Iterator<U>> iteratorTransform) {
    // ...
    this.dependencies = List.of(new OneToOneDependency<>(parent));
}

@Override
public List<Dependency<?>> dependencies() {
    return dependencies;
}

7.3.2 Dependency：窄依赖与 ShuffleDependency

对于 reduceByKey，情况不同。ShuffledRDD 的每个 Reduce 分区都要读多个 Map 分区写出的文件，所以它返回的是 ShuffleDependency：

private final ShuffleDependency<K, V> shuffleDependency;
private final List<Dependency<?>> dependencies;

public ShuffledRDD(RDD<KeyValuePair<K, V>> parent,
                   int numReducePartitions,
                   BinaryOperator<V> reduceFunc) {
    // ...
    this.shuffleDependency = new ShuffleDependency<>(
            parent,
            numReducePartitions,
            shuffleDir,
            reduceFunc);
    this.dependencies = List.of(shuffleDependency);
}

@Override
public List<Dependency<?>> dependencies() {
    return dependencies;
}

dependencies() 返回 ShuffleDependency，调度器沿 RDD 血缘回溯时，就能从这里发现 Shuffle 边界。至于 Map 输出该怎样写，则由这条依赖保存所需信息：

public final class ShuffleDependency<K, V>
        implements Dependency<KeyValuePair<K, V>> {

    private final RDD<KeyValuePair<K, V>> rdd;
    private final int numReducePartitions;
    private final File shuffleDir;
    private final BinaryOperator<V> reduceFunc;

    public ShuffleDependency(
            RDD<KeyValuePair<K, V>> rdd,
            int numReducePartitions,
            File shuffleDir,
            BinaryOperator<V> reduceFunc) {
        this.rdd = Objects.requireNonNull(rdd, "rdd");
        this.numReducePartitions = numReducePartitions;
        this.shuffleDir = Objects.requireNonNull(shuffleDir, "shuffleDir");
        this.reduceFunc = Objects.requireNonNull(reduceFunc, "reduceFunc");
    }

    //...

    public int partition(Object key) {
        return (key.hashCode() & Integer.MAX_VALUE) % numReducePartitions;
    }

    public File mapOutputFile(int mapId, int reduceId) {
        return new File(shuffleDir,
                String.format("map_%d_reduce_%d", mapId, reduceId));
    }
}

ShuffleDependency 的 rdd() 方法返回父 RDD，调度器借此知道上游是谁。numReducePartitions 决定每个 Map 分区要写几个文件。shuffleDir 决定这些文件放在哪里。reduceFunc 用来在 Map 端先做一次本地 combine。partition() 方法根据 key 计算目标 Reduce 分区编号。mapOutputFile() 方法构造 Map 输出文件的路径。

ShuffleDependency 保存了 Shuffle 所需的全部信息，但它本身不执行任何操作。如何利用这些信息来调度执行，是 DAGScheduler 的职责。

7.3.3 Stage：切出来的一段计算

第三个对象是 Stage。Stage 不是 RDD 的替代品，它只是调度器根据 RDD 依赖切出来的执行段。完整实现见 Stage.java。

public record Stage(
        int id,
        RDD<?> rdd,
        boolean shuffleMap,
        List<Stage> parents,
        Optional<ShuffleDependency<?, ?>> shuffleDependency) {
}

id 是 Stage 的编号，由调度器按创建顺序递增分配（0、1、2……）。它给每个 Stage 一个唯一身份，让调度器能凭编号识别、区分每个 Stage。

rdd 表示这个 Stage 要算到哪里。对于最终 Stage，它就是用户提交的最终 RDD；对于 ShuffleMapStage，它是宽依赖的父 RDD，也就是写 shuffle 文件之前那段窄依赖链的末端。

shuffleMap 区分两种 Stage：

值	含义
true	这是 ShuffleMapStage，输出是 shuffle 中间文件
false	这是 ResultStage，输出是最终结果

parents 是必须先完成的父 Stage。

shuffleDependency 只在 ShuffleMapStage 上有值。因为只有 ShuffleMapStage 需要写出 shuffle 文件，ResultStage 只需要读取最终结果。

"Stage" 在中文里常被译作"阶段"。这个命名很贴切：一个 Stage 就是一段中间不需要写盘、可以连续算完的计算。窄依赖链上的 map、filter 可以流水线执行，直到遇到 Shuffle 边界才停下来——停下来的那个边界，就是两个 Stage 的分隔线。每一个 Stage 内部，数据像流水一样流过；Stage 和 Stage 之间，数据必须先写入磁盘，再重新读取。

三个对象介绍完毕：RDD 用 dependencies() 报告自己的依赖类型，ShuffleDependency 标出宽依赖，Stage 表示切分后的执行段。沿 RDD 血缘回溯、遇到宽依赖就切 Stage 的那个调度器，就是 DAGScheduler。下一节看它具体怎么切。

7.4 DAGScheduler 怎么切 Stage

DAGScheduler 划分 Stage 的依据是整条 RDD 血缘，所以先把它画成一张图。示例程序里的 RDD 血缘如下：

这张 RDD 血缘图有两个特点：箭头代表依赖方向（下游依赖上游），且不会绕成环（RDD 不可变，依赖只指向已存在的父 RDD）。这样的图就是 DAG（有向无环图），DAGScheduler 的名字正源于此——它要在这样的血缘图上排定执行顺序。

DAGScheduler 划分 Stage 的方法如下：

1. 从最终 RDD（比如 ShuffledRDD）出发，调用它的 dependencies()，看看它依赖谁。
2. 如果依赖是 ShuffleDependency，就停下来——这是一个 Stage 边界。把父 RDD 装进一个 ShuffleMapStage，当前 RDD 装进 ResultStage。
3. 如果依赖是 OneToOneDependency（窄依赖），就不停——继续往父 RDD 走，直到遇到下一个 ShuffleDependency 或到达源头。

Stage 是在 RDD DAG 上沿 Shuffle 边界、把一段窄依赖链合并成的节点。图中，RDD 之间的窄依赖链合并成了 Stage 节点，每个节点内部可以流水线执行，节点之间靠 Shuffle 文件衔接。

示例血缘是一条直线，所以划分出来的是一棵只有两层的 Stage 树。

7.4.1 从 action 到结果的全景

把 Stage 的创建和执行放到同一条时间线上，关系会更直观：

sequenceDiagram
    participant U as 用户代码
    participant R as ShuffledRDD
    participant SC as SparkContext
    participant D as DAGScheduler
    participant TS as TaskScheduler
    participant T as 分区任务
    participant F as shuffle 文件

    U->>R: collect()
    R->>SC: runJob(this, partitionFunction)
    SC->>D: runJob(finalRdd, taskScheduler, partitionFunction)

    Note over R,D: 先沿 RDD 血缘创建 Stage
    D->>R: dependencies()
    R-->>D: ShuffleDependency
    Note over D: 创建 ShuffleMapStage<br/>并保存 ShuffleDependency
    Note over D: 创建下游 ResultStage

    Note over D,F: 先执行 ShuffleMapStage
    loop 每个 Map 分区
        D->>D: 创建 ShuffleMapTask(mapId)
    end
    D->>TS: submitTasks(shuffleMapTasks)
    loop 每个 ShuffleMapTask
        TS->>T: 提交到线程池
        T->>F: 写出当前 Map 分区的各个桶文件
    end
    TS-->>D: 所有 Map 任务完成

    Note over D,F: 再执行 ResultStage
    loop 每个 Reduce 分区
        D->>D: 创建 ResultTask(reduceId)
    end
    D->>TS: submitTasks(resultTasks)
    loop 每个 ResultTask
        TS->>T: 提交到线程池
        T->>R: compute(reduceId)
        R->>F: 读取属于 reduceId 的 Map 输出
        F-->>R: 返回 Map 输出
        R-->>T: 返回当前 Reduce 分区结果
    end
    TS-->>D: 返回所有分区结果
    D-->>SC: 返回 action 结果
    SC-->>R: 返回各分区结果
    R-->>U: collect() 合并并返回

时序图中，先执行 ShuffleMapStage 之前的交互只是创建 Stage，尚未提交任何任务。实际的执行从 loop 每个 Map 分区 开始：DAGScheduler 把 ShuffleMapStage 展开成一批 ShuffleMapTask，把 ResultStage 展开成一批 ResultTask。TaskScheduler 不再判断应该创建哪种任务，只负责把收到的任务提交到线程池。

为了保持 transformation 的惰性语义，ShuffledRDD 的构造方式没有改成“构造时就写文件”。它仍然只记录配方：

ShuffledRDD<String, Integer> shuffled = mapped.reduceByKey(
        Integer::sum, 2);

此时 shuffle 目录里仍然没有文件。只有 collect() 这样的 action 实际执行时，文件才会出现。

时序图里能分出两段：前面创建 Stage（沿血缘切），后面执行 Stage（先写文件、再读文件）。接下来先看"创建"这段——看 DAGScheduler 怎么从最终 RDD 找到 Shuffle 边界、切出 Stage（7.4.2、7.4.3）；7.5 再看"执行"这段，看父 Stage 怎么先写完文件、子 Stage 才开始读。

7.4.2 从最终 RDD 开始：createResultStage 与 getParentStages

action（比如 collect()）触发后，调用链从 SparkContext 进入 DAGScheduler：SparkContext.runJob 把最终 RDD 转交给 DAGScheduler.runJob，后者先调用 createResultStage 把 Stage 树建立出来（verbose 时会打印划分结果，7.6 会看到），再交给私有 runJob 去执行。代码里用注释标明了每个方法所属的文件：

//SparkContext.java
public <T, U> List<U> runJob(
        RDD<T> rdd,
        Function<Iterator<T>, U> partitionFunction) {
    return dagScheduler.runJob(rdd, taskScheduler, partitionFunction);
}

//DAGScheduler.java
public <T, U> List<U> runJob(
        RDD<T> finalRdd,
        TaskScheduler taskScheduler,
        Function<Iterator<T>, U> partitionFunction) {
    Stage finalStage = createResultStage(finalRdd);
    if (verbose) {
        System.out.println("Stage 划分结果:");
        printStage(finalStage, 0);
    }
    return runJob(finalStage, taskScheduler, partitionFunction);
}

Stage createResultStage(RDD<?> finalRdd) {
    return newResultStage(finalRdd);
}

newResultStage 做两件事：先沿血缘找出这个 RDD 的父 Stage，再把它们和当前 RDD 组合成一个 Stage。因为是最终 RDD，shuffleMap 传 false，所以生成的是 ResultStage：

private Stage newResultStage(RDD<?> rdd) {
    List<Stage> parents = getParentStages(rdd);
    return new Stage(
            nextStageId.getAndIncrement(),
            rdd,
            false,
            parents,
            Optional.empty());
}

找父 Stage 的工作落在 getParentStages 上：

private List<Stage> getParentStages(RDD<?> rdd) {
    Set<Stage> parents = new LinkedHashSet<>();
    Set<RDD<?>> visited = new HashSet<>();
    visit(rdd, parents, visited);
    return List.copyOf(parents);
}

它准备两个集合：parents 收集切出来的父 Stage，visited 记录已经走过的 RDD、避免重复遍历。沿血缘回溯、判定哪里切 Stage 的，是 visit()，下一段展开。

7.4.3 沿血缘回溯：visit

visit() 沿 RDD 血缘往回走，只做一个判断：

private void visit(RDD<?> rdd, Set<Stage> parents, Set<RDD<?>> visited) {
    if (!visited.add(rdd)) {
        return;
    }

    for (Dependency<?> dependency : rdd.dependencies()) {
        if (dependency instanceof ShuffleDependency<?, ?> shuffleDependency) {
            parents.add(newShuffleMapStage(shuffleDependency));
        } else {
            visit(dependency.rdd(), parents, visited);
        }
    }
}

判断规则只有两条：

看到 OneToOneDependency 这类窄依赖：继续往父 RDD 走。
看到 ShuffleDependency：停下来，创建一个父 Stage。

注意，划分 Stage 的判断依据始终是依赖类型，不是算子名称。reduceByKey 之所以会划分出新的 Stage，是因为它产生的是 ShuffleDependency；map 之所以不划分，是因为它产生的是窄依赖。如果你自己实现一个新算子，只要它产生窄依赖，它也不会切出新的 Stage。

创建 ShuffleMapStage 时，调度器把 ShuffleDependency 放进 Stage：

private Stage newShuffleMapStage(ShuffleDependency<?, ?> dependency) {
    List<Stage> parents = getParentStages(dependency.rdd());
    return new Stage(
            nextStageId.getAndIncrement(),
            dependency.rdd(),
            true,
            parents,
            Optional.of(dependency));
}

以示例程序的血缘为例：

ListRDD -> MapPartitionsRDD -> ShuffledRDD

从 ShuffledRDD 往回看，第一条依赖是 ShuffleDependency。所以这里切开：

Stage 0: ListRDD -> MapPartitionsRDD
Stage 1: ShuffledRDD

ListRDD -> MapPartitionsRDD 之间是 OneToOneDependency，属于窄依赖，不再继续切。

这就是从 RDD DAG 到 Stage 图的转换。RDD 图保留每一步 transformation，Stage 图只保留必须分段执行的边界。

7.5 先执行父 Stage，再执行子 Stage

上一节已经划分出了 Stage 树：

ResultStage（计算 ShuffledRDD）
  └─ ShuffleMapStage（计算 MapPartitionsRDD）

父 Stage 负责写 shuffle 文件，子 Stage 负责读这些文件。如果子 Stage 先执行，它会发现文件还没写出来，随即失败。所以执行的先后顺序由 Shuffle 文件的依赖关系强制决定——父 Stage 必须先执行完，子 Stage 才能开始。

现在只看一个问题：调用 submitStage(ResultStage) 后，代码怎样保证父 Stage 完成以前，ResultStage 的任务不会开始？

7.5.1 submitStage：先递归父 Stage

入口就是 submitStage()：

private <T, U> List<U> submitStage(
        Stage stage,
        TaskScheduler taskScheduler,
        Function<Iterator<T>, U> partitionFunction,
        Set<Integer> finishedStages) {
    for (Stage parent : stage.parents()) {
        submitShuffleMapStage(parent, taskScheduler, finishedStages);
    }

    if (verbose) {
        System.out.println("提交 " + stage);
    }
    return submitMissingTasks(stage, taskScheduler, partitionFunction);
}

这段代码的执行顺序如下：

sequenceDiagram
    participant D as DAGScheduler
    participant TS as TaskScheduler
    participant M as ShuffleMapTask
    participant R as ResultTask

    D->>D: submitStage(ResultStage)
    Note over D: for 循环：对每个父 Stage 调 submitShuffleMapStage
    D->>TS: submitTasks(所有 ShuffleMapTask)
    TS->>M: 线程池并行写 map 输出
    TS->>TS: 逐个 await(future)，等全部 Map 任务完成
    TS-->>D: submitTasks() 返回
    Note over D: shuffle 文件已全部写好，submitShuffleMapStage 才返回
    Note over D: for 循环结束，才开始提交 ResultStage
    D->>TS: submitTasks(所有 ResultTask)
    TS->>R: 线程池并行读 map 输出、跑分区函数
    TS->>TS: 逐个 await(future)，收集各分区结果
    TS-->>D: 返回所有分区结果

submitShuffleMapStage(...) 是同步调用——它返回时，父 Stage 的所有任务已经执行完毕，Map 输出文件已经全部写完。for 循环结束后，代码才继续提交 ResultStage。

return submitMissingTasks(stage, taskScheduler, partitionFunction);

DAGScheduler 负责保证这个先后顺序。把每个分区任务交给线程池的，仍然是 TaskScheduler。

7.5.2 ShuffleMapStage：写端

上一节已经看到 submitShuffleMapStage() 是阻塞等待父 Stage 全部完成的入口。这一节深入看它的内部实现：

private void submitShuffleMapStage(
        Stage stage,
        TaskScheduler taskScheduler,
        Set<Integer> finishedStages) {
    if (!finishedStages.add(stage.id())) {
        return;
    }

    for (Stage parent : stage.parents()) {
        submitShuffleMapStage(parent, taskScheduler, finishedStages);
    }
    if (!stage.shuffleMap()) {
        throw new IllegalArgumentException("stage must be a ShuffleMapStage");
    }

    if (verbose) {
        System.out.println("提交 " + stage);
    }
    submitMissingTasks(stage, taskScheduler);
    if (verbose) {
        System.out.println("  shuffle map 输出已写入磁盘");
    }
}

finishedStages 用来避免同一个 Stage 被重复提交。

中间的 for 循环处理更早的父 Stage。当前例子里只有两段，所以父 Stage 没有再往上的 Stage。

最后的 submitMissingTasks(stage, taskScheduler) 取得这条 Stage 保存的 ShuffleDependency：

private void submitMissingTasks(Stage stage, TaskScheduler taskScheduler) {
    ShuffleDependency<?, ?> dependency = stage.shuffleDependency()
            .orElseThrow(() -> new IllegalStateException("missing shuffle dependency"));
    submitShuffleMapTasks(stage, taskScheduler, dependency);
}

接着，DAGScheduler 为父 RDD 的每个 Map 分区创建一个 ShuffleMapTask：

private <K, V> void submitShuffleMapTasks(
        Stage stage,
        TaskScheduler taskScheduler,
        ShuffleDependency<K, V> dependency) {
    RDD<KeyValuePair<K, V>> rdd =
            (RDD<KeyValuePair<K, V>>) stage.rdd();
    List<ShuffleMapTask<K, V>> tasks = new ArrayList<>();
    for (Partition partition : rdd.partitions()) {
        tasks.add(new ShuffleMapTask<>(rdd, partition, dependency));
    }
    taskScheduler.submitTasks(tasks);
}

ShuffleMapTask.call() 分三步：先准备一组空桶，再把父分区的数据逐条分到桶里，最后把每个桶写成文件。

先准备桶。一个 Map 分区要为每个 Reduce 分区各写一个文件，所以先按 Reduce 分区数创建一组空 HashMap，每个桶负责存放这个 Map 分区里属于同一个 Reduce 分区的 key/value：

List<Map<K, V>> buckets = new ArrayList<>();
for (int i = 0; i < dependency.numReducePartitions(); i++) {
    buckets.add(new HashMap<>());
}

再分桶。rdd.iterator(partition) 是上一章那条窄依赖流水线的入口——ListRDD -> MapPartitionsRDD 的 compute 会在这里一层套一层地启动，逐条输出 key/value。每读到一条，先用 dependency.partition(key) 算出它该去哪个 Reduce 分区（也就是桶号），再把这条合并进对应的桶：

Iterator<KeyValuePair<K, V>> iterator = rdd.iterator(partition);
while (iterator.hasNext()) {
    KeyValuePair<K, V> kv = iterator.next();
    int bucketId = dependency.partition(kv.key());
    buckets.get(bucketId).merge(
            kv.key(), kv.value(), dependency.reduceFunc());
}

这里的 merge 是 Map 端的本地合并：同一个 key 在这一个 Map 分区里出现多次，会先被 reduceFunc 合并为一条再写出，从而减少写入的数据量。

最后写入磁盘。桶号就是 Reduce 分区号，所以第 i 个桶写成 map_当前分区号_reduce_i：

for (int reduceId = 0;
     reduceId < dependency.numReducePartitions();
     reduceId++) {
    writeMapOutput(partition.index(), reduceId, buckets.get(reduceId));
}
return null;

写文件的方法和上一章的文件格式保持一致：先写条数，再写每个 key/value。

private void writeMapOutput(int mapId, int reduceId, Map<K, V> data) {
    File file = dependency.mapOutputFile(mapId, reduceId);
    try (ObjectOutputStream out = new ObjectOutputStream(
            new BufferedOutputStream(new FileOutputStream(file)))) {
        out.writeInt(data.size());
        for (var entry : data.entrySet()) {
            out.writeObject(entry.getKey());
            out.writeObject(entry.getValue());
        }
    } catch (IOException e) {
        throw new UncheckedIOException("写入 shuffle 文件失败: " + file, e);
    }
}

到这里，写端三个对象的职责就分开了：

对象	职责
ShuffleDependency	保存写 Map 输出所需的信息：父 RDD、Reduce 分区数、shuffle 目录、合并函数、分区方法、文件路径构造
ShuffleMapTask	计算一个 Map 分区，并把这个分区的数据写成多个 Reduce 桶文件
ShuffledRDD	表示 shuffle 之后的结果 RDD，在 Reduce 阶段读取已经写出的 Map 输出

也就是说，写文件属于 ShuffleMapTask，读文件属于 ShuffledRDD。二者通过同一个 ShuffleDependency 连接到同一批 shuffle 文件。

7.5.3 ResultStage：读端

父 Stage 已经把 shuffle 文件写完。子 Stage 的 ResultTask 现在要读这些文件，再把结果交给用户调用的 action。

ResultStage 只记录最终要计算哪个 RDD。创建 ResultTask 时，还要回答另一个问题：计算完一个分区以后，当前 action 想得到什么结果？

同一个 RDD 可以调用不同的 action：

rdd.collect();
rdd.count();
rdd.reduce(Integer::sum);

三种 action 计算的是同一条 RDD 血缘，区别只发生在最后一步：

action	拿到一个最终分区后做什么	单个分区的返回值
collect()	把迭代器中的元素放进 List	List<T>
count()	一边遍历，一边计数	Long
reduce()	一边遍历，一边用 operator 合并	U（局部合并结果）

因此，SparkContext.runJob(...) 除了接收最终 RDD，还要接收一个分区函数：

public <T, U> List<U> runJob(
        RDD<T> rdd,
        Function<Iterator<T>, U> partitionFunction) {
    return dagScheduler.runJob(rdd, taskScheduler, partitionFunction);
}

它的输入是最终 RDD 一个分区的 Iterator<T>，输出是这个分区交给 action 的结果 U。三个 action 只是传入不同的分区函数：

// collect：当前分区 -> 元素列表
sparkContext.runJob(this, iterator -> {
    List<T> list = new ArrayList<>();
    iterator.forEachRemaining(list::add);
    return list;
});

// count：当前分区 -> 元素个数
sparkContext.runJob(this, iterator -> {
    long count = 0;
    while (iterator.hasNext()) {
        iterator.next();
        count++;
    }
    return count;
});

// reduce：当前分区 -> 局部合并结果
sparkContext.runJob(
        this, iterator -> reducePartition(iterator, operator));

submitMissingTasks() 是 Stage 变成 Task 的地方。先看 ResultStage：

private <T, U> List<U> submitMissingTasks(
        Stage stage,
        TaskScheduler taskScheduler,
        Function<Iterator<T>, U> partitionFunction) {
    RDD<T> rdd = (RDD<T>) stage.rdd();
    List<ResultTask<T, U>> tasks = new ArrayList<>();
    for (Partition partition : rdd.partitions()) {
        tasks.add(new ResultTask<>(
                rdd, partition, partitionFunction, verbose));
    }
    return taskScheduler.submitTasks(tasks);
}

这里的 stage.rdd() 就是最终的 ShuffledRDD。循环每遍历到一个 Reduce 分区，就创建一个 ResultTask。因此，两个 Reduce 分区会得到两个 ResultTask。

ResultTask 计算一个最终 RDD 分区，再执行当前 action 传入的 partitionFunction：

public U call() {
    return partitionFunction.apply(rdd.iterator(partition));
}

这行代码很关键：ResultTask 不关心当前 action 是 collect()、count() 还是 reduce()。它只负责把当前分区的迭代器交给分区函数。

对 collect() 来说，分区函数会把元素收入 List，所以每个 ResultTask 返回当前分区的元素列表；最后 RDD.collect() 再把这些列表合并起来。count() 和 reduce() 也使用同一个 ResultTask，但它们会遍历迭代器并返回计数或局部合并结果。

rdd.iterator(partition) 最终进入的是 ShuffledRDD.compute()。它不再补做 Map 端写盘——写文件这一步已经被 Stage 调度提前完成。它的职责变窄了：只读取当前 Reduce 分区需要的 map_*_reduce_i 文件。

@Override
public Iterator<KeyValuePair<K, V>> compute(Partition partition) {
    Objects.requireNonNull(partition, "partition");
    if (partition.index() < 0 || partition.index() >= partitions.size()) {
        throw new IllegalArgumentException("unknown partition: " + partition);
    }

    return toKeyValuePairs(readAndMergeReducePartition(partition.index()));
}

读取时，它把所有 Map 分区写给当前 Reduce 分区的文件合并起来，相同 key 用 reduceFunc 再合并一次：

private Map<K, V> readAndMergeReducePartition(int reduceId) {
    Map<K, V> merged = new HashMap<>();
    for (int mapId = 0; mapId < numMapPartitions; mapId++) {
        Map<K, V> mapOutput = readMapOutput(mapId, reduceId);
        for (var entry : mapOutput.entrySet()) {
            merged.merge(entry.getKey(), entry.getValue(), reduceFunc);
        }
    }
    return merged;
}

7.5.4 TaskScheduler：统一的提交入口

注意，TaskScheduler 的接口在本章发生了一个变化：它不再接收 RDD 并自行创建任务，而是接收一个已经创建的任务列表（List<? extends Callable<T>>）。因为 Stage 划分后，ShuffleMapTask 和 ResultTask 由 DAGScheduler 负责创建，TaskScheduler 只需要专注于"提交并等待执行"这一件事。

两类 Task 都创建完成以后，进入的是同一个 TaskScheduler.submitTasks(...)：

public <T> List<T> submitTasks(List<? extends Callable<T>> tasks) {
    List<Future<T>> futures = new ArrayList<>();
    for (Callable<T> task : tasks) {
        futures.add(executor.submit(task));
    }

    List<T> result = new ArrayList<>();
    for (Future<T> future : futures) {
        result.add(await(future));
    }
    return result;
}

TaskScheduler 不需要知道这些任务属于哪个 Stage，也不需要判断它们是写 Shuffle 文件还是计算最终结果。它只提交任务、等待 Future，并按提交顺序返回结果。

现在再总结 Stage 和 Task 的关系：

ShuffleMapStage
  -> DAGScheduler 为每个 Map 分区创建一个 ShuffleMapTask
  -> TaskScheduler 提交这些任务

ResultStage
  -> DAGScheduler 为每个结果分区创建一个 ResultTask
  -> TaskScheduler 提交这些任务

到这里，两个调度器的分工已经清晰：

调度器	关心的问题
DAGScheduler	作业要切成几个 Stage，每个 Stage 应该创建哪些 Task
TaskScheduler	把收到的 Task 提交给线程池，并等待执行结果

两个调度器处理的是不同层次的问题：先按 Shuffle 边界安排大步骤，再把具体的分区任务交给线程池执行。

7.6 跑一次，看 Stage 树

把 7.1 的示例程序运行起来（完整实现见 Main.java）。

运行：

mvn -pl ch07-stage-dag package
java -Dfile.encoding=UTF-8 -cp ch07-stage-dag/target/classes com.sparklearn.Main

可以看到 Stage 划分结果：

ResultStage 1 (rdd=ShuffledRDD, parents=[0])
  ShuffleMapStage 0 (rdd=MapPartitionsRDD, parents=[])

这棵树要从下面往上读。

ShuffleMapStage 0 负责把 MapPartitionsRDD 之前的窄依赖链算完，并写出 shuffle 文件。

ResultStage 1 表示最终要提交的 ShuffledRDD 计算。它会在父 Stage 完成后，交给 TaskScheduler 启动 Reduce 分区任务。

执行时会先出现下面这个顺序：

提交 ShuffleMapStage 0 (rdd=MapPartitionsRDD, parents=[])
  shuffle map 输出已写入磁盘
提交 ResultStage 1 (rdd=ShuffledRDD, parents=[0])
  提交 ResultStage 的分区任务

随后底层的 TaskScheduler 启动 ShuffledRDD 的 Reduce 分区任务，读文件并返回结果。用户侧仅写了 shuffled.collect()。

结果里 key 的顺序仍然不重要，因为最后仍然是从 HashMap 转成结果列表：

private Iterator<KeyValuePair<K, V>> toKeyValuePairs(Map<K, V> merged) {
    List<KeyValuePair<K, V>> result = new ArrayList<>();
    for (var entry : merged.entrySet()) {
        result.add(new KeyValuePair<>(entry.getKey(), entry.getValue()));
    }
    return result.iterator();
}

只要计数相符即可：

hello -> 4
world -> 2
spark -> 2
java  -> 1

7.7 本章小结

到这里，Stage 和 DAGScheduler 已经能够正常工作。

核心规则只有一句：窄依赖留在同一个 Stage，宽依赖划分出新的 Stage。

这条规则把前面的几个概念串联起来：

1. 迭代器流水线，解释了窄依赖为什么可以放在同一个 Stage。
2. Dependency 分类，提供了切 Stage 的类型判据。
3. Shuffle 文件，解释了宽依赖为什么必须成为物理边界。

执行 reduceByKey 的 action 时，DAGScheduler 会先沿血缘找到 Shuffle 边界，创建父 ShuffleMapStage，主动写好中间文件；随后最终 ResultStage 才会创建 ResultTask，读取这些文件并返回结果。

本章之前，TaskScheduler 已经可以并行计算 RDD 的各分区，Shuffle 文件读写也已经实现，但 Map 端写盘和 Reduce 端读取混在同一个 compute() 里。本章把 Shuffle 边界显式化为 Stage，让调度器先完成 Map 端写文件，再执行 Reduce 端读文件——这是调度层面上的一次关键重组。