Optimizing PySpark Jobs

When to use

• A Spark/PySpark job is slow, spills to disk, or OOMs (driver or executor).
• One or a few tasks straggle while the rest finish (skew).
• Huge shuffles, wide stages, or exploding output.
• Do NOT use for pure SQL warehouse tuning (use optimizing-sql-queries).

Workflow

- [ ] Read the Spark UI: find the slow stage and its shuffle/skew
- [ ] Confirm the cause: skew, too many/few partitions, wide shuffle, or driver pull
- [ ] Fix joins (broadcast small side; salt skewed keys)
- [ ] Right-size partitions; enable AQE
- [ ] Re-run and compare stage time/shuffle bytes

1. Read the Spark UI (Stages/SQL tab). Find the stage dominating wall-clock; look at shuffle read/write and the task-duration distribution (a long tail = skew).
2. Diagnose the dominant cause before changing config.
3. Fix joins first — broadcast the small side; handle skewed keys.
4. Right-size partitions and let Adaptive Query Execution coalesce them.
5. Re-measure in the UI; confirm shuffle bytes / stage time dropped.

Patterns

Broadcast the small side to avoid a shuffle join:

from pyspark.sql.functions import broadcast
fact.join(broadcast(small_dim), "dim_id")

Enable Adaptive Query Execution (coalesces partitions, converts to broadcast, handles skew automatically):

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

Salt a skewed join key when one key value dominates (AQE off or extreme skew): add a random suffix to the hot key on both sides so it spreads across tasks.

Right-size partitions — aim for ~128MB per partition. Use repartition(n) to increase parallelism before a wide operation; use coalesce(n) to reduce small output files without a shuffle.

Cache only reused DataFrames, and unpersist when done:

df.cache(); df.count()   # materialize once if used multiple times
...
df.unpersist()

Common pitfalls

• collect() / toPandas() on large data — pulls everything to the driver and OOMs it. Aggregate first, or write to storage.
• Python UDFs in hot paths — serialize row-by-row; use built-in/SQL functions, or pandas/vectorized UDFs when unavoidable.
• Default 200 shuffle partitions on big or tiny data — too few causes spills, too many causes overhead; enable AQE or set spark.sql.shuffle.partitions to fit the data.
• repartition() when coalesce() suffices — repartition forces a full shuffle; use coalesce to only shrink partition count.
• Ignoring skew — a few straggler tasks dominate runtime; salt or use AQE skew join.
• Caching everything — wastes memory and forces eviction; cache only reused frames.
• Many small output files — coalesce before write or use adaptive file sizing.