How do I tune spark.sql.shuffle.partitions?

Summary

• spark.sql.shuffle.partitions sets how many partitions Spark uses when it shuffles data for joins or aggregations. The default is 200.
• Adaptive Query Execution (AQE) has been on by default since Spark 3.2.0 and coalesces small shuffle partitions at runtime, so you need a large enough starting number rather than an exact one.
• The Spark UI shows which way the value is off: many tiny tasks suggest too many partitions, while shuffle spill or an OutOfMemoryError in reduce tasks suggests too few.
• On Databricks, setting the value to auto turns on auto-optimized shuffle, and auto is the default on serverless compute.

How do I tune spark.sql.shuffle.partitions?
spark.sql.shuffle.partitions is the number of partitions Spark uses when it shuffles data for joins or aggregations. The default is 200. A task works on a single partition, so the value sets how many tasks the stage after a shuffle starts with.

What goes wrong when the value is off?

Too many partitions produce many tiny tasks, and the Databricks documentation notes that very small tasks suffer more from scheduling and task setup overhead. Too few partitions make each task's input large. The Spark tuning guide describes one result: an OutOfMemoryError can come from the working set of a single reduce task being too large, and the simplest fix is to increase the level of parallelism so that each task's input set is smaller.

How do I see which way it is off?

Open the Spark UI. A job's detail page in the Jobs tab lists each stage with Shuffle read and Shuffle write columns (bytes and records). A stage's detail page in the Stages tab shows task duration, GC time, Shuffle Read Size / Records, and Shuffle spill (memory) and (disk). In the SQL tab, an Exchange node reports metrics such as shuffle bytes written and shuffle records written. Shuffle spill means shuffled data did not fit in memory (the Spark RDD guide says Spark spills to disk when data does not fit in memory), and more partitions make each task's input smaller. The Details section of the SQL tab also shows runtime statistics while a query runs: look for Statistics(..., isRuntime=true) in the plan.

How do I pick and set a value?

The Spark tuning guide says: "In general, we recommend 2-3 tasks per CPU core in your cluster." A stage runs one task per partition, so applying that guidance to a cluster with 40 cores gives roughly 80 to 120 partitions as a first value. Adjust from what the Spark UI shows.

spark.conf.set("spark.sql.shuffle.partitions", "400")  # example value

SET spark.sql.shuffle.partitions=400;

In conf/spark-defaults.conf, each line is a key and a value separated by whitespace:

spark.sql.shuffle.partitions 400

spark-submit also accepts --conf spark.sql.shuffle.partitions=400. To read the value a session is using, run SET spark.sql.shuffle.partitions;. Values set in spark-defaults.conf, SparkConf or on the command line also appear in the Environment tab of the Spark UI.

How does AQE change the tuning?

With AQE on, the Spark performance tuning guide says that coalescing post-shuffle partitions "simplifies the tuning of shuffle partition number". Spark can pick the number at runtime once you set a large enough initial number of shuffle partitions. That initial number comes from spark.sql.adaptive.coalescePartitions.initialPartitionNum, which falls back to spark.sql.shuffle.partitions.

Setting	Default	Role
spark.sql.adaptive.enabled	true	Turns AQE on
spark.sql.adaptive.coalescePartitions.enabled	true	Coalesces contiguous shuffle partitions to avoid too many small tasks
spark.sql.adaptive.coalescePartitions.initialPartitionNum	none	Initial partition count before coalescing
spark.sql.adaptive.advisoryPartitionSizeInBytes	64 MB	Advisory size when coalescing or splitting skewed partitions
spark.sql.adaptive.coalescePartitions.parallelismFirst	true	Sizes coalesced partitions from the cluster's parallelism instead of the advisory size; the docs recommend false on a busy cluster
spark.sql.adaptive.skewJoin.enabled	true	Splits skewed partitions in shuffled joins

A partition counts as skewed when it is larger than spark.sql.adaptive.skewJoin.skewedPartitionFactor (5.0) times the median partition size and larger than spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes (256 MB).

What changes on Databricks?

• AQE is enabled by default. The docs say it applies to non-streaming queries that contain at least one exchange (usually a join, aggregate or window) or a sub-query.
• Setting spark.sql.shuffle.partitions to auto enables auto-optimized shuffle, which "automatically determines this number based on the query plan and the query input data size".
• On serverless compute, spark.sql.shuffle.partitions is one of the few Spark properties you can set, and its default there is auto.
• Databricks generally recommends against configuring most Spark properties.

FAQs

Does AQE make this setting irrelevant?

No. The initial number is the starting point AQE coalesces from, which is why the docs ask for a large enough one.

How do streaming queries differ?

The Databricks docs note that for Structured Streaming, the value cannot be changed between query restarts from the same checkpoint location. The Apache Spark 4.1.0 release notes list SPARK-53942 (support changing shuffle partitions in stateless streaming workloads) and SPARK-53941 (support AQE in stateless streaming workloads), so check the notes for your version.

How do I set it in a Spark Declarative Pipelines pipeline?

The SDP programming guide's example pipeline spec sets spark.sql.shuffle.partitions: "1000" under configuration. The dataset decorators also accept a spark_conf argument that can override pipeline or cluster settings for one dataset.