AWS Training
Modules Listen Certification
0:00 0:00

← PySpark

Starts this lesson and continues through 6 more to the end of the course.

Partitions, skew, and the shuffle you caused

The symptom that defines this lesson

199 tasks finish in 20 seconds. One task runs for 40 minutes. Then it fails with an out-of-memory error.

That's skew, and it is the single most common reason a Spark job is slow rather than merely big. Adding executors does not help — 200 more idle machines watching one task is not faster.

Engine mechanics are elsewhere. How the shuffle is written and read, how join strategies are chosen, how adaptive query execution intervenes — module E7 of the sibling db-training repo. This lesson is the PySpark-level diagnosis and the code you actually change.

Two kinds of partition — don't conflate them

What it is Controlled by
Spark partitions Units of parallelism in memory. One task per partition. repartition, coalesce, shuffle partition count, input splits
Storage partitions Directory structure in S3, e.g. dt=2026-08-12/ partitionBy on write

They interact — reading a partitioned dataset with a matching filter prunes files, which changes how many Spark partitions you start with — but they are different things and people mix them up constantly in interviews.

A useful phrasing: storage partitions decide what you read; Spark partitions decide how you process it.

Diagnosing skew

Before tuning anything, look at the data:

# 1. How skewed is the join/group key?
(df.groupBy("customer_id").count()
   .orderBy(desc("count"))
   .show(20))

# 2. How lumpy are the current partitions?
from pyspark.sql.functions import spark_partition_id
(df.groupBy(spark_partition_id().alias("pid")).count()
   .orderBy(desc("count"))
   .show(20))

The second one is underused and tells you directly whether your parallelism is real. If one partition holds 60% of the rows, you have one task doing 60% of the work regardless of cluster size.

In the Spark UI, the tell is the task duration distribution for a stage: compare max against median. A max that is orders of magnitude above the median is skew, not slowness.

Causes, in order of frequency

1. A dominant key. NULL, 0, -1, "UNKNOWN", or a genuine whale customer. Nulls are the classic — every unmatched row gets the same key and lands in one partition.

2. Joining on a low-cardinality column. Joining on country when 70% of rows are one country gives you a partition holding 70% of the data.

3. Unsplittable input files. A single large GZIP file cannot be split, so it becomes one partition and one task no matter what you configure. This is a genuinely common cause of "why is my job single-threaded".

4. groupBy on something enormous where one group dominates — see the Series-to-Scalar and applyInPandas memory constraints from lesson 3, which turn skew into an OOM rather than a delay.

Fixes

Filter out the junk key first. Often the entire fix:

# Nulls aren't joinable anyway — don't drag them through the shuffle
df.filter(col("customer_id").isNotNull()).join(dim, "customer_id")

Broadcast the small side. If one side fits in executor memory, there is no shuffle at all and skew stops mattering:

from pyspark.sql.functions import broadcast
df.join(broadcast(small_dim), "customer_id")

⚠️ Broadcasting something too large moves the failure to the driver, which must collect it first. Broadcast deliberately, based on the actual size, not hopefully.

Salt the hot key. The standard technique when you can't avoid the join:

from pyspark.sql.functions import rand, floor, lit, explode, array

N = 10
# Spread the hot side across N sub-keys
left  = df.withColumn("salt", floor(rand() * N))
# Replicate the small side N times, once per salt value
right = dim.withColumn("salt", explode(array(*[lit(i) for i in range(N)])))

out = left.join(right, ["customer_id", "salt"]).drop("salt")

The cost is real: you have multiplied the right side by N. Salt only the keys that are actually hot, and only when broadcasting isn't available.

Repartition on a better key, or on a composite key with higher cardinality.

Let adaptive query execution help. Modern Spark can detect and split skewed partitions at runtime. ⚠️ I have not verified the specific configuration keys and defaults for AQE skew handling in your Spark version — check the docs for the version your Glue/EMR release pins, and confirm AQE is actually enabled there rather than assuming it.

repartition vs coalesce

Shuffle? Can increase partitions? Use for
repartition(n) Yes — full shuffle Yes Rebalancing, fixing skew, increasing parallelism
repartition("col") Yes — Co-locating by key before a write or join
coalesce(n) No No — reduce only Cheaply reducing partition count before a write

⚠️ coalesce has a non-obvious side effect. Because it avoids a shuffle, it merges partitions by narrowing the upstream parallelism too. df.map(...).coalesce(1) doesn't run the map with full parallelism and then merge — it can run the whole upstream computation with a single task. If you need one output file and full parallelism upstream, use repartition(1) and accept the shuffle.

That distinction is a frequent interview question and the "no shuffle, so it's always cheaper" answer is the wrong one.

The small files problem

Writing with high parallelism produces one file per partition. 2,000 partitions → 2,000 files. On S3 with a downstream reader, thousands of small files are slow to list and slow to open.

# Target a sensible file size on write
df.repartition(50).write.parquet(path)

# Or, when writing partitioned output, control files per storage partition
df.repartition("dt").write.partitionBy("dt").parquet(path)

⚠️ I'm deliberately not quoting an optimal file size. Recommendations vary by engine and by reader, and the right number depends on your query pattern. The principle: many tiny files and a few huge files are both bad; aim for a middle you've measured. Module DL2 covers this properly for lakes.

Check yourself

  1. 199 tasks fast, one slow, then OOM. What is it, and does a bigger cluster help?
  2. Difference between a Spark partition and a storage partition?
  3. Why can a single GZIP file make your job effectively single-threaded?
  4. What does salting cost you?
  5. Why isn't coalesce(1) always cheaper than repartition(1)?
Answers
  1. Skew — one partition holds far more data than the rest. A bigger cluster doesn't help; the extra executors sit idle while one task does the work. Fix the distribution instead.
  2. Spark partitions are units of in-memory parallelism (one task each); storage partitions are the directory layout on S3 created by partitionBy. Storage partitions decide what you read; Spark partitions decide how you process it.
  3. GZIP isn't splittable, so the whole file must be read by one task. Parallelism configuration can't override that.
  4. It multiplies the replicated side by the salt factor N — real extra shuffle and compute. Salt only genuinely hot keys, and only when you can't broadcast.
  5. Because coalesce avoids the shuffle by narrowing upstream parallelism as well — it can force the whole upstream computation into a single task. repartition(1) pays for a shuffle but keeps upstream parallelism.

Teaching this section

← PreviousUDFs, Arrow, and pandas UDFsNext →Production PySpark — testing, structure, and AWS