Reading the shuffle in `.explain()` and the Spark UI

Description

Lesson 3 showed the Exchange node over forty rows — real, but invisible in cost. This lesson repeats the exact same experiment over kiosko_orders_at_scale — ten million rows, synthetic, declared as such — and this time the Exchange moves data you can actually count: bytes, records, and a partition count that changes based on how Spark decides to coalesce them. You're also going to see, for the first time in this guide, a .join() that genuinely triggers a real shuffle — something lesson 3 left open on purpose — and you're going to read the Spark UI directly through its REST API, the same mechanism you already used in module 2 to confirm who the executor is.

Connection to the module. This lesson introduces no new concept about what a shuffle is — lesson 3 already built that; it measures, with real evidence, the same phenomenon at the scale this module exists for.

An analogy: the same reordering, now with thousands of real boxes

In lesson 3, the order "reorganize by zip code" landed on seven trucks carrying forty boxes total — work so small the time it takes to reorganize is indistinguishable from zero. This lesson gives the same order to twelve trucks carrying ten million boxes. The instruction is identical — "group by store_id" — but now the reordering moves real data: each truck has to decide, box by box, which new truck each one belongs to, based on its zip code, and that decision and that movement leave a measurable trace: how many bytes traveled, how many records got reorganized, how many stops (partitions) resulted at the end.

Worked example, part 1: the same groupBy, now at scale

# shuffle_at_scale.py
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import (
    StructType, StructField, StringType, IntegerType, DoubleType, TimestampType,
)

spark = SparkSession.builder.appName("kiosko-spark").master("local[*]").getOrCreate()

scale_schema = StructType([
    StructField("order_id", StringType(), False),
    StructField("franchise_id", IntegerType(), False),
    StructField("store_id", StringType(), False),
    StructField("product_id", StringType(), False),
    StructField("quantity", IntegerType(), False),
    StructField("unit_price", DoubleType(), False),
    StructField("order_ts", TimestampType(), False),
])
orders_at_scale_df = spark.read.csv(
    "kiosko_orders_at_scale.csv", schema=scale_schema, header=True, enforceSchema=False,
)

# fact_orders_at_scale: the same data, with revenue calculated -- with no need
# for any JOIN against dim_store or dim_product, because quantity and unit_price
# already live in the same row (exactly as in module 3).
fact_orders_at_scale_df = orders_at_scale_df.withColumn(
    "revenue", F.col("quantity") * F.col("unit_price")
)

print("=== groupBy('store_id').sum('revenue') over 10,000,000 rows ===")
revenue_by_store = fact_orders_at_scale_df.groupBy("store_id").sum("revenue")
revenue_by_store.explain()
revenue_by_store.show()

spark.stop()

What to expect. Running python3 shuffle_at_scale.py, the output is exactly this (executed in this run, over kiosko_orders_at_scale.csv, 10,000,000 rows):

=== groupBy('store_id').sum('revenue') over 10,000,000 rows ===
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[store_id#2], functions=[sum(revenue#8)])
   +- Exchange hashpartitioning(store_id#2, 200), ENSURE_REQUIREMENTS, [plan_id=23]
      +- HashAggregate(keys=[store_id#2], functions=[partial_sum(revenue#8)])
         +- Project [store_id#2, (cast(quantity#4 as double) * unit_price#5) AS revenue#8]
            +- FileScan csv [store_id#2,quantity#4,unit_price#5] Batched: false, DataFilters: [], Format: CSV, ...

+--------+-----------------+
|store_id|     sum(revenue)|
+--------+-----------------+
|     S02|9699999.999998083|
|     S01|9574999.999993788|
|     S03|7262500.000003301|
+--------+-----------------+

The plan is, in structure, identical to lesson 3's: partial HashAggregate, Exchange hashpartitioning(store_id, 200), final HashAggregate. The difference isn't in the plan's shape — it's in what that Exchange actually moves, and that's what part 3 of this lesson is going to measure with the Spark UI. Notice also the sum(revenue) values: 9699999.999998083 instead of a clean 9700000.0 — the same floating-point phenomenon data-engineering-foundations-guide already documented, now visible in a sum of ten million multiplication-and-addition operations instead of forty. Rounded to two decimals (F.round(F.sum("revenue"), 2), the pattern you already used in module 3), this number is exactly 9,700,000.00 — matching the figure lesson 4 verified with assert over the raw data, before Spark touched anything.

Worked example, part 2: the first join that genuinely shuffles

Lesson 3 left a question open: when does a join stop resolving as a BroadcastHashJoin and start paying a real shuffle? The full answer — Spark's decision criterion, with the configurable threshold — is module 5's topic. But this lesson shows you the first concrete case: when both sides of a JOIN are too large to be copied whole to every partition, Spark has no alternative — it has to shuffle both sides by the JOIN key.

print("=== join of orders_at_scale_df against itself, by product_id ===")
left = orders_at_scale_df.select("order_id", "product_id", "store_id")
right = orders_at_scale_df.select("product_id").distinct()
joined = left.join(right, "product_id")
joined.explain()

What to expect (executed in this run):

=== join of orders_at_scale_df against itself, by product_id ===
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Project [product_id#3, order_id#0, store_id#2]
   +- SortMergeJoin [product_id#3], [product_id#11], Inner
      :- Sort [product_id#3 ASC NULLS FIRST], false, 0
      :  +- Exchange hashpartitioning(product_id#3, 200), ENSURE_REQUIREMENTS, [plan_id=39]
      :     +- Filter isnotnull(product_id#3)
      :        +- FileScan csv [order_id#0,store_id#2,product_id#3] Batched: false, DataFilters: [isnotnull(product_id#3)], Format: CSV, ...
      +- Sort [product_id#11 ASC NULLS FIRST], false, 0
         +- HashAggregate(keys=[product_id#11], functions=[])
            +- Exchange hashpartitioning(product_id#11, 200), ENSURE_REQUIREMENTS, [plan_id=35]
               +- HashAggregate(keys=[product_id#11], functions=[])
                  +- Filter isnotnull(product_id#11)
                     +- FileScan csv [product_id#11] Batched: false, DataFilters: [IsNotNull(product_id)], Format: CSV, ...

This plan has two Exchange nodes, one for each side of the JOIN — both with hashpartitioning(product_id, 200), the same key, to guarantee rows with the same product_id end up in the same partition on both sides, regardless of which file or original partition they came from. The strategy's name changed too: it's no longer BroadcastHashJoin — it's SortMergeJoin, the strategy Spark uses when neither side of the JOIN is small enough for broadcast. Over orders_at_scale_df (ten million rows on both sides), copying either side whole to every partition would be more expensive than the shuffle itself — so the optimizer picks the strategy that actually scales.

Diagram: two ways of resolving a JOIN, depending on size

flowchart TD
    A{"Does the small side\nfit in memory, to be copied\nto every partition?"}
    A -->|"Yes -- e.g. dim_store,\n3 rows -- lesson 3"| B["BroadcastHashJoin\nno shuffle Exchange,\njust BroadcastExchange"]
    A -->|"No -- both sides\nlarge, this lesson"| C["SortMergeJoin\nwith an Exchange on BOTH sides,\nreorganized by the JOIN key"]

    style B fill:#9c6,stroke:#333
    style C fill:#f96,stroke:#333

Worked example, part 3: reading the shuffle in the Spark UI, through its own REST API

Module 1 mentioned the Spark UI at localhost:4040; module 2 (lesson 2) showed that whole interface is powered by a public REST API, at http://localhost:4040/api/v1/. This lesson uses that same API, now over the /stages endpoint, to read precisely what this lesson's part 1 Exchange moved. The pattern: leave a SparkSession running in the background after running the query, and query its REST API while it's still alive.

# ui_shuffle_bg.py
import time
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DoubleType, TimestampType

spark = SparkSession.builder.appName("kiosko-spark").master("local[*]").getOrCreate()
scale_schema = StructType([
    StructField("order_id", StringType(), False),
    StructField("franchise_id", IntegerType(), False),
    StructField("store_id", StringType(), False),
    StructField("product_id", StringType(), False),
    StructField("quantity", IntegerType(), False),
    StructField("unit_price", DoubleType(), False),
    StructField("order_ts", TimestampType(), False),
])
orders_at_scale_df = spark.read.csv("kiosko_orders_at_scale.csv", schema=scale_schema, header=True, enforceSchema=False)
revenue_by_store = (
    orders_at_scale_df
    .withColumn("revenue", F.col("quantity") * F.col("unit_price"))
    .groupBy("store_id")
    .sum("revenue")
)
print("APP_ID", spark.sparkContext.applicationId)
revenue_by_store.collect()
time.sleep(45)  # keeps the SparkSession alive to query its REST API
spark.stop()

Run in the background, and queried with curl while it's still alive:

python3 ui_shuffle_bg.py &
curl -s "http://localhost:4040/api/v1/applications/<app-id>/stages"

What to expect (real JSON response from this run, summarized to each stage's relevant fields):

stageId=0  status=COMPLETE   numTasks=12  inputRecords=10000000  inputBytes=602074847
           shuffleWriteBytes=2736  shuffleWriteRecords=36  shuffleReadBytes=0

stageId=1  status=SKIPPED    numTasks=12  (alternate plan AQE decided not to use)

stageId=2  status=COMPLETE   numTasks=1   inputRecords=0  inputBytes=0
           shuffleWriteBytes=0  shuffleReadBytes=2736  shuffleReadRecords=36

This is the definitive proof that the physical plan's Exchange is a real event, not just an annotation: groupBy("store_id").sum("revenue") ran as two stages, split by the shuffle. Stage 0 read the 10,000,000 rows from the CSV (inputRecords=10000000, inputBytes=602074847 — matching the file's real size on disk) using its 12 partitions (numTasks=12, the same number lesson 2 already measured with .rdd.getNumPartitions()), calculated a partial sum of revenue by store_id within each partition, and wrote those partial results to the shuffle: shuffleWriteBytes=2736, shuffleWriteRecords=36. That 36 isn't arbitrary: it's 12 partitions × 3 stores — each partition produces, at most, one partial sum for each of the three stores it can contain. Stage 2 read exactly those same 2736 bytes and 36 records (shuffleReadBytes=2736, shuffleReadRecords=36 — the write and read numbers match, byte for byte) and combined them into the three-row final result you already saw in part 1.

Diagram: two stages, separated by the Exchange

flowchart LR
    subgraph Stage0["Stage 0 -- 12 tasks in parallel"]
        A["FileScan:\n10,000,000 rows read\n602,074,847 bytes"]
        B["Partial HashAggregate:\nup to 3 sums per partition"]
        C["Shuffle write:\n2,736 bytes, 36 records"]
        A --> B --> C
    end

    C -.->|"Exchange\nhashpartitioning(store_id, 200)"| D

    subgraph Stage2["Stage 2 -- 1 task (AQE coalesced)"]
        D["Shuffle read:\n2,736 bytes, 36 records"]
        E["Final HashAggregate:\n3 rows, one per store"]
        D --> E
    end

    style C fill:#f96,stroke:#333
    style D fill:#f96,stroke:#333

Notice something this module's lesson 7 is going to explain in depth: stage 2 ran with a single task (numTasks=1), not 200 — the number spark.sql.shuffle.partitions indicates. That's Adaptive Query Execution acting in real time: since the combined result (three rows, 2736 bytes) is tiny, AQE decided it wasn't worth spreading it across 200 empty or near-empty partitions, and coalesced them into one. This lesson only flags the phenomenon with evidence; lesson 7 explains it end to end.

Common mistakes

Looking for the Exchange only in .explain() and never confirming with the Spark UI (or the other way around). What happens: someone reads .explain(), sees Exchange hashpartitioning(...), and assumes they understood the shuffle's full cost without ever checking how many bytes or records actually moved. Why it happens: .explain() confirms a shuffle is going to happen, but doesn't say how much it's going to cost in practice — that magnitude only shows up after execution, in the Spark UI or its REST API. How to spot it: if you've never queried /api/v1/applications/.../stages (or the Spark UI's "Stages" tab) for a query you already know shuffles, you're missing half the evidence — you know that there's a cost, but not how much. How to fix it: use .explain() to predict the plan's structure before running it, and the Spark UI (or its REST API, as in this lesson) to confirm the real magnitude after running it — the two sources complement each other, neither replaces the other.

Reading numTasks=1 in stage 2 as "the shuffle was pointless." What happens: someone sees the final stage ran with a single task and concludes all the work of partitioning into 200 hash groups was wasted, since it all ended up in one task anyway. Why it happens: it seems contradictory to talk about "partitioning by hash into 200 groups" and end up with a single task. How to spot it: if you conclude the Exchange hashpartitioning(store_id, 200) was "unnecessary" because the final result used a single task, you're missing a distinction between two things: the Exchange did move the data correctly grouped by store_id — that work was necessary for the final sum to be correct; what AQE optimized afterward was how many parallel tasks were needed to process that already-grouped result, given it was small enough that a single task sufficed. How to fix it: always separate "did the shuffle reorganize the data correctly?" (yes, whenever an Exchange shows up with the correct key) from "how many partitions does the result have after the shuffle?" (a separate decision, affected by AQE, that lesson 7 develops).

Assuming shuffleWriteRecords=36 means only 36 rows of the original dataset moved. What happens: someone reads shuffleWriteRecords=36 and concludes Spark only had to move thirty-six of the ten million original rows during the shuffle. Why it happens: 36 is such a small number compared to 10,000,000 that it seems to describe how many original rows crossed the network. How to spot it: if you can't explain why 36 is exactly 12 × 3 (partitions × stores), you're missing the key piece: those 36 records aren't original orders_at_scale_df rows — they're partial sums, already aggregated within each partition by the first HashAggregate, before the Exchange moves them. The ten million original rows really did get read and processed (that work is in stage 0's inputRecords=10000000), but the shuffle itself only moves that first pass's already-reduced result, not the whole dataset — that prior reduction (the partial per-partition aggregation) is precisely what makes a groupBy far cheaper than moving all ten million full rows across the network.

Exercises

Exercise 1 — Confirm .distinct() over store_id, at scale, also produces an Exchange with the same structural figures. Run orders_at_scale_df.select("store_id").distinct().explain() over the dataset at scale and compare the plan against this lesson's groupBy.

distinct_at_scale = orders_at_scale_df.select("store_id").distinct()
distinct_at_scale.explain()
distinct_at_scale.show()

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[store_id#2], functions=[])
   +- Exchange hashpartitioning(store_id#2, 200), ENSURE_REQUIREMENTS, [plan_id=...]
      +- HashAggregate(keys=[store_id#2], functions=[])
         +- FileScan csv [store_id#2] Batched: false, DataFilters: [], Format: CSV, ...

+--------+
|store_id|
+--------+
|     S02|
|     S01|
|     S03|
+--------+

Exercise 2 — Calculate how many partitions this lesson's shuffle would have if AQE were disabled. Without running anything yet, using the value of spark.sql.shuffle.partitions you already know (200, the default), predict how many partitions groupBy("store_id").sum("revenue")'s final stage would have if spark.sql.adaptive.enabled were False. Verify your prediction.

spark.conf.set("spark.sql.adaptive.enabled", False)
revenue_no_aqe = orders_at_scale_df.withColumn("revenue", F.col("quantity") * F.col("unit_price")).groupBy("store_id").sum("revenue")
revenue_no_aqe.collect()
print(f"post-shuffle partitions, without AQE = {revenue_no_aqe.rdd.getNumPartitions()}")

post-shuffle partitions, without AQE = 200

Exercise 3 — Explain, without code, why this lesson's self-join had to shuffle both sides, not just one. In 2-3 sentences, explain why this lesson's SortMergeJoin plan includes an Exchange on each of the two sides of the JOIN, instead of just one.

Summary and next step

This lesson measured, with real, measurable evidence, exactly what lesson 3 only showed in outline: a groupBy("store_id").sum("revenue") over ten million rows produces the same Exchange hashpartitioning(store_id, 200) as over forty, but this time the shuffle moves 2,736 bytes and 36 records — confirmed with the Spark UI's own REST API — splitting execution into two real stages: one that reads and partially aggregates (stage 0, 12 tasks, 10,000,000 input rows), and another that combines the result (stage 2, coalesced by AQE into a single task). You also saw, for the first time, a join that genuinely pays the shuffle's full cost — SortMergeJoin, with an Exchange on both sides — when neither table is small enough for broadcast.

Before moving on you should be able to: read a .explain() plan and predict how many stages execution is going to split into; explain why shuffleWriteRecords=36 doesn't mean "only 36 rows moved"; and tell a BroadcastHashJoin apart from a SortMergeJoin just by looking at whether the plan has one or two Exchange nodes.

Lesson 6 takes the same orders_at_scale_df and answers a practical question: if a DataFrame's partition count isn't the one you need, how do you change it — and does that operation itself trigger another shuffle?

Resources

• Apache Spark — Monitoring and Instrumentation (the Spark UI's full REST API, including the /applications/[app-id]/stages endpoint used in this lesson, with the complete list of per-stage metrics: shuffleWriteBytes, shuffleReadBytes, inputRecords, among others). spark.apache.org/docs/latest/monitoring.html.
• Apache Spark — SQL Performance Tuning, JOIN strategy sections (SortMergeJoin versus BroadcastHashJoin — the full reference module 5 develops). spark.apache.org/docs/latest/sql-performance-tuning.html.
• This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — this lesson's exact specification: .explain() showing Exchange, verified over kiosko_orders_at_scale.