Module 4: Partitions And The Cost Of Shuffle

Project: Kiosko at scale, partitioned

Description

This project closes the module by pulling the previous six lessons together into a single script: you read kiosko_orders_at_scale (generated in lesson 4) with an explicit schema, measure its partitions, calculate revenue and group by store — seeing the shuffle's Exchange in .explain() — compare repartition() against coalesce(), and close by measuring Adaptive Query Execution's real effect on the post-shuffle partition count. All of it, verified with assert, over the full ten million rows.

Connection to the module. This project introduces no new concept — it's the final integration of lessons 2 through 7, in the same order you built them, with a single script run end to end over the complete synthetic dataset.

An analogy: the full banquet audit

Pick back up lesson 4's analogy: the recipe multiplied for a banquet of two hundred fifty thousand people. This project is that banquet's full audit, already served: how many tables (partitions) got set up to serve it, what happened when someone asked to reorganize the food by a new criterion (the shuffle), what happened when you decided to use more tables or merge the existing ones (repartition() versus coalesce()), and whether the final number of tables made sense for the real size of what got served, or whether an automatic system (AQE) had to correct it on the fly. An audit that doesn't check every one of these points, with evidence, isn't a full audit — it's just someone's word that "it went fine."

The material: everything this module built, in a single flow

You need: kiosko_orders_at_scale.csv, the 10,000,000-row file generated in lesson 4 (generate_orders_at_scale(250_000), written with step1_generate.py), in the same folder where you're going to run this script. If you haven't generated it yet, lesson 4 has the full code to produce it.

The reference solution, verified

Parts 1 and 2 — Open the session, read the dataset at scale

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

print("=== Kiosko at scale, partitioned: module 4 final delivery ===\n")

print("Part 1 -- opening the SparkSession")
spark = (
    SparkSession.builder
    .appName("kiosko-spark")
    .master("local[*]")
    .getOrCreate()
)
print(f"Spark version: {spark.version}\n")

print("Part 2 (L2, L4) -- reading kiosko_orders_at_scale.csv, generated in lesson 4")
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,
)
row_count = orders_at_scale_df.count()
num_partitions = orders_at_scale_df.rdd.getNumPartitions()
print(f"orders_at_scale_df.count() = {row_count}")
print(f"orders_at_scale_df.rdd.getNumPartitions() = {num_partitions}")
assert row_count == 10_000_000
print("Verification: 10,000,000 rows read -> OK\n")

Part 3 — revenue, the groupBy, and the Exchange in .explain()

print("Part 3 (L3, L5) -- revenue calculated, groupBy with shuffle, verified with assert")
fact_orders_at_scale_df = orders_at_scale_df.withColumn(
    "revenue", F.col("quantity") * F.col("unit_price")
)
revenue_by_store = (
    fact_orders_at_scale_df
    .groupBy("store_id")
    .agg(F.round(F.sum("revenue"), 2).alias("total_revenue"))
    .orderBy("store_id")
)
print("revenue_by_store's physical plan (notice the Exchange nodes -- the shuffle):")
revenue_by_store.explain()
revenue_by_store.show()

by_store = {r["store_id"]: r["total_revenue"] for r in revenue_by_store.collect()}
total_revenue = round(sum(by_store.values()), 2)
print(f"by_store = {by_store}")
print(f"total_revenue = {total_revenue}")
assert by_store == {"S01": 9575000.0, "S02": 9700000.0, "S03": 7262500.0}
assert total_revenue == 26_537_500.00
print("Verification: 26,537,500.00 in total revenue, exact breakdown by store -> OK\n")

Part 4 — repartition() vs coalesce()

print("Part 4 (L6) -- repartition() vs coalesce(), compared by partition count")
print(f"orders_at_scale_df partitions (base)        = {orders_at_scale_df.rdd.getNumPartitions()}")
repartitioned_50 = orders_at_scale_df.repartition(50)
print(f"repartition(50) partitions                  = {repartitioned_50.rdd.getNumPartitions()}")
coalesced_2 = orders_at_scale_df.coalesce(2)
print(f"coalesce(2) partitions                      = {coalesced_2.rdd.getNumPartitions()}")
coalesced_over = orders_at_scale_df.coalesce(1000)
print(f"coalesce(1000) partitions (can't grow)      = {coalesced_over.rdd.getNumPartitions()}")
assert repartitioned_50.rdd.getNumPartitions() == 50
assert coalesced_2.rdd.getNumPartitions() == 2
assert coalesced_over.rdd.getNumPartitions() == num_partitions
print("Verification: repartition freely grows and shrinks, coalesce only shrinks -> OK\n")

Part 5 — spark.sql.shuffle.partitions and Adaptive Query Execution

print("Part 5 (L7) -- spark.sql.shuffle.partitions and Adaptive Query Execution")
print(f"spark.sql.shuffle.partitions = {spark.conf.get('spark.sql.shuffle.partitions')}")
print(f"spark.sql.adaptive.enabled (default) = {spark.conf.get('spark.sql.adaptive.enabled')}")

def build_groupby():
    return (
        orders_at_scale_df
        .withColumn("revenue", F.col("quantity") * F.col("unit_price"))
        .groupBy("store_id")
        .agg(F.sum("revenue").alias("total_revenue"))
    )

spark.conf.set("spark.sql.adaptive.enabled", False)
q_no_aqe = build_groupby()
q_no_aqe.collect()
partitions_no_aqe = q_no_aqe.rdd.getNumPartitions()
print(f"post-shuffle partitions, AQE off = {partitions_no_aqe}")

spark.conf.set("spark.sql.adaptive.enabled", True)
q_aqe = build_groupby()
q_aqe.collect()
partitions_with_aqe = q_aqe.rdd.getNumPartitions()
print(f"post-shuffle partitions, AQE on  = {partitions_with_aqe}")

assert partitions_no_aqe == 200
assert partitions_with_aqe == 1
print("Verification: 200 without AQE (fixed value), 1 with AQE (coalesced by real size) -> OK\n")

Part 6 — Final summary

print("Part 6 -- module 4 final summary")
print(f"Rows processed: {row_count:,}")
print(f"Verified total revenue: {total_revenue:,}")
print(f"Read partitions (this machine's cores): {num_partitions}")

spark.stop()
print("=== spark.stop() -- module 4 closed, partitions and shuffle verified over 10M rows ===")

What to expect. Running python3 kiosko_orders_at_scale_partitioned.py in full (all six parts together), the output is exactly this (executed in this run, PySpark 4.2.0, on a machine with 12 logical cores):

=== Kiosko at scale, partitioned: module 4 final delivery ===

Part 1 -- opening the SparkSession
Spark version: 4.2.0

Part 2 (L2, L4) -- reading kiosko_orders_at_scale.csv, generated in lesson 4
orders_at_scale_df.count() = 10000000
orders_at_scale_df.rdd.getNumPartitions() = 12
Verification: 10,000,000 rows read -> OK

Part 3 (L3, L5) -- revenue calculated, groupBy with shuffle, verified with assert
revenue_by_store's physical plan (notice the Exchange nodes -- the shuffle):
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Sort [store_id#2 ASC NULLS FIRST], true, 0
   +- Exchange rangepartitioning(store_id#2 ASC NULLS FIRST, 200), ENSURE_REQUIREMENTS, [plan_id=54]
      +- HashAggregate(keys=[store_id#2], functions=[sum(revenue#26)])
         +- Exchange hashpartitioning(store_id#2, 200), ENSURE_REQUIREMENTS, [plan_id=51]
            +- HashAggregate(keys=[store_id#2], functions=[partial_sum(revenue#26)])
               +- Project [store_id#2, (cast(quantity#4 as double) * unit_price#5) AS revenue#26]
                  +- FileScan csv [store_id#2,quantity#4,unit_price#5] Batched: false, ...

+--------+-------------+
|store_id|total_revenue|
+--------+-------------+
|     S01|    9575000.0|
|     S02|    9700000.0|
|     S03|    7262500.0|
+--------+-------------+

by_store = {'S01': 9575000.0, 'S02': 9700000.0, 'S03': 7262500.0}
total_revenue = 26537500.0
Verification: 26,537,500.00 in total revenue, exact breakdown by store -> OK

Part 4 (L6) -- repartition() vs coalesce(), compared by partition count
orders_at_scale_df partitions (base)        = 12
repartition(50) partitions                  = 50
coalesce(2) partitions                      = 2
coalesce(1000) partitions (can't grow)      = 12
Verification: repartition freely grows and shrinks, coalesce only shrinks -> OK

Part 5 (L7) -- spark.sql.shuffle.partitions and Adaptive Query Execution
spark.sql.shuffle.partitions = 200
spark.sql.adaptive.enabled (default) = true
post-shuffle partitions, AQE off = 200
post-shuffle partitions, AQE on  = 1
Verification: 200 without AQE (fixed value), 1 with AQE (coalesced by real size) -> OK

Part 6 -- module 4 final summary
Rows processed: 10,000,000
Verified total revenue: 26537500.0
Read partitions (this machine's cores): 12
=== spark.stop() -- module 4 closed, partitions and shuffle verified over 10M rows ===

Six parts, six checks, and the same 26,537,500.00 lesson 4 already confirmed with assert over the raw data, now reproduced again by Spark, with a real .groupBy() that went through an Exchange. Notice a new detail in Part 3's plan you didn't see in previous lessons in this exact form: two Exchange nodes show up, not one. The first, Exchange hashpartitioning(store_id#2, 200), is the groupBy's shuffle you already know from lessons 3 and 5. The second, Exchange rangepartitioning(store_id#2 ASC NULLS FIRST, 200), is the extra shuffle added by the .orderBy("store_id") at the end of the chain — the same phenomenon you confirmed in lesson 3's exercise 1: globally sorting by a key also requires reorganizing data across partitions, even though it's a different operation from grouping.

Diagram: the six parts, closing out the whole module

flowchart TD
    A["Part 1: SparkSession opened"] --> B
    B["Part 2 (L2, L4): kiosko_orders_at_scale.csv\nread -- 10,000,000 rows, 12 partitions"] --> C
    C["Part 3 (L3, L5): revenue calculated,\ngroupBy with Exchange, 26,537,500.00 verified"] --> D
    D["Part 4 (L6): repartition(50)=50,\ncoalesce(2)=2, coalesce(1000)=12 (no growth)"] --> E
    E["Part 5 (L7): AQE off=200 partitions,\nAQE on=1 post-shuffle partition"] --> F["Module 4 closed:\npartitions and shuffle\nverified over 10M rows"]

Closing out this module's checklist, piece by piece

Module pieceStatus at the end of this project
What a partition is, measured over real dataResolved — lesson 2, 7 partitions over orders_df, 12 over the dataset at scale
Why groupBy/join/distinct trigger a shuffleResolved — lesson 3, evidence with .explain() and Spark's official quote
kiosko_orders_at_scale generated, deterministic, verifiedResolved — lesson 4, 10,000,000 rows, 26,537,500.00, no random
The shuffle read in .explain() and in the Spark UI, at real scaleResolved — lesson 5, 2,736 bytes / 36 records confirmed by the REST API
repartition() vs coalesce()Resolved — lesson 6, Exchange in one, none in the other
spark.sql.shuffle.partitions and AQEResolved — lesson 7, 200 without AQE, 1 with AQE, measured
Broadcast join vs sort-merge join, window functionsPending — module 5
Catalyst, .explain() in its phases, cachingPending — module 6
Partitioned Parquet at scale, vectorized UDFsPending — module 7
Distributed capstone, full decision treePending — module 8

Four of eight modules resolved — halfway through this guide. And with this module closed, you have something no previous module gave you: real evidence, measured in bytes and partitions, of why distributing a calculation costs something — the foundation module 5 is going to build joins and windows at scale on top of.

Common mistakes

Running this project without having generated kiosko_orders_at_scale.csv first. What happens: someone jumps straight into this project without having run lesson 4's step1_generate.py, and the script fails with a file-not-found error in Part 2. Why it happens: it's tempting to treat the closing project as an independent starting point, instead of as the integration of what you already built. How to spot it: if spark.read.csv("kiosko_orders_at_scale.csv", ...) fails with AnalysisException: Path does not exist, you're missing lesson 4's file in the same folder where you run this script. How to fix it: this project, like module 3's, reuses artifacts generated in earlier lessons of the same module — run lesson 4's step1_generate.py first (or confirm the file already exists in your working folder) before running this full script.

Turning in the project without Part 3's asserts, trusting only .show()'s eyeball check. What happens: someone runs the six parts, sees the numbers "look right" in revenue_by_store.show()'s output, and considers the project done without checking whether Part 3's asserts actually passed. Why it happens: 9575000.0, 9700000.0, and 7262500.0 on screen look correct at a glance, and formally verifying them seems like an extra step over a result that already looks fine. How to spot it: if your final delivery didn't run Part 3's asserts all the way through with no AssertionError, you have no real guarantee the breakdown by store is exact — the same trap module 3 already warned about with fact_orders. How to fix it: this project's asserts — over the row count, the breakdown by store, the partition count in every comparison — aren't decorative: they're the proof that the whole script, not just a fragment, produces the correct result end to end.

Reading post-shuffle partitions, AQE on = 1 as the expected behavior for any query. What happens: someone memorizes "with AQE, post-shuffle partitions always end up at 1" from this single result, without remembering that 1 depends on the specific size of this groupBy's result (three rows). Why it happens: seeing a concrete number repeated across several lessons in this module (lesson 5, lesson 7, and now this project) can feel like a universal rule, instead of a specific consequence of the result's size. How to spot it: if you predict 1 post-shuffle partition for any groupBy you see from here on in this guide, without considering how many rows that particular aggregation produces, you're missing the nuance lesson 7's exercise 2 already explored (with franchise_id, 250,000 distinct values, the result is much bigger). How to fix it: always remember AQE decides the post-shuffle partition count based on the result's real size, measured at runtime — there's no fixed number that applies to every case, and that's exactly AQE's advantage over spark.sql.shuffle.partitions's static value.

Exercises

Exercise 1 — Extend the project with a Part 7: the partition count by franchise_id. Add a seventh part that calculates orders_at_scale_df.groupBy("franchise_id").count(), measures how many partitions the post-shuffle result has with AQE active, and confirms with assert that the total count of distinct franchises is 250,000.

See solution
print("Part 7 -- grouping by franchise_id, 250,000 distinct values")
franchise_counts = orders_at_scale_df.groupBy("franchise_id").count()
franchise_counts.collect()
num_franchises = franchise_counts.count()
result_partitions = franchise_counts.rdd.getNumPartitions()
print(f"distinct num_franchises = {num_franchises}")
print(f"post-shuffle partitions (franchise_id, AQE on) = {result_partitions}")
assert num_franchises == 250_000
print("Verification: 250,000 distinct franchises -> OK")

Expected output (executed in this run):

Part 7 -- grouping by franchise_id, 250,000 distinct values
distinct num_franchises = 250000
post-shuffle partitions (franchise_id, AQE on) = 1
Verification: 250,000 distinct franchises -> OK

Surprise, if you expected a number greater than 1: it's still a single partition, just like Part 3's groupBy("store_id"), even though this result has 250,000 rows instead of 3. The reason, developed in more detail in lesson 7's exercise 2: AQE coalesces partitions based on how many bytes the real result weighs, not how many rows it has — and 250,000 rows of two small columns (franchise_id, count) still weigh just a few megabytes, well below the 64 MB per-partition target (spark.sql.adaptive.advisoryPartitionSizeInBytes). "More rows" isn't the same as "more bytes," and AQE decides on bytes.

Exercise 2 — Confirm the parts' order matters: move Part 4 before Part 3, and explain what breaks. Without running anything, explain in 2-3 sentences whether reordering the script — calculating repartition()/coalesce() (Part 4) before calculating revenue and verifying it (Part 3) — would change the asserts' final result.

See solution

It wouldn't change the asserts' result, because Part 4 operates on orders_at_scale_df (the DataFrame read directly from the CSV, with no revenue calculated), not on fact_orders_at_scale_df or on revenue_by_store — the two parts are, in this script, independent of each other in terms of what data they consume. However, reordering them would break the project's pedagogical narrative: Part 4 makes more sense read after having seen, in Part 3, a real shuffle in action (groupBy with Exchange) — without that foundation, "compare partitions before and after repartition()" loses the context of why that partition change matters in the first place. The script's order, just like module 3's project, isn't arbitrary: it follows the same order you built the knowledge in throughout the module.

Exercise 3 — Explain, without looking at the guide's design, what this pipeline is missing to become module 8's full capstone. In a 4-6 sentence paragraph, describe what transformations, checks, or decisions kiosko_orders_at_scale_partitioned.py is missing to become module 8's full distributed pipeline.

See solution

Today, this script reads kiosko_orders_at_scale and calculates revenue directly, with no JOIN against dim_store or dim_product at all — those dimensions never show up — and with no window function for rankings or running totals by store. It's also missing any conscious decision between BroadcastHashJoin and SortMergeJoin (there's no JOIN here at all, just a direct groupBy over an already-calculated table) and window functions for top product by store and by day (module 5); reading the plan with Catalyst through its full phases and explicit .cache() decisions over DataFrames reused several times (module 6); and partitioned writes with partitionBy("store_id") plus a vectorized pandas_udf, instead of the direct revenue calculation used today (module 7). The module 8 capstone assembles all those pieces — broadcast joins, windows, judgment-driven caching, partitioned Parquet, vectorized UDFs — into a single pipeline that runs end to end over the ten million synthetic rows, and closes with the full "do I need Spark?" decision tree, applied both to real forty-row Kiosko (the answer is still no) and to a much larger hypothetical Kiosko.

Summary and next step: the end of module 4

With this mini-project you close out module 4 in full. You read kiosko_orders_at_scale — this guide's first synthetic data, 10,000,000 rows generated deterministically in lesson 4 — and confirmed Spark assigns it 12 partitions on this machine. You calculated revenue and grouped by store, seeing the shuffle's Exchange node written into .explain(), and verified with three asserts that the result — 26,537,500.00 in total revenue, with the same proportional breakdown by store as always — matches exactly the arithmetic lesson 4 already confirmed over the raw data. You compared repartition() against coalesce() by resulting partition count, and closed by measuring Adaptive Query Execution's real effect: 200 post-shuffle partitions without AQE, versus just 1 with AQE active, over the same groupBy.

You took this module's central step: you stopped treating "partition" and "shuffle" as abstract words, and measured them with concrete numbers — bytes, records, partition count — over a data volume that actually demands something of Spark. What you still haven't done is decide, with judgment, between two JOIN strategies when both are possible, or use a window function over data at scale.

Where you're headed. Module 5 — joins-and-window-functions-at-scale — takes exactly kiosko_orders_at_scale and answers the question this module deliberately left open in lessons 3 and 5: when does Spark choose BroadcastHashJoin (no shuffle, copying the small table to every partition) and when does it choose SortMergeJoin (with a shuffle, when neither table fits for broadcast) — with spark.sql.autoBroadcastJoinThreshold's exact threshold, forced and compared with evidence. Then, window functions — Window.partitionBy().orderBy() — for cumulative revenue by store and a top-product ranking, the same class of question data-modeling-for-analytics-guide already solved with arrays in DuckDB, now with a distributed engine's native window.

Resources

  • Apache Spark — SQL Performance Tuning (Catalyst, shuffle partitions, Adaptive Query Execution — the central reference for this entire module, and in particular for Parts 3 and 5 of this project). spark.apache.org/docs/latest/sql-performance-tuning.html.
  • Apache Spark — RDD Programming Guide, "Shuffle operations" section (the official shuffle definition backing this whole project's Part 3). spark.apache.org/docs/latest/rdd-programming-guide.html.
  • This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — module 4's full objective and its place in the eight-module plan.
  • data-modeling-for-analytics-guide DESIGN doc — the source of the windows-and-running-totals question module 5 picks back up with Spark's native API. src/guides/data-modeling-for-analytics-guide/DISENO.md