---
title: "How to Count Spark Stages by Sight and Optimize Shuffles"
description: "Want to predict execution stages before running a job? Learn how to count Spark stages by sight with code breakdowns, shuffle rules, and optimization traps."
url: https://srinimf.com/2026/09/27/how-to-count-spark-stages-by-sight-and-optimize-shuffles/
date: 2026-09-27
modified: 2026-10-11
author: "Srini"
image: https://i0.wp.com/srinimf.com/wp-content/uploads/2026/09/pexels-photo-14314638.jpeg?fit=1880%2C1253&ssl=1
categories: ["PySpark"]
tags: ["Apache Spark", "Data Engineering", "Spark Stages"]
type: post
lang: en
---

# How to Count Spark Stages by Sight and Optimize Shuffles

Learning to count Spark stages is essential for optimizing your Spark jobs. This guide teaches you [how to](https://srinimf.com/2024/09/27/how-to-resolve-pyspark-sql-puzzle-merchant-transaction-data/) 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](https://spark.apache.org/docs/latest/).

Here is a practical [guide to](https://srinimf.com/2024/11/22/a-comprehensive-guide-to-databricks-log-types-and-access/) calculating Spark stages by sight, complete with rules, edge cases, and code examples that demonstrate how Spark stages form at shuffle boundaries.

Related[PySpark Code: Calculate Click Rates and Salary Matches](https://srinimf.com/2024/09/30/how-to-calculate-click-percentage-dept-wise-salary-matches-in-pyspark/)

![How to count Spark stages: visual guide showing shuffle boundaries and Exchange operators in Apache Spark execution plans](https://i0.wp.com/srinimf.com/wp-content/uploads/2026/09/srinimf-6ab91fd8ec56c.png?resize=1184%2C864&ssl=1)

## 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 codedf.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](https://srinimf.com/about/) page, or feel free to [contact us](https://srinimf.com/contact-us/) with any questions about Spark optimization.

### Share this:

-
- [Tweet](https://twitter.com/share)
- [![](https://i0.wp.com/assets.pinterest.com/images/pidgets/pinit_fg_en_rect_gray_20.png)](https://www.pinterest.com/pin/create/link/?url=https%3A%2F%2Fsrinimf.com%2F2026%2F09%2F27%2Fhow-to-count-spark-stages-by-sight-and-optimize-shuffles%2F&media=https%3A%2F%2Fi0.wp.com%2Fsrinimf.com%2Fwp-content%2Fuploads%2F2026%2F09%2Fpexels-photo-14314638.jpeg%3Ffit%3D1200%252C800%26ssl%3D1&description=How%20to%20Count%20Spark%20Stages%20by%20Sight%20and%20Optimize%20Shuffles)
-
- [Share on Tumblr](https://www.tumblr.com/share)
- [Share on WhatsApp (Opens in new window) WhatsApp](https://srinimf.com/2026/09/27/how-to-count-spark-stages-by-sight-and-optimize-shuffles/?share=jetpack-whatsapp)
- [More](#)
-

- [Print (Opens in new window) Print](https://srinimf.com/2026/09/27/how-to-count-spark-stages-by-sight-and-optimize-shuffles/?share=print)
- [Share on Reddit (Opens in new window) Reddit](https://srinimf.com/2026/09/27/how-to-count-spark-stages-by-sight-and-optimize-shuffles/?share=reddit)
- [Share on Telegram (Opens in new window) Telegram](https://srinimf.com/2026/09/27/how-to-count-spark-stages-by-sight-and-optimize-shuffles/?share=telegram)
- [Email a link to a friend (Opens in new window) Email](mailto:?subject=%5BShared%20Post%5D%20How%20to%20Count%20Spark%20Stages%20by%20Sight%20and%20Optimize%20Shuffles&body=https%3A%2F%2Fsrinimf.com%2F2026%2F09%2F27%2Fhow-to-count-spark-stages-by-sight-and-optimize-shuffles%2F&share=email)
-

### Like this:

Like Loading…