Spark shuffle partitions: how to choose the number

Set spark.sql.shuffle.partitions so each partition holds at most about 128 MB and every core gets at least 2 tasks, then round up to a multiple of the core count. A shuffle that moves 200 GB on 64 cores needs 1,600 partitions of 128 MB; the default of 200 gives each task 1 GB.

The rest of this article is how to measure the data, what adaptive query execution does to the number you set, and when the rule changes.

What the setting controls

A join or a group by has to bring every row with the same key to the same place. Spark does that with a shuffle: each task of the stage before writes its rows out sorted by destination, and the next stage reads them back as spark.sql.shuffle.partitions partitions, 1 task each.

Each core runs 1 task at a time, so a stage runs in rounds: 1,600 tasks on 64 cores take 25 rounds. The number decides 2 things at once:

• How much each task reads. Too few partitions and each task holds more than its memory can sort, so it spills to disk, and a failed task takes long to redo.
• How busy the cores are. Fewer tasks than cores leaves cores idle, and a last round of 8 tasks leaves 56 of 64 cores waiting.

The default is 200 in Apache Spark and on Databricks, whatever the size of the data or the cluster.

Measure what the shuffle moves

Run the job once and open the stage in the Spark UI. Shuffle Write on the stage before the shuffle is the data that crosses it. That is the figure to divide, not the size of the input files: a filter can make it far smaller, a join that multiplies rows far bigger.

Then count the cores: executors times cores per executor, or the total under Executors in the UI.

The rule

1. Divide Shuffle Write by 128 MB. 128 MB is the size Spark already uses for each partition when it reads files (spark.sql.files.maxPartitionBytes).
2. Take at least 2 tasks per core. Spark's tuning guide recommends 2 to 3, so a slow task does not hold up a whole round.
3. Round up to a multiple of the cores, so the last round is full.

For 200 GB on 64 cores: 200 GB ÷ 128 MB is 1,600, well above 2 × 64 = 128, and already a multiple of 64. 1,600 partitions, in 25 rounds. Try your own sizes, and the default, below.

Try itSpark shuffle partitions, in rounds

Data the shuffle moves

200GB

Shuffle Write in the Spark UI

1101001,000

Cores in the cluster

64cores

each runs 1 task at a time

16642561,024

Shuffle partitions

200partitions

spark.sql.shuffle.partitions: 200 unless you set it

2002,00020,000

Adaptive query execution

OnOff

Use the suggested 1,600

Each task reads1 GB

Tasks200

Rounds4

Each round, the share of cores busy

What adaptive query execution fixes, and what it cannot

Adaptive query execution has been on by default since Spark 3.2.0. After the shuffle writes its data, it reads the real size of every partition and merges small neighbours into fewer tasks. So a number set too high costs little: 20,000 partitions of 200 GB on 64 cores are about 10 MB each, and it merges them into about 3,300 tasks of 61 MB.

It does not work the other way. It never splits a partition that is too big, except a skewed one in a join, so at the default of 200 those 200 GB still run as 200 tasks of 1 GB. That is the case for setting the number high and letting adaptive query execution merge.

One default surprises people. With spark.sql.adaptive.coalescePartitions.parallelismFirst at true, as it is in Apache Spark, merging stops at 1 task per core, not at the 64 MB target: a 2 GB shuffle on 64 cores runs as 67 tasks of 31 MB. Spark's own docs recommend setting it to false on a busy cluster shared by many jobs, which gives 34 tasks of 60 MB. Databricks merges toward 64 MB but keeps at least 2 tasks per core by default.

When one key holds most of the data

The rule assumes keys are spread evenly. When 1 customer or 1 null value holds 20% of the rows, 1 partition is many times the median, and the stage waits on that 1 task.

In a sort-merge join, adaptive query execution splits a partition over 5 times the median and over 256 MB into pieces. In a group by it cannot, and more partitions do not help either: every row of the key still lands in 1. The Spark shuffle partition calculator shows the split, the longest task and how much longer the stage takes for your share of the hot key, and how to find and fix data skew covers salting and the other fixes. The concept page on data skewShuffle and Data SkewA shuffle moves rows between machines so rows with the same key meet, for a join or group by; skew is when one key has far more rows, so one task does most of the work. has the basics.

On Databricks

Databricks has the same default of 200, and a setting Apache Spark does not: spark.sql.shuffle.partitions set to auto turns on auto-optimized shuffle, which picks the number from the query plan and the size of the input. Try it before tuning by hand.

Structured Streaming is the exception: the number cannot change between restarts from the same checkpoint, so set it for the size the stream will reach before its first run.

Set it per job, not per cluster

1 cluster runs shuffles of very different sizes, so set the number in the job, before the query that needs it:

spark.conf.set("spark.sql.shuffle.partitions", 1600)

For the cluster around it, the Spark sizing calculator picks executors and memory from the data a job reads and shuffles, and gives this number for the cores it chooses.