DEV Community

Cover image for Spark Performance Deep Dive on Databricks: Shuffle Tuning, Skew Handling, and Z-Ordering with Delta Lake + Unity Catalog
Jubin Soni
Jubin Soni Subscriber

Posted on

Spark Performance Deep Dive on Databricks: Shuffle Tuning, Skew Handling, and Z-Ordering with Delta Lake + Unity Catalog

The problem with "just add more workers"

Most Spark performance issues on Databricks aren't solved by scaling the cluster — they're caused by shuffle and skew, and no amount of extra nodes fixes a badly partitioned join. This post builds a realistic pipeline (order events joined against a small dimension table, aggregated, and written to Delta Lake) from the ground up, and uses it to work through:

  1. How Spark's shuffle actually behaves during a wide transformation
  2. Diagnosing and fixing data skew with salting and adaptive query execution (AQE)
  3. Laying out the resulting Delta table with Z-Ordering so downstream queries skip irrelevant files
  4. Governing access to the whole pipeline with Unity Catalog

Architecture overview

Pipeline shape — a batch job reading raw events, joining against a dimension table, aggregating, and writing to a governed Delta table:

Pipeline description

What happens inside a shuffle stage — this is the part most tutorials skip, and it's the key to understanding why skew hurts:

Shuffle c  description

Step 1 — Set up governed tables in Unity Catalog

Everything downstream depends on tables being registered under Unity Catalog, which gives you centralized access control and lineage instead of per-workspace table grants.

-- setup.sql, run in a Databricks SQL or notebook cell
CREATE CATALOG IF NOT EXISTS retail_analytics;
CREATE SCHEMA IF NOT EXISTS retail_analytics.events;

CREATE TABLE IF NOT EXISTS retail_analytics.events.raw_orders (
  order_id STRING,
  customer_id STRING,
  product_id STRING,
  quantity INT,
  event_ts TIMESTAMP
)
USING DELTA
LOCATION 'abfss://data@<storage-account>.dfs.core.windows.net/raw_orders';

CREATE TABLE IF NOT EXISTS retail_analytics.events.dim_products (
  product_id STRING,
  category STRING,
  unit_cost DOUBLE
)
USING DELTA
LOCATION 'abfss://data@<storage-account>.dfs.core.windows.net/dim_products';

GRANT SELECT ON TABLE retail_analytics.events.raw_orders TO `analysts`;
Enter fullscreen mode Exit fullscreen mode

Step 2 — Read and force a broadcast join for the small dimension table

dim_products is small, so let Spark broadcast it rather than shuffle both sides of the join.

# pipeline.py
from pyspark.sql import functions as F

orders = spark.table("retail_analytics.events.raw_orders")
products = spark.table("retail_analytics.events.dim_products")

joined = orders.join(F.broadcast(products), on="product_id", how="left")
Enter fullscreen mode Exit fullscreen mode

Without the explicit broadcast hint, Spark's cost-based optimizer usually picks a broadcast join automatically for small tables, but being explicit avoids surprises when the dimension table grows past spark.sql.autoBroadcastJoinThreshold (default 10MB) without anyone noticing.

Step 3 — The aggregation that triggers a shuffle

groupBy on customer_id is a wide transformation — Spark must shuffle rows so all records for a given key land on the same reducer.

agg = (
    joined
    .groupBy("customer_id", "category")
    .agg(
        F.sum(F.col("quantity") * F.col("unit_cost")).alias("total_spend"),
        F.count("order_id").alias("order_count"),
    )
)
Enter fullscreen mode Exit fullscreen mode

If one customer_id (say, a test account or a large B2B buyer) accounts for a disproportionate share of rows, this is where skew shows up: one reducer task runs for minutes while the rest of the stage finishes in seconds. You'll see this in the Spark UI as a single long-running task in an otherwise short stage.

Step 4 — Fixing skew with salting

Adaptive Query Execution (AQE) handles a lot of skew automatically in modern Databricks Runtime, but for known hot keys, explicit salting is still the most predictable fix.

from pyspark.sql import functions as F
import random

SALT_BUCKETS = 20

# Add a salt column to spread the hot key across multiple reducers
salted = joined.withColumn("salt", (F.rand() * SALT_BUCKETS).cast("int"))

partial_agg = (
    salted
    .groupBy("customer_id", "category", "salt")
    .agg(
        F.sum(F.col("quantity") * F.col("unit_cost")).alias("total_spend"),
        F.count("order_id").alias("order_count"),
    )
)

# Second pass: combine the salted partial aggregates into the final result
final_agg = (
    partial_agg
    .groupBy("customer_id", "category")
    .agg(
        F.sum("total_spend").alias("total_spend"),
        F.sum("order_count").alias("order_count"),
    )
)
Enter fullscreen mode Exit fullscreen mode

This two-phase pattern — pre-aggregate on a salted key, then combine — is the same trick used inside combiners in older MapReduce systems. It trades a bit of extra shuffle for eliminating the single-reducer bottleneck.

Also worth setting explicitly rather than relying on the default of 200:

spark.conf.set("spark.sql.shuffle.partitions", "auto")  # let AQE size it dynamically
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
Enter fullscreen mode Exit fullscreen mode

Step 5 — Write to Delta Lake with optimized writes

(
    final_agg.write
    .format("delta")
    .mode("overwrite")
    .option("delta.autoOptimize.optimizeWrite", "true")
    .saveAsTable("retail_analytics.events.customer_spend_summary")
)
Enter fullscreen mode Exit fullscreen mode

Optimized writes shuffle data before writing so you get fewer, better-sized files instead of many small ones — this costs some write-time latency in exchange for faster reads later.

Step 6 — Z-Order the table for the queries that matter

If most downstream queries filter on customer_id, co-locate related rows physically so Spark can skip files that can't match the filter.

OPTIMIZE retail_analytics.events.customer_spend_summary
ZORDER BY (customer_id);
Enter fullscreen mode Exit fullscreen mode

Z-ordering aims to produce evenly balanced data files by row count rather than raw size, and its effectiveness depends on the column having reasonably high cardinality — Z-ordering by a low-cardinality column like category alone gives little benefit. On newer Delta Lake / Databricks Runtime versions, Liquid Clustering is generally the preferred choice for new tables since it adapts as query patterns change, while ZORDER remains relevant mainly for existing tables not yet migrated.

-- Preferred on new tables (Databricks Runtime supporting Liquid Clustering):
CREATE TABLE retail_analytics.events.customer_spend_summary
CLUSTER BY (customer_id);
Enter fullscreen mode Exit fullscreen mode

Comparing shuffle-mitigation techniques

Technique Fixes Cost When to use
Broadcast join Shuffle on the large side of a join Extra memory on executors Small (<~10MB by default) dimension/lookup tables
AQE skew join handling Automatic detection of skewed partitions Minor planning overhead Default-on; good general safety net
Manual salting Known, severe hot keys Extra shuffle for the two-phase aggregate Hot keys that AQE doesn't fully resolve
Repartition by key Uneven task distribution before a shuffle stage One extra shuffle Pre-shaping data before multiple downstream joins
Z-Ordering Slow reads due to unnecessary file scans Shuffle + rewrite during OPTIMIZE Existing tables, high-cardinality filter columns
Liquid Clustering Same as Z-Order, plus evolving query patterns Ongoing background clustering New tables on supporting runtime versions

Governance: tying it back to Unity Catalog

Because every table above was created under retail_analytics, access control, audit logging, and lineage are handled centrally rather than per-cluster:

-- Restrict PII-adjacent columns without duplicating the table
CREATE VIEW retail_analytics.events.customer_spend_summary_masked AS
SELECT
  customer_id,
  category,
  total_spend,
  order_count
FROM retail_analytics.events.customer_spend_summary;

GRANT SELECT ON VIEW retail_analytics.events.customer_spend_summary_masked TO `bi_readers`;
Enter fullscreen mode Exit fullscreen mode

Production considerations

  • Diagnose before tuning. Check the Spark UI's stage view for one long-running task among many short ones — that's the signature of skew, not just "the job is slow."
  • Predictive optimization can run OPTIMIZE and ANALYZE automatically on Unity Catalog managed tables, which reduces the need for scheduled maintenance jobs for many workloads.
  • Don't Z-Order everything. It requires shuffling and rewriting the table (or partition), so reserve it for columns that are actually filtered on frequently downstream.

References

  1. Best practices: Delta Lake — Azure Databricks / Microsoft Learn
  2. Data skipping (Z-ordering) — Azure Databricks / Microsoft Learn
  3. Optimizations — Delta Lake documentation
  4. Delta Lake Under the Hood: What Every Data Engineer Should Know — Databricks Community
  5. Mastering Delta Lake Performance: Z-Ordering vs Liquid Clustering — Medium

Top comments (0)