Data Skew in PySpark: How to Find It and Fix It

Data skew is when a few values of a key hold most of the rows, so after a shuffle one partition is far larger than the rest and one task runs long after the others have finished. In Spark 3.2 and later, adaptive query execution (AQE) splits skewed partitions in sort-merge joins on its own, so first check that it is on and that it can see your skew. When it cannot, for instance in an aggregation, salting the key spreads the hot value by hand.

Every example below was run on PySpark 4.2 with 4 million events, 60% of which belong to one customer, guest, and the rest to 100,000 others.

How to tell your job is skewed

In the Spark UI, open the slow stage and sort its tasks by duration. If one or a few tasks take far longer than the median, or read far more shuffle data, the stage is skewed. On Databricks the same shows in the query profile. The cause is usually one of:

• a placeholder value such as guest, unknown, 0 or NULL used for many rows;
• a few very large customers, products or devices;
• a date column where one day holds a backfill.

To find the hot keys, count rows per key:

(orders.groupBy("customer_id").count()
       .orderBy(F.desc("count"))
       .show(5))

+-----------+------+
|customer_id| count|
+-----------+------+
|      guest|599949|
|       c253|   462|

Why it hurts

A join or groupBy shuffles rows so that all rows with the same key land in the same partition. With 8 partitions, our 4 million rows split like this:

All of guest has to go to one partition, so one task does more than 2.5 million rows of work while the others do a few hundred thousand. Adding more partitions does not help: repartition(n, key) still sends every row of one key to the same partition. The job finishes only when that last task does, and the task may also spill to disk or run out of memory.

Fix 1: let AQE split skewed joins

Adaptive query execution re-plans a query while it runs, using the real sizes of the shuffled data. With spark.sql.adaptive.skewJoin.enabled, it finds partitions of a sort-merge join that are much larger than the rest and splits them into smaller tasks, copying the matching rows of the other side to each. Both settings are on by default since Spark 3.2.0, and Databricks turns them on too.

A partition counts as skewed when it is both more than 5 times the median partition size (spark.sql.adaptive.skewJoin.skewedPartitionFactor, default 5.0) and larger than 256 MB (skewedPartitionThresholdInBytes). On our test data, which is smaller than that, we lowered the threshold to 1 MB and the final plan showed the split:

+- SortMergeJoin(skew=true) [customer_id#1], [customer_id#70], Inner
   :- Sort [customer_id#1 ASC NULLS FIRST], false, 0
   :  +- AQEShuffleRead skewed

SortMergeJoin(skew=true) and AQEShuffleRead skewed in df.explain() or the SQL tab of the Spark UI, after the query has run, tell you AQE handled it. If they are missing, check that AQE is on, and whether your hot partitions are below the thresholds. Lowering them makes AQE act on smaller skew:

spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "64MB")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "3")

AQE's skew handling applies to joins, not to aggregations, and not to every join type, so the next fixes still matter.

Fix 2: broadcast the small side

If the other side of the join is small, broadcast it. Each executor gets a full copy, the large table is not shuffled at all, and skew in it no longer matters:

from pyspark.sql import functions as F

joined = orders.join(F.broadcast(customers), "customer_id")

Spark does this by itself when it estimates a side to be under spark.sql.autoBroadcastJoinThreshold, 10 MB by default. The broadcast() hint forces it for a table you know fits in executor memory. Our guide to PySpark joins explains how Spark chooses a join strategy.

Fix 3: salt the key

When both sides are large, or the skew is in an aggregation, spread the hot key by hand. Add a random number, the salt, to the key on the large side, and copy each row of the other side once for every salt value:

SALT = 16

orders_salted = orders.withColumn("salt", (F.rand() * SALT).cast("int"))
customers_salted = customers.crossJoin(
    spark.range(SALT).withColumnRenamed("id", "salt"))

joined = orders_salted.join(customers_salted, ["customer_id", "salt"]).drop("salt")

The rows of guest now go to 16 different partitions. On a million orders, the salted join gave exactly the same row count and per-tier totals as the plain join. The cost is that the small side grows 16 times, so choose the smallest salt that brings the largest partition down to the size of the others, and salt only the hot keys if the other side is large.

For an aggregation, salt and aggregate twice: first by key and salt, then by key alone.

totals = (orders.withColumn("salt", (F.rand() * SALT).cast("int"))
    .groupBy("customer_id", "salt").agg(F.sum("amount").alias("part"))
    .groupBy("customer_id").agg(F.sum("part").alias("total")))

This works for sums, counts, minimums and maximums. An average must be computed as a sum divided by a count, and a distinct count cannot be split this way.

Fix 4: deal with the hot value itself

Often the skewed key is not real data. Rows with customer_id = 'guest' or NULL may not need the join at all:

known = orders.filter("customer_id != 'guest'").join(customers, "customer_id")
guests = orders.filter("customer_id = 'guest'")   # handled on their own

Process them separately and unionByName the results. This is often the simplest fix and the fastest.

What not to rely on

• repartition(n, key) does not fix skew. Every row of one key still lands in one partition, as the table above shows.
• The skew join hint was a Databricks feature. Databricks has retired its documentation and says the hints are not required because AQE handles skew. Open-source Spark ignores the hint without an error.
• More memory keeps a skewed task from failing, but it still runs alone while the rest of the cluster waits.

Tecyfy Takeaway

Find the hot keys first with a groupBy().count(). Then let AQE split skewed joins, which it does by default, and check the plan for skew=true. Broadcast the small side when it fits, filter out placeholder keys such as guest or NULL and handle them on their own, and salt the key when both sides are large or the skew is in an aggregation. Do not count on repartition or extra memory to fix it.