Use Adaptive Query Execution to Reoptimize Spark 3.x Joins and Shuffle Partitions at Runtime

The problem is static estimates, not bad queries

Spark builds a physical plan before any data moves. When the optimizer guesses row counts or partition sizes incorrectly, you get too many tiny shuffle tasks, a sort‑merge join where a broadcast would be cheaper, or a single skewed key that stalls a stage. Adaptive Query Execution lets Spark 3.x observe real shuffle sizes and replan mid‑execution instead of trusting the initial estimates.

What AQE changes at runtime

AQE operates after shuffle map stages finish and actual statistics are available. From those measurements it can:

• Coalesce small shuffle partitions into fewer, larger tasks to reduce scheduling overhead.
• Switch a planned sort‑merge join to a broadcast hash join when the post‑shuffle size of one side is below the broadcast threshold.
• Split skewed join keys into multiple sub‑tasks for adaptive skew join.

The initial plan shown by explain() is therefore not the executed plan when AQE is on. The final plan is materialized after the shuffle.

Enable AQE with a focused configuration

AQE requires Spark 3.0 or later. The core flag is spark.sql.adaptive.enabled. Two commonly used sub‑options control the behaviors above.

# Run on the driver before creating the query. No cluster privileges required, only driver session scope.
spark.conf.set('spark.sql.adaptive.enabled', 'true')
spark.conf.set('spark.sql.adaptive.coalescePartitions.enabled', 'true')
spark.conf.set('spark.sql.adaptive.adaptiveSkewJoin.enabled', 'true')

For a persistent setting, add the same keys to spark-defaults.conf or the cluster manager configuration. Check that a managed service is not overriding them at the cluster level; some services disable AQE by default.

Worked example: join with uncertain sizes

Assume two DataFrames, fact and dim, joined on dim.key. Static statistics suggest both sides are large, so the planner picks a sort‑merge join with spark.sql.shuffle.partitions tasks.

With AQE enabled, Spark first executes the shuffle map stage for both sides, measures the actual output bytes per partition, then decides:

• If the measured size of dim after shuffle is below spark.sql.autoBroadcastJoinThreshold, the physical plan is replaced with BroadcastHashJoin.
• If many shuffle partitions are tiny, they are coalesced to a target size governed by spark.sql.adaptive.advisoryPartitionSizeInBytes.
• If a key shows disproportionate size, adaptive skew join splits it.

Run the query, then open the Spark UI SQL tab for that query id. Inspect the Summary and Details sections for the executed physical plan and shuffle write metrics. The plan should show the adapted operators, and the stage task count should reflect coalescing.

How to verify AQE is actually taking effect

Do not rely on explain() alone. Compare plans with AQE on and off on the same session.

1. With AQE enabled, run df.explain('extended') to see the initial logical and physical plan.
2. Execute the query and open the Spark UI SQL tab. Check the final physical plan displayed for the query. Look for operators like AdaptiveSparkPlan, BroadcastExchange, and indicators of coalesced partitions.
3. Compare stage metrics in the History Server: task count, shuffle read/write bytes, and task duration distribution between a run with spark.sql.adaptive.enabled=true and spark.sql.adaptive.enabled=false.
4. Confirm configuration is not overridden by checking spark.conf.get('spark.sql.adaptive.enabled') on the driver after session start and reviewing spark-defaults.conf and cluster manager configs.

These checks confirm whether replanning occurred and which optimization was applied.

Limits and when AQE does not help

AQE works best when statistics are collected during shuffle. It does not apply to Structured Streaming continuous mode and has limited effect on pipelines dominated by Python UDFs where Spark cannot accurately measure sizes.

Flag names and defaults changed across Spark 3.0 through 3.5. Confirm the exact names for your distribution and version.

Adaptive skew join depends on detecting skew during execution. If skew keys are not identified or broadcast size limits are hit, performance may not improve.

Coalescing partitions creates larger tasks. With extreme input skew, a coalesced task can exceed executor memory and cause OOM errors. Monitor executor memory and consider lowering the advisory partition size or disabling coalescing for known skewed workloads.

Common mistakes

Manually tuning spark.sql.shuffle.partitions while AQE is enabled is a frequent conflict. AQE overrides static partition targets based on observed shuffle file sizes, so manual values can be ignored or cause unexpected task counts.

Expecting explain() to show the final plan. The initial plan is static; the adapted plan only appears in the UI after execution.

Enabling AQE without checking service defaults. Some managed Spark services disable AQE or set different thresholds at the cluster level, silently overriding session settings.

Practical decision rule

Enable AQE for ad‑hoc SQL and ETL jobs on Spark 3.x where input sizes are variable and join selectivity is uncertain. Keep it off for workloads that require deterministic physical plans, such as tests that assert exact task counts, or when you have already tuned partition sizes for a stable, well‑characterized dataset.