Diagnosing Data Skew in PySpark for Azure Databricks AutoBroadcast Joins Salting and AQE
A PySpark job can look healthy for 95% of its runtime, then sit frozen on a handful of tasks while one executor keeps spilling, retrying, and finally dying with an out-of-memory error. That pattern is often not a cluster sizing problem. It is usually data skew.
In Azure Databricks, skew shows up most clearly during shuffle-heavy stages: joins, aggregations, window functions, `dropDuplicates`, and wide transformations. One or a few keys carry far more rows than the rest, so Spark assigns too much data to a small number of partitions. The result is uneven task duration, large shuffle reads, executor memory pressure, and sometimes stranded tasks that keep the job alive long after most executors are idle.
This post focuses on practical PySpark Data Skew Tuning: Salting, Broadcast Joins & AQE for engineers tuning production workloads on Azure Databricks.
How data skew causes shuffle OOM crashes and stranded tasks
Spark distributes work by partition. During a shuffle, records with the same join key or aggregation key usually move to the same reducer partition. If a key is extremely common, one shuffle partition receives a much larger share of data.
A non-skewed shuffle might look like this:
Partition | Shuffle read | Task duration |
0 | 480 MB | 42 s |
1 | 510 MB | 45 s |
2 | 495 MB | 43 s |
3 | 505 MB | 44 s |
A skewed shuffle can look like this:
Partition | Shuffle read | Task duration |
0 | 420 MB | 39 s |
1 | 390 MB | 36 s |
2 | 8.7 GB | 18 min |
3 | 410 MB | 38 s |
That large partition becomes dangerous because one task must process it. Spark can add executors, but it cannot split that one task unless the query plan or runtime feature allows it. The oversized task may:
Spill heavily to local disk.
Spend a long time in garbage collection.
Hit executor memory limits during sort, hash join, or aggregation.
Retry on another executor and fail again.
Hold up downstream stages while other tasks finish early.
This is how skew creates stranded tasks. Most of the stage completes, but one or two tasks keep running. In the Spark UI, the progress bar sits at something like 199/200 tasks for a long time. Cluster utilisation drops, yet the job is still not complete.
A common failure pattern in Databricks logs looks like this:
```text
ExecutorLostFailure
Container killed by YARN for exceeding memory limits
or
java.lang.OutOfMemoryError: Java heap space
or
org.apache.spark.shuffle.FetchFailedException
```
The exact message depends on cluster mode, runtime, and where the pressure occurs. The tuning process is the same: find the skew, then reduce the size of the largest per-task unit of work.
Start diagnosis in the Spark UI
In Azure Databricks, open the job run, then inspect the Spark UI.
Check these screens first:
Stages Look for stages where most tasks finish quickly but the max task duration is much higher than the median.
Tasks Sort by shuffle read size, spill, or duration. One task reading many times more data than the median is a strong skew signal.
SQL/DataFrame tab Find the physical operator causing the shuffle, such as `SortMergeJoin`, `HashAggregate`, `Exchange`, or `Window`.
Executors Check whether failures concentrate on specific stages rather than across the whole job.
You can also run quick profiling from PySpark before tuning the join.
```python
from pyspark.sql import functions as F
fact = spark.table("sales.fact_orders")
key_profile = (
fact.groupBy("customer_id")
.count()
.orderBy(F.desc("count"))
)
key_profile.show(20, truncate=False)
summary = key_profile.agg(
F.max("count").alias("max_key_rows"),
F.expr("percentile_approx(count, 0.50)").alias("median_key_rows"),
F.avg("count").alias("avg_key_rows")
)
summary.show()
```
If the biggest key is hundreds or thousands of times larger than the median, repartitioning alone will not fix the problem. Spark will still send the same hot key to the same reducer partition.
When raising autoBroadcastJoinThreshold helps
Broadcast joins avoid shuffle on the large side by sending a small table to every executor. In a typical fact-to-dimension join, this can remove the shuffle that exposes skew.
Spark controls automatic broadcast join selection with:
```text
spark.sql.autoBroadcastJoinThreshold
```
If a table is smaller than this threshold, Spark may choose a broadcast hash join instead of a shuffle join. Apache Spark’s common default is small, and Databricks Runtime settings can vary, so check the active value rather than assuming it.
```python
spark.conf.get("spark.sql.autoBroadcastJoinThreshold")
```
Raise the threshold when the dimension table is slightly too large
Raising `autoBroadcastJoinThreshold` is useful when all of these are true:
One side of the join is clearly small, usually a dimension or lookup table.
The small side is only slightly above the current threshold.
The small side fits safely in executor memory after serialisation.
The join currently uses `SortMergeJoin` and causes a large shuffle.
The row count and table size are stable enough for production use.
For example, suppose a customer dimension table is larger than the threshold, while the fact table has a skewed `customer_id`. A shuffle join will move fact rows by `customer_id`, pushing the hot customer into one reducer. A broadcast join sends the dimension to executors, letting each fact partition join locally.
Set the threshold at the session level for testing:
```python
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", 100 1024 1024) # 100 MB
fact = spark.table("sales.fact_orders")
dim_customer = spark.table("sales.dim_customer")
joined = fact.join(dim_customer, on="customer_id", how="left")
joined.explain("formatted")
```
Look for `BroadcastHashJoin` in the physical plan.
```text
BroadcastHashJoin [customer_id], [customer_id], LeftOuter, BuildRight
```
You can also use an explicit broadcast hint when testing a specific join:
```python
from pyspark.sql.functions import broadcast
joined = fact.join(
broadcast(dim_customer),
on="customer_id",
how="left"
)
joined.explain("formatted")
```
Do not raise it blindly
Broadcasting is not free. Every executor needs memory to hold the broadcasted relation. If the table has wide rows, nested columns, or poor compression, a threshold that looks safe from storage size may still hurt executor memory.
Avoid raising the threshold when:
Both sides are large fact tables.
The small table grows unpredictably.
Executors already run close to memory limits.
Many broadcast joins run in the same query.
The join condition is not selective and creates row explosion.
A good Databricks workflow is to compare plans before and after the change, then test with production-like data. Watch executor memory, broadcast time, and task duration. If broadcast removes the skewed shuffle stage, it is often the cleanest fix.
How to implement key salting on skewed fact tables
Broadcast joins work only when one side is small enough. If the skewed join is between a very large fact table and another sizeable table, you need to split hot keys across multiple partitions. Key salting does exactly that.
Salting adds an extra value to the join key. Instead of all rows for a hot key going to one reducer, Spark spreads them across several salted keys.
For a skewed fact table, the usual pattern is:
Add a salt column to the skewed fact table.
Expand the matching dimension or right-side table for the hot keys across the same salt range.
Join on the original key plus the salt.
Detect the hot keys first
Do not salt every key unless needed. Start by identifying the keys that cause skew.
```python
from pyspark.sql import functions as F
fact = spark.table("sales.fact_orders")
dim = spark.table("sales.customer_pricing")
hot_keys = (
fact.groupBy("customer_id")
.count()
.where(F.col("count") > 1_000_000) # Example threshold
.select("customer_id")
)
hot_keys.cache()
hot_keys.show(20, truncate=False)
```
The threshold should come from your data distribution. A simple rule is to mark keys far above the median or above a size that creates unsafe task memory.
Salt only the skewed keys
The next example salts only hot keys. Non-hot keys keep salt `0`, which avoids unnecessary data expansion.
```python
from pyspark.sql import functions as F
salt_buckets = 16
fact_with_hot_flag = (
fact.join(hot_keys.withColumn("is_hot", F.lit(True)), "customer_id", "left")
)
salted_fact = (
fact_with_hot_flag
.withColumn(
"salt",
F.when(
F.col("is_hot") == True,
(F.rand(seed=42) * salt_buckets).cast("int")
).otherwise(F.lit(0))
)
.drop("is_hot")
)
```
This spreads rows for each hot `customer_id` across 16 salt buckets. The right value depends on the skew size and cluster capacity. Too few buckets leave skew behind. Too many buckets add overhead.
Now expand the matching right-side rows for hot keys.
```python
hot_dim = (
dim.join(hot_keys, "customer_id", "inner")
.withColumn("salt", F.explode(F.sequence(F.lit(0), F.lit(salt_buckets - 1))))
)
normal_dim = (
dim.join(hot_keys, "customer_id", "left_anti")
.withColumn("salt", F.lit(0))
)
salted_dim = hot_dim.unionByName(normal_dim)
```
Then join using both columns.
```python
joined = salted_fact.join(
salted_dim,
on=["customer_id", "salt"],
how="left"
)
joined.explain("formatted")
```
Salt aggregations with care
Salting is also useful for skewed aggregations. Use a two-stage aggregation.
```python
salt_buckets = 32
salted = (
fact.withColumn(
"salt",
(F.rand(seed=99) * salt_buckets).cast("int")
)
)
partial = (
salted.groupBy("customer_id", "salt")
.agg(F.sum("order_amount").alias("partial_amount"))
)
final = (
partial.groupBy("customer_id")
.agg(F.sum("partial_amount").alias("total_amount"))
)
```
This reduces the amount of data one reducer handles for the hot key. The final aggregation still groups by `customer_id`, but it processes far fewer rows because the partial stage has already reduced them.
Salting has trade-offs. It adds columns, changes join logic, and may expand the right-side data. Use it where the skew cost is clearly higher than the extra processing.
How AQE dynamically coalesces partitions
Adaptive Query Execution, usually called AQE, lets Spark adjust parts of the query plan at runtime. In Azure Databricks, AQE is commonly used for SQL and DataFrame workloads, but always check your cluster or SQL warehouse configuration.
```python
spark.conf.get("spark.sql.adaptive.enabled")
```
Enable it for a test session if needed:
```python
spark.conf.set("spark.sql.adaptive.enabled", "true")
```
AQE uses runtime statistics from completed shuffle stages. One of its key features is post-shuffle partition coalescing. Spark may start with many shuffle partitions, then combine small adjacent partitions after it knows their actual sizes.
The setting below controls the initial shuffle partition count:
```python
spark.conf.set("spark.sql.shuffle.partitions", "800")
```
Without AQE, Spark would create 800 reduce partitions for many shuffle operations. That may produce too many tiny tasks when the data is smaller than expected.
With AQE coalescing enabled, Spark can combine those small partitions into fewer larger tasks.
```python
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", 128 1024 1024)
orders = spark.table("sales.fact_orders")
daily_sales = (
orders.groupBy("order_date")
.agg(F.sum("order_amount").alias("sales_amount"))
)
daily_sales.explain("formatted")
```
In the Spark UI, AQE plans may show adaptive plan changes. You may see operators such as:
```text
AdaptiveSparkPlan
AQEShuffleRead coalesced
```
Coalescing helps when shuffle partitions are too small. It reduces scheduling overhead and can make stages more efficient.
It does not automatically solve every hot-key problem. If one partition is huge because one key dominates, coalescing small partitions around it will not make that partition smaller. For skewed joins, AQE may also apply skew join handling when enabled and when the plan qualifies, but salting or broadcast may still be needed for severe or complex skew.
A useful pattern on Databricks is to combine a higher initial `spark.sql.shuffle.partitions` value with AQE coalescing. The high initial value gives Spark room to split work finer. AQE then reduces the cost of tiny partitions after it observes real shuffle sizes.
```python
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.shuffle.partitions", "1000")
result = (
spark.table("sales.fact_orders")
.join(spark.table("sales.dim_product"), "product_id")
.groupBy("category_id")
.agg(F.sum("order_amount").alias("amount"))
)
result.write.mode("overwrite").format("delta").saveAsTable("sales.category_sales")
```
For production, keep these settings close to the workload, not scattered across notebooks. Cluster policies, job parameters, or shared utility modules make changes easier to review.
A practical tuning flow for Azure Databricks jobs
Start with evidence. Skew tuning becomes messy when every setting changes at once.
Use this sequence for production investigations:
Find the skewed stage
Use the Spark UI. Sort tasks by duration, shuffle read, spill, and failed attempts.
Find the skewed key
Profile likely join or aggregation keys with `groupBy().count()` and compare max, median, and average counts.
Check the physical plan
Run `explain("formatted")`. Confirm whether Spark is using `SortMergeJoin`, `BroadcastHashJoin`, `HashAggregate`, or another expensive operator.
Try broadcast when one side is small
Raise `autoBroadcastJoinThreshold` only when the smaller side safely fits in executor memory. Prefer testing with a broadcast hint before changing broader settings.
Use salting for severe hot keys
Salt the skewed fact table and expand the matching right side only for hot keys. Keep the salt bucket count measurable and reviewed.
Enable and verify AQE
AQE coalescing helps with too many small shuffle partitions. It pairs well with higher initial shuffle partition counts, but it does not replace skew-specific fixes.
Compare before and after
Measure stage duration, max task duration, shuffle spill, failed tasks, and executor memory. The best fix reduces the worst task, not only the average task.
A simple logging pattern helps compare runs:
```python
def print_spark_settings():
keys = [
"spark.sql.autoBroadcastJoinThreshold",
"spark.sql.shuffle.partitions",
"spark.sql.adaptive.enabled",
"spark.sql.adaptive.coalescePartitions.enabled",
"spark.sql.adaptive.advisoryPartitionSizeInBytes",
]
for key in keys:
try:
print(f"{key} = {spark.conf.get(key)}")
except Exception:
print(f"{key} is not set")
print_spark_settings()
```
Also keep the tuned query small enough to reason about. If a notebook contains many joins, cache or write intermediate Delta tables during diagnosis. That makes the Spark UI easier to map back to the code.
The takeaway
Data skew is not fixed by adding more workers alone. If one key creates one oversized shuffle partition, Spark still has to run that oversized task somewhere. That is why skew leads to executor shuffle OOM crashes, repeated retries, and stranded tasks.
For Azure Databricks performance tuning, use the right fix for the shape of the problem:
Raise `autoBroadcastJoinThreshold` when a slightly larger dimension table can safely be broadcast and doing so removes a skewed shuffle.
Use key salting when a hot key in a large fact table must be split across multiple reducers.
Use AQE coalescing to reduce waste from many small shuffle partitions after Spark observes runtime data sizes.
The strongest workflow is evidence-led. Find the skew in the Spark UI, prove it with key distribution checks, inspect the physical plan, then apply the smallest change that reduces the largest task.
Comments