Answer in Stages

Concepts covered: Shuffle Optimization

An answer ordered by component (driver, then cluster manager, then executors) describes the org chart of a Spark application. An answer ordered by stage describes the job's timeline, and the timeline is where the cost lives, because every stage boundary is a point where data is written, moved and read again. Reordering the same facts around stages is the change the rest of this lesson builds on.

Component order against stage order

Ordered by component

• The driver creates a SparkSession and requests resources
• The cluster manager launches executors
• The DAG scheduler splits the work into stages
• Executors run tasks and return results

Ordered by stage

• Plan: the write triggers a physical plan with a broadcast join and 2 shuffles
• Stage 1: about 4000 scan tasks filter, join and pre-sum, then write shuffle files
• Stage 2: 200 reducers fetch those files, finish the sum, write range output
• Stage 3: sorted ranges are written out as the report

The stage-ordered version still mentions the driver and the executors, but each appears at the moment it does work, attached to a count. The interviewer can interrupt at any stage and ask "why that many?", and the answer is already in the sentence.

Where the boundaries fall

Spark cuts a stage at every shuffle Exchange in the physical plan. Operators between exchanges (scan, filter, project, the map side of an aggregate, a broadcast hash join) run pipelined inside the same task over the same partition, with no intermediate result written anywhere. That pipelining is why a filter adds almost no cost: it rides inside a task that is already reading the data. An Exchange is where a task stops, writes its output grouped by destination partition, and hands off to tasks in the next stage that start by fetching it.

In the SQL / DataFrame tab, pipelined operators sit inside boxes labelled WholeStageCodegen with a number, and the Exchange nodes sit between those boxes, so the stage boundaries are visible in the graph.

Watch a boundary form

This run filters and groups 9 trips read as 3 partitions, with 3 reducers on the other side. Watch the first stage do the scan, the filter and a partial sum inside each task, then watch the partial sums cross to the reducers.

city_totals = (
    trips
    .filter("fare > 5")
    .groupBy("city")
    .sum("fare")
)
city_totals.collect()

Two trips fail the filter (a 4.0 fare in SF and a 3.5 fare in LA) and never leave their tasks. Each map task sends at most 1 partial sum per city across the boundary, so this shuffle carries at most 9 rows however many trips each partition held. The driver collects 3 totals: SF 39.5, NYC 34.5 and LA 22.0. The same property explains stage 1 of the March job: heavy on reading, light on shuffling, because each task holds far more trips than there are cities.

The spoken shape

Four beats, a number in each

• Plan: The write or collect makes the driver build the physical plan. Name its shuffle exchanges and the join strategy it chose.
• Stages: Count them out loud as shuffles plus 1. Say which stage reads the most bytes and which writes the most.
• Tasks: Input splits of 128MB set the scan's task count. The 200 shuffle partitions are the planned count after each exchange, which adaptive execution may merge when they turn out small. Divide by cores to get waves.
• Shuffle and result: Say roughly how many bytes cross each boundary, and whether the result returns to the driver or lands in storage.

About this lesson

This section is part of the The Architecture Recital lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.