Learning to count Spark stages is essential for optimizing your Spark jobs. This guide teaches you how to count Spark stages by identifying shuffle boundaries and understanding transformation dependencies, so you can estimate stage counts, spot expensive operations, and tune your workloads for better performance.

Spark stages are the execution units that Apache Spark uses to run a job. By understanding how transformations and shuffles create stage boundaries, developers can estimate stage counts, identify expensive operations, and optimize Spark workloads.

Knowing your Spark stages count helps you identify expensive shuffles, optimize bottlenecks, and navigate the Spark UI without guessing. For a deeper understanding of the underlying architecture, refer to the official Apache Spark documentation.

Here is a practical guide to calculating Spark stages by sight, complete with rules, edge cases, and code examples that demonstrate how Spark stages form at shuffle boundaries.

The Golden Rule: How Shuffles Create Spark Stages and Boundaries

Spark breaks down execution into a hierarchy: Jobs → Stages → Tasks.

• Actions (count(), collect(), write(), show()) trigger Jobs.
• Stages are determined strictly by Wide Dependencies (Shuffles).

Total Stages in a Linear Branch = (Number of Shuffles) + 1

1. Narrow Transformations in Spark (Zero New Stages)

In a narrow transformation, each partition of the parent RDD/DataFrame is processed independently without network exchange. Spark pipelines these together inside the same stage.

• Common examples: select(), filter(), withColumn(), drop(), map(), flatMap().

2. Wide Transformations in Spark (Add a Stage Boundary)

In a wide transformation, data must be grouped, rearranged, or aggregated across cluster nodes. This triggers a shuffle, splitting the DAG into a write phase and a read phase.

• Common examples: groupBy(), reduceByKey(), join(), distinct(), repartition(), cube(), rollup(), orderBy() / sort().

4 Code Examples: Counting Spark Stages by Sight

Example 1: Purely Narrow Pipeline (1 Stage)

Python

df = spark.read.parquet("s3://bucket/data/")
result = (
    df.filter("age > 25")
      .select("name", "salary")
      .withColumn("annual_bonus", col("salary") * 0.1)
)
result.write.mode("overwrite").parquet("s3://bucket/output/")

• Breakdown:
  • read, filter, select, withColumn: All narrow dependencies.
  • write: Triggers 1 Action.
  • Shuffles: 0
• Total Stages: $0 + 1 =$ 1 Stage

Example 2: Aggregation with a Shuffle (2 Stages)

Python

df = spark.read.csv("s3://bucket/sales.csv", header=True)
aggregated = (
    df.filter("region = 'EMEA'")
      .select("category", "amount")
      .groupBy("category")
      .sum("amount")
)
aggregated.write.format("parquet").save("s3://bucket/sales_summary/")

• Breakdown:
  • Stage 1 (Map Side): Reads the data, applies filters/projections, runs local pre-aggregation (partial sum), and writes shuffle files.
  • Shuffle Boundary: Network redistribution of rows grouped by category.
  • Stage 2 (Reduce Side): Reads shuffled blocks, calculates final sums, and writes output to storage.
• Total Stages: $1 \text{ (Shuffle)} + 1 =$ 2 Stages

Example 3: Mixing Multiple Wide Dependencies (3 Stages)

Python

df = spark.read.json("s3://bucket/web_logs/")
processed = (
    df.repartition(100)                     # Wide (Shuffle #1)
      .filter("status = 200")                # Narrow
      .groupBy("endpoint")                  # Wide (Shuffle #2)
      .count()
)
processed.collect()                         # Action

• Breakdown:
  • Stage 1: Scan data $\rightarrow$ Partition hash calculation for repartition(100).
  • Shuffle 1: Exchange data across 100 partitions.
  • Stage 2: Apply filter on local partitions $\rightarrow$ Map-side group aggregation for endpoint.
  • Shuffle 2: Exchange data by endpoint key.
  • Stage 3: Aggregate counts $\rightarrow$ send final result to driver via collect().
• Total Stages: $2 \text{ (Shuffles)} + 1 =$ 3 Stages

Example 4: Joining Two Tables (3 Stages)

When two unrelated datasets are shuffled into a standard Sort-Merge Join, Spark executes two independent upstream stages in parallel before joining them in a third.

Python

orders = spark.read.parquet("s3://bucket/orders/")
customers = spark.read.parquet("s3://bucket/customers/")
joined = (
    orders.filter("order_status = 'COMPLETED'")
          .join(customers, on="customer_id", how="inner")
          .select("order_id", "customer_name", "amount")
)
joined.write.parquet("s3://bucket/joined_orders/")

• Breakdown:
  • Stage 1: Reads orders, applies filter, hashes/partitions by customer_id for shuffle.
  • Stage 2: Reads customers, hashes/partitions by customer_id for shuffle.
  • Stage 3: Reads shuffled data from both Stage 1 and Stage 2, performs the merge join, executes select, and writes output.
• Total Stages: 3 Stages (2 map stages running in parallel + 1 join/reduce stage).

3 Common Traps That Change the Expected Stage Count

Real-world Spark optimization can alter the naive stage count:

1. Broadcast Joins: If customers is small enough (or hint broadcast(customers) is used), Spark skips shuffling entirely. It broadcasts the small table to executors, reducing the join from 3 stages to 1 stage. For more on optimization strategies, see PySpark Performance Anti-Patterns.
2. coalesce() vs repartition(): repartition() forces a shuffle (+1 stage), but coalesce(n) collapses existing partitions without network shuffling, keeping it inside the same stage (0 added stages).
3. Adaptive Query Execution (AQE): In modern Spark 3.x/4.x with AQE enabled (spark.sql.adaptive.enabled = true), Spark dynamically re-evaluates the physical plan at runtime. It can combine stages or convert a sort-merge join to a broadcast join on the fly based on intermediate shuffle file statistics.

How to Verify the Actual Execution Plan in the Spark UI

To verify your Spark stages calculation directly without opening the Spark UI:

Python

# Check the execution plan in code
df.explain(True)
# For RDDs, count the lineage depth:
print(rdd.toDebugString().decode("utf-8"))

In explain(), look for Exchange operators. Each Exchange node represents a shuffle and therefore marks a stage boundary. For more insights into optimizing Apache Spark performance, explore our About Srini page, or feel free to contact us with any questions about Spark optimization.