When One Key Has Everyone

Nine of every ten Daily Grind orders are large. Group by size on a cluster and the shuffle sends nine-tenths of the rows to one place. Ninety-nine tasks finish in a minute; one runs for an hour; the job is "almost done" the whole time. That is skew, and it is the commonest reason a Spark job is slow when the code is fine.

Watch it happen

A thousand orders, with every tenth one small:

orders = spark.createDataFrame(
    [(i, "large" if i % 10 else "small") for i in range(1, 1001)],
    "order_id INT, size STRING")

result = orders.groupBy("size").count()

Now shuffle them by size into four pieces, the way a groupBy("size") would, and count the rows in each:

by_key = (orders.repartition(4, "size")
                .withColumn("partition", F.spark_partition_id()))
result = by_key.groupBy("partition").count()

Two of four pieces are empty, and one holds 90% of the rows. On a cluster, that is the task that runs for an hour.

Salting

Add a second key that spreads the hot one. A salt is a small number derived from something that varies, here a hash of the order id, so large becomes large·0 … large·3 and the shuffle has four places to send it:

salted = (orders
    .withColumn("salt", F.pmod(F.hash("order_id"), F.lit(4)))
    .repartition(4, "size", "salt")
    .withColumn("partition", F.spark_partition_id()))
result = salted.groupBy("partition").count()

Not even. A hash is not a dealer, and 508 is still twice 240. But no piece holds 900, and on a cluster the hour becomes a few minutes. The cost is a second pass: an aggregate over a salted key has to be done per salt and then combined, and a salted join has to explode the small side so it carries every salt value (the Long or Wide lesson's explode), one copy per salt, before the keys match.

What Spark does on its own

Since Spark 3.2, adaptive query execution is on by default, and one of its three jobs is this one: "This feature dynamically handles skew in sort-merge join by splitting (and replicating if needed) skewed tasks into roughly evenly sized tasks." That covers joins. A skewed groupBy is still yours to salt.

The symptom, and the fix

The Stages tab shows a stage with every task finished but one, and that task's input is a hundred times the median. Find the key: groupBy(key).count() sorted descending says which value is hot. Then: for a join, let adaptive query execution split it, or broadcast the small side so there is no shuffle at all (How a Join Is Planned, next); for an aggregate, salt the key and aggregate twice; and ask whether the hot value is a null or a placeholder that should have been filtered out first, which it often is.

Self-check

groupBy("country") is slow, and the Stages tab shows one task with 80% of the input. What is the first thing to check?

Which country, with a count per key, and whether the hot value is a real country or an unknown or null placeholder that should not be grouped at all. Salting a key that should have been filtered is work spent on rows nobody wanted.

Nothing to grade here. Next, How a Join Is Planned: the four jobs one join line can become, and the one with no shuffle at all.