Adaptive Query Execution

Description

Every physical plan in this guide, since module 2, started with the same line: AdaptiveSparkPlan isFinalPlan=false. It was never explained what that means — only mentioned, in passing, that the plan "could still change." This lesson closes that story: Adaptive Query Execution (AQE) is the mechanism, active by default since Spark 3.2, that revises the physical plan after a shuffle stage finishes running, using real statistics — bytes, rows — instead of the estimates lesson 3 saw with mode="cost". You already saw one of its two main capabilities without the full name, in module 4, lesson 7: coalescing empty shuffle partitions. This lesson repeats that evidence in a paragraph, and adds the second capability, which no earlier lesson has shown yet: switching a complete JOIN's strategy, from SortMergeJoin to BroadcastHashJoin, after seeing that one side turned out much smaller than Catalyst had estimated.

Connection to the module. Lessons 2 and 3 of this module showed the physical plan as a fixed object, wrapped in AdaptiveSparkPlan. This lesson opens that wrapper: what AQE revises, when it does it, and what evidence it leaves in .explain() when it decides to change something.

An analogy: the GPS that only recalculates at stops, not all the time

The GPS from this module's introduction doesn't continuously recalculate the route, frame by frame — that would be very expensive, and most of the time unnecessary. It recalculates at specific moments: when you reach an intersection where it had to make a decision anyway. Adaptive Query Execution works with the same discipline: it doesn't constantly review the plan while running — it reviews it at the points where it has to stop and coordinate anyway: the end of every shuffle stage, the exact moment Spark has already written that stage's complete result and knows, precisely, how many bytes it weighs. There, and only there, AQE compares that real weight against what it had assumed before running, and decides whether the rest of the plan still makes sense as it was, or whether it's worth rewriting.

Worked example, part 1: partition coalescing, in a paragraph with its own evidence

You already measured this first capability with real evidence in module 4, lesson 7 — here it's repeated in its shortest form, to have on hand before seeing the second capability.

# aqe_partitions_recap.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)

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

print("spark.sql.shuffle.partitions =", spark.conf.get("spark.sql.shuffle.partitions"))

spark.conf.set("spark.sql.adaptive.enabled", False)
q_off = build_query()
q_off.collect()
print(f"post-shuffle partitions, AQE OFF = {q_off.rdd.getNumPartitions()}")

spark.conf.set("spark.sql.adaptive.enabled", True)
q_on = build_query()
q_on.collect()
print(f"post-shuffle partitions, AQE ON  = {q_on.rdd.getNumPartitions()}")

spark.stop()

What to expect (executed in this run, over fact_orders_at_scale, ten million rows):

spark.sql.shuffle.partitions = 200
post-shuffle partitions, AQE OFF = 200
post-shuffle partitions, AQE ON  = 1

Identical to module 4's result: groupBy("store_id") produces a result with only three rows, and AQE, seeing that real size after the shuffle finished, coalesces the 200 potential partitions into one. If you need the complete mechanism — spark.sql.adaptive.advisoryPartitionSizeInBytes, the AQEShuffleRead coalesced node, the Initial Plan versus the Final Plan — it's developed in depth in module 4, lesson 7; this lesson doesn't repeat it twice.

Worked example, part 2: switching a complete JOIN's strategy

This is AQE's new capability no earlier lesson has shown: when Catalyst can't know, before executing, that one side of the JOIN's result is going to be tiny — because that side comes from an aggregation, not from a file with a size known ahead of time — it plans with the safer strategy for the general case: SortMergeJoin. But if, upon execution, that side turns out much smaller than the initial plan assumed, AQE can replace that strategy with BroadcastHashJoin, mid-execution, with nothing forced on your part.

# aqe_join_switch.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[*]").config("spark.driver.memory", "4g").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_df.withColumn("revenue", F.col("quantity") * F.col("unit_price"))

# top_products: an aggregation's result (only 4 rows -- one product for P001..P004),
# but Catalyst CANNOT know that ahead of time -- it comes from a shuffle, not a file.
top_products = fact.groupBy("product_id").agg(F.sum("revenue").alias("total_revenue_by_product"))

# Join against fact itself (10M rows) by product_id -- projecting little so
# collect() doesn't exhaust driver memory; what matters is the join's STRATEGY.
joined = fact.select("product_id").join(top_products, "product_id")

print("=== BEFORE running: only the INITIAL plan exists (isFinalPlan=false) ===")
joined.explain(mode="formatted")

rows = joined.collect()
print(f"\ncollect() -> {len(rows)} rows")

print("\n=== AFTER collect(), on the SAME object: Final Plan vs Initial Plan ===")
joined.explain(mode="formatted")

spark.stop()

What to expect, before running (executed in this run):

=== BEFORE running: only the INITIAL plan exists (isFinalPlan=false) ===
== Physical Plan ==
AdaptiveSparkPlan (14)
+- Project (13)
   +- SortMergeJoin Inner (12)
      :- Sort (4)
      :  +- Exchange (3)
      :     +- Filter (2)
      :        +- Scan csv  (1)
      +- Sort (11)
         +- HashAggregate (10)
            +- Exchange (9)
               +- HashAggregate (8)
                  +- Project (7)
                     +- Filter (6)
                        +- Scan csv  (5)

... (detail blocks omitted -- confirm product_id as both sides' key)

(14) AdaptiveSparkPlan
Output [2]: [product_id#3, total_revenue_by_product#9]
Arguments: isFinalPlan=false

Before anything runs, Catalyst chooses SortMergeJoin — the safe option, because top_products is a HashAggregate's result over a shuffle, and Catalyst has no way to anticipate, without running anything, that this result is going to have only four rows.

What to expect, after .collect(), on the same joined object (executed in this run):

=== AFTER collect(), on the SAME object: Final Plan vs Initial Plan ===
== Physical Plan ==
AdaptiveSparkPlan (30)
+- == Final Plan ==
   ResultQueryStage (18)
   +- * Project (17)
      +- * BroadcastHashJoin Inner BuildRight (16)
         :- AQEShuffleRead (5), local
         :  +- ShuffleQueryStage (4), Statistics(sizeInBytes=228.9 MiB, rowCount=1.00E+7)
         :     +- Exchange (3)
         :        +- * Filter (2)
         :           +- Scan csv  (1)
         +- BroadcastQueryStage (15), Statistics(sizeInBytes=8.0 MiB, rowCount=4)
            +- BroadcastExchange (14)
               +- * HashAggregate (13)
                  +- AQEShuffleRead (12), coalesced
                     +- ShuffleQueryStage (11), Statistics(sizeInBytes=1536.0 B, rowCount=48)
                        +- Exchange (10)
                           +- * HashAggregate (9)
                              +- * Project (8)
                                 +- * Filter (7)
                                    +- Scan csv  (6)
+- == Initial Plan ==
   Project (29)
   +- SortMergeJoin Inner (28)
      :- Sort (21)
      :  +- Exchange (20)
      :     +- Filter (19)
      :        +- Scan csv  (1)
      +- Sort (27)
         +- HashAggregate (26)
            +- Exchange (25)
               +- HashAggregate (24)
                  +- Project (23)
                     +- Filter (22)
                        +- Scan csv  (6)

...

(30) AdaptiveSparkPlan
Output [2]: [product_id#3, total_revenue_by_product#9]
Arguments: isFinalPlan=true

This is the complete evidence. The Final Plan replaced the Initial Plan's SortMergeJoin with a * BroadcastHashJoin Inner BuildRight — and the reason is in the line BroadcastQueryStage (15), Statistics(sizeInBytes=8.0 MiB, rowCount=4): once top_products's shuffle finished writing, AQE measured its real size — 8.0 MiB, 4 rows, one for each real Kiosko product (P001 through P004) — and confirmed it sat far below the broadcast threshold (10 MB by default, the same criterion from module 5). With that real evidence in hand, AQE rewrote the rest of the plan, still before executing, to use the cheaper strategy. The Initial Plan never disappears from .explain() — it stays there, as a record of what Catalyst would have done without this correction.

Diagram: the two ways AQE rewrites a plan

flowchart TD
    subgraph Coalesce["Partition coalescing -- already seen in M4L7 and in part 1"]
        A1["Exchange hashpartitioning(store_id, 200)"] --> A2["AQE measures: small real result\n(3 rows, a few KB)"]
        A2 --> A3["AQEShuffleRead coalesced:\n1 read partition, not 200"]
    end

    subgraph Switch["JOIN strategy switch -- new in this lesson"]
        B1["Initial plan: SortMergeJoin\n(Catalyst doesn't know top_products's real size)"] --> B2["AQE measures: top_products's\nreal shuffle = 8.0 MiB, 4 rows"]
        B2 --> B3["Final Plan: BroadcastHashJoin\n(8.0 MiB is below the 10 MB threshold)"]
    end

    style A3 fill:#9c6,stroke:#333
    style B3 fill:#9c6,stroke:#333

Going deeper: AQE's three documented optimizations, and which one this guide skips

Spark's official documentation describes Adaptive Query Execution precisely:

"Adaptive Query Execution (AQE) is an optimization technique in Spark SQL that makes use of the runtime statistics to choose the most efficient query execution plan, which is enabled by default since Apache Spark 3.2.0."

And it explicitly names three features under that same mechanism: coalescing post-shuffle partitions (this lesson's part 1, and module 4, lesson 7), converting a sort-merge join into a broadcast join (this lesson's part 2, with real evidence), and a third one this guide does not run a dedicated example for: optimizing a skewed join, where a single key concentrates a disproportionate number of rows compared to the others, and AQE can split that overloaded partition into several smaller ones to process in parallel. It's named here, with its official citation, so you know it exists and where to look for it — kiosko_orders_at_scale has no key skewed that way (every store and every product receives a proportional, stable fraction of the 250,000 synthetic franchises), so this guide doesn't build executed evidence for that third optimization.

All three share the same underlying principle, and it's worth stating plainly: AQE never changes what a query computes — it only changes how it executes it, based on information that only exists after part of the work has already happened. The final result — 26,537,500.00 in total revenue, verified in every module of this guide — is exactly the same with AQE active or disabled; the only thing that changes is the cost of getting to that result.

Common mistakes

Looking for the strategy switch in a .explain() of a DataFrame that never ran. What happens: someone builds this lesson's part 2 query, calls .explain() just once, sees SortMergeJoin, and concludes AQE didn't work. Why it happens: it's easy to forget a Spark DataFrame is lazy — nothing runs until there's a real action, and AQE can only re-optimize something that already ran, at least in part. How to spot it: if your .explain() shows isFinalPlan=false and just one section (no == Final Plan == or == Initial Plan ==), you haven't run any action on that object yet. How to fix it: call an action (.collect(), .count() through .rdd.count(), or any operation that materializes the result) on the same DataFrame object, and call .explain() again on that same object — not on a copy, nor on a new query built from scratch.

Assuming any action, like .count(), updates the same object's cached plan. What happens: someone calls joined.count() expecting that to be enough for joined.explain() to later show the Final Plan, and is surprised it still shows only the initial plan. Why it happens: it seems reasonable that any action "counts" as having run the DataFrame. How to spot it: if you used .count() as your trigger action and the following .explain() still shows no Final Plan, that's exactly the expected behavior — not a bug in your code. How to fix it: .count() internally builds and runs a different plan (one designed specifically to count), not the original object's plan — so running .count() leaves no trace in the joined object's own planning cache. Use .collect() (or any action operating directly on the DataFrame without rebuilding a new plan behind it) when you need the same object to reflect its final plan.

Reading BroadcastQueryStage (15), Statistics(sizeInBytes=8.0 MiB, rowCount=4) as an estimate, the same way as mode="cost"'s figures. What happens: someone sees these figures inside the Final Plan and treats them with the same caution as lesson 3's estimates — "this is just a prediction, don't confirm anything with this." Why it happens: both look similar in the text — the word Statistics next to a size in bytes — and lesson 3 explicitly taught you to distrust those figures as real measurements. How to spot it: if your explanation doesn't distinguish which phase of the plan each figure shows up in, check again: mode="cost"'s figures live in the Optimized Logical Plan, before running anything; a BroadcastQueryStage's figures inside a Final Plan live after that stage has actually already run. How to fix it: the figures inside a QueryStage inside a Final Plan (with isFinalPlan=true) are real measurements, taken after running that stage — it's exactly the evidence AQE used to decide the strategy switch, not a prediction.

Exercises

Exercise 1 — Repeat part 2's experiment with dim_store_df instead of top_products. dim_store_df (dim_store.csv) is a real table, with a size Catalyst knows ahead of time — it doesn't come from a shuffle. Build fact.select("store_id").join(dim_store_df, "store_id"), call .explain(mode="formatted") before any action, and predict whether you're going to see SortMergeJoin or BroadcastHashJoin from the initial plan, with no AQE switch needed.

joined_dim = fact.select("store_id").join(dim_store_df, "store_id")
print("=== explain BEFORE any action ===")
joined_dim.explain(mode="formatted")

=== explain BEFORE any action ===
== Physical Plan ==
AdaptiveSparkPlan (8)
+- Project (7)
   +- BroadcastHashJoin Inner BuildRight (6)
      ...
(8) AdaptiveSparkPlan
Arguments: isFinalPlan=false

Exercise 2 — Explain, without code, why part 2's Initial Plan never disappears from .explain(), even though you already know it isn't what ran. In 2-3 sentences, justify why Spark keeps that section instead of only showing the Final Plan.

Exercise 3 — Explain, without code, why this lesson never uses the word "faster" to describe the switch from SortMergeJoin to BroadcastHashJoin. In 2-3 sentences, explain what evidence this lesson does use to describe the switch's benefit, and why that evidence is more reliable than a measured time.

Summary and next step

This lesson finished explaining AdaptiveSparkPlan, the wrapper that accompanied every physical plan since module 2: Adaptive Query Execution revises the plan at every shuffle boundary, with real instead of estimated statistics, and can rewrite the rest of the plan in two distinct ways — coalescing empty partitions (already seen in module 4, repeated here in a paragraph) and switching a JOIN's complete strategy, from SortMergeJoin to BroadcastHashJoin, when one side turns out much smaller than estimated. You verified the latter with real evidence: an Initial Plan with SortMergeJoin, a Final Plan with BroadcastHashJoin, and the exact figure (8.0 MiB, 4 rows) that justified the switch.

Before moving on you should be able to: explain exactly when AQE revises a plan (at the end of every shuffle stage, not continuously); read a .explain() with Final Plan/Initial Plan and explain what changed and why; and name AQE's three documented optimizations, even though this guide only runs evidence for two.

Lessons 5 through 7 leave Catalyst and AQE behind to resolve this module's second half: when .cache() saves real work, and when it only takes up memory with no benefit at all.

Resources

• Apache Spark — SQL Performance Tuning, "Adaptive Query Execution" section (the three officially documented optimizations: partition coalescing, sort-merge-to-broadcast join conversion, and skewed-join optimization — the exact quote used in this lesson). 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: AQE compared True vs False over the same groupBy at scale, reading the post-shuffle partition count.
• This guide's module 4, lesson 7 (shuffle-partitions-and-adaptive-query-execution.md) — the complete development of partition coalescing, including spark.sql.adaptive.advisoryPartitionSizeInBytes and the AQEShuffleRead node, which this lesson doesn't repeat in detail.