Adaptive Query Execution in Apache Spark 3.x

AQE rewrites Spark physical plans at runtime using actual statistics, often reducing shuffle data by ~30%. This post explains the mechanism, a concrete example, trade-offs, and verification steps.

How AQE Works in Two Phases

In Spark 3.0+, the Adaptive Query Execution (AQE) feature rewrites the physical plan after the first shuffle has produced runtime statistics. Before AQE, the planner decides partition counts and join strategies at compile time, which can lead to under- or over-partitioned shuffles when data distribution is unexpected.

Initial plan generation

When a job starts, Spark generates a deterministic physical plan based on catalog statistics and cost estimates. This plan includes a fixed number of shuffle partitions, chosen join strategies, and partition pruning rules. The plan is static. It does not change during execution.

Runtime re-planning

After the first shuffle completes, Spark collects metrics such as partition sizes, data skew, and join cardinality. Using these metrics, AQE may trigger one or more of the following optimizations: dynamic partition pruning (skipping irrelevant partitions), shuffle aggregation (combining map output before the second shuffle), and broadcast join selection (choosing a broadcast join when one side is small enough). The planner then rewrites the remaining stages, often reducing the total shuffle volume and improving executor memory usage.

Concrete Example – groupBy Aggregation on a Large CSV

Imagine a 120 GB CSV file with 10 million rows, and you need a group-by aggregation on a non-key column. Without AQE, Spark creates a ShuffleHashAggregate stage with a fixed partition count (typically spark.sql.shuffle.partitions, default 200). The shuffle moves all row groups to executors, and the final aggregation happens in reducers.

Enable AQE by adding the configuration when submitting your job or starting a shell:

\nspark-submit --conf spark.sql.adaptive.enabled=true my-job.jar\n

or in a Spark shell:

\nspark.sql("set spark.sql.adaptive.enabled=true")\n

In the Spark UI Stages tab, the shuffle stage plan often changes from ShuffleHashAggregate to a DynamicShuffleReducer. You may also see a reduction in shuffle read bytes of roughly 20-30%, and a modest decrease in job duration because fewer records are shuffled across the network.

To compare, run the same job with spark.sql.adaptive.enabled=false and note the UI metrics. The difference confirms AQE's effect without needing external profiling tools.

Trade-offs, Limits and Verification

Plan non-determinism: The same query may produce different physical plans across runs when AQE is enabled, which can trip up unit tests that assert exact stage names or task counts.
CPU overhead during re-planning: Analyzing runtime statistics and re-optimizing consumes executor CPU. On very small jobs (under a few hundred megabytes) the overhead may outweigh the shuffle-reduction benefit.
Custom UDF constraints: Non-deterministic UDFs or those with side effects may not survive the plan rewrite, causing runtime errors if not annotated or restructured.
Small-dataset penalty: When the data volume is low, the cost of collecting metrics and re-planning can add latency compared to a static, well-tuned plan.

Practical verification: Run the job twice - once with AQE disabled and once enabled - and compare the Shuffle read bytes and Executor run time values in the Spark UI. Additionally, set spark.sql.adaptive.shuffle.targetPostShuffleInputSize to a target size (e.g., 128 MB) and observe the resulting partition count in the UI; lowering the target typically increases partition count, which can help with skew but may increase coordination overhead.

Getting the Most Out of AQE

For production clusters, keep AQE enabled and monitor the Spark UI after each major query pattern change. Start with the defaults, then enable spark.sql.adaptive.coalescePartitions.enabled=true to further reduce shuffle partition count for skewed joins. Use spark.sql.adaptive.shuffle.targetPostShuffleInputSize to set a per-partition input size goal and watch how the UI's partition count adjusts. Disable AQE when debugging plan-specific issues, running deterministic test suites, or when using connectors that do not expose runtime statistics (e.g., some JDBC drivers).

Adaptive Query Execution in Apache Spark 3.x: Dynamic Optimization Explained