TopolearnPartition counts and shuffles determine parallelism, memory per task and most of a job's runtime.
A DataFrame is split into partitions, and each partition is processed by one task. Partition count sets the maximum parallelism: 4 partitions on a 100-core cluster use only 4 cores.
Where partitions come from:
- Reading files: roughly one partition per 128 MB of input (
spark.sql.files.maxPartitionBytes). - After a shuffle (from groupBy, join, distinct, orderBy, window):
spark.sql.shuffle.partitions, default 200, which adaptive query execution then coalesces based on actual data size.
A shuffle writes each task's output to local disk, bucketed by key, and the next stage's tasks fetch their buckets over the network. That costs disk, network and serialization, so it's the main expense in most jobs.
Controlling partitions:
df.repartition(n)orrepartition(n, "key"): full shuffle to exactly n partitions, by key if given. Use it to increase parallelism or co-locate keys.df.coalesce(n): merges partitions without a shuffle; only reduces the count. Good before writing to avoid many small files. But coalescing to 1 makes the whole preceding stage run as a single task.df.rdd.getNumPartitions()shows the current count.
A healthy target is tasks of roughly 100 to 200 MB of data and at least two to three times as many partitions as cores.
Example
from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.master("local[2]").getOrCreate()
df = spark.range(0, 1_000_000, numPartitions=8)
print(df.rdd.getNumPartitions()) # 8
spark.conf.set("spark.sql.adaptive.enabled", "false") # show raw settings
spark.conf.set("spark.sql.shuffle.partitions", "50")
grouped = df.groupBy((F.col("id") % 10).alias("k")).count()
print(grouped.rdd.getNumPartitions()) # 50 (set by the shuffle)
print(df.coalesce(2).rdd.getNumPartitions()) # 2 (no shuffle)
by_key = df.repartition(4, (F.col("id") % 4).alias("k"))
print(by_key.rdd.getNumPartitions()) # 4 (full shuffle by key)
sizes = by_key.rdd.glom().map(len).collect()
print(sorted(sizes)) # [0, 250000, 250000, 500000] (hash collisions: 2 keys share one)Remember
- One task per partition: partition count caps parallelism
- Wide operations shuffle; shuffles dominate runtime
repartitionshuffles to any count;coalesceonly merges, without a shuffle
Watch out for
coalesce(1)before writing a big result, turning the last stage into one slow task.