DEV Community

pawan deore
pawan deore

Posted on

Apache Spark Partitioning, Skew, and Broadcast Joins: A Practical Guide

Apache Spark Partitioning, Skew, and Broadcast Joins: A Practical Guide

Apache Spark can process large datasets in parallel because it divides work into partitions. When those partitions are balanced, executors finish work at roughly the same time. When one partition contains far more data than the others, a single task becomes a bottleneck. This is called data skew.

Partitioning, skew handling, and join strategy are tightly connected:

  • Partitioning determines where records are processed.
  • Joins often trigger expensive data shuffles.
  • Skew causes a small number of join or aggregation tasks to run much longer than the rest.
  • Broadcast joins can avoid shuffling a large fact table when one side is small enough.

This tutorial assumes Spark DataFrames, PySpark, and a typical analytics workload with a large fact table and smaller lookup tables.

Partitioning: Spark’s Unit of Parallel Work

A Spark DataFrame is split into partitions. Each partition is processed by one task at a time, typically on an executor core.

For example, a DataFrame with 200 partitions can create up to 200 parallel tasks for a stage. The actual concurrency depends on available executor cores.

There are several different meanings of “partition” in data engineering, and confusing them causes design mistakes:

Term Meaning
Spark partition An in-memory or execution-time chunk of a DataFrame or RDD
Shuffle partition A partition created after Spark redistributes data for joins, aggregations, or repartition()
File partition A physical output directory such as date=2025-03-08/ in Parquet storage
Table partition A logical storage partition defined in a metastore or table format

A Parquet table partitioned by event_date can still have poorly balanced Spark partitions during a join. Storage layout and execution layout are related, but they are not the same thing.

The most important operations are:

df.repartition(200)
df.repartition(200, "customer_id")
df.repartitionByRange(200, "event_date")
df.coalesce(50)
Enter fullscreen mode Exit fullscreen mode

repartition() creates a full shuffle. Spark moves rows across executors to produce the requested partitioning.

coalesce() usually reduces partitions without a full shuffle. It is useful after filtering data down, but it does not rebalance uneven partitions. If one partition was already huge, coalesce() can preserve that imbalance.

How Joins Create Shuffles

Consider a common pipeline:

Raw order files
    |
    v
Filter and clean orders
    |
    +--------------------+
    |                    |
    v                    v
Large orders table   Small customers table
    |                    |
    +------ join --------+
              |
              v
Aggregate revenue by customer segment
              |
              v
Write Parquet or Delta output
Enter fullscreen mode Exit fullscreen mode

A join between two large DataFrames usually requires Spark to shuffle both sides by the join key:

orders                          customers
  |                                  |
  v                                  v
Exchange hash(customer_id)      Exchange hash(customer_id)
  |                                  |
  +------------- join -------------+
Enter fullscreen mode Exit fullscreen mode

If both tables are distributed by customer_id, rows for the same customer are placed in the same shuffle partition. That makes the join possible, but it creates a problem when one customer appears far more often than others.

For example:

customer_id = C_HOT      700 million orders
customer_id = C_001        3 million orders
customer_id = C_002        2 million orders
...
Enter fullscreen mode Exit fullscreen mode

All rows for C_HOT generally land in one hash partition. One task now processes hundreds of millions of rows while other tasks finish quickly.

This is why adding more executors does not necessarily solve skew. The skewed key still belongs to one partition unless the execution strategy changes.

A Runnable PySpark Example

The following example creates a skewed orders dataset:

  • One customer, C_HOT, has 700 orders.
  • One hundred normal customers have 3 orders each.
  • The customer dimension is small enough to broadcast.
  • The expected joined row count is 1000.
  • The expected revenue is 13000.

Run it with a local Spark installation:

pip install pyspark
Enter fullscreen mode Exit fullscreen mode
from pyspark.sql import SparkSession, functions as F

spark = (
    SparkSession.builder
    .master("local[*]")
    .appName("partitioning-skew-broadcast-demo")
    .config("spark.sql.shuffle.partitions", "8")
    .config("spark.sql.adaptive.enabled", "true")
    .getOrCreate()
)

# 700 orders from one unusually active customer.
hot_orders = [
    (f"HOT-{i:04d}", "C_HOT", 10)
    for i in range(700)
]

# 100 normal customers, each with 3 orders.
normal_orders = [
    (f"NORMAL-{customer_num:03d}-{order_num}", f"C_{customer_num:03d}", 20)
    for customer_num in range(1, 101)
    for order_num in range(3)
]

orders = spark.createDataFrame(
    hot_orders + normal_orders,
    ["order_id", "customer_id", "amount"]
)

customers = spark.createDataFrame(
    [("C_HOT", "enterprise")]
    + [(f"C_{i:03d}", "standard") for i in range(1, 101)],
    ["customer_id", "segment"]
)

# Force a hash shuffle by customer_id so we can observe skew.
orders_by_customer = orders.repartition(8, "customer_id").cache()
orders_by_customer.count()  # Materialize the cached shuffled DataFrame.

partition_sizes = (
    orders_by_customer.rdd
    .mapPartitions(lambda rows: [sum(1 for _ in rows)])
    .collect()
)

print("Rows per partition:", partition_sizes)
print("Largest partition:", max(partition_sizes))

# Broadcast the small dimension table.
joined = (
    orders_by_customer
    .join(F.broadcast(customers), on="customer_id", how="inner")
)

result = (
    joined
    .agg(
        F.count("*").alias("joined_rows"),
        F.sum("amount").alias("total_revenue")
    )
    .first()
)

print("Joined rows:", result["joined_rows"])
print("Total revenue:", result["total_revenue"])

assert result["joined_rows"] == 1000
assert result["total_revenue"] == 13000
assert max(partition_sizes) >= 700

joined.explain(mode="formatted")

spark.stop()
Enter fullscreen mode Exit fullscreen mode

The exact list of partition sizes varies because Spark’s hash partitioning determines where each customer key lands. However, one partition must contain at least the 700 C_HOT rows.

The important result is that the join uses a broadcast strategy for customers. In the physical plan, look for an operator similar to:

BroadcastHashJoin
Enter fullscreen mode Exit fullscreen mode

Spark sends the small customers table to each executor. The large orders DataFrame does not need to be shuffled just to perform this join.

When Broadcast Joins Help—and When They Do Not

A broadcast join is appropriate when one join input is genuinely small after filtering and selecting only needed columns.

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

Conceptually, Spark does this:

Large orders partitions        Small customers table
        |                              |
        |                         Broadcast to all executors
        v                              v
Local join in each executor partition
Enter fullscreen mode Exit fullscreen mode

The main benefit is avoiding a shuffle of the large table for the join.

However, broadcast joins do not solve every skew problem:

  1. A skewed scan remains skewed.

    If one input file is much larger than the rest, one scan task may still lag.

  2. A later aggregation can still skew.

    Joining orders to customers may be fast, but this can still produce a skewed shuffle:

   joined.groupBy("customer_id").sum("amount")
Enter fullscreen mode Exit fullscreen mode

All C_HOT rows must still be brought together for that aggregation.

  1. Broadcasting a table that is too large can fail. Every executor needs memory for the broadcast relation. A dimension that looks small in compressed Parquet can expand significantly after deserialization and hashing.

Spark can automatically broadcast relations below spark.sql.autoBroadcastJoinThreshold, but the threshold is configuration-dependent. Explicit broadcast() is useful when you understand the size and want to communicate intent, not as a blind optimization.

If you want a broader learning path after practicing these concepts, this optional data engineering roadmap for beginners can help place Spark optimization topics in context.

Detecting and Fixing Data Skew

Start by confirming that skew exists. In the Spark UI, look for a stage where:

  • Most tasks finish quickly.
  • One or a few tasks run much longer.
  • A few tasks read or write substantially more shuffle data.
  • Executor memory pressure or spill occurs only in a few tasks.

You can also inspect key frequency directly:

orders.groupBy("customer_id") \
    .count() \
    .orderBy(F.desc("count")) \
    .show(20, truncate=False)
Enter fullscreen mode Exit fullscreen mode

Let Adaptive Query Execution Help First

Adaptive Query Execution (AQE) allows Spark to make some decisions after it observes runtime shuffle statistics.

Useful settings include:

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
Enter fullscreen mode Exit fullscreen mode

For eligible shuffle joins, Spark can split unusually large shuffle partitions into smaller pieces. AQE can also coalesce tiny shuffle partitions.

AQE is often the first mitigation to try because it requires no business-logic changes. But it cannot always split work for a single extremely hot key, especially when all records for that key must be processed together.

Use Key Salting for Severe Hot Keys

Salting splits one hot key into multiple artificial keys.

For example, instead of joining only on:

C_HOT
Enter fullscreen mode Exit fullscreen mode

you create:

C_HOT_0
C_HOT_1
...
C_HOT_7
Enter fullscreen mode Exit fullscreen mode

The fact table assigns a deterministic salt:

salt_count = 8

salted_orders = orders.withColumn(
    "salt",
    F.when(
        F.col("customer_id") == "C_HOT",
        F.pmod(F.hash("order_id"), F.lit(salt_count))
    ).otherwise(F.lit(0))
)
Enter fullscreen mode Exit fullscreen mode

The dimension table must duplicate only the hot key across the same salt values. Normal customer rows keep salt 0. Then join on both customer_id and salt.

Salting is most useful when both join inputs are too large to broadcast. It adds complexity and duplicates dimension rows, so do not apply it to every key. Identify the small set of problematic keys first.

Production Pitfalls

A few common choices create avoidable Spark performance problems.

Using repartition() everywhere.

Every repartition is a shuffle. Repartition only when it improves a downstream operation, such as balancing data before a large write or aligning work before a non-broadcast join.

Using too many shuffle partitions.

Thousands of tiny tasks add scheduler overhead and can produce small output files. Too few partitions create large tasks, memory pressure, and poor parallelism. Tune based on observed input size, shuffle volume, task duration, and cluster capacity.

Assuming a broadcast hint is always safe.

A forced broadcast can cause executor out-of-memory failures. Filter dimensions early, select only necessary columns, and verify size under realistic production data volumes.

Writing files partitioned by high-cardinality keys.

Writing data partitioned by customer_id can create massive directory and small-file counts. Storage partitions should usually reflect common pruning filters, such as date, region, or source system.

Ignoring post-join cardinality.

A duplicate key in a dimension table can turn a one-to-many join into a many-to-many join. This both inflates results and amplifies skew.

Testing Partitioning and Join Behavior

Test correctness before tuning performance.

For joins, create assertions for:

  • Expected row counts
  • Expected aggregate totals
  • Null handling for unmatched keys
  • Duplicate dimension keys
  • Known hot-key behavior

For example, the runnable example asserts both row count and total revenue. In a real pipeline, add a test where a dimension intentionally contains duplicate customer_id values and verify that the pipeline rejects or deduplicates them according to your business rule.

For performance regression testing:

  1. Build a representative sample containing hot keys.
  2. Run the job with Spark event logging enabled.
  3. Compare shuffle read/write volume, task duration distribution, spill metrics, and output row counts.
  4. Validate the physical plan with explain("formatted").
  5. Test with production-like executor memory and partition counts; local mode cannot reveal all distributed failure modes.

Takeaway

Spark partitioning determines how work is distributed, while joins and aggregations often force data to shuffle. A skewed key can make one partition dominate an entire stage.

Use broadcast joins when one side is safely small, because they avoid shuffling the large side. Use AQE to handle many skewed shuffle joins automatically. For persistent, severe hot keys that cannot be broadcast, use targeted salting and validate both correctness and resource usage.

The best optimization is not “more partitions” by default—it is choosing a partitioning and join strategy that matches the actual shape of your data.

Top comments (0)

⌂ 换网址