Spark Performance Tuning: Shuffles, Joins and Caching

Most slow Spark jobs are slow for a handful of reasons: too much data moved between machines in shuffles, the wrong join strategy, one partition far bigger than the rest, or the same data recomputed repeatedly. Guessing at configuration rarely helps; reading the query plan does. This guide shows each problem in real plans produced by local Spark on January 2024 New York yellow taxi trips (NYC TLC: trip record data), then what to change.

Shuffles, and why they cost so much

Spark divides each job into stages; a stage boundary is where data must be redistributed across executors, a shuffle (Spark docs: cluster overview). Grouping, joining two large tables, distinct and repartition all shuffle: rows are written out, sent over the network and read back. In a plan, a shuffle appears as an Exchange. Here is an aggregation, with a small helper that captures the plan as text:

import io
import re
import urllib.request
from contextlib import redirect_stdout
from pathlib import Path
from pyspark.sql import SparkSession, functions as F

base = "https://d37ci6vzurychx.cloudfront.net"
for name, url in {"yellow_tripdata_2024-01.parquet": f"{base}/trip-data/yellow_tripdata_2024-01.parquet",
                  "taxi_zone_lookup.csv": f"{base}/misc/taxi_zone_lookup.csv"}.items():
    if not Path(name).exists():
        urllib.request.urlretrieve(url, name)

spark = SparkSession.builder.master("local[4]").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")
trips = spark.read.parquet("yellow_tripdata_2024-01.parquet")
zones = spark.read.csv("taxi_zone_lookup.csv", header=True, inferSchema=True)

def plan(df):
    """The physical plan as text, without Spark's internal column ids such as #7."""
    buf = io.StringIO()
    with redirect_stdout(buf):
        df.explain()
    return re.sub(r"#\d+L?", "", buf.getvalue())

print("shuffle partitions:", spark.conf.get("spark.sql.shuffle.partitions"),
      "| adaptive execution:", spark.conf.get("spark.sql.adaptive.enabled"))
by_zone = trips.groupBy("PULocationID").count()
print(re.search(r"Exchange hashpartitioning\([^)]*\)", plan(by_zone)).group())
by_zone.collect()
print("partitions after the shuffle, adaptive on: ", by_zone.rdd.getNumPartitions())
spark.conf.set("spark.sql.adaptive.enabled", "false")
print("partitions after the shuffle, adaptive off:", trips.groupBy("PULocationID").count().rdd.getNumPartitions())
spark.conf.set("spark.sql.adaptive.enabled", "true")

shuffle partitions: 200 | adaptive execution: true
Exchange hashpartitioning(PULocationID, 200)
partitions after the shuffle, adaptive on:  1
partitions after the shuffle, adaptive off: 200

The plan hashes rows by pickup zone into 200 partitions, the default for spark.sql.shuffle.partitions (Spark docs: performance tuning). But the aggregated result is only a few hundred rows, so 200 partitions means 200 mostly empty tasks. Adaptive query execution, on by default since Spark 3.2.0, measured the real shuffle output at runtime and coalesced it into one partition. With adaptive execution off, all 200 remain. The lesson: keep adaptive execution on, and treat the 200 default as a starting point, raising it for very large shuffles so each partition stays a manageable size.

Joins: broadcast when one side is small

Joining two large tables requires shuffling both by the join key and sorting, a sort-merge join. If one side is small, Spark can instead send a full copy of it to every executor, a broadcast hash join, and the large side never moves:

strategy = lambda df: re.search(r"BroadcastHashJoin|SortMergeJoin|ShuffledHashJoin", plan(df)).group()
on = trips.PULocationID == zones.LocationID
print("default:", strategy(trips.join(zones, on)))
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
print("broadcast threshold disabled:", strategy(trips.join(zones, on)))
print("threshold disabled, broadcast hint:", strategy(trips.join(F.broadcast(zones), on)))
spark.conf.unset("spark.sql.autoBroadcastJoinThreshold")

default: BroadcastHashJoin
broadcast threshold disabled: SortMergeJoin
threshold disabled, broadcast hint: BroadcastHashJoin

The 265-row zone table is far below the automatic broadcast threshold of 10 MB, so Spark broadcast it. Disabling the threshold forced a sort-merge join, shuffling three million trips to join a lookup table. A broadcast hint overrides the setting. In practice:

• Check the join strategy of every expensive join in the plan.
• Hint a broadcast when you know a table is small but Spark's size estimate is wrong, common after filters or with formats lacking statistics.
• Never broadcast a large table: every executor must hold it in memory.

Adaptive execution can also convert a sort-merge join into a broadcast join at runtime when one side turns out small.

Skew: one partition doing most of the work

A stage finishes only when its slowest task does. If one partition holds far more data than the others, one task runs long while the rest sit idle. Taxi pickups are skewed by zone, and hash partitioning by zone shows it:

sizes = [r["count"] for r in trips.repartition(8, "PULocationID")
                                 .groupBy(F.spark_partition_id().alias("p")).count().orderBy("p").collect()]
print("rows per partition:", sizes)
print(f"largest / smallest: {max(sizes) / min(sizes):.1f}x")
top = trips.groupBy("PULocationID").count().orderBy(F.desc("count")).first()
print(f"busiest zone {top['PULocationID']}: {top['count']} trips, {top['count'] / trips.count():.1%} of all")

rows per partition: [479399, 445765, 96742, 359505, 211262, 574147, 378257, 419547]
largest / smallest: 5.9x
busiest zone 132: 145240 trips, 4.9% of all

The largest partition holds almost six times the rows of the smallest, and a single zone, 132, JFK Airport, accounts for 4.9% of all trips; every JFK row must land in the same partition. At cluster scale, keys like this, a default customer ID, a null, a dominant product, produce tasks that run for hours. Remedies:

• Let adaptive execution split skewed joins, which it does by default.
• Filter or handle nulls and placeholder keys before joining or grouping.
• Salt hot keys: add a random suffix to spread one key over several partitions, aggregate, then combine.

Caching: only when data is reused

Every action recomputes its plan from the source. If several actions reuse the same intermediate result, caching it saves the recomputation:

clean = trips.where(F.col("fare_amount") > 0).select("PULocationID", "payment_type", "fare_amount").cache()
clean.count()                                   # an action fills the cache
print("storage level:", clean.storageLevel)
print("later queries read the cache:", "InMemoryTableScan" in plan(clean.groupBy("payment_type").count()))
clean.unpersist()
spark.stop()

storage level: Disk Memory Deserialized 1x Replicated
later queries read the cache: True

The cache is lazy, filled by the first action, and stored in memory, spilling to disk if needed; later plans read it through an InMemoryTableScan. Cache after filtering and selecting, as here, so the cached copy is small. Do not cache data used once: it costs memory and time for no benefit. Unpersist when finished.

Use the Spark UI alongside the plan

The plan shows what Spark intends to do; the Spark UI shows what actually happened. While an application runs, the UI, served by the driver on port 4040 by default, lists every job and stage. Two views repay attention. The stage detail page shows task durations as a distribution, so a maximum far above the median is skew made visible. The SQL tab shows each query's plan with runtime metrics such as rows and bytes per operator, so you can see which Exchange moved the most data. Tune the stage that dominates the time, not the one that looks untidy.

A tuning checklist

1. Read the plan with explain() before changing settings.
2. Select columns and filter early, so Parquet scans read less.
3. Check join strategies; broadcast small tables.
4. Keep adaptive query execution on, and size shuffle partitions for large jobs.
5. Look for skewed keys and nulls in join and group columns.
6. Cache only data reused across actions.
7. Write output partitioned sensibly; partitioning and bucketing in Spark covers layout.
8. Replace Python UDFs with built-in functions, as the PySpark tutorial recommends.

If the architecture behind stages and executors is unfamiliar, start with Apache Spark explained for beginners.

Performance optimization is part of the Apache Spark module, and performance tuning and cost optimization of the monitoring module, in the Data Engineering course.

Sources

• Performance tuning (Spark SQL documentation) — Apache Spark
• Cluster mode overview (Apache Spark documentation) — Apache Spark
• TLC Trip Record Data (NYC Taxi and Limousine Commission) — NYC Taxi and Limousine Commission