Spark Partitioning, Shuffles, and Data Skew

Find the work behind the slow task

A partition is a logical chunk of a dataset processed by a task in a stage. It is not necessarily one file or one machine: an executor can run several tasks, and later stages can use a different partition layout. Spill means writing intermediate data to disk when the operation cannot keep it in memory; garbage collection (GC) reclaims unused objects. A long task tail can delay a Spark stage even when most tasks have finished. Data skew means an uneven distribution of data or work across processing units. A hot customer key is one cause, but row width, expensive functions, join fan-out, spill, garbage collection, and slow hosts can also create uneven durations. Code and query design influence the distribution; skew is not a property of source data alone.

Compare the slow task with its peers using input and shuffle bytes, records, spill, GC time, duration, and the executed plan. A task taking forty times the median does not prove its partition has forty times the rows. The median may even be zero for sparse partitions. Inspect the complete distribution and count the keys used by the expensive operator, with an appropriate sampling or profiling budget.

This article follows partition selection, hash redistribution, hot-key handling, broadcast joins, salting, and file layout. The lab uses actual local Spark execution and checks results and partition or file counts. It does not estimate production elapsed time or network bytes from a fixed bytes-per-row assumption.

Partition count is different from key distribution

A Spark shuffle redistributes records to satisfy an operator’s required distribution. For ordinary hash partitioning by customer_id, equal keys choose the same bucket. Raising the bucket count cannot divide the rows of that one key across buckets, although it can reduce collisions with other keys. This describes that hash exchange, not every wide operation: range, round-robin, and adaptive execution have different rules.

Shuffle writes and reads can involve local disk, remote fetches, serialization, memory, and spill. Some reads are local and aggregation may send partial states instead of all raw rows. Input row count times a guessed row width is therefore neither actual shuffle size nor measured network traffic. Filtering or projecting unneeded data can help when it preserves the query’s meaning.

The lab makes exactly 6,000 of 10,000 rows belong to customer 1. With either eight or sixteen hash partitions, that key has one destination. A round-robin repartition can distribute those rows, but a later groupBy on customer_id must bring their contributions together again. Inspect the distribution at the operator that matters rather than treating an earlier even layout as a permanent property.

Choose a starting count and measure it

Spark SQL documents 200 as the default shuffle partition count; it is a default to inspect, not a universal mistake or a workload target. Consider expected shuffle bytes, decoded working memory, available concurrent task capacity, skew, and scheduling overhead. A fixed two-to-three-tasks-per-slot rule cannot determine every workload’s best count. Input scan partitions, shuffle partitions, and output directory partitions are separate choices.

With equal-duration tasks and eight continuously available slots, 32 tasks need four ideal waves. Real slots become available as tasks finish rather than waiting for a synchronized wave barrier. Summing an assumed 20 ms over 2,000 tasks gives 40 seconds of aggregate task cost, not automatically 40 seconds of stage wall time. More memory can reduce spill or GC and improve runtime in some workloads; it is not always merely a way to keep a hot task alive.

AQE can coalesce supported post-shuffle partitions using runtime statistics. Its advisory size and parallelism settings are inputs to a policy, not a promise about each task’s memory footprint. Test a reasonable initial count against actual task-size and duration distributions, preserve result equality, and change one consequential choice at a time.

Select a join strategy without changing the answer

The build side is the relation prepared for matching, often as a hash lookup; the streamed side supplies rows that probe it. A broadcast join distributes a build side so streamed-side tasks can look up matches without that join requiring a key shuffle of the streamed side. It still transfers data, and earlier input imbalance or a later aggregation can remain skewed. Repeated build keys can create large output even when the build side itself is small.

Check the estimated size, the built hash representation, driver-side construction, executor memory, timeout, and supported join type. An on-disk size threshold is not a safe memory limit. A hint requests a strategy and may not be applicable. In the lab, automatic broadcast is disabled and explicit merge and broadcast hints make the chosen operators visible. Full multiset equality checks the answer; no single-run speed ratio is claimed.

Null treatment is part of join semantics. Ordinary equality does not match nulls, while null-safe equality can. Inner joins may already filter null keys; left joins must retain unmatched left rows. Splitting a null branch and unioning it back requires the correct right-side null columns and bag-preserving union. Assigning arbitrary keys, excluding unknown categories, or deduplicating a dimension can change the answer and is not a generic performance repair.

Salt only when the merge rule is valid

Key salting adds a secondary key to divide a hot logical group into subgroups. Use a stable per-row identifier to derive a reproducible salt; hashing the hot key alone gives every hot row the same salt. Eight salt values do not guarantee eight physical destinations because their hashes can collide or AQE can alter subsequent tasks.

For a sum and count, first group by original key plus salt, then sum partial sums and partial counts by the original key. A mean must combine sum and count, not average the subgroup means. Exact distinct counts and medians cannot generally be merged by adding partial results. Even a sum may change floating-point rounding order or overflow finite decimal/integer types; verify the numeric domain.

Spark already uses partial aggregation for suitable built-in aggregates, as the plain sum plan in the lab shows. A hot input key may therefore emit only a small partial state per upstream partition. Manual salting can add a second exchange and more state without improving the bottleneck. Apply it to an observed operator and compare complete results before measuring performance.

For an inner or left join with facts on the streamed side, assign each hot fact one salt and replicate every matching build row across the salt values. Cold keys keep one salt. The lab deliberately includes duplicate dimension keys, an unmatched key, and a null fact key, and verifies all seven output rows against the unsalted left join. This recipe must be reconsidered for right or full outer joins because replicated unmatched build rows can multiply output.

Repartition, coalesce, and adaptive skew handling

repartition(n, key) requests a hash redistribution by that key; repartition(n) uses a round-robin distribution in the tested Spark SQL path. Neither promises equal execution cost per task. repartitionByRange has another distribution rule. Repartitioning before an operation that immediately requires a different distribution can add a redundant exchange.

DataFrame coalesce reduces partitions through a narrow dependency without adding a full shuffle. It cannot split one existing partition to increase parallelism and need not merge only partitions already on the same executor. A drastic reduction can reduce upstream parallelism and change the size distribution, rather than preserve a skew ratio exactly. The lab contrasts coalesce(8) and repartition(8) from a single input partition.

AQE skew-join optimization applies only to eligible plans and sufficiently large uneven shuffle partitions. Spark uses both a relative-to-median condition and an absolute byte threshold; enabling AQE alone does not establish that a split happened. Inspect the final adaptive plan and task metrics. Split/replication also adds work. It is not a general repair for an indivisible slow source task, nor does it automatically perform manual salting of every aggregation. The lab disables AQE to isolate the explicit operations and does not claim to test adaptive skew splitting.

Count committed files rather than predicting them from tasks

Writer tasks, table partition directories, and data files have different meanings. A task touching three directory values can write three files. A file rollover limit can make one task write several files within one directory. Format, writer implementation, empty tasks, commit behavior, and table optimizations affect the final count. Two hundred tasks times thirty directories is a conditional task-directory count, not a universal 6,000-file upper bound.

The lab writes 120 rows from one task into three directories: three Parquet data files. With maxRecordsPerFile=10, the same rows produce twelve data files. It lists committed part files and reads every row back; it does not infer compressed file sizes from row counts.

Repartitioning by date concentrates each date’s rows but can create a new bottleneck for a large date. Additional within-date distribution may help when parallel writing is necessary. Coalesce can reduce tasks when existing distribution is acceptable. Choose the layout from the reader’s needs and measured file sizes, then verify counts, content, and any compaction or replacement procedure. Fewer files alone does not establish a correct or faster dataset.

Lab: actual local Spark execution

Use Python 3.12, pyspark==4.2.0, and Java 17 with JAVA_HOME set to that installation. Driver and Python workers use the same Python minor version. Run the setup, then each example in a standalone process without another Spark session. local[2] uses local worker threads and loopback communication, not a multi-node cluster. The helper stops Spark and removes its temporary files. SQL settings explicitly fix UTC, TIMESTAMP_LTZ, ANSI mode, and eight shuffle partitions; AQE is disabled to keep the inspected plan simple. Spark startup warnings about logging, incubator modules, and the native Hadoop fallback are separate from recorded stdout. No runtime speedup or large-data memory claim is tested.

import os, sys, tempfile
from pathlib import Path
from contextlib import contextmanager, redirect_stdout
from io import StringIO
from decimal import Decimal
from pyspark.sql import SparkSession, functions as F, types as T

os.environ["PYSPARK_PYTHON"] = sys.executable
os.environ["SPARK_LOCAL_IP"] = "127.0.0.1"

@contextmanager
def local_spark():
    with tempfile.TemporaryDirectory(prefix="de022-") as work:
        spark = (SparkSession.builder.master("local[2]").appName("skew-lab")
                 .config("spark.driver.host", "127.0.0.1")
                 .config("spark.driver.bindAddress", "127.0.0.1")
                 .config("spark.ui.enabled", "false")
                 .config("spark.ui.showConsoleProgress", "false")
                 .config("spark.sql.warehouse.dir", work)
                 .config("spark.sql.session.timeZone", "UTC")
                 .config("spark.sql.timestampType", "TIMESTAMP_LTZ")
                 .config("spark.sql.ansi.enabled", "true")
                 .config("spark.sql.shuffle.partitions", "8")
                 .config("spark.sql.autoBroadcastJoinThreshold", "-1")
                 .config("spark.sql.adaptive.enabled", "false")
                 .getOrCreate())
        spark.sparkContext.setLogLevel("ERROR")
        try:
            yield spark, Path(work)
        finally:
            spark.stop()

def explain_text(df):
    out = StringIO()
    with redirect_stdout(out):
        df.explain(mode="extended")
    return out.getvalue()

def sorted_rows(df):
    return sorted(tuple(r) for r in df.collect())

def orders_fixture(spark):
    return spark.range(0, 10_000, numPartitions=4).select(
        F.col("id").alias("order_id"),
        F.when(F.col("id") < 6000, F.lit(1)).otherwise(2 + F.col("id") % 100).alias("customer_id"),
        (F.col("id") % 97 + 1).alias("amount"))

def partition_sizes(df):
    return sorted(df.rdd.mapPartitions(lambda rows: [sum(1 for _ in rows)]).collect())

def salt_hot(df, buckets):
    if type(buckets) is not int or buckets <= 0:
        raise ValueError("buckets must be a positive integer")
    return df.withColumn("salt", F.when(F.col("customer_id") == 1,
        F.pmod(F.xxhash64("order_id"), F.lit(buckets))).otherwise(F.lit(0)))

1. Compare hot-key destinations with round-robin distribution.

with local_spark() as (spark, work):
    orders = orders_fixture(spark)
    assert orders.filter("customer_id = 1").count() == 6000
    for n in (8, 16):
        by_key = orders.repartition(n, "customer_id")
        sizes = partition_sizes(by_key)
        hot = by_key.filter("customer_id = 1").select(F.spark_partition_id().alias("pid"))
        assert len(sizes) == n and sum(sizes) == 10_000
        assert hot.distinct().count() == 1 and max(sizes) >= 6000
        print("key partitions:", n, "hot-key rows:", 6000, "hot-key destinations:", hot.distinct().count())
    scattered = orders.repartition(8)
    assert sum(partition_sizes(scattered)) == 10_000
    print("round-robin largest partition below hot-key size:", max(partition_sizes(scattered)) < 6000)
# key partitions: 8 hot-key rows: 6000 hot-key destinations: 1
# key partitions: 16 hot-key rows: 6000 hot-key destinations: 1
# round-robin largest partition below hot-key size: True

2. Verify salted sum/count and existing partial aggregation.

with local_spark() as (spark, work):
    orders = orders_fixture(spark)
    plain = orders.groupBy("customer_id").agg(F.sum("amount").alias("total"), F.count("amount").alias("n"))
    salted = salt_hot(orders, 8)
    partial = salted.groupBy("customer_id", "salt").agg(F.sum("amount").alias("subtotal"), F.count("amount").alias("n"))
    merged = partial.groupBy("customer_id").agg(F.sum("subtotal").alias("total"), F.sum("n").alias("n"))
    assert sorted_rows(plain) == sorted_rows(merged)
    assert salted.filter("customer_id != 1 AND salt != 0").count() == 0
    assert salted.filter("customer_id = 1").select("salt").distinct().count() == 8
    hot_parts = salted.repartition(8, "customer_id", "salt").filter("customer_id = 1").select(F.spark_partition_id().alias("pid"))
    print("plain and two-step sum/count equal:", sorted_rows(plain) == sorted_rows(merged))
    print("hot salt values:", 8, "physical destinations at most eight:", 1 <= hot_parts.distinct().count() <= 8)
    print("cold keys remain unsalted:", salted.filter("customer_id != 1 AND salt != 0").count() == 0)
    assert "partial_sum" in explain_text(plain)
    print("plain sum already has partial aggregation:", True)
# plain and two-step sum/count equal: True
# hot salt values: 8 physical destinations at most eight: True
# cold keys remain unsalted: True
# plain sum already has partial aggregation: True

3. Preserve duplicate, unmatched, and null rows in a salted left join.

with local_spark() as (spark, work):
    facts = spark.createDataFrame([(1,1,10),(2,1,20),(3,2,30),(4,9,40),(5,None,50)],
                                  "order_id long, customer_id long, amount long")
    dimension = spark.createDataFrame([(1,"a"),(1,"b"),(2,"c")], "customer_id long, label string")
    reference = facts.join(dimension, "customer_id", "left").select("order_id","amount","label")
    salted = salt_hot(facts, 4)
    expanded = dimension.withColumn("salt", F.explode(F.when(F.col("customer_id") == 1,
        F.sequence(F.lit(0), F.lit(3))).otherwise(F.array(F.lit(0)))))
    result = salted.join(expanded, ["customer_id","salt"], "left").select("order_id","amount","label")
    expected = [(1,10,"a"),(1,10,"b"),(2,20,"a"),(2,20,"b"),(3,30,"c"),(4,40,None),(5,50,None)]
    assert sorted_rows(reference) == sorted_rows(result) == expected
    print("salted left join preserves all seven output rows:", sorted_rows(result) == expected)
    print("small-side rows before and after hot-key expansion:", dimension.count(), expanded.count())
    assert expanded.count() == 9
# salted left join preserves all seven output rows: True
# small-side rows before and after hot-key expansion: 3 9

4. Compare merge and broadcast plans with complete result equality.

with local_spark() as (spark, work):
    orders = orders_fixture(spark)
    dimension = spark.range(1,102).select(F.col("id").alias("customer_id"), (F.col("id") % 3).alias("segment"))
    shuffled = orders.join(dimension.hint("merge"), "customer_id")
    broadcast = orders.join(F.broadcast(dimension), "customer_id")
    a = shuffled.select("order_id","amount","segment")
    b = broadcast.select("order_id","amount","segment")
    assert a.exceptAll(b).limit(1).count() == b.exceptAll(a).limit(1).count() == 0
    assert a.count() == b.count() == 10_000
    sp, bp = explain_text(shuffled), explain_text(broadcast)
    assert "SortMergeJoin" in sp and "BroadcastHashJoin" in bp
    print("same complete join rows:", True)
    print("merge and broadcast operators observed:", "SortMergeJoin" in sp, "BroadcastHashJoin" in bp)
    print("broadcast exchange still transfers the build side:", "BroadcastExchange" in bp)
# same complete join rows: True
# merge and broadcast operators observed: True True
# broadcast exchange still transfers the build side: True

5. Try coalesce and repartition from one input partition.

with local_spark() as (spark, work):
    hot_input = spark.range(0,1000,numPartitions=1)
    reduced = hot_input.coalesce(8)
    spread = hot_input.repartition(8)
    assert reduced.rdd.getNumPartitions() == 1 and spread.rdd.getNumPartitions() == 8
    assert sum(partition_sizes(spread)) == 1000
    merged = spread.coalesce(2)
    assert merged.rdd.getNumPartitions() == 2 and merged.count() == 1000
    print("coalesce cannot increase one input partition:", reduced.rdd.getNumPartitions())
    print("repartition and subsequent coalesce counts:", spread.rdd.getNumPartitions(), merged.rdd.getNumPartitions())
    print("row total preserved:", merged.count() == 1000)
# coalesce cannot increase one input partition: 1
# repartition and subsequent coalesce counts: 8 2
# row total preserved: True

6. Count committed Parquet files with and without rollover.

with local_spark() as (spark, work):
    rows = spark.range(0,120,numPartitions=4).select("id", (F.col("id") % 3).alias("day"))
    path1, path2 = work / "one-task", work / "rolled"
    rows.coalesce(1).write.partitionBy("day").parquet(str(path1))
    rows.coalesce(1).write.option("maxRecordsPerFile", 10).partitionBy("day").parquet(str(path2))
    files1 = list(path1.rglob("part-*.parquet"))
    files2 = list(path2.rglob("part-*.parquet"))
    assert len(files1) == 3 and len(files2) == 12
    restored = spark.read.parquet(str(path2))
    assert restored.count() == 120
    assert sorted_rows(restored.select("id","day")) == sorted_rows(rows)
    print("one task, three directories, data files:", len(files1))
    print("same task with ten-row rollover, data files:", len(files2))
    print("all written rows restored:", restored.count() == 120)
# one task, three directories, data files: 3
# same task with ten-row rollover, data files: 12
# all written rows restored: True

References: Spark performance tuning, DataFrame coalesce, DataFrame repartition, Null semantics, Spark configuration.