Module 4: Partitions And The Cost Of Shuffle

`repartition()` vs `coalesce()`

Description

Lesson 2 measured that orders_at_scale_df has 12 partitions — the core count of the machine that read it — and that number isn't always the one you need: sometimes more parallelism helps before an expensive operation, and sometimes less helps, so you don't end up writing hundreds of tiny files. Spark offers two ways to change the partition count, repartition() and coalesce(), and this lesson demonstrates, with .explain(), that they aren't interchangeable: one pays for a full shuffle, the other doesn't.

Connection to the module. This lesson directly applies everything lessons 2, 3, and 5 built: you already know what a partition is, you already know what triggers a shuffle, and now you're going to see that even the operation of "changing the partition count" may or may not be a shuffle, depending on which of the two tools you use.

An analogy: redistributing from scratch versus merging what you already have

Imagine you have twelve loaded trucks, each with a portion of the order, and you need to end up with exactly fifty trucks instead of twelve — more trucks, each with less cargo. The only way to pull this off is to unload all the contents of the twelve original trucks and redistribute them among the fifty new ones, box by box — a complete reordering, exactly like lesson 3's. That's repartition().

Now imagine the opposite case: you have twelve trucks and need to end up with only two, to simplify the final delivery. Here you don't need to unload anything and redistribute from scratch — it's enough to combine the cargo from several existing trucks into a few, without touching each box's contents or redistributing them by any new criterion. The first six trucks merge into truck 1, the other six into truck 2, done. That's coalesce(): it combines adjacent partitions, with no need for every executor to agree on a global distribution criterion — which is why, generally, it's cheaper.

Worked example, part 1: the numbers, compared

# repartition_vs_coalesce.py
from pyspark.sql import SparkSession
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,
)

baseline = orders_at_scale_df.rdd.getNumPartitions()
print(f"baseline (freshly read)          = {baseline}")

repartitioned_50 = orders_at_scale_df.repartition(50)
print(f"repartition(50)                 = {repartitioned_50.rdd.getNumPartitions()}")

coalesced_2 = orders_at_scale_df.coalesce(2)
print(f"coalesce(2)                     = {coalesced_2.rdd.getNumPartitions()}")

coalesced_over = orders_at_scale_df.coalesce(1000)
print(f"coalesce(1000) (can't grow)     = {coalesced_over.rdd.getNumPartitions()}")

spark.stop()

What to expect. Running python3 repartition_vs_coalesce.py, the output is exactly this (executed in this run, starting from orders_at_scale_df with 12 partitions):

baseline (freshly read)          = 12
repartition(50)                 = 50
coalesce(2)                     = 2
coalesce(1000) (can't grow)     = 12

Three results, three different behaviors. repartition(50) took the DataFrame from 12 to 50 partitions — more than it had, with no issue at all. coalesce(2) reduced it to 2. But coalesce(1000) — asking for more partitions than already exist — stayed at exactly 12, not growing by a single extra partition. This isn't a code bug: it's a deliberate limitation of coalesce(), and part 2 of this lesson explains exactly why.

Worked example, part 2: the real difference, in .explain()

print("=== explain of repartition(50) ===")
repartitioned_50.explain()

print("\n=== explain of coalesce(2) ===")
coalesced_2.explain()

What to expect (executed in this run):

=== explain of repartition(50) ===
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
      +- Exchange RoundRobinPartitioning(50), REPARTITION_BY_NUM, [plan_id=9]
         +- FileScan csv [...]
+- == Initial Plan ==
   Exchange RoundRobinPartitioning(50), REPARTITION_BY_NUM, [plan_id=9]
   +- FileScan csv [...]

=== explain of coalesce(2) ===
== Physical Plan ==
Coalesce 2
+- FileScan csv [...]

There's the difference, written into the physical plan. repartition(50) produces an Exchange RoundRobinPartitioning(50) node — an Exchange, the same kind of node you already identified in lessons 3 and 5 as the mark of a real shuffle. The strategy is called RoundRobinPartitioning because, unlike hashpartitioning (which groups by a column's value), it distributes rows evenly across the new partitions with no grouping criterion at all — each row goes, on a rotating turn, to the next available partition — with the sole goal of balancing load. coalesce(2), by contrast, produces a single-line plan: Coalesce 2 directly over the FileScan. There's no Exchange at all. coalesce() works by merging existing partitions within the same stage, with no need for data to cross the network between executors — which is why, generally, it's the cheaper option when what you need is to reduce the partition count.

Diagram: why coalesce() can't grow

flowchart TD
    subgraph Repartition["repartition(n) -- full shuffle"]
        A1["12 original\npartitions"] -->|"Exchange:\nevery row gets\nredistributed"| A2["n new,\nbalanced partitions"]
    end

    subgraph Coalesce["coalesce(n) -- no shuffle, just merges"]
        B1["12 original\npartitions"] -->|"merges adjacent\npartitions, no data\nmoves between executors"| B2["n partitions\n(only if n <= 12)"]
        B1 -.->|"coalesce(1000):\nno way to SPLIT\na partition without a shuffle"| B3["stays at 12\n-- can't grow"]
    end

    style A2 fill:#69c,stroke:#333
    style B2 fill:#9c6,stroke:#333
    style B3 fill:#c66,stroke:#333

The technical reason coalesce(1000) stays at 12 is direct: merging adjacent partitions (combining 0 and 1 into one, for example) doesn't require moving data between executors — each merged partition simply inherits the content of the ones that formed it — but splitting a partition into several smaller ones would require deciding, row by row, which of the new partitions each one goes to, and that's exactly a shuffle's job. Since coalesce() is deliberately designed to avoid the shuffle, it simply doesn't offer that capability — if you ask for more partitions than already exist, it stays at the maximum it can achieve without shuffling: the partition count you already had.

Going deeper: when to use each one

The choice between repartition() and coalesce() depends on the direction of the change and on how much the uniform distribution of data matters:

SituationToolWhy
You need more partitions (more parallelism before an expensive operation)repartition(n)coalesce() can't grow without shuffling; repartition() is the only option
You need fewer partitions, and don't mind losing load balance between themcoalesce(n)Avoids the shuffle — cheaper — but can leave partitions very uneven in size
You need fewer partitions, with balanced load between themrepartition(n)Pays the shuffle's cost in exchange for uniformly sized partitions
You're going to write the result to disk and want fewer output filescoalesce(n) (typically)Module 7 of this guide uses exactly this pattern before partitionBy(...).parquet(...), to avoid ending up with hundreds of tiny files per partition

It's worth noting the nuance in the third row: coalesce() merges adjacent partitions as they are, without redistributing their content — if your twelve original partitions have very uneven sizes (for example, if the source file wasn't split evenly), coalesce(2) can end up with one partition much heavier than the other, because it simply merged the ones that already existed with no rebalancing at all. repartition(), by going through an Exchange with RoundRobinPartitioning, does distribute rows evenly across the new partitions, regardless of how they were distributed before — that balance is exactly what you're paying for with the shuffle's cost.

Common mistakes

Using repartition() when coalesce() would have been enough, paying an unnecessary shuffle. What happens: someone needs to reduce the partition count from 1000 to 10, before writing an output file, and uses repartition(10) out of habit or because it's the name they remember better. Why it happens: repartition() is, in practice, the better-known name of the two operations, and it's easy not to think of the alternative when the goal is simply "fewer partitions." How to spot it: if your .explain() shows an Exchange right before an operation that only needed to reduce the partition count (not redistribute them in a balanced way), you probably paid for a shuffle coalesce() would have spared you. How to fix it: when the goal is reducing partitions and you don't care about losing exact balance between them — the most common case before writing output files — use coalesce(); save repartition() for when you need more partitions, or when you need the resulting partitions to end up evenly balanced regardless of how they were distributed before.

Trying to increase the partition count with coalesce(), and not understanding why it doesn't work. What happens: someone calls coalesce(200) on a DataFrame with 12 partitions, expecting to end up with 200, and is surprised to see the partition count didn't change. Why it happens: the name coalesce doesn't, by itself, suggest any directional limitation — it seems reasonable to expect it to accept any number, larger or smaller. How to spot it: if your coalesce(n) with n larger than the current partition count doesn't produce the number you asked for, it isn't a bug in your code — it's documented behavior: coalesce() can only reduce. How to fix it: if you need more partitions than you already have, the only correct tool is repartition(n), which does pay the shuffle cost needed to split existing partitions into more pieces.

Assuming repartition() with no numeric argument, using a column instead, produces the same partition count as spark.sql.shuffle.partitions. What happens: someone uses orders_at_scale_df.repartition("store_id") (repartition by a column's value, without specifying how many partitions), and expects a predictable partition count ahead of time, without checking it. Why it happens: repartition(n) with an explicit number is predictable by definition; it's easy to assume the column-name variant behaves the same way, with spark.sql.shuffle.partitions's value as the final number. How to spot it: over a DataFrame with very few distinct values in the repartitioning column (like store_id, with only three possible values), the real partition count with no explicit number can end up determined by the number of unique values, not by a fixed, predictable-without-checking number. How to fix it: when using repartition()'s column-name variant instead of an explicit number, always confirm the result with .rdd.getNumPartitions() instead of assuming a value — this is exactly the kind of behavior module 5 explores in more depth, when franchise_id (with 250,000 distinct values, very different from store_id) comes into play.

Exercises

Exercise 1 — Confirm repartition(12) (the same number it already had) still triggers a shuffle. Without running anything yet, predict whether orders_at_scale_df.repartition(12).explain() shows an Exchange, even though the input and output partition counts are the same (12). Verify your prediction.

See solution

Prediction: yes, it still shows an Exchange. repartition(n) doesn't compare the current partition count against n to decide whether a shuffle is needed or not — it always redistributes the rows again with RoundRobinPartitioning, regardless of whether the resulting number matches the one it already had.

repartitioned_12 = orders_at_scale_df.repartition(12)
print(f"repartition(12) partitions = {repartitioned_12.rdd.getNumPartitions()}")
repartitioned_12.explain()

Expected output:

repartition(12) partitions = 12
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
      +- Exchange RoundRobinPartitioning(12), REPARTITION_BY_NUM, [plan_id=...]
         +- FileScan csv [...]
+- == Initial Plan ==
   Exchange RoundRobinPartitioning(12), REPARTITION_BY_NUM, [plan_id=...]
   +- FileScan csv [...]

Confirmed: the output partition count (12) equals the input count, but the Exchange is still there — repartition() always pays the shuffle's cost, even when, in this particular case, the final number turns out to be the same as the starting one. This is exactly the opposite of what coalesce(12) would do on this same DataFrame: since the requested number (12) matches the current one, coalesce() wouldn't need to merge anything, and its plan would show no Exchange at all.

Exercise 2 — Measure how many rows land in each partition after coalesce(2), and compare them. Using lesson 2's mapPartitions pattern, confirm whether the two resulting partitions from coalesce(2) are similar in size or very different from each other.

See solution
from pyspark.sql import functions as F

counts = (
    coalesced_2
    .withColumn("pid", F.spark_partition_id())
    .groupBy("pid")
    .count()
    .orderBy("pid")
    .collect()
)
rows_per_partition = [r["count"] for r in counts]
print(f"Rows per partition after coalesce(2): {rows_per_partition}")
print(f"Sum: {sum(rows_per_partition)}")
assert sum(rows_per_partition) == 10_000_000

Expected output (executed in this run — the exact split depends on how the 12 original partitions were merged into 2 groups):

Rows per partition after coalesce(2): [5071306, 4928694]
Sum: 10000000

Confirmed: the two resulting partitions aren't exactly equal (5,071,306 versus 4,928,694, a difference of almost one hundred fifty thousand rows) — coalesce() merged groups of adjacent partitions as they were (the first six of the twelve originals into one, the last six into the other), with no rebalancing of content row by row. Compare this against the per-partition counts you already measured in lesson 2 ([854430, 846972, 846972, 846971, 842971, 832990, 832989, 832989, 832990, 832989, 832990, 763747]): the first coalesce(2) partition is exactly the sum of the first six (5,071,306), the second is the sum of the last six (4,928,694) — no row changed groups relative to where it already was, only the complete groups got merged. If this imbalance mattered for your case — for example, over a dataset where the source files are much more unevenly sized — repartition(2) would produce more even partitions, in exchange for paying the cost of a full Exchange.

Exercise 3 — Explain, without code, why RoundRobinPartitioning doesn't work for grouping by store_id. In 2-3 sentences, explain why repartition(50) — which uses RoundRobinPartitioning — wouldn't be a valid alternative for achieving what groupBy("store_id") achieves with hashpartitioning.

See solution

RoundRobinPartitioning distributes rows across the new partitions in a simple rotating order, without looking at any column's value — its only goal is to balance the row count per partition, not to group related rows together. After a repartition(50) with this strategy, S01, S02, and S03 rows would end up mixed across the fifty new partitions, exactly as they were before (or worse, more scattered). hashpartitioning(store_id, ...), by contrast, calculates the destination partition from store_id's value, guaranteeing every row with the same value ends up in the same partition — the exact property a groupBy needs to be able to correctly sum without comparing across partitions again.

Summary and next step

This lesson compared two ways of changing a DataFrame's partition count: repartition(n), which always pays for a full shuffle (Exchange RoundRobinPartitioning) but can grow or shrink the partition count and leaves them balanced; and coalesce(n), which can only reduce the partition count, merging existing ones with no Exchange at all, cheaper but with no balance guarantee. You verified with real numbers: repartition(50) took 12 partitions to 50; coalesce(2) reduced them to 2; and coalesce(1000) stayed at exactly 12, because it can't split partitions without shuffling.

Before moving on you should be able to: explain why coalesce() can't increase the partition count; read a .explain() and say, unambiguously, whether a repartitioning operation paid for a shuffle or not; and decide, for a concrete case (more parallelism before an expensive join, or fewer output files before writing to disk), which of the two tools to use.

Lesson 7 closes this module's loop with the question you already brushed against in lesson 5: spark.sql.shuffle.partitions is fixed at 200 by default — what happens when that number doesn't make sense for your data's real size, and how does Adaptive Query Execution fix it without you having to adjust it by hand?

Resources