Spark DAG, Stages and Tasks

Every Spark action turns into a job. Spark splits each job into stages and each stage into tasks. Knowing this chain helps you read the Spark UI, find slow steps, and understand why some operations cost far more than others.

The Layman View

A restaurant receives an order for a pizza. The kitchen breaks the order into steps: prepare dough, add toppings, bake, and box. Some steps run side by side, and some wait for the previous step to finish. Spark breaks your code into steps the same way and runs independent steps in parallel.

The Chain from Code to Tasks

Your code
   |  action such as count() or write()
   v
Job
   |  split at every shuffle
   v
Stages
   |  one task per partition
   v
Tasks  -->  run on executors

What Is a DAG

DAG stands for directed acyclic graph. Spark records every transformation as a node and links each node to the node it depends on. Arrows run in one direction and never loop back. The graph describes the full recipe for the result.

read file --> filter --> map --> groupBy --> sum --> write

Narrow and Wide Dependencies

Transformations fall into two groups based on how data moves.

  • Narrow transformations, such as filter, map, and select, work on each partition alone. No data travels between machines.
  • Wide transformations, such as groupBy, join, and distinct, need rows from many partitions. Spark shuffles data across the network.
Narrow:   Partition 1 --> Partition 1'      each partition stays on its machine
          Partition 2 --> Partition 2'

Wide:     Partition 1 --\  /--> Partition A   every output partition
          Partition 2 ----X----> Partition B   reads from many inputs
          Partition 3 --/  \--> Partition C

Stages Split at Shuffles

Spark groups consecutive narrow steps into a single stage and runs them in one pass. Each wide step ends a stage and begins another.

df = spark.read.parquet("/data/sales")
result = (df.filter("amount > 0")
            .select("region", "amount")
            .groupBy("region")
            .sum("amount"))
result.write.parquet("/data/report")
Stage 1: read --> filter --> select --> write shuffle files
------------------------------------------------------ shuffle
Stage 2: read shuffle files --> sum --> write result

Stage 2 cannot start until every task of Stage 1 finishes, because the sums need all rows of each region.

Tasks Match Partitions

A stage creates one task for each partition of its data. A file split into 8 partitions produces 8 tasks. Each executor core runs one task at a time.

Stage with 8 tasks, cluster with 4 cores:

Round 1: Task 1  Task 2  Task 3  Task 4
Round 2: Task 5  Task 6  Task 7  Task 8

More cores shorten the number of rounds. Tasks that vary in size lengthen the last round.

Jobs Come from Actions

Transformations only add nodes to the graph. An action triggers execution.

total  = df.count()        # job 1
sample = df.take(5)        # job 2
df.write.parquet("/out")   # job 3

Each line above creates its own job. The Jobs tab of the Spark UI lists them in order.

Fault Tolerance Through Lineage

The graph records how every partition was built. Spark loses an executor now and then. Spark replays only the lost partition from its recorded steps and leaves the rest alone.

Partition 3 lost  -->  re-read source chunk --> filter --> map --> partition 3 rebuilt

Seeing the Plan

result.explain()

Each Exchange line in the printed plan marks a shuffle and a stage boundary. Count the Exchange lines to predict the number of stages.

Practical Lessons

  • Fewer shuffles mean fewer stages and faster jobs.
  • Filter before wide operations to shrink shuffle data.
  • Tune partition counts so tasks run for seconds, not milliseconds or hours.
  • Check the Stages tab when a job runs slowly. The slow stage names the wide operation to fix.