Spark's DAG scheduler turns a chain of RDD transformations into a
physical execution plan by splitting it into STAGES. Narrow
transformations (map, filter) can be pipelined -- fused together and
run record-by-record with nothing materialized in between. A SHUFFLE
(data movement across partitions, needed by groupBy, join,
reduceByKey) can't be pipelined that way: its output has to be fully
written out before the next stage can start reading it. So a new stage
begins at every shuffle-requiring operation.
Implement compute_stages(dag):
dag: a list of(node_id, parents, is_shuffle)tuples, given in topological order (every parent appears earlier in the list than the node that references it).- A node with no parents starts a new stage (it's a data source).
- A node flagged
is_shuffle=Truealways starts a new stage, regardless of how many parents it has (a join has two -- one per side being joined). - A node flagged
is_shuffle=False(narrow) joins the SAME stage as its one parent -- it must have at most one parent; more than one is an error (ValueError) since this model doesn't represent multi-parent narrow operations likeunion. - A stage
Sdepends on stageTif any node inShas a parent whose stage isT(andT != S).
Return (stage_of, stage_order, stage_deps):
stage_of:{node_id: stage_id},stage_idan int assigned in creation order (starting at 0).stage_order: everystage_id, sorted ascending -- a valid execution order (a dependency's id is always smaller than its dependent's, since it was necessarily created first).stage_deps:{stage_id: sorted list of stage_ids it depends on}.
Also raise ValueError if a node's parent hasn't been processed yet
(the input isn't actually topologically sorted).
Example
dag = [
("read", [], False),
("map", ["read"], False),
("groupBy", ["map"], True),
("mapAfter", ["groupBy"], False),
]
stage_of, stage_order, stage_deps = compute_stages(dag)
# stage_of -> {"read": 0, "map": 0, "groupBy": 1, "mapAfter": 1}
# stage_order -> [0, 1]
# stage_deps -> {0: [], 1: [0]}
A join of two independently-produced sides gets its own stage depending on BOTH upstream stages:
dag = [
("a_src", [], False), ("a_map", ["a_src"], False),
("b_src", [], False), ("b_map", ["b_src"], False),
("join", ["a_map", "b_map"], True),
]
# "a_src"/"a_map" share one stage, "b_src"/"b_map" share a different
# one, and "join" is a third stage depending on both.
Real-world grounding
- Spark's
DAGScheduler, which walks the RDD dependency graph backward from the final action and creates a newShuffleMapStageat everyShuffleDependencyboundary:core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
Real Spark's stage graph also tracks partitioning info, handles stage retries on fetch failure, and can skip recomputing a stage whose shuffle output is still available on disk from a prior run -- this exercise implements the core "split at shuffle boundaries, track inter-stage dependencies" structure both are built on.