Skip to main content

Module 8 — Tuning: partitions, memory, shuffle

Everything so far has been about writing correct code. This module is about making that same code run in five minutes instead of fifty. The four tuning levers that actually matter on a Spark ML job are the number of partitions, the choice between repartition and coalesce, the shape of the data at the shuffle boundary, and the memory each executor gets. Read the web UI first, tune second.

The partition count is the first thing to check​

spark.sql.shuffle.partitions controls how many partitions a wide transformation (module 3) produces. Its default is 200, chosen for a cluster of moderate size ten years ago and rarely appropriate today.

Two heuristics work in practice. Aim for partitions of 128 MB to 512 MB after the shuffle: too small wastes cores in scheduling, too large starves parallelism and risks running out of memory. And aim for a partition count that is two to four times the total number of cores in the cluster, so a slow task does not idle the fastest ones.

spark.conf.set("spark.sql.shuffle.partitions", 400)

On the flight-delays project running on a 64-core cluster, moving from the default 200 to 400 typically cuts a groupBy-heavy job by 30 %. On a laptop with 8 cores, 32 is closer to right.

repartition vs coalesce​

Both change the number of partitions of a DataFrame, but they do very different things.

repartition(n) shuffles the data into exactly n partitions with a hash of the specified column, or round-robin if no column is given. It is expensive — every row moves — but the result is well-balanced.

coalesce(n) merges existing partitions without a shuffle, only ever reducing the count. It is nearly free because no data crosses the network, but the result can be uneven if the input was already unbalanced.

Use repartition(n, "carrier") before a groupBy("carrier") when you want a specific parallelism, especially when data is skewed. Use coalesce(n) before writing to reduce the number of output files, once the computation is done. Doing the opposite — coalesce before an aggregation, repartition before a write — is a classic reason for slow jobs.

Key skew is the quiet killer​

A groupBy or a join shuffles rows so that all rows with the same key land on the same task. When one key holds 50 % of the rows — a single airline that dominates the data, an "UNKNOWN" bucket that swallowed every failed lookup — one task ends up with 50 % of the work while every other task finishes and waits.

The Spark UI shows this exactly. Open the stage's task list and sort by duration. If the ninety-ninth-percentile task is ten times slower than the median, you have skew.

Two techniques handle it. The first, and the one to try first, is salting: pre-append a small random integer to the skewed key so that the heavy key is spread across a few tasks, then aggregate in two passes.

salted = df.withColumn("key_salted",
    F.concat_ws("-", F.col("carrier"), (F.rand() * 8).cast("int")))
partial = salted.groupBy("key_salted").agg(F.avg("arrival_delay").alias("part"))
# then aggregate away the salt on the small result

The second is adaptive query execution, on by default since Spark 3.0. spark.sql.adaptive.enabled=true and spark.sql.adaptive.skewJoin.enabled=true let the optimizer split skewed partitions at runtime. Leave it on. It handles the common cases well and turns manual salting into a fallback.

Executor memory, in three parts​

An executor's memory is divided by Spark into three regions. Reserved memory (300 MB) is used by Spark itself. Execution memory is for shuffles, joins, sorts, and aggregations. Storage memory is for cached DataFrames.

Two failure modes matter. The first, and the one every Spark user meets, is OutOfMemory: Java heap space on an executor during a shuffle: the intermediate data does not fit in execution memory. The response is to give executors more memory (spark.executor.memory=8g), to increase the partition count so each partition is smaller, or to fix the skew that is inflating one partition disproportionately. The second is silent slowness: cached DataFrames evict each other because storage memory is undersized. Cache only what you actively reuse (module 3), and monitor the "Storage" tab of the UI.

Reading the shuffle stage​

Open the Spark UI on a running job, click the stage that says "Exchange" in its description, then look at the summary metrics.

  • Shuffle Read and Shuffle Write should be of the same order across tasks. A 10:1 ratio between the max and the median is skew.
  • Task duration — see above.
  • Spill (Memory) and Spill (Disk) columns non-zero mean execution memory was tight and Spark wrote intermediate data to disk. Small spill is fine; multi-gigabyte spill is a call for more memory or fewer rows per partition.

This diagnostic loop — read the UI, form a hypothesis, change one parameter, re-run — replaces every attempt to tune Spark by reading configuration guides. The configuration reference is a menu; the UI is the diagnosis.

collect() is not a debugging tool at this scale

.collect() brings every row back to the driver. On a job that will run on tens of millions of rows this is a driver crash waiting to happen — the JVM heap on the driver is a fraction of what the whole cluster has. Debug with .show(n), .limit(n).toPandas() or a small sample; use collect() only on aggregates known to be tiny.

Summary​

  • Set spark.sql.shuffle.partitions to 2 to 4 times the total cores and aim for 128–512 MB partitions.
  • repartition shuffles and balances, coalesce merges without shuffle; use coalesce before writing, repartition before wide operations.
  • Key skew is the top cause of slow jobs; enable adaptive query execution and reach for salting on hot keys.
  • Read the Spark UI — shuffle read/write balance, task duration percentiles, spill — before you touch executor memory or partition counts.

Next module: writing the results out so someone else's job can read them.