DEV Community

Karan Pulyala
Karan Pulyala

Posted on

Diagnosing and Fixing Data Skew in PySpark on Azure Databricks (2026 Edition)

A familiar failure: the stage hits 99% and stalls. 199 tasks are done, one is still running, and adding nodes doesn't help, because the bottleneck is one partition, not the cluster.

That's data skew. In this post I'll go through how I diagnose it and how I fix it on Azure Databricks as of 2026, roughly in the order I'd try each fix. Every snippet runs on Databricks Free Edition, which is serverless-only, so I'll flag the places where serverless behaves differently.

Not affiliated with or endorsed by Databricks. Config names and defaults are from the Apache Spark and Databricks docs as of October 2026. Check them against your runtime version.

What skew actually is

A shuffle sends every row with the same key to the same partition. That's what makes joins and groupBy work. If one key holds 30% of your rows, one task gets 30% of the work, and the stage waits for it.

The usual culprits:

  • NULL or placeholder keys ('UNKNOWN', -1, ''). In my experience this is the cause more often than anything else.
  • A handful of "whale" entities: one huge customer, one default store, one bot user agent.
  • Time skew: one day (Black Friday, a backfill date) is 50× the others.
  • File/read skew, which isn't about keys at all. A few huge or unsplittable files (gzipped CSV, for example) make the first stage lopsided.

Step 1: Build a skewed dataset to practice on

Paste this into a Free Edition notebook. It creates 20M orders: ~30% with a NULL customer, ~20% for one whale customer (42):

from pyspark.sql import functions as F

N = 20_000_000
orders = (
    spark.range(N)
    .withColumn("r", F.rand(seed=7))
    .withColumn(
        "customer_id",
        F.when(F.col("r") < 0.30, F.lit(None).cast("bigint"))   # ~30% NULL
         .when(F.col("r") < 0.50, F.lit(42).cast("bigint"))     # ~20% whale
         .otherwise((F.rand(seed=11) * 1_000_000).cast("bigint"))
    )
    .withColumn("amount", F.round(F.rand(seed=3) * 100, 2))
    .drop("r")
)

customers = (
    spark.range(1_000_000).withColumnRenamed("id", "customer_id")
    .withColumn("segment", F.when(F.col("customer_id") % 10 == 0, "enterprise").otherwise("smb"))
)
Enter fullscreen mode Exit fullscreen mode

Step 2: Confirm it. Don't guess.

Key distribution. This is the cheapest check and the one I always run first:

total = orders.count()
(orders.groupBy("customer_id").count()
    .withColumn("pct_of_rows", F.round(100 * F.col("count") / total, 2))
    .orderBy(F.desc("count"))
    .show(10))
Enter fullscreen mode Exit fullscreen mode

Double-digit percentages at the top mean you've found it.

Partition-level distribution. This shows what a shuffle on that key does to your tasks:

parts = (orders.repartition(200, "customer_id")
    .groupBy(F.spark_partition_id().alias("pid")).count())

parts.agg(
    F.max("count").alias("max_rows"),
    F.percentile_approx("count", 0.5).alias("median_rows"),
).withColumn("max_over_median", F.col("max_rows") / F.col("median_rows")).show()
Enter fullscreen mode Exit fullscreen mode

A max/median ratio in the tens or hundreds means one task will do most of the work, no matter how big your cluster is.

In the Spark UI (classic compute). Open the slow stage and look at Summary Metrics. Compare Max with Median for Duration, Shuffle Read Size / Records, and Spill (Disk). A healthy stage has a fairly flat distribution. A skewed stage has a median of seconds and a max of minutes, and the Event Timeline shows one long bar.

On serverless (including Free Edition) there's no Spark UI. Click See performance under the cell and open the query profile. Look for the operator with a disproportionate share of time and rows. It's operator-level, not a per-task histogram, so the key-distribution query above is even more useful there.

Step 3: Let AQE do its job (and know when it can't)

Adaptive Query Execution re-plans at stage boundaries using real shuffle statistics. It's on by default in Spark 3.2+ and on Databricks. Its skew-join feature splits an oversized partition into smaller ones and replicates the matching partition on the other side. A partition counts as skewed only when both of these are true:

  • it's larger than spark.sql.adaptive.skewJoin.skewedPartitionFactor (default 5) × the median partition size, and
  • it's larger than spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes (default 256MB).

Split sizes aim at spark.sql.adaptive.advisoryPartitionSizeInBytes (default 64MB). When it works, the final plan shows a SortMergeJoin with isSkew=true.

joined = orders.join(customers, "customer_id", "left")
joined.write.mode("overwrite").format("noop").save()   # run it without writing anything
Enter fullscreen mode Exit fullscreen mode

noop runs the whole job without writing output. Then check the plan in the query profile (serverless) or the SQL tab (classic).

What AQE won't handle, according to the docs:

  • Broadcast joins. Skew handling only applies to sort-merge and shuffle-hash joins.
  • Some join sides. In a LEFT OUTER join only the left side can be split. FULL OUTER gets no skew handling at all.
  • Partitions under 256MB. If your "skew" is 200MB against a 2MB median, the threshold blocks it. On classic compute you can lower the thresholds. On serverless you can't, because only a short allowlist of Spark confs (like spark.sql.shuffle.partitions, which defaults to auto there) can be set.
  • Scan-stage skew, aggregations and window functions. The skew-join rule is about joins.

One honest caveat about the demo: on a small Free Edition dataset the skewed partition may not cross 256MB, so you may not see isSkew=true. That's the threshold working as designed, not a broken demo.

Step 4: Broadcast the small side

If the other side fits in memory, a broadcast hash join removes the shuffle completely. With no shuffle there are no hot partitions, so skew stops mattering:

joined = orders.join(F.broadcast(customers), "customer_id", "left")
Enter fullscreen mode Exit fullscreen mode

Spark auto-broadcasts below spark.sql.autoBroadcastJoinThreshold (10MB in OSS). Databricks AQE can also switch to a broadcast at runtime below spark.databricks.adaptive.autoBroadcastJoinThreshold (documented at 30MB). Databricks' own guidance is still to hint the broadcast when you know a side is small, because a planned broadcast beats one that AQE only chooses after both sides have already shuffled.

Two cautions: compressed Parquet can expand a lot in memory, and you can't broadcast the preserved side of an outer join (the left side of a LEFT JOIN).

Step 5: Treat hot keys separately (especially NULLs)

For NULL keys this is free. NULLs never match in an equi-join, so pushing them through the shuffle is pure waste:

is_null = F.col("customer_id").isNull()

matched = orders.filter(~is_null).join(customers, "customer_id", "left")
nulls   = orders.filter(is_null)            # no partner rows to find anyway

result = matched.unionByName(nulls, allowMissingColumns=True)
Enter fullscreen mode Exit fullscreen mode

For an inner join, just filter the NULLs out. The same split works for a known whale key: send it through a broadcast path and the rest through the normal join.

Step 6: Salting, the targeted version

Salting spreads one key across N partitions by adding a random suffix on the big side and replicating the small side N times. The common version salts every key, which multiplies the whole dimension table by N. I prefer to salt only the hot keys:

SALT = 16
hot_keys = [42]  # take these from your Step 2 profiling, not from a hard-coded guess

orders_s = orders.withColumn(
    "salt",
    F.when(F.col("customer_id").isin(hot_keys), (F.rand(seed=1) * SALT).cast("int"))
     .otherwise(F.lit(0))
)

customers_s = customers.withColumn(
    "salt",
    F.explode(
        F.when(F.col("customer_id").isin(hot_keys), F.sequence(F.lit(0), F.lit(SALT - 1)))
         .otherwise(F.array(F.lit(0)))
    )
)

salted = orders_s.join(customers_s, ["customer_id", "salt"], "left").drop("salt")
Enter fullscreen mode Exit fullscreen mode

Only the whale's dimension row gets replicated (16 rows, not 16 million). Size SALT so each slice of the hot key is close to a normal partition.

For aggregations, be selective too. Spark already does partial aggregation on the map side, so plain sum/count on a skewed key often hurts less than you'd expect. Salting pays off for aggregations that can't be combined that way, like collect_list, exact countDistinct and percentiles, and for window functions partitioned by a hot key. Two-phase aggregation looks like this:

partial = (orders.withColumn("salt", (F.rand(seed=5) * SALT).cast("int"))
    .groupBy("customer_id", "salt")
    .agg(F.sum("amount").alias("s"), F.count("*").alias("c")))

final = partial.groupBy("customer_id").agg(F.sum("s").alias("total_amount"),
                                           F.sum("c").alias("order_count"))
Enter fullscreen mode Exit fullscreen mode

Step 7: Repartitioning (and its trap)

repartition(400, "customer_id") does not fix key skew. Every row for key 42 still hashes to the same partition, and you've paid for a shuffle. What helps:

  • repartition(400) with no key (round-robin) when you just need even work for narrow, row-by-row processing.
  • The REBALANCE hint before a write. AQE can split skewed partitions in a rebalance (spark.sql.adaptive.optimizeSkewsInRebalancePartitions.enabled, default true), which avoids one giant output file:
orders.hint("rebalance").write.mode("overwrite").saveAsTable("workspace.default.orders_demo")
Enter fullscreen mode Exit fullscreen mode

(workspace.default is the default catalog and schema on Free Edition. Use your own elsewhere.)

Step 8: Fix the layout when the skew is on read

If the first stage is the lopsided one, AQE can't help, because there's no shuffle statistic to re-plan from yet. That's a file-layout problem: a few giant files, thousands of tiny ones, unsplittable compression, or Hive-style partitions where one date= folder is huge.

On Delta, the current recommendation for new tables is liquid clustering. Databricks lists "tables with heavy data skew" among its target use cases:

spark.sql("""
  CREATE OR REPLACE TABLE workspace.default.orders_clustered
  CLUSTER BY (customer_id)
  AS SELECT * FROM workspace.default.orders_demo
""")
spark.sql("OPTIMIZE workspace.default.orders_clustered")
Enter fullscreen mode Exit fullscreen mode

Be clear about what this buys you: evenly sized files and better data skipping. It does not remove skew from a later join on customer_id. The shuffle still sends key 42 to one partition. With predictive optimization on Unity Catalog managed tables, CLUSTER BY AUTO lets Databricks choose the keys.

How to explain this in an interview

Interviewers are checking whether you have a process, not a definition. A structure that works:

  1. Symptom: "A stage stuck at 99% with one long task. Max task duration far above the median, a big shuffle read on that task, often spill."
  2. Confirm: "I profile the key distribution, and the first thing I check is NULL and default keys."
  3. Fix in order of cost: "Check that AQE skew join is applicable (join type, thresholds), broadcast the small side if I can, route NULL or hot keys around the join, and use targeted salting as a last resort."
  4. Know the limits: "AQE only handles skew in shuffle joins. It won't touch broadcast joins or scan-stage skew, and on serverless I can't tune its thresholds."
  5. Close the loop: "Then I verify. I check for isSkew=true in the plan, compare max against median again, and add a skew check to the pipeline so a new whale customer doesn't bring it back."

Point 5 is where you stand out. Most candidates stop at "I'd use salting."

Checklist

  • [ ] Max vs. median task duration and shuffle read (Spark UI), or the slowest operator (query profile)
  • [ ] Top-10 key distribution, including NULL and placeholder values
  • [ ] Is it a join, an aggregation, a window function, or the scan stage?
  • [ ] AQE on and applicable: shuffle join, a join type that allows splitting that side, partitions over the thresholds
  • [ ] Small side broadcastable? Hint it.
  • [ ] NULLs filtered or routed around the join
  • [ ] Salt only the hot keys, sized from real numbers
  • [ ] Scan-stage skew → fix the files (compaction, liquid clustering), not the join
  • [ ] Re-measure, then add monitoring

If you're prepping for Databricks interviews, I put 10 questions like this with model answers into a free sample: https://pulyala.gumroad.com/l/vonemy (the full 61-question pack is here: https://pulyala.gumroad.com/l/ivjsaf).

Not affiliated with Databricks. Databricks and Azure Databricks are trademarks of their respective owners.

Top comments (0)