Adaptive Query Execution

Before Spark 3.0, the Catalyst Optimizer decided the physical execution plan (e.g. join strategy, partition counts) before the job started, relying on static database catalog statistics. If data was skewed or sizes were estimated incorrectly, Spark executed sub-optimal, slow plans.

Adaptive Query Execution (AQE) solves this by dynamically re-optimizing and modifying the physical plan at runtime, using real-time statistics collected during stage completions.

The 3 Pillars of AQE

AQE optimizes your queries during execution through three main features:

graph TD
    subgraph AQEFeatures["Adaptive Query Execution (AQE)"]
        direction TB
        F1["1. Dynamic Partition Coalescing<br>- Merges tiny post-shuffle partitions<br>- Prevents high task-scheduler overhead"]
        F2["2. Dynamic Join Switching<br>- Converts Sort-Merge to Broadcast Join<br>- If table size is smaller than expected"]
        F3["3. Dynamic Skew Join Handling<br>- Detects heavy-key partitions<br>- Splits skewed rows into parallel tasks"]
    end
    style AQEFeatures fill:#f0fdf4,stroke:#16a34a,stroke-width:2px;

1. Dynamic Coalescing of Shuffle Partitions

• The Problem: If you shuffle data, Spark creates spark.sql.shuffle.partitions (default 200). For small datasets, this launches 200 tiny tasks, causing huge scheduler overhead.
• AQE Solution: At the end of a map stage, Spark checks the actual size of the shuffled partitions. If they are tiny, AQE merges adjacent small partitions into a few sensible partitions, reducing task counts and accelerating execution.

2. Dynamic Join Selection

• The Problem: A table estimated at 15 MB gets compiled as a Sort-Merge Join because it exceeds the 10 MB broadcast threshold. But a filter WHERE age > 60 actually reduces the table to 2 MB.
• AQE Solution: When the scan stage completes, AQE checks the actual filtered size. Seeing it is only 2 MB, AQE dynamically changes the physical plan from a Sort-Merge Join to a high-speed Broadcast Hash Join at runtime!

3. Dynamic Skew Join Handling

• The Problem: Data is skewed. One partition is 5 GB while others are 10 MB. The job gets stuck at 99% waiting for the single 5 GB partition task to complete.
• AQE Solution: AQE monitors task sizes. When it detects a skewed partition, it splits that partition into multiple sub-partitions and joins them in parallel using a duplicated key strategy, completely eliminating the long-tail latency!

Configuring AQE in PySpark

AQE is enabled by default in Spark 3.0+. You can fine-tune its behavior using the following configurations:

from pyspark.sql import SparkSession

# 1. Setup Spark enabling all AQE features
spark = SparkSession.builder \
    .appName("Adaptive Query Execution") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .master("local[*]") \
    .getOrCreate()

# 2. Key configurations explained:
# - 'spark.sql.adaptive.enabled': Master switch to activate AQE.
# - 'spark.sql.adaptive.coalescePartitions.enabled': Merges small post-shuffle partitions.
# - 'spark.sql.adaptive.skewJoin.enabled': Automatically splits skewed partitions during joins.
# - 'spark.sql.adaptive.advisoryPartitionSizeInBytes': The target size for coalesced partitions (default 64MB).

# 3. Running a query with AQE
# Spark will automatically analyze partition statistics between execution stages 
# and update join paths and partition layouts dynamically!
df1 = spark.range(1, 1000000).repartition(10)
df2 = spark.range(1, 100).repartition(5)

result_df = df1.join(df2, "id")
result_df.show()