Definition
- Stage – a set of transformations that can run without any data shuffle, so all tasks execute independently on their partitions and can be pipelined together.
- DAGScheduler – the component in the Driver that splits the logical DAG into stages, cutting a new boundary at every wide (shuffle) dependency.
- One task per partition – within a stage, Spark creates exactly one task per partition; the task count equals the partition count.
- Job – the work triggered by one action, divided into stages at shuffle boundaries.
- Shuffle-map stage – produces shuffle output for the next stage; a result stage produces the action's output.
Picture a multi-phase reporting project at the National Bank of India. A stage is one phase where every clerk can work entirely on their own box of files without passing anything to a neighbour. The DAGScheduler is the planner who draws the phase boundaries, and he draws a new boundary exactly wherever the files must be re-sorted and exchanged between teams – that re-sort is the shuffle. Inside any one phase, the work is split so that one clerk handles one box: one task per box of files.
Understanding the Concepts
What is a Spark stage?
A stage is defined by the absence of a shuffle inside it. When transformations have only narrow dependencies – map, filter, flatMap – each partition's computation depends solely on one upstream partition, so Spark can pipeline the whole chain inside a single task without moving data across the network.
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum as spark_sum, upper
spark = SparkSession.builder.appName("StageDemo").getOrCreate()
# These narrow transformations all pipeline into ONE stage
txns = spark.read.parquet("/data/transactions")
filtered = txns.filter(col("amount") > 1000) # narrow
renamed = filtered.withColumn("acct", upper(col("account_id"))) # narrow
selected = renamed.select("acct", "amount") # narrow
# All three transformations run in the same stage (Stage 0)
# This groupBy introduces a shuffle, creating a new stage boundary
totals = selected.groupBy("acct").agg(spark_sum("amount")) # wide
totals.count() # Action triggers execution
# Spark UI shows:
# Stage 0: scan + filter + rename + select (pipelined, no shuffle)
# Stage 1: groupBy shuffle + aggregation
How does the DAGScheduler determine stage boundaries?
The DAGScheduler walks the RDD or DataFrame lineage backwards from the action. Every time it encounters a wide dependency (shuffle), it marks a stage boundary. The result is a directed acyclic graph of stages.
# Example: two shuffles create three stages
txns = spark.read.parquet("/data/transactions")
# Stage 0: scan + filter (narrow)
filtered = txns.filter(col("amount") > 1000)
# Shuffle 1: groupBy creates Stage 1
account_totals = filtered.groupBy("account_id") \
.agg(spark_sum("amount").alias("total"))
# Read another dataset
accounts = spark.read.parquet("/data/accounts")
# Shuffle 2: join creates Stage 2
enriched = account_totals.join(accounts, "account_id")
enriched.count()
# Physical plan shows:
# Stage 0: Scan transactions + filter
# Stage 1: Scan accounts (parallel with Stage 0)
# Stage 2: Shuffle for groupBy
# Stage 3: Shuffle for join + result
enriched.explain()
What is the relationship between stages, tasks, and partitions?
The relationship is precise and fundamental to understanding Spark performance.
Job (triggered by one action)
└── Stage 0 (scan, 100 partitions)
│ └── 100 tasks (one per partition)
└── Stage 1 (groupBy, 200 partitions)
│ └── 200 tasks (one per partition)
└── Stage 2 (result, 1 partition)
└── 1 task
# See partition and task counts
print(f"Input partitions: {txns.rdd.getNumPartitions()}")
# After groupBy, partition count is controlled by:
print(spark.conf.get("spark.sql.shuffle.partitions")) # Default: 200
# Tune for your data size
spark.conf.set("spark.sql.shuffle.partitions", "100")
# Now Stage 1 will have 100 tasks instead of 200
What are narrow vs wide dependencies?
The type of dependency determines whether a shuffle (and thus a stage boundary) is needed.
| Dependency | Examples | Shuffle? | Stage Boundary? |
|---|---|---|---|
| Narrow | map, filter, select, withColumn, flatMap | No | No |
| Wide | groupBy, join, repartition, distinct, sort | Yes | Yes |
# Narrow: no shuffle, same stage
df1 = txns.filter(col("amount") > 0) # narrow
df2 = df1.withColumn("doubled", col("amount") * 2) # narrow
df3 = df2.select("account_id", "doubled") # narrow
# All in one stage
# Wide: shuffle, new stage
df4 = df3.groupBy("account_id").sum("doubled") # wide -> new stage
df5 = df4.orderBy("account_id") # wide -> new stage
How do you control the number of tasks in a shuffle stage?
The number of post-shuffle partitions (and therefore tasks) is controlled by spark.sql.shuffle.partitions. Tuning this is critical for performance.
# Too few partitions: tasks are large, underusing the cluster
spark.conf.set("spark.sql.shuffle.partitions", "10")
txns.groupBy("account_id").sum("amount").count()
# 10 tasks, each processing a lot of data
# Too many partitions: tasks are tiny, scheduling overhead dominates
spark.conf.set("spark.sql.shuffle.partitions", "10000")
txns.groupBy("account_id").sum("amount").count()
# 10000 tasks, most finish instantly, driver wastes time scheduling
# Right-sized: aim for tasks of 100-200 MB each
spark.conf.set("spark.sql.shuffle.partitions", "200")
# Or enable AQE to auto-tune:
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
How do you read stages in the Spark UI for debugging?
The Spark UI Stages tab is the primary tool for understanding job performance. The number of stages tells you how many shuffles occurred, and the task metrics within each stage reveal bottlenecks.
# After running a job, navigate to the Spark UI at localhost:4040
# Jobs tab: see which job is slow and how many stages it has
# Stages tab: drill into the slow stage
# Task metrics to check:
# - Duration spread (max vs median): skew if max >> median
# - Shuffle read/write: how much data moved
# - GC time: memory pressure if GC is a large fraction
Spark UI Stages Tab: ┌──────────────────────────────────────────────┐ │ Stage 0: Scan + Filter (narrow) │ │ Tasks: 100, Duration: 0.5s median │ │ No shuffle (all local reads) │ ├──────────────────────────────────────────────┤ │ Stage 1: GroupBy (after shuffle) │ │ Tasks: 200, Duration: 2.1s median │ │ Shuffle Read: 500 MB total │ │ Max task: 45s <-- possible skew! │ └──────────────────────────────────────────────┘
What is the difference between shuffle-map stages and result stages?
A shuffle-map stage produces shuffle output that the next stage consumes. A result stage produces the final output of the action (the number for count, the rows for collect, the files for write).
# In this job:
txns.groupBy("account_id").sum("amount").count()
# Stage 0 is a shuffle-map stage: produces shuffle output
# Stage 1 is a result stage: reads shuffle output and returns the count
What should a senior engineer know?
- The number of stages directly reflects the number of shuffles, so minimizing unnecessary wide operations reduces stages and improves performance.
- Task count per stage equals partition count, so tuning partitions tunes parallelism.
- AQE can dynamically coalesce partitions after a shuffle, reducing the effective task count if many partitions are small.
- At the National Bank of India, reading the stage and task view in the Spark UI was the standard diagnostic for slow nightly jobs: the stage count revealed unnecessary shuffles, and task metrics revealed skew.
- A handful of straggler tasks in a stage almost always means data skew in a shuffle partition, not a code bug.
Visual Overview
Logical DAG (before DAGScheduler)
┌──────────┐
│ Scan │ narrow ┌──────────┐ narrow ┌──────────┐
│ txns │────────>│ Filter │────────>│ Select │
└──────────┘ └──────────┘ └─────┬────┘
│
wide (shuffle)
│
┌─────▼────┐
│ GroupBy │
│ + Sum │
└─────┬────┘
│
┌─────▼────┐
│ Count │
│ (action) │
└──────────┘
Physical Stages (after DAGScheduler)
┌─────────────────────────────────┐
│ Stage 0 (shuffle-map) │
│ Scan + Filter + Select │ Pipelined (narrow)
│ 100 tasks (= 100 partitions) │
└────────────────┬────────────────┘
│ shuffle
┌────────────────▼────────────────┐
│ Stage 1 (result) │
│ GroupBy + Sum + Count │
│ 200 tasks (= 200 partitions) │
└─────────────────────────────────┘
Key Takeaway
A Spark stage is a set of transformations that can run without any shuffle. All its tasks operate independently on their own partitions and Spark pipelines the whole narrow chain inside one task. The reason a stage ends is a wide dependency: when a shuffle is needed, the work on the other side cannot begin until every task before it has produced its shuffle output, creating a synchronization barrier.
The DAGScheduler walks the lineage backwards from the action and cuts a stage boundary at every wide dependency. The result is a directed acyclic graph of stages: shuffle-map stages that produce shuffle output and a final result stage. The relationship between stages, tasks, and partitions is precise: a job is divided into stages at shuffle boundaries, each stage is divided into tasks, and there is exactly one task per partition per stage.
At the National Bank of India, this structure was the primary diagnostic tool because it is exactly what the Spark UI shows. The number of stages told how many shuffles a job had, the task count per stage told the partition count and parallelism, and straggler tasks running much longer than the rest almost always indicated data skew. Tuning the post-shuffle partition count so tasks were neither too few nor too many was a regular optimization.
The senior perspective is that minimizing unnecessary wide operations reduces stages and improves performance, that task count is tuned through partition count, and that AQE can dynamically coalesce small partitions after a shuffle.
Remember one sentence: the DAGScheduler cuts stages at every shuffle boundary, and within each stage there is exactly one task per partition, which is the foundation of Spark's parallelism.
---