DEV Community

Cover image for PySpark Performance Tuning: A Checklist You Can Actually Follow
DatanestDigital
DatanestDigital

Posted on

PySpark Performance Tuning: A Checklist You Can Actually Follow

Most slow Spark jobs aren't slow because Spark is slow. The data landed badly, the shuffle is enormous, or the physical plan is doing something expensive nobody asked for. The good news: the usual-suspect list is short, and you can work through it in order.

This is a checklist, not a theory lecture. Every section says what to look at, how to look at it, and when it's worth your time. No fake benchmark tables here — performance claims are deliberate ranges, because your data, cluster, and budget differ. Treat everything as a hypothesis and measure on your own workload.

0. Read the plan before you touch a knob

df.explain("cost") shows the physical plan: which joins Spark picked, where the Exchange nodes (shuffles) are, and what Catalyst decided to do with your query. The Spark UI's SQL tab shows the same thing per stage, with timings and row counts.

If you don't know what Spark is actually doing, every tuning decision is a guess. A SortMergeJoin where you expected a broadcast, or a shuffle between two trivial filters, is the kind of thing the plan shows in ten seconds and profiling takes an hour to find.

1. Partitioning and skew

Two separate problems hide under "partitioning":

  • Too few or too many shuffle partitions. The default is 200 (spark.sql.shuffle.partitions), which is wrong for almost every job — too many for a small cluster, too few for a big one. A common starting point is 2-4 partitions per executor core, but that's a heuristic, not a law. Measure stage times, then adjust.
  • Skew. One task runs for 40 minutes while siblings finish in two — obvious in the UI stage view. Skew usually comes from a hot key (a few users, customers, or device IDs owning most rows). Classic fixes: salt the join key with a random suffix, or let AQE handle it on a recent Spark (section 5).

Also check read-side partitioning: spark.sql.files.maxPartitionBytes defaults to 128 MB, but thousands of tiny source files mean thousands of tiny tasks regardless. Fix the source before tuning the knob.

2. Broadcast joins

If one side of a join is small enough, Spark can copy it to every executor and avoid shuffling the big side entirely. The default threshold is 10 MB (spark.sql.autoBroadcastJoinThreshold); the plan shows BroadcastHashJoin when it applies.

You can hint manually: big_df.join(small_df, "key").hint("broadcast"). In practice this is the cheapest win on dimension-style joins — often an order of magnitude faster than the sort-merge equivalent, depending on your data. It can backfire: raise the threshold too high and the driver ships a huge table to every executor. Keep the broadcast side genuinely small.

3. Avoid the shuffle before you optimize it

A shuffle is not just network traffic — it's serialization, disk writes, and a stage boundary that stops pipelining. Cheapest fix: don't create it.

  • Filter before join. Fewer rows through the shuffle.
  • Aggregate before join. Fewer rows, same answer, if the aggregate is valid at that point.
  • coalesce(n) vs repartition(n). Coalesce reduces partitions without a full shuffle (only safe when reducing); repartition always shuffles. Use coalesce on write, repartition only when you genuinely need to rebalance.
  • reduceByKey over groupByKey. Reduce combines on the map side before anything crosses the network.
  • Bucketing. For tables joined repeatedly on the same key, pre-bucketing can make joins shuffle-free. Plan it early — it's a schema decision.

4. Cache with intent, not habit

df.cache() is lazy — nothing is cached until the first action forces materialization. It pays off only when the same DataFrame is reused across multiple actions or iterations (training loops, iterative algorithms). If you read it once and move on, caching just adds a serialization round-trip.

When you do cache, use the default MEMORY_AND_DISK and unpersist() when done. Caching "everything, just in case" evicts what you actually need later — Spark's cache is a shared pool competing with shuffle buffers and execution memory.

5. Let AQE do the boring stuff

Adaptive Query Execution has been on by default since Spark 3.2 (on 3.0/3.1 you enable it manually via spark.sql.adaptive.enabled). At runtime AQE can:

  • coalesce post-shuffle partitions down to something sane,
  • detect skew and split the hot partitions,
  • convert a sort-merge join into a broadcast join when stats say the table is actually small.

AQE removes a surprising amount of hand-tuning — especially the "guess the shuffle partition count" game from section 1. Not magic, but the closest thing Spark has to a self-tuning default. Check spark.sql.adaptive.coalescePartitions.enabled and spark.sql.adaptive.skewJoin.enabled before building your own workarounds.

6. Diagnose spills before blaming the cluster

Spill means a task ran out of execution memory and wrote intermediate data to disk. The UI shows it as shuffle spill or storage memory spill. Spill is a symptom, not a disease: too much data per task, too little memory, or a fat shuffle.

Fix in this order: shrink the shuffle (sections 2-3), then tune executor memory and spark.memory.fraction, and only then add nodes. Throwing machines at a spill caused by a fat shuffle pays for a symptom.

7. The small-files problem

Write a batch job with 200 shuffle partitions and you get 200 output files. A streaming job writing per micro-batch adds a new batch of tiny files every few minutes. Over weeks, a table becomes thousands of files that slow every read, file listing, and catalog operation. It compounds — fix it early.

Practical fixes: coalesce() on write for batch jobs, and for Delta Lake, OPTIMIZE with ZORDER BY on the columns you filter by. A reasonable target is files in the 64-256 MB range — but like every number in this post, verify against your own workload.

8. Keep Python out of the hot path

Python UDFs are the classic "one function that makes everything crawl" culprit: each row serializes through PySpark's bridge and breaks Catalyst's whole-stage code generation. Replacing a per-row UDF with built-in SQL expressions or a vectorized pandas UDF routinely moves the needle more than any cluster knob.

Prefer selectExpr, when/otherwise, and built-ins. For genuinely custom logic, use vectorized UDFs so serialization is amortized over batches, not paid per row.

The checklist

Run top to bottom; each step is cheaper than the next:

  1. Read the physical plan. Know your shuffles.
  2. Fix partitioning and skew (2-4 partitions per core as a starting point; salt or AQE for hot keys).
  3. Broadcast the small side of joins.
  4. Eliminate shuffles: filter, aggregate, coalesce, bucket.
  5. Cache only reused DataFrames; unpersist when done.
  6. Confirm AQE is on; let it tune partition counts and skew.
  7. Diagnose spills — shrink the shuffle before buying nodes.
  8. Kill per-row Python UDFs; use vectorized or built-ins.

Where to go deeper

The Data Engineering store at Datanest carries the tools behind this checklist: the Spark Performance Masterclass (25+ optimization patterns for Databricks), plus the Spark ETL Framework, PySpark Utils Library, and Delta Lake Patterns for the pipeline around it. If you want the whole set, the Data Engineering Bundle includes all 17 tools — $523 bought separately, yours for $199 (save $324, 62%).

Browse the Data Engineering store

Top comments (0)