Why Stages Exist

Concepts covered: Spark Execution Model

A task sees only its own partition

Go back to stage 0 of the revenue report. Task 0 holds some orders and task 1 holds the rest, and either can hold US rows, so neither can produce the final US total alone. Spark settles this by ending the stage. Each task writes its partial sums to its executor's local disk, split into 1 piece per destination partition, and a fresh set of tasks fetches the pieces it owns from every executor. The Shuffles lessons take that write and fetch apart; this lesson needs only the timing it forces.

The barrier between stages

A stage 1 task that owns the US piece cannot know it has every US row until every stage 0 task has finished writing, since the last one to finish might hold US rows too. So Spark starts no task in a stage until all tasks of the stage it reads from have succeeded. A stage lasts as long as its slowest task, and everything downstream waits for it.

That is the 2 a.m. dashboard from the top of this lesson. With 199 of 200 tasks done, the next stage is waiting on 1 task and will wait 9 minutes as readily as 9 seconds. Adding executors at that point does not speed that task up, because 1 task runs on 1 core. A partition that received far more rows than its neighbours is the usual cause, and the Data Skew lessons deal with it.

Inside a stage

• Each task carries its rows through every step back to back, in memory
• Tasks never wait on each other; each finishes on its own schedule
• A filter or a new column costs CPU time and nothing else

Across a boundary

• Every task writes its output to local disk, split by destination partition
• The next stage's tasks fetch their pieces from every executor over the network
• The next stage launches only after the slowest task here has written its output

That disk and network work is why boundaries are worth counting. In the 50GB run, each of the 400 tasks in the first stage writes 1 output file holding 200 pieces, 1 for each task in the next stage, and every one of those 200 tasks fetches its piece from all 400 files.

• 80,000

What creates a boundary

Calls that need rows from other partitions

• groupBy(...).agg(...): rows sharing a key must meet in 1 task.
• distinct() and dropDuplicates(): 2 copies of a row can start in different partitions.
• join of 2 large tables: matching keys must land in the same task. A join to a small table can avoid it, as the Joins lessons show.
• orderBy: rows move into value ranges, so each task sorts 1 range and the ranges line up.
• repartition(n) or repartition("col"): moving rows is the whole purpose of the call.
• A window with partitionBy: each window's rows must sit in 1 task.

Everything else you have used so far, filter, select, withColumn, union, runs inside the stage already in progress.

With adaptive execution off, how many stages does this job run?

(orders
    .repartition("status")
    .groupBy("region")
    .agg(F.sum("profit").alias("profit"))
    .collect())

• It runs 3: repartition makes 1 boundary and groupBy needs a second for its own key
• It runs 2: the groupBy reuses the layout that repartition already built
• It runs 1: neither call needs rows from other partitions to finish its work
• It runs 2: repartition runs inside the reading stage and only groupBy is a boundary

A job with 2 boundaries

This program on the practice data counts line items per order and sorts the orders by that count. Trace it before running: a groupBy boundary, then an orderBy boundary, so 3 stages. The rows come back in descending order across the whole result because each task in the last stage sorted 1 contiguous range of counts, and the driver lines the ranges up.

Now add .limit(5) to the end of the chain and run it again. The 5 rows match the top of the full result, yet that version runs 2 stages, for the top-N reason in the tip above.

Stage 3 of a job has 200 tasks. 199 of them finish within 2 minutes, and the last finishes at 11 minutes. Stage 4 reads stage 3's output. When does stage 4's first task start?

• At about 2 minutes, reading the 199 finished outputs while the last task runs
• At about 2 minutes, with the missing task's rows merged in later by the driver
• Right away, because stages placed on different executors run at the same time
• At about 11 minutes, once all 200 tasks have written their output pieces

About this lesson

This section is part of the Execution Model: Intermediate 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.