Skip to main content
PySpark advanced Lesson 9 of 10

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:

  1. explain() — count Exchange nodes. Can any join be broadcast?
  2. Stage timings in the UI — which stage dominates? One slow stage is a different problem from many.
  3. Task duration spread within that stage — max far above median means skew.
  4. Spill columns — non-zero disk spill means partitions are too large.
  5. Input file count — thousands of small files means fix the writer, not the reader.
  6. 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.

Frequently Asked Questions

What should I tune first?
Nothing. Read the query plan and the stage timings first, because the fix is almost always structural — a missing broadcast, a skewed key, filtering after a UDF — rather than a config value. Changing spark.executor.memory before you know which stage is slow is guessing.
Is spark.sql.shuffle.partitions still worth setting?
Less than it used to be. Adaptive Query Execution coalesces shuffle partitions at runtime in Spark 3, so the old advice to tune it per job is mostly obsolete. Set it when AQE is disabled, or when you know the output size precisely.
How do I know if a job is spilling?
The Spark UI's stage detail shows Spill (Memory) and Spill (Disk) columns. Non-zero disk spill means a task's working set exceeded its share of executor memory and was written out. Fix it with more partitions so each task handles less, before reaching for more memory.
Why are small files a problem?
Each file becomes at least one task, and task scheduling overhead is a few milliseconds regardless of how little data the task reads. Ten thousand 1 MB files take far longer than a hundred 100 MB files holding the same data, and the metadata listing alone can dominate a cloud-storage read.