Spark partitions: how many to read, shuffle partitions, repartition vs coalesce

Part 1 ran a job whose first stage had 4 tasks and whose shuffle was planned with 200 partitions, then run with 3. The number of tasks is the number of partitions, and it decides how much work runs at once: too few and cores sit idle, too many and Spark spends its time starting tiny tasks. This part shows where each number comes from.

How many partitions Spark reads

Save this as partitions.py and run it:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[4]").appName("partitions").getOrCreate()
orders = spark.read.parquet("lake/orders")

orders.groupBy(F.spark_partition_id().alias("partition")).count().orderBy("partition").show()

+---------+------+
|partition| count|
+---------+------+
|        0|500000|
|        1|500000|
|        2|500000|
|        3|500000|
+---------+------+

spark_partition_id() gives the partition each row is in, so this counts the rows of each. The 8 files, 22 MB in all, became 4 partitions of 2 files each. Spark gets there in 3 steps:

1. It counts each file as 4 MB bigger than it is, spark.sql.files.openCostInBytes, because opening a file costs time whatever its size. 8 files of about 2.75 MB count as about 7 MB each, 55 MB in all.
2. It aims for at least 1 partition per core. spark.sql.files.minPartitionNum defaults to the parallelism, 4 here, so the target is 55 / 4, about 14 MB a partition.
3. It caps a partition at spark.sql.files.maxPartitionBytes, 128 MB by default, and packs files into each partition until the next would pass the target. 2 files fit, 3 do not.

So small data reads as 1 partition per core, and large data as partitions of 128 MB. Set maxPartitionBytes lower and the read splits finer, up to a limit:

for size in ["128m", "4m", "1m"]:
    spark.conf.set("spark.sql.files.maxPartitionBytes", size)
    print(size, spark.read.parquet("lake/orders").select(F.spark_partition_id()).distinct().count())
spark.conf.unset("spark.sql.files.maxPartitionBytes")

128m 4
4m 8
1m 8

The last line puts the setting back. Spark reads it when a query runs, not when you call spark.read, so leaving it at 1 MB would change every query after it. At 4 MB, each file is its own partition. At 1 MB, the Spark UI shows 24 tasks in the read stage, 3 per file, but only 8 of them hold rows. A Parquet file is stored as row groups, blocks of rows kept together, and a reader can only split a file between them. Each of these files is 1 row group, so 2 of every 3 tasks open the file and find nothing to read.

The same arithmetic explains the small files problem. 10,000 files of 100 KB count as 10,000 × 4.1 MB, so Spark gives them far more partitions, and tasks, than 1 GB of data needs, and each task spends its time opening files. Module 4 compacts small files into large ones.

How many partitions a shuffle makes

After a shuffle, the number of partitions is spark.sql.shuffle.partitions, 200 unless you set it, whatever the size of the data. 3.3 MB of partial sums in 200 partitions is about 16 KB each. With adaptive query execution off, part 1's write ran 200 tasks and made 200 files.

With it on, the default, Spark merges neighbouring small partitions once the shuffle has written and their sizes are known. Its default, parallelismFirst, ignores the 64 MB target size and tries to keep every core busy, but makes no partition smaller than 1 MB, minPartitionSize. 3.3 MB became 3 partitions, so 3 tasks and 3 files.

Merging only works in 1 direction. A shuffle of 2 TB at 200 partitions is 10 GB a partition, and adaptive query execution will not split them, outside a skewed join. So set the number high enough for your largest shuffle and let Spark merge it down where the data is small. On a cluster the number decides how many rounds of tasks the cores run, which this demo shows:

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

ToolSpark shuffle partition calculatorThe data a shuffle moves, your cores and one heavy key in; the partitions, the tasks adaptive query execution makes and the settings out.

On Databricks serverless, spark.sql.shuffle.partitions is one of the few settings you can change, and its default there is auto: Databricks picks the number from the plan and the size of the data.

repartition or coalesce

You can also set the partitions yourself, in 2 ways:

for name, df in [
    ("repartition(16)", orders.repartition(16)),
    ("coalesce(2)", orders.coalesce(2)),
    ("coalesce(16)", orders.coalesce(16)),
]:
    print(name, df.select(F.spark_partition_id()).distinct().count())

repartition(16) 16
coalesce(2) 2
coalesce(16) 4

repartition(16) is a shuffle: its plan has Exchange RoundRobinPartitioning(16), and every row moves, so the 16 partitions are even. Give it columns, repartition("customer_id"), and it groups the rows by them instead. coalesce(2) has no Exchange in its plan, just Coalesce 2: each new partition takes over some of the old ones, so nothing moves. It can only lower the number, which is why coalesce(16) stays at 4.

No shuffle sounds like a bargain, but it has a cost. Because coalesce starts no new stage, the stage it is in runs with the lower number of tasks from its first step. df.coalesce(1).write to get 1 file makes the whole stage, reading and transforming included, run as 1 task on 1 core. If the stage does real work, use repartition(1): the work before the shuffle stays parallel, and only the write is 1 task.

So, as a rule:

• coalesce to write fewer files from data that is already small, such as a daily summary.
• repartition to spread out data that is unevenly split, or to group it by a key before writing.
• Neither, most of the time. With adaptive query execution on, the defaults are usually right, and every repartition is a shuffle you pay for.

These partitions live in memory while a job runs. Partitioning a table on disk by a column, so a query reads only some folders, is a different idea, Data Partitioning and ClusteringData Partitioning and ClusteringTwo ways to lay out a large table so a query reads only the files it needs: partitioning splits it by a column such as the date, clustering keeps similar values in the same files., which module 4 covers.

Part 3 looks at the other 2 things adaptive query execution changes as a job runs: the kind of join, and a partition that is far bigger than the rest.