Tuning Spark Shuffle for Large Spatial Joins

Read the task-time distribution before touching a single configuration value. If the slowest task in the join stage takes fifty times the median, the problem is skew and no amount of parallelism fixes it; if the distribution is tight and the stage is simply long, more partitions genuinely help. Almost all wasted tuning effort on spatial joins comes from applying the second remedy to the first problem. This page is the diagnostic procedure behind Apache Sedona for distributed spatial joins.

Quick Reference

Symptom Diagnosis Fix Primary Use Case
max ≈ 2–3 × median Even distribution, stage just long Raise shuffle.partitions Genuinely large joins
max > 10 × median Skew — one dense cell Density-aware partitioner Urban-concentrated data
High spill bytes Partitions exceed execution memory More partitions or more memory Wide rows, large candidate sets
Thousands of tiny tasks Over-partitioned Fewer partitions Small joins on large clusters
One task never finishes Coincident geometry in one cell Secondary key or capacity cap Address and utility data

The Distribution Is the Diagnosis

A Spark stage finishes when its slowest task finishes. That single sentence explains why the median task time is almost useless on its own and why the maximum is the number that matters.

Geographic data guarantees an uneven distribution. Points, buildings, addresses, and observations concentrate where people are, so a uniform grid over a country produces cells whose populations differ by three orders of magnitude. Partition on that grid and one task receives a thousand times the work of another. The stage takes as long as that task, and every other executor idles waiting for it.

Raising spark.sql.shuffle.partitions divides the already small partitions into smaller ones. The dense cell is still one cell, still one partition, still one task. The stage time does not move, the cluster looks busier, and the tuning appears to have done nothing — because it did.

Task durations before and after replacing a uniform grid partitioner Two task-duration histograms for the same join stage. Before, almost every task finishes in under ten seconds but a single task runs for four hundred seconds, so the stage takes four hundred seconds and the cluster is idle for most of that time. After switching to a density-aware partitioner, task durations cluster between eleven and nineteen seconds with no outlier, and the stage finishes in nineteen seconds. The total work done is the same in both cases. Uniform grid — one task holds the stage 400 s median 8 s task duration → stage: 400 s · cluster idle for 98% of it max / median = 50 → skew, not parallelism Density-aware — no outlier left median 14 s task duration → stage: 19 s · cluster busy throughout max / median = 1.4 → evenly distributed

Diagnosing and Configuring

python
# Requires: pyspark>=3.5, apache-sedona[spark]>=1.6  (Python 3.10+)
from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class StageProfile:
    tasks: int
    median_ms: float
    max_ms: float
    shuffle_written: int
    spill_bytes: int

    @property
    def skew_ratio(self) -> float:
        return self.max_ms / max(self.median_ms, 1.0)

    def diagnosis(self) -> str:
        if self.skew_ratio > 10:
            return (
                f"skew: slowest task is {self.skew_ratio:.0f}× the median. Switch to a "
                f"density-aware partitioner; raising parallelism will not help."
            )
        if self.spill_bytes > self.shuffle_written * 0.1:
            return (
                f"spill: {self.spill_bytes/1e9:.1f} GB spilled. Raise the partition "
                f"count or the executor memory fraction."
            )
        if self.skew_ratio < 3 and self.tasks < 200:
            return "even but under-parallelised — raise shuffle.partitions"
        return "healthy"


def profile_stage(event_log: Path, stage_id: int) -> StageProfile:
    """Read a stage's task-time distribution from the Spark event log."""
    durations: list[float] = []
    shuffle = spill = 0
    with open(event_log, encoding="utf-8") as handle:
        for line in handle:
            event = json.loads(line)
            if event.get("Event") != "SparkListenerTaskEnd":
                continue
            if event.get("Stage ID") != stage_id:
                continue
            info = event.get("Task Info", {})
            durations.append(float(info.get("Finish Time", 0) - info.get("Launch Time", 0)))
            metrics = event.get("Task Metrics", {})
            shuffle += metrics.get("Shuffle Write Metrics", {}).get("Shuffle Bytes Written", 0)
            spill += metrics.get("Disk Bytes Spilled", 0)

    if not durations:
        raise ValueError(f"no tasks found for stage {stage_id} in {event_log}")
    durations.sort()
    return StageProfile(
        tasks=len(durations),
        median_ms=durations[len(durations) // 2],
        max_ms=durations[-1],
        shuffle_written=shuffle,
        spill_bytes=spill,
    )


def partitions_for(shuffle_bytes: int, *, target_mb: int = 150) -> int:
    """Partition count from shuffle volume, not from core count.

    Core count tells you how many tasks can run at once; it says nothing about
    how much work one task should hold.
    """
    if shuffle_bytes <= 0:
        raise ValueError("shuffle volume must be positive")
    return max(16, int(shuffle_bytes / (target_mb * 1024**2)))

Validation

Change one thing, re-measure the distribution, and confirm the ratio moved. A stage time that improves without the ratio improving means something else changed and the skew is still there, waiting.

python
# Requires: the module above — before-and-after on the same stage
before = profile_stage(Path("/tmp/spark-events/app_uniform"), stage_id=7)
after = profile_stage(Path("/tmp/spark-events/app_kdbtree"), stage_id=7)

for label, profile in (("uniform grid", before), ("kdb-tree", after)):
    print(f"{label:<14} {profile.tasks:>5} tasks · median {profile.median_ms/1000:6.1f}s "
          f"· max {profile.max_ms/1000:7.1f}s · skew {profile.skew_ratio:5.1f}× "
          f"· spill {profile.spill_bytes/1e9:5.1f} GB")
    print(f"{'':<14} {profile.diagnosis()}")

Expected results after fixing the partitioner: skew ratio falling from double digits to under about 2, median task time rising slightly (each task now does more), maximum falling dramatically, and spill unchanged unless the partition count also changed.

Fix in this order, and re-measure between each step A four-step sequence. First, measure the task-time distribution for the join stage. Second, if the maximum exceeds ten times the median, replace the partitioner with a density-aware one and re-measure before doing anything else. Third, once the distribution is even, set the partition count from shuffle volume rather than core count. Fourth, only then adjust memory to control spill. A note warns that steps three and four applied before step two produce no improvement and consume a great deal of time. Apply the remedies in this order, re-measuring between each 1 · Measure median and max task time for the join stage 2 · Fix skew density-aware partitioner if max > 10 × median 3 · Size partitions from shuffle volume, not from core count 4 · Control spill memory fraction, only once even re-measure before step 3 Steps 3 and 4 applied before step 2 produce no improvement and consume a great deal of time. That is the single most common way tuning effort is wasted on spatial joins. Adaptive skew handling helps at the margin but does not replace a partitioner that understands density. Spill is usually a partition-count problem, not a memory problem Disk spill plotted against shuffle partition count for the same join on the same cluster. At low partition counts each partition exceeds execution memory and spills heavily. Raising the partition count reduces per-partition size and spill falls to zero without any change to the memory configuration. Beyond that point more partitions only add scheduling overhead. Disk spill against shuffle partition count spill reaches zero — no memory change needed 1.9 TB spilled 128 512 4,096 shuffle partitions → · beyond the elbow, more partitions buy only scheduling overhead

Edge Cases and Caveats

Adaptive execution masking the problem. Spark’s adaptive skew handling splits oversized shuffle partitions automatically, which improves the symptom enough that the underlying partitioning never gets fixed. That is fine until the dense cell exceeds what splitting can help with. Check the skew ratio even when adaptive execution is enabled and the stage looks acceptable.

Partition counts inherited from another job. A spark.sql.shuffle.partitions value tuned for a different workload is one of the most common causes of both over- and under-partitioning. Derive it per job from the measured shuffle volume, as partitions_for does.

Spill treated as a memory problem when it is a partition problem. Disk spill means a partition did not fit in execution memory; raising memory works, and raising the partition count usually works better and costs nothing. Try the partition count first, especially when the stage is already skew-free.

Frequently Asked Questions

Why doesn’t raising shuffle partitions fix a slow spatial join?

Because the problem is usually skew, and skew is about how work is distributed rather than how finely it is divided. If one grid cell over a city holds sixty per cent of the candidate pairs, that cell remains one partition however many partitions exist in total, so the stage still waits for it. More partitions divides the already-small partitions further and leaves the large one untouched. The fix is a density-aware partitioner that subdivides the dense region itself.

How do I tell skew apart from simply having too little parallelism?

Compare the maximum task duration against the median for the stage. If the maximum is within about two or three times the median, the stage is evenly distributed and more parallelism will help. If the maximum is ten or fifty times the median, one or two tasks are holding the stage and adding executors changes nothing — those extra executors sit idle waiting. The ratio is visible in the Spark UI’s stage summary without any instrumentation.

How many shuffle partitions should a spatial join use?

Enough that each task handles a few seconds of work, which usually means two to four times the total core count rather than exactly the core count. Too few and a straggler dominates; too many and scheduling overhead plus tiny shuffle files start to cost more than the work. Derive it from the shuffle volume — target roughly one hundred to two hundred megabytes of shuffle data per partition — and adjust from the measured task times.


← Back to Apache Sedona for Distributed Spatial Joins