Let Spark Coalesce Shuffle Partitions at Runtime with AQE

A batch job with uneven shuffle tasks

You run a nightly Spark batch job that joins a large events table to a small dimension table. You set 200 shuffle partitions—the tasks that read and redistribute data across the network before a join—because that worked on a smaller dataset. In production, most tasks finish in under two minutes, but a few tasks run for 40 minutes and hold up the whole stage. The cause is usually data skew: one or a few join keys appear far more often than others, so the executors holding those keys do disproportionate work.

Manually guessing the right shuffle partition count is fragile. Add more partitions and small tasks waste time on task overhead; add too few and stragglers dominate. Apache Spark 3.x offers Adaptive Query Execution (AQE), which can coalesce shuffle partitions at runtime and replicate skewed join keys so the work spreads more evenly across executors.

This post shows how to enable AQE for a batch job, what to check in the Spark UI, and where the feature can create new problems instead of solving them.

What AQE changes at runtime

AQE is enabled by default in Spark 3.x when spark.sql.adaptive.enabled is true. It works with the Spark SQL optimizer and the DataFrame and Dataset APIs. When AQE runs, Spark can adjust the physical plan after it sees actual shuffle file sizes and data distribution.

The two optimizations most relevant to shuffle tuning are:

• Coalescing post-shuffle partitions. Spark merges small shuffle output files into fewer tasks, reducing task launch overhead when the initial partition count was too high.
• Skew join handling. When a join key is skewed, Spark replicates the rows for that key across multiple tasks instead of sending all matching rows to one task.

AQE can also perform dynamic partition pruning, which reduces scanned data when joining a large fact table to a small dimension table, but that optimization depends on additional settings and the join pattern.

Version assumption

This guidance targets Spark 3.2 or later. AQE defaults and supported optimizations changed across Spark 3.1, 3.2, 3.3, and 3.4, especially for Structured Streaming, so confirm spark.version before relying on a behavior described here.

Enable AQE for a batch job

Set these properties when you submit the job. Run spark-submit on a gateway or edge node where you have permission to submit applications and read the input data:

spark-submit \
  --master CLUSTER_MASTER_URL \
  --deploy-mode cluster \
  --conf spark.sql.adaptive.enabled=true \
  --conf spark.sql.adaptive.coalescePartitions.enabled=true \
  --conf spark.sql.adaptive.skewJoin.enabled=true \
  --conf spark.sql.adaptive.advisoryPartitionSizeInBytes=128MB \
  --class com.example.batch.JoinJob \
  APPLICATION_JAR.jar \
  INPUT_PATH OUTPUT_PATH

spark.sql.adaptive.advisoryPartitionSizeInBytes suggests a target size for coalesced shuffle partitions. Spark uses it as guidance, not a hard limit.

Required permissions: submit access to the Spark cluster, read access to the input path, and write access to the output path. If you run in client mode, the driver runs on the submitting node; in cluster mode, the driver runs on the cluster.

Worked example: skewed join

Suppose you join events to users on user_id. A small number of user_id values appear millions of times. Without AQE, one shuffle task may read all events for those hot keys while other tasks sit idle.

// Scala / Spark Shell example. Run spark-shell with AQE enabled.
val events = spark.read.parquet("INPUT_PATH/events")
val users = spark.read.parquet("INPUT_PATH/users")

val joined = events.join(users, Seq("user_id"))
joined.write.mode("overwrite").parquet("OUTPUT_PATH/joined")

After the job runs, open the Spark UI and go to the SQL tab for the completed job. Look at the physical plan for the join stage.

• If coalescing worked, the number of shuffle write tasks should be lower than the initial shuffle partition count, and task durations should be more uniform.
• If skew join handling worked, the Spark UI stage detail should show the skewed key split across multiple tasks rather than one long-running task.

Check the Stages tab for shuffle read bytes and task duration variance. High variance across tasks in the same stage is the classic sign that skew remains.

Do not copy expected numbers from another job. Shuffle bytes and task times depend on your data distribution, file format, and cluster hardware.

Trade-offs and limitations

AQE removes some manual tuning, but it introduces operational trade-offs:

• Fewer output files. Coalescing can produce fewer output files than expected. If a downstream process expects a specific partitioning or file-size target, check the output directory after the job.
• Memory pressure. Skew join replication increases shuffle read for the replicated keys and can raise memory pressure on executors handling those partitions.
• Plans change at runtime. An EXPLAIN plan captured before execution may not match the final executed plan because AQE decides partitions after seeing shuffle statistics.
• Streaming support is limited. AQE is primarily a batch and Spark SQL feature. Structured Streaming support is limited and version-dependent, so do not assume streaming jobs receive the same optimizations.

Verify the result and roll back if needed

Confirm the Spark version:

spark-shell --conf spark.sql.adaptive.enabled=true
scala> spark.version

Run the job once with AQE enabled and once with it disabled, comparing stage duration and output file count in the Spark UI. If AQE causes slower stages or too few output files, disable it for that job:

--conf spark.sql.adaptive.enabled=false

Disabling AQE is the rollback path because enabling it changes the runtime physical plan and output partitioning.

Closing

AQE is useful when your shuffle partition count is a guess and your data has uneven keys. It will not replace right-sized executors, appropriate file formats, or careful memory configuration. Enable it, compare the Spark UI stage details before and after, and watch output file counts and executor memory.