Tune Spark 3.x Batch ETL Shuffle Partitions with Adaptive Query Execution

The problem: a static shuffle-partition guess

You inherit a nightly Spark batch ETL job that writes to a data lake. The previous engineer set spark.sql.shuffle.partitions to 200 and left it there. On small days the stage is over-partitioned and task overhead dominates; on big days some partitions skew and straggler tasks hold up the whole job. The output directory fills with thousands of small files, and downstream consumers complain about slow reads.

Picking one fixed partition count for every dataset and every day is the core issue. Spark 3.x includes Adaptive Query Execution (AQE), a Spark SQL feature that replans parts of a query while it runs, using statistics collected during execution. For batch ETL, the most relevant behaviors are shuffle-partition coalescing and dynamic partition pruning.

What AQE changes at runtime

AQE is controlled by the master switch spark.sql.adaptive.enabled. When it is on, Spark can adjust the physical plan after seeing actual data sizes:

• Shuffle partition coalescing: after a shuffle, Spark merges small output partitions into larger ones, up to spark.sql.adaptive.advisoryPartitionSizeInBytes. This reduces per-task overhead and small-file writes without requiring you to guess the right partition count up front.
• Dynamic partition pruning (DPP): when a large fact table joins a smaller filtered dimension table, Spark uses values from the dimension side to skip reading irrelevant partitions of the fact table.
• Broadcast conversion: a sort-merge join can become a broadcast hash join at runtime if statistics show one side is smaller than the broadcast threshold.

Exact defaults for these sub-features differ between Spark 3.x releases, so confirm behavior against your target version before relying on them.

Enable it and verify the effective configuration

You can set AQE in three common places. Cluster-wide in spark-defaults.conf affects every application, so test on a single pipeline first. Per-job submission is safer for an initial rollout.

spark-submit --master  --conf spark.sql.adaptive.enabled=true --conf spark.sql.adaptive.coalescePartitions.enabled=true --conf spark.sql.adaptive.advisoryPartitionSizeInBytes=134217728 --class com.example.ETLJob /opt/jobs/etl.jar

Run spark-submit from a gateway or edge node where the Spark client is installed and you have permission to submit to the target queue or cluster. Replace placeholders in angle brackets with your cluster master and application coordinates.

After the driver starts, confirm the effective flags inside your Spark application or through spark-shell:

scala> spark.conf.get("spark.sql.adaptive.enabled")
scala> spark.conf.get("spark.sql.adaptive.coalescePartitions.enabled")
scala> spark.conf.get("spark.sql.adaptive.advisoryPartitionSizeInBytes")

Expected check: each returns true or the byte value you set. An empty string or false means AQE is off for that session.

Diagnose the plan before and after

Use explain('extended') on a representative query to compare physical plans. With AQE enabled, look for adaptive plan nodes such as AdaptiveSparkPlan, coalesced ShuffleExchange operators, BroadcastExchange replacing sort-merge shuffle, or pruning predicates on partitioned scans. Do not assume a specific operator name; print the plan for your own query and read the actual output.

In the Spark UI, open the SQL tab for the completed application and compare shuffle write task counts against your static spark.sql.shuffle.partitions setting. Fewer tasks with larger input bytes per task is the signal that coalescing happened.

Worked example: filtered dimension join

Consider a join between a sales fact table partitioned by event_date and a stores dimension table filtered to one region. Without dynamic partition pruning, Spark may scan every sales partition. With AQE and DPP enabled, Spark reads the filtered stores side first, collects the relevant join keys, and uses them to skip sales partitions that cannot match.

To verify, run the join with AQE disabled, note the number of sales partitions read in the Spark UI, then enable AQE and run the same query. A reduction in scanned partitions or input bytes for sales indicates pruning took effect. Confirm with explain('extended') that the plan includes a pruning predicate rather than a full table scan.

Trade-offs and limitations

• Unpredictable output partitions: coalescing changes the number of output files based on data size. Downstream jobs expecting a fixed partition count may fail or produce unexpected layouts. Add an explicit repartition(n) before writing if the consumer requires a stable count.
• Executor memory pressure: runtime broadcast conversion depends on statistics. If the broadcast side is larger than estimated or skewed, executors can run out of memory. Keep spark.sql.autoBroadcastJoinThreshold appropriate for your executor memory.
• Local shuffle reader blast radius: spark.sql.adaptive.localShuffleReader.enabled improves data locality by reading shuffle data on the same executor, but unstable executors increase the failure blast radius because local shuffle data is tied to that executor.
• Skew handling is separate: skewed join optimization is an AQE sub-option and may not be enabled by default in your release.

Practical check and rollback

After enabling AQE on one pipeline, compare the Spark UI shuffle write stage before and after: task count, input size per task, and output file count. If output partitions vary too much for downstream consumers, pin the count with df.repartition(n) before write.

Because AQE is a configuration change, rollback means removing the --conf flags or reverting spark-defaults.conf and redeploying the job. Test on a non-production queue first.

Closing

AQE removes the need to hard-code shuffle partitions for every batch, but it is not a set-and-forget switch. Enable it on a single ETL job, verify with explain and the Spark UI, and treat the resulting partition count as a runtime result rather than a contract.