Back to Advanced Data Engineering
Q6Question 6 of 24

AQE's three major optimizations

Key Concepts
  • Post-shuffle partition coalescing merges many tiny shuffle partitions into fewer balanced ones so the next stage has appropriately sized tasks.
  • Skew join optimization detects oversized shuffle partitions and splits them into sub-partitions, replicating the other join side to eliminate straggler tasks.
  • Dynamic join strategy change replaces a planned shuffle join with a broadcast join when AQE discovers at runtime that one side is small enough.
  • All three optimizations key off real statistics gathered at shuffle boundaries and are on by default in modern Spark.
  • AQE is not a replacement for good data design; extreme single-key skew still needs manual salting.
  • These optimizations work together and can be verified by checking for AdaptiveSparkPlan nodes in the physical plan.

Definition

AQE coalesces tiny partitions into balanced ones, splits oversized skewed partitions to kill stragglers, and switches a shuffle join to a broadcast join when a side turns out small.
  • Post-shuffle partition coalescing – merges many small shuffle partitions into fewer, larger partitions so the next stage runs a sensible number of tasks instead of hundreds of tiny ones.
  • Skew join optimization – detects shuffle partitions that are abnormally large and splits them into sub-partitions, replicating the matching side, so a single hot key no longer produces one straggler task.
  • Dynamic join strategy change – replaces a planned shuffle join (like sort-merge) with a broadcast join when AQE sees at runtime that one side is small enough to broadcast.
  • All three are driven by real statistics at shuffle boundaries and are on by default in Spark 3.x.

Picture clerks at the National Bank of India dividing up branch files after a big re-sort. Coalescing is when the supervisor sees many clerks each holding a nearly empty folder and combines them so fewer clerks carry fuller, balanced folders. Skew handling is when one clerk is buried under a mountain of files for one huge branch while others sit idle, so the supervisor splits that mountain across several clerks. Dynamic join change is when two teams were going to do a slow full re-sort to combine their lists, but the supervisor notices one team's list is tiny and photocopies it to everyone, skipping the re-sort entirely.

Understanding the Concepts

What is Adaptive Query Execution?

Adaptive Query Execution (AQE) is a Spark SQL feature that re-optimizes the query plan at runtime using real statistics collected at shuffle boundaries. Instead of committing to a single static plan before any data is read, AQE adjusts the plan between stages based on the actual data it sees.

PYTHON
from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("AQEDemo") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

# Confirm AQE is active
print(spark.conf.get("spark.sql.adaptive.enabled"))  # true

When AQE is on, the physical plan wraps stages in AdaptiveSparkPlan nodes that can be re-planned.

How does post-shuffle partition coalescing work?

By default a shuffle produces a fixed number of partitions (often 200) regardless of how much data there actually is. After a filter or aggregation, many partitions end up nearly empty, so the next stage launches hundreds of tiny tasks with more scheduling overhead than useful work.

PYTHON
from pyspark.sql import functions as F

txn = spark.read.parquet("/lake/silver/transactions")

# After filtering, most of the 200 shuffle partitions are nearly empty
small_branch = txn.filter(F.col("branch_code") == "NBI-DEL-042") \
    .groupBy("account_id") \
    .agg(F.sum("amount").alias("total"))

small_branch.explain()
TEXT
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   *(2) HashAggregate(keys=[account_id])
   +- AQEShuffleRead coalesced
      +- ShuffleQueryStage
         +- Exchange hashpartitioning(account_id, 200)

The AQEShuffleRead coalesced line shows AQE merged the tiny partitions. You control this with spark.sql.adaptive.coalescePartitions.enabled and the advisory partition size target.

How does skew join optimization work?

Data skew happens when one key has vastly more records than others. After the shuffle, that key's partition is enormous and produces a straggler task.

PYTHON
# The mega-account ACC-9999 has 10 million transactions
# while all others have a few thousand
txn = spark.read.parquet("/lake/silver/transactions")
accounts = spark.read.parquet("/lake/silver/accounts")

joined = txn.join(accounts, "account_id")
joined.explain()
TEXT
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   *(5) SortMergeJoin [account_id], [account_id]
   :- AQEShuffleRead
   :  +- ShuffleQueryStage (skewedShuffle)

AQE detects the skewed partition when it exceeds both a factor of the median and an absolute byte threshold, then splits it into sub-partitions and replicates the matching rows from the other side.

PYTHON
# Key AQE skew settings
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256m")

How does dynamic join strategy change work?

Before the job runs, Catalyst must choose a join algorithm from size estimates. If it thinks both sides are large, it picks a shuffle join like sort-merge. But after an upstream filter, one side may actually be tiny.

PYTHON
# Before filtering, the branch table looks large enough for sort-merge
branches = spark.read.parquet("/lake/silver/branches")

# After filtering to one region, it shrinks below the broadcast threshold
region_branches = branches.filter(F.col("region") == "WEST")

# AQE sees the real post-filter size and switches to broadcast
result = txn.join(region_branches, "branch_code")
result.explain()
TEXT
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   *(2) BroadcastHashJoin [branch_code], [branch_code]
   :- *(2) FileScan parquet transactions
   +- BroadcastExchange
      +- *(1) Filter (region = WEST)
         +- *(1) FileScan parquet branches

AQE converted the planned sort-merge join to a broadcast join at runtime, avoiding the expensive shuffle of the large transaction table.

How do the three optimizations work together?

All three key off the same runtime statistics gathered at shuffle boundaries. In a single query, AQE might coalesce small partitions from an initial aggregation, detect skew in a subsequent join, and switch that join to broadcast, all without the developer touching the plan.

PYTHON
# A realistic pipeline that benefits from all three
result = (
    txn
    .filter(F.col("txn_date") >= "2026-10-01")
    .join(region_branches, "branch_code")       # dynamic broadcast
    .groupBy("account_id")
    .agg(F.sum("amount").alias("total_spend"))   # coalescing
)
result.write.format("delta").save("/lake/gold/account_spend")

What are the limitations of AQE?

AQE is powerful but not magic.

PYTHON
# Limitation 1: extreme single-key skew may still need manual salting
from pyspark.sql import functions as F

# If ACC-9999 has 100 million rows, even split sub-partitions may be too large
salted = txn.withColumn(
    "salted_key",
    F.concat(F.col("account_id"), F.lit("_"), (F.rand() * 10).cast("int"))
)

# Limitation 2: AQE only helps jobs that have shuffles
# A pure narrow pipeline (read -> filter -> write) has no shuffle boundary
# for AQE to act on

# Limitation 3: the broadcast threshold must be set sensibly
# Broadcasting a table that is too large blows executor memory
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50m")  # be careful

How do you verify AQE is active?

Check the physical plan for AdaptiveSparkPlan and look at the Spark UI's SQL tab for the adaptive plan details.

PYTHON
result.explain(True)
# Look for:
#   AdaptiveSparkPlan isFinalPlan=true
#   AQEShuffleRead coalesced
#   skewedShuffle
#   BroadcastHashJoin (when sort-merge was planned)
TEXT
# In the Spark UI SQL tab, the plan shows:
# - Number of partitions before and after coalescing
# - Skew-split partitions highlighted
# - Original vs final join strategy

What should a senior engineer know?

  • AQE removes most everyday tuning but complements good data design rather than replacing it.
  • Coalescing is almost always a free win on every aggregation.
  • Skew handling removes the most common cause of mysterious straggler tasks.
  • Dynamic join switching captures the frequent case where filtering shrinks a table below the broadcast limit.
  • For extreme skew, combine AQE with manual salting for the few keys it cannot fully absorb.
PYTHON
# Senior pattern: verify AQE improvements in the Spark UI
# Before AQE: 200 tasks, 190 finish in 2s, 1 straggler at 45 min
# After AQE:  200 tasks coalesced to 20, skew split into 5 sub-tasks,
#             all finish within 30s

What are common interview follow-up questions?

  • Does AQE work with streaming? AQE is for batch SQL plans; structured streaming has its own adaptive mechanisms.
  • Can you disable individual AQE features? Yes, each has its own sub-switch under the master spark.sql.adaptive.enabled flag.
  • Does AQE change the result? No, it only changes the execution strategy. The logical result is identical.

Visual Overview

DIAGRAM
  AQE's Three Optimizations
  ─────────────────────────

  1. Post-Shuffle Coalescing
  ┌───┬───┬───┬───┬───┬───┬───┬───┐
  │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ 8 │  8 tiny partitions
  └───┴───┴───┴───┴───┴───┴───┴───┘
           AQE merges ▼
  ┌───────────┬───────────┐
  │   1+2+3+4 │   5+6+7+8 │  2 balanced partitions
  └───────────┴───────────┘

  2. Skew Join Optimization
  Before:  [===]  [===]  [========================]  [===]
                          ▲ skewed partition
  After:   [===]  [===]  [======] [======] [======]  [===]
                         split + replicate other side

  3. Dynamic Join Strategy Change
  Planned:  Sort-Merge Join (both sides shuffled)
  Runtime:  one side shrinks after filter
  AQE:      switches to Broadcast Hash Join
            (small side broadcast, large side no shuffle)

Key Takeaway

Adaptive Query Execution delivers three runtime optimizations, all driven by real statistics collected at shuffle boundaries. Post-shuffle partition coalescing merges many tiny shuffle partitions into fewer balanced ones so the next stage runs appropriately sized tasks instead of hundreds of nearly empty ones. Skew join optimization detects oversized shuffle partitions and splits them into sub-partitions, replicating the other join side, so a single hot key no longer produces one straggler task that runs long after everything else finishes.

Dynamic join strategy change is the third optimization. Before the job runs, Catalyst picks a join algorithm from size estimates, but after an upstream filter one side may actually be tiny. At the shuffle boundary AQE sees the real size and converts a planned sort-merge join into a broadcast join, shipping the small side everywhere and skipping the expensive shuffle of the large side entirely.

At the National Bank of India these three together turned several fragile, skew-prone nightly joins into stable, fast jobs without rewriting a single query. Coalescing was a free win on every aggregation, skew splitting eliminated the straggler tasks we saw on month-end spikes from one mega-branch, and dynamic join switching repeatedly caught the case where filtering shrank a dimension table below the broadcast limit.

The caution is that AQE is not magic. Extreme single-key skew sometimes still needs manual salting, the broadcast threshold must be set sensibly for executor memory, and none of these help a job that has no shuffle in the first place. AQE is a powerful default that removes most everyday tuning, but it complements good data design rather than replacing it.

Remember one sentence: AQE coalesces tiny partitions, splits skewed ones, and switches to broadcast joins at runtime using real statistics, making most Spark jobs faster without changing a line of query code.

---

Share
Use ← → arrow keys to navigate