Sizing Shuffle Partitions
Concepts covered: Configuration Tuning
Every Exchange that groupBy, distinct, orderBy or a count-less repartition(col) adds splits its output into spark.sql.shuffle.partitions pieces, 200 unless you change it. It is 1 value for the whole session, so the same 200 applies to a 2GB job and a 2TB job. Each piece becomes 1 task in the next stage, and when that stage ends in a write, each of those tasks that holds rows writes at least 1 file.
Below, the program sets the count to 2 before grouping 9 orders by region. Each of the 3 input partitions combines its own rows first, then hashes the partial counts into 2 reduce tasks, and the region totals come out of those 2. The board draws the plan as set, with no runtime merging, so the count you set is the task count you see.
spark.conf.set("spark.sql.shuffle.partitions", 2)
orders.groupBy("region").count().collect()| order_id | region |
|---|---|
| 1 | US |
| 2 | EU |
| 3 | US |
| 4 | APAC |
| 5 | EU |
| 6 | LATAM |
| 7 | US |
| 8 | APAC |
| 9 | EU |
Sizing from shuffle bytes
Size the count from the bytes the stage writes to the shuffle, divided by a target per partition. A target of 100 to 200MB keeps each task's slice well inside its memory while leaving enough tasks to occupy every core. When the count exceeds the core count, round it to a multiple of the cores so the last round of tasks fills every slot. Take a cluster with 100 cores and a 150MB target. Job A writes 75GB to the shuffle and job B writes 1.2GB:
- 500
- 8
Job A finishes in 5 rounds of 100 tasks. At the default it would hand each task 375MB, 2.5 times the target. Job B fits in a single round. With runtime partition merging off, the default gives each of its 200 tasks 6MB, and each task spends more time being scheduled than computing. On Spark 3.2 and later that merging is on by default and folds small reduce partitions together after the shuffle runs, so 200 becomes a ceiling and job B runs on fewer tasks. To check a running job, divide the stage's total shuffle read by its task count, both shown for the stage in the Spark UI.
Scoping the setting to one job
Set the count for the job that needs it and leave other jobs at theirs. At launch, --conf spark.sql.shuffle.partitions=500 on spark-submit scopes it to that application. Inside 1 session that runs several queries, call spark.conf.set before each action, since the plan reads the setting when the action triggers planning:
spark.conf.set("spark.sql.shuffle.partitions", 500)
per_user = events.groupBy("event_date", "user_id").count()
per_user.write.parquet("/lake/daily_users")
spark.conf.set("spark.sql.shuffle.partitions", 8)
per_country = events.groupBy("country").count()
per_country.write.parquet("/lake/country_counts")Where 1 step needs a specific count, repartition(n) or repartition(n, col) states it in the code at that step, where a teammate reading the job will see it.
That runtime merging is covered in a later lesson. It only combines partitions, so a setting that is too small stays too small, and a count passed to repartition(n) is kept as written.
Controlling the files a write emits
A write leaves 1 file per final-stage task that holds rows, multiplied by the partition values each task holds when partitionBy is on. A small aggregate written at the default count leaves up to 200 files of a few kilobytes each. These settings change that:
| Before the write | Files | Extra shuffle | Use when |
|---|---|---|---|
| Nothing | 1 per final-stage task holding rows | No | That task count already gives the file size you want |
| coalesce(16) | Up to 16 | No | The output is small and the merged stage's work is light |
| repartition(16) | 16, near equal in size | Yes | The final stage is heavy or its partitions are lopsided |
| repartition("click_date") plus partitionBy("click_date") | 1 per date | Yes | Writing many dates in one run |
| option("maxRecordsPerFile", n) on the writer | A new file every n rows | No | A single task would otherwise write an oversized file |
coalesce(16) after a groupBy merges the reduce side, so the aggregation itself then runs on 16 tasks. That suits a light aggregate and slows a heavy one, where repartition(16) keeps the aggregation at full width and pays 1 extra shuffle for 16 even files:
- events (source, Partitioned files)
- Scan and partial aggregate (transform, Stage 1); parallelism: 400 partitions
- Final aggregate (transform, Stage 2); parallelism: 200 partitions
- repartition(16) and write (transform, Stage 3); parallelism: 16 partitions
- daily_rollup (storage, 16 files)
- events -> Scan and partial aggregate
- Scan and partial aggregate -> Final aggregate
- Final aggregate -> repartition(16) and write
- repartition(16) and write -> daily_rollup
A nightly job's aggregation stage reads 60GB of shuffle data on a cluster with 40 cores. Which shuffle partition count fits a 150MB target per partition?
- Set it to 400, giving each task about 150MB in 10 full rounds on 40 cores
- Keep the default 200, which already gives each of the 40 cores 5 tasks
- Set it to 40, 1 per core, so the stage finishes in a single round
- Set it to 4,000, since more partitions always shorten the stage
In code, the count is set on the session before the action that runs the aggregation. With spark.sql.shuffle.partitions at 8, explain() on a groupBy shows Exchange hashpartitioning(key, 8) between the partial and final aggregates, and the reduce stage is planned as 8 tasks. Set it after the action and that job keeps whatever count was in force when it planned.
About this lesson
This section is part of the Partitioning: Intermediate lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.