How post-shuffle partition coalescing works

• Problem – spark.sql.shuffle.partitions is a single static number (default 200) that applies to every shuffle, regardless of how much data actually flows through it.
• Solution – AQE post-shuffle coalescing observes the actual size of each partition after the shuffle write and merges adjacent small partitions into fewer groups that approach the advisory target size.
• Effect – the next stage runs fewer, right-sized tasks instead of hundreds of nearly empty ones, cutting per-task scheduling overhead.
• Config – controlled by spark.sql.adaptive.coalescePartitions.enabled (on by default), tuned by spark.sql.adaptive.advisoryPartitionSizeInBytes and spark.sql.adaptive.coalescePartitions.minPartitionSize.

What is post-shuffle partition coalescing?

Post-shuffle partition coalescing is an Adaptive Query Execution (AQE) optimization that merges small shuffle partitions into fewer, larger ones at runtime. It addresses the fundamental problem that spark.sql.shuffle.partitions is a static number that applies to every shuffle in the job, regardless of how much data actually flows through each one.

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder \
    .appName("NBI-CoalesceDemo") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.shuffle.partitions", "200") \
    .getOrCreate()

# Read the bank's daily transactions
txns = spark.read.parquet("s3://nbi-data-lake/transactions/2024-01-15/")

Why is the static shuffle partition count a problem?

The default 200 shuffle partitions apply to every shuffle. For a large dataset, 200 might be too few. But after a heavy filter or aggregation, most of those 200 partitions end up nearly empty. Since Spark runs one task per partition, you launch 200 tasks that each process almost nothing, and the fixed per-task overhead dominates.

# Filter to a small branch: most of the 200 partitions will be nearly empty
small_branch = txns.filter(F.col("branch") == "Shimla")

# This groupBy creates 200 shuffle partitions, but Shimla has very few rows
branch_total = small_branch.groupBy("accountId") \
    .agg(F.sum("amount").alias("total"))

# Without coalescing: 200 tasks, most doing nothing
# With coalescing: AQE merges them into ~5 well-sized tasks
branch_total.explain("formatted")

How does coalescing work internally?

Spark still writes the initial large number of shuffle partitions. But before the next stage reads them, AQE examines the actual byte size of each partition. It then groups contiguous partitions together so that each coalesced group approaches the advisory target size. The next stage reads these merged groups.

Coalescing only merges contiguous partitions to preserve any ordering guarantees from the shuffle. It does not reorder or reshuffle data.

# AQE coalescing in action:
# Step 1: Spark writes 200 shuffle partitions
# Step 2: AQE checks sizes: [1KB, 2KB, 0KB, 1KB, ..., 0KB]
# Step 3: AQE merges contiguous small partitions toward target (64MB)
# Step 4: Result: 5 coalesced partitions instead of 200

# You can see this in the Spark UI SQL tab:
# "CustomShuffleReader" or "AQEShuffleRead" with coalesced partitions
branch_total.write.mode("overwrite") \
    .parquet("s3://nbi-data-lake/shimla-totals/")

What configurations control coalescing?

Several settings work together to control the behavior.

| Configuration | Default | Purpose |
|---|---|---|
| spark.sql.adaptive.enabled | true | Master switch for all AQE optimizations |
| spark.sql.adaptive.coalescePartitions.enabled | true | Enable/disable partition coalescing specifically |
| spark.sql.adaptive.advisoryPartitionSizeInBytes | 64MB | Target size for each coalesced partition |
| spark.sql.adaptive.coalescePartitions.minPartitionSize | 1MB | Minimum size to prevent over-coalescing |
| spark.sql.adaptive.coalescePartitions.parallelismFirst | true | Prefer more partitions to maximize parallelism |

# Tune coalescing for the bank's workload
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionSize", "4m")

# Now coalesced partitions target 128MB and never go below 4MB

How is coalescing different from repartition and coalesce?

Post-shuffle coalescing is automatic and happens at runtime between stages. repartition() and coalesce() are explicit DataFrame operations you write in code.

| Aspect | AQE Coalescing | repartition(N) | coalesce(N) |
|---|---|---|---|
| When | Runtime, between stages | Code-time, triggers shuffle | Code-time, no shuffle |
| Direction | Only reduces partition count | Can increase or decrease | Only reduces |
| Trigger | Automatic by AQE | Explicit in code | Explicit in code |
| Shuffle | No extra shuffle | Full shuffle | No shuffle |
| Scope | Post-shuffle partitions only | Any DataFrame | Any DataFrame |

# AQE coalescing: automatic, no code change
result = txns.groupBy("branch").agg(F.sum("amount"))
# AQE merges small partitions after the groupBy shuffle

# Explicit coalesce: manual, for controlling output files
result.coalesce(10).write.parquet("s3://nbi-data-lake/output/")
# Forces 10 output files, no shuffle

# Explicit repartition: manual, triggers a full shuffle
result.repartition(50, "branch").write.parquet("s3://nbi-data-lake/output/")
# 50 partitions distributed by branch key

How does coalescing interact with skew handling?

Coalescing reduces the partition count by merging small partitions. AQE's skew handling goes the opposite direction by splitting oversized partitions. Together, they drive partitions toward a balanced target from both sides.

# Skewed data: one branch has 90% of transactions
skewed = txns.groupBy("branch").agg(F.sum("amount"))

# AQE does both:
# 1. Splits the oversized Mumbai partition (skew handling)
# 2. Merges the tiny partitions for small branches (coalescing)
# Result: balanced partition sizes across the stage

spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
skewed.explain("formatted")

When does coalescing NOT help?

Coalescing only reduces the task count. It does not increase it. So it fixes the too-many-small-partitions problem but not the too-few-huge-partitions problem. For that, you need a higher initial shuffle-partitions value or explicit repartitioning.

# Problem: 200 partitions but each is 2GB (too few, too large)
# Coalescing will NOT split them. You need more partitions.
spark.conf.set("spark.sql.shuffle.partitions", "2000")

# Now Spark starts with 2000 partitions.
# If some are too small after filtering, AQE coalesces them back down.
# This is the "set high and let AQE reduce" strategy.

What should a senior engineer know?

A senior engineer knows that coalescing is why you can set spark.sql.shuffle.partitions to a high value and let AQE adaptively reduce it per stage. They verify coalescing is working by checking the Spark UI SQL tab for "AQEShuffleRead" nodes that show the coalesced partition count.

# Senior practice: set partitions high, let AQE coalesce
spark.conf.set("spark.sql.shuffle.partitions", "1000")

# For large shuffles, 1000 partitions are used.
# For small post-filter shuffles, AQE coalesces to ~10.
# One setting handles both cases.

# Verify in the Spark UI SQL tab:
# Look for "coalesced" in the AQEShuffleRead node