Job, Stage, Task
Concepts covered: Spark Execution Model
3 nested units
Before any data moves, Spark cuts your program into 3 nested units. A job is the work 1 action asks for. A stage is a stretch of that job in which no row has to leave the task holding it. A task is 1 stage applied to 1 partition, and it is the only unit an executor ever runs.
| Unit | Created by | How many | Where the UI shows it |
|---|---|---|---|
| Job | An action: `collect()`, `count()`, a write | At least 1 per action | Jobs tab, 1 row each |
| Stage | A point where rows must regroup across partitions | Boundaries in the job, plus 1 | Stages tab, 1 row each |
| Task | A stage applied to 1 partition | The stage's partition count | The task table inside a stage |
The first row says "at least" because a few operations, orderBy among them, submit an extra job ahead of yours. To pick its value ranges, the sort samples the rows arriving at it, so that sampling job runs every stage upstream of the sort to produce those rows. Spark keeps the shuffle files that job writes, and your job then lists the stages already done as skipped instead of running them again. The stage counts in this section describe the plan of the job your action submits, and the Lazy Evaluation lessons trace the extra job in full.
Counting the stages in one program
Regrouping rows so that every key lands in 1 task is called a shuffle. The report below sums net revenue per region over 6 orders stored in 2 partitions, and each shuffle splits its output into 2 partitions. Before you press play, mark each call where you expect a stage to end.
(orders.filter(F.col("status") == "completed")
.withColumn("net", F.col("amount") - F.col("discount"))
.groupBy("region").agg(F.sum("net").alias("revenue"))
.orderBy(F.col("revenue").desc())
.collect())| order_id | region | status | amount | discount |
|---|---|---|---|---|
| 1 | EU | completed | 120 | 10 |
| 2 | US | completed | 80 | 0 |
| 3 | US | cancelled | 60 | 0 |
| 4 | APAC | completed | 90 | 15 |
| 5 | EU | completed | 40 | 5 |
| 6 | US | completed | 150 | 20 |
| Stage | What each task does | Tasks | Why the stage ends |
|---|---|---|---|
| 0 | Reads its partition, drops rows that are not completed, computes `net`, keeps a running sum per region | 2, 1 per input partition | `groupBy` needs all rows of a region in 1 task, and a region's rows can sit in any partition |
| 1 | Adds up the partial sums it receives into final revenue per region | 2, 1 per shuffle partition | `orderBy` needs rows split into revenue ranges, so each task can sort 1 range |
| 2 | Sorts its range and sends its rows to the driver | 2, 1 per range | The job is done once `collect()` holds every row |
That plan has 3 stages and 6 tasks for 1 action, and the driver ends up with 210 for US and 145 for EU, with APAC at 75. The Jobs tab shows 2 jobs, though. The sampling job runs stage 0, then 2 tasks that aggregate the regions and sample the result. The main job finds stage 0's shuffle files already written, lists that stage as skipped, and runs stages 1 and 2, so the per-region sums are computed twice and 8 tasks run in all. filter and withColumn never get a stage of their own. Each output row depends on 1 input row, so Spark runs both inside the task that reads the partition and builds the partial sums in the same pass over its rows. A stage ends only where a task would need rows another task holds.
Scale the input to 50GB and the shape stays the same while the numbers grow. Spark reads files in pieces of up to 128MB by default, so the read stage gets around 400 tasks, and it runs once, inside the sampling job. Every stage after a boundary, the sampling pass included, gets 200 tasks, the default value of spark.sql.shuffle.partitions.
- 1,000
These counts assume adaptive execution is off. Spark 3.2 and later turn it on by default and can merge small partitions mid-job, a behaviour with lessons of its own.
In the program below, which call ends the first stage, and which call submits the job?
(orders
.filter(F.col("profit") > 0)
.withColumn("profit_k", F.col("profit") / 1000)
.groupBy("status")
.agg(F.sum("profit_k").alias("profit_k"))
.collect())- filter ends the first stage, and collect() submits the job
- groupBy ends the first stage, and agg submits the job
- groupBy ends the first stage, and collect() submits the job
- withColumn ends the first stage, and groupBy submits the job
Why a count() shows 2 stages
Run this on the practice data. It has no groupBy anywhere, only a filter and a new column before the count, yet the Stages tab lists 2 stages for it. Spark counts in 2 halves: each task counts the rows in its own partition, then those per-partition numbers move to a single partition, where 1 task adds them up. That move is a boundary like any other, so the second stage has exactly 1 task.
Predict before you edit: change the last line to heavy.groupBy("order_id").count().count() and run it. The job grows to 3 stages: the read, the per-order totals, and the single task that counts orders. The number returned changes from line items to orders holding at least 1 such line.
A job reads events as 120 partitions and runs the code below, with spark.sql.shuffle.partitions at 200 and adaptive execution off. How many stages and tasks does it run?
(events
.filter(F.col("country") == "BR")
.select("user_id", "event_type")
.groupBy("event_type")
.agg(F.count("*").alias("events"))
.collect())- It runs 3 stages and 440 tasks, because filter and select get a stage each
- It runs 2 stages and 320 tasks: 120 that read and pre-aggregate, 200 that finish
- It runs 2 stages and 240 tasks, because both stages run 1 task per input partition
- It runs 1 stage and 120 tasks, because the groupBy runs inside the reading tasks
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.