1. Scheduling

Spark DAG Scheduler: Stage Boundary at Shuffle

seniorbig-data-processingscheduling

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=True always 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 like union.
  • A stage S depends on stage T if any node in S has a parent whose stage is T (and T != S).

Return (stage_of, stage_order, stage_deps):

  • stage_of: {node_id: stage_id}, stage_id an int assigned in creation order (starting at 0).
  • stage_order: every stage_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

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.

hints
stage_scheduler.pysign in to run
Loading...
test result

This is a preview. Sign in to run this exercise against the real tests.

Spark DAG Scheduler: Stage Boundary at Shuffle — nossletter.tech