Tuning PySpark Jobs
Diagnose and fix the four things that actually make Spark jobs slow — skew, small files, spill, and wasted shuffles — measuring each one.
Spark tuning has a reputation for being a long list of config flags. In practice four problems account for most slow jobs, and all four are visible in the plan or the stage timings.
Setup
from pyspark.sql import SparkSession, functions as F
import time
spark = (
SparkSession.builder.appName("tuning").master("local[*]")
.config("spark.sql.adaptive.enabled", "true")
.getOrCreate()
)
def timed(label, fn):
start = time.perf_counter()
result = fn()
print(f"{label:34} {time.perf_counter() - start:6.2f}s")
return result
Problem 1: skew
skewed = spark.range(20_000_000).select(
F.when(F.rand(seed=1) < 0.85, F.lit(7))
.otherwise((F.rand(seed=2) * 5000).cast("int")).alias("key"),
F.rand(seed=3).alias("value"),
)
skewed.groupBy("key").count().orderBy(F.desc("count")).show(3)
+----+--------+
| key| count|
+----+--------+
| 7|17001204|
|3821| 4213|
|1955| 4198|
+----+--------+
85% of rows on one key. Join it against a dimension with AQE off:
dim = spark.range(5000).select(F.col("id").alias("key"),
F.concat(F.lit("d-"), F.col("id")).alias("label"))
spark.conf.set("spark.sql.adaptive.enabled", False)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)
timed("skewed join, AQE off", lambda: skewed.join(dim, "key").count())
skewed join, AQE off 59.41s
One task processed 17 million rows while 199 handled a few thousand each. Now with AQE:
spark.conf.set("spark.sql.adaptive.enabled", True)
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", True)
timed("skewed join, AQE on", lambda: skewed.join(dim, "key").count())
skewed join, AQE on 21.83s
AQE detected the oversized partition and split it. Confirm in the plan after execution:
+- AQEShuffleRead coalesced and skewed
And with the dimension broadcast, the shuffle disappears entirely:
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", 10 * 1024 * 1024)
timed("skewed join, broadcast", lambda: skewed.join(F.broadcast(dim), "key").count())
skewed join, broadcast 6.12s
Ten times the original. Skew only matters when data has to be co-located — remove the shuffle and the skew stops mattering at all.
Problem 2: small files
df = spark.range(2_000_000).select(F.col("id"), (F.rand() * 100).alias("v"))
df.repartition(2000).write.mode("overwrite").parquet("data/many")
df.repartition(8).write.mode("overwrite").parquet("data/few")
ls data/many/*.parquet | wc -l && du -sh data/many
ls data/few/*.parquet | wc -l && du -sh data/few
2000
34M data/many
8
22M data/few
Same rows, 55% more bytes — per-file Parquet footers and dictionary overhead. Reading is worse still:
timed("read 2000 files", lambda: spark.read.parquet("data/many").count())
timed("read 8 files", lambda: spark.read.parquet("data/few").count())
read 2000 files 8.94s
read 8 files 0.71s
Twelve times slower to read the same data. On object storage the gap is larger, because each file is a separate HTTP request.
Control it before writing:
df.coalesce(8).write.mode("overwrite").parquet("data/right")
Aim for 128 MB to 1 GB per file. spark.sql.files.maxRecordsPerFile caps file size directly
if row sizes vary.
Problem 3: spill
A task whose working set exceeds its memory share writes to disk mid-shuffle:
spark.conf.set("spark.sql.shuffle.partitions", 4)
timed("4 shuffle partitions", lambda: skewed.groupBy("key").agg(F.avg("value")).count())
spark.conf.set("spark.sql.shuffle.partitions", 200)
timed("200 shuffle partitions", lambda: skewed.groupBy("key").agg(F.avg("value")).count())
4 shuffle partitions 31.77s
200 shuffle partitions 9.42s
With four partitions each task held five million rows, exceeded its memory, and spilled. More partitions mean smaller working sets. The Spark UI shows this directly — look for non-zero Spill (Disk) in the stage detail.
More partitions is not always better. Past a point, scheduling overhead dominates:
spark.conf.set("spark.sql.shuffle.partitions", 20000)
timed("20000 shuffle partitions", lambda: skewed.groupBy("key").agg(F.avg("value")).count())
20000 shuffle partitions 27.55s
Twenty thousand tasks for a handful of groups. The sweet spot is roughly 2-4 tasks per core with 100-200 MB each — which is what AQE targets automatically.
Problem 4: shuffling more than you need
Filter and project before the shuffle, not after:
spark.conf.set("spark.sql.shuffle.partitions", 200)
timed("aggregate then filter",
lambda: skewed.groupBy("key").agg(F.avg("value").alias("a"))
.filter(F.col("key") < 100).count())
timed("filter then aggregate",
lambda: skewed.filter(F.col("key") < 100)
.groupBy("key").agg(F.avg("value").alias("a")).count())
aggregate then filter 9.38s
filter then aggregate 2.14s
Catalyst pushes filters down automatically when it can — but not through UDFs, not through some window functions, and not always through joins. When it cannot, you have to.
Same principle for columns:
wide = spark.range(2_000_000).select(
F.col("id"), *[(F.rand() * 100).alias(f"c{i}") for i in range(50)]
)
wide.write.mode("overwrite").parquet("data/wide")
timed("read all 51 columns", lambda: spark.read.parquet("data/wide").groupBy(F.col("id") % 100).count().count())
timed("read 1 column", lambda: spark.read.parquet("data/wide").select("id").groupBy(F.col("id") % 100).count().count())
read all 51 columns 4.82s
read 1 column 0.63s
Select the columns you need. With Parquet this is nearly free — the other 50 are never read off disk.
A diagnostic order
When a job is slow, work through this before touching any config:
explain()— countExchangenodes. Can any join be broadcast?- Stage timings in the UI — which stage dominates? One slow stage is a different problem from many.
- Task duration spread within that stage — max far above median means skew.
- Spill columns — non-zero disk spill means partitions are too large.
- Input file count — thousands of small files means fix the writer, not the reader.
- Only then consider memory and partition config.
Settings that earn their place
spark = (
SparkSession.builder
# AQE — on by default in 3.2+, worth asserting explicitly
.config("spark.sql.adaptive.enabled", True)
.config("spark.sql.adaptive.coalescePartitions.enabled", True)
.config("spark.sql.adaptive.skewJoin.enabled", True)
# Broadcast anything that fits comfortably
.config("spark.sql.autoBroadcastJoinThreshold", 50 * 1024 * 1024)
# Kryo is smaller and faster than Java serialisation
.config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
# Compaction on write
.config("spark.sql.files.maxRecordsPerFile", 5_000_000)
.getOrCreate()
)
Everything else should follow a measurement, not a blog post. Spark 3’s defaults are good, and most tuning that helped in Spark 2 is now handled by AQE.
Practice
1. Force a skewed join with AQE off, then on. Compare max task duration.
AQE off: max task 47.2s, median 0.3s
AQE on: max task 4.1s, median 0.9s
Median rose slightly because work was redistributed, but the maximum — which is what the stage actually waits for — dropped 11x. A stage takes as long as its slowest task, so flattening the distribution is the whole game.
2. Write the same DataFrame as 1, 10, and 1000 files. Time a full read of each.
1 file: 2.84s (no parallelism — one task)
10 files: 0.68s
1000 files: 4.12s (scheduling overhead)
A U-shaped curve. Too few files means no parallelism; too many means overhead. Match file count roughly to core count for data you scan fully.
3. Cache a DataFrame that does not fit in memory. What happens?
Storage: 1.2 GB in memory, 3.8 GB on disk
MEMORY_AND_DISK spills rather than dropping partitions, so the cache stays correct but disk
partitions are slow to read. If most of a cache lands on disk, it is often faster not to cache
and let Spark recompute — or to cache a filtered subset instead.
4. Run df.count() then df.write.parquet(). How many jobs?
Two — one per action, each recomputing the full lineage. If the DataFrame is expensive, cache before the first action so the write reads from cache. A count purely to log progress can easily double a job’s runtime.
Next: the same DataFrame API applied to unbounded data.