Spark performance: AQE, skew and spills

A Spark job slows down, and its slowest stage shows 7 of its 8 tasks finished while one is still running. That one task holds 820,000 of the stage’s 1,000,000 rows, and more cores would not take a single row off it. So when a job is slow, look first at the tasks of its slowest stage and count the rows each one holds: one task far above the rest is skew, every task too large for its memory spills, and thousands of near-empty tasks mean a stage split too finely.

The row and task counts below were printed by PySpark 4.2.0 in local mode on fixed generated data: counts you can reproduce, never durations. Some of what they show, adaptive execution repairs by itself.

Three shapes of a slow stage, and the count behind each

A stage runs one task per partition and ends with its last task; where partitions come from is covered in Spark joins, shuffles and partitions. Given the partitions, three spreads of rows explain most slow stages:

• Skew. One key value holds far more rows than the rest, so one partition does too, and its task is the stragglerstraggler: A task that runs far longer than the other tasks of its stage, so the whole stage, and every stage after it, waits for that one task to finish. the stage waits for.
• Spills. Every task holds more than its memory takes, so its sort or hash table goes to local disk, a spillIn Spark, a task writing part of its sort buffer or hash table to local disk because it could not get more execution memory, then reading that part back to finish its work., and is read back.
• Too many tiny tasks. Thousands of partitions hold a few rows each, and starting a task costs more than its work.

Data skewdata skew: An uneven spread of rows over the partitions of a stage, usually because one key value is far more common than the others, so the task that holds that key does most of the stage's work. is the shape more machines never fix, because one task runs on one core. Spills are quieter: the Spark UI’s stage page reports them per task as shuffle spill (memory) and shuffle spill (disk). A fourth cost shows in no task count: a DataFrame that two actions read is computed twice, unless caching keeps it.

Adaptive query executionadaptive query execution: A Spark feature, on by default, that re-plans parts of a query at runtime from the actual rows and bytes it has seen, under an AdaptiveSparkPlan root. (AQE) re-plans a query between stages from the sizes finished stages report, and AQE is on by default in Spark 4.2. Its skew handling is documented for the two joins that shuffle both sides, sort-merge and shuffled hash. The 4.2 UI shows a query’s initial and final adaptive plans side by side, where a split shows. Read a slow stage in this order:

Only the last check needs the query itself.

Skew as a number: the largest partition against the mean

A ratio says more than a list of tasks. The run below builds 1,000,000 orders over 41 cities, sends four orders in five to city 0, and spreads them over 8 partitions twice, by a hash of the order id and of the city id, counting each partition’s rows with spark_partition_id().

"""Session driver: run skew_counts.py on PySpark 4.2.0 in local mode, inside a private lab folder.

Spark needs a Java runtime and the pyspark package. The driver points at the pinned ones (Java 21 and PySpark
4.2.0), gives Spark a home and a temporary folder inside the lab, binds it to 127.0.0.1 and prints only what the
program prints: configuration values and row counts, never Spark's own log lines or paths.
"""

import os
import pwd
import subprocess
import sys
import tempfile
from pathlib import Path

JAVA_HOME = os.path.join(pwd.getpwuid(os.getuid()).pw_dir, ".local", "share", "mise", "installs", "java",
                         "temurin-21.0.12+101.0.LTS")
HERE = Path(__file__).resolve().parent
PROGRAM = "skew_counts.py"


def pinned_pyspark():
    """The folder that holds the pinned pyspark 4.2.0 package, or None to use the one installed for this Python."""
    for folder in [Path.cwd(), *Path.cwd().parents, *HERE.parents]:
        for found in folder.glob("*/toolchains/pyspark-4.2.0/python"):
            return str(found)
    return None


with tempfile.TemporaryDirectory(prefix="spt-skew-") as lab:
    for name in ("home", "tmp", "out"):
        os.makedirs(os.path.join(lab, name))
    java_home = JAVA_HOME if os.path.isdir(JAVA_HOME) else os.environ.get("JAVA_HOME", "")
    env = {
        "PATH": os.pathsep.join([os.path.join(java_home, "bin"), "/usr/bin", "/bin"]),
        "JAVA_HOME": java_home,
        "HOME": os.path.join(lab, "home"),
        "TMPDIR": os.path.join(lab, "tmp"),
        "WORK": os.path.join(lab, "out"),
        "SPARK_LOCAL_IP": "127.0.0.1",
        "PYSPARK_PYTHON": sys.executable,
        "TZ": "UTC",
        "LANG": "en_US.UTF-8",
    }
    package = pinned_pyspark()
    if package:
        env["PYTHONPATH"] = package
    done = subprocess.run([sys.executable, str(HERE / PROGRAM)], cwd=lab, env=env,
                          capture_output=True, text=True, check=False)
    print(done.stdout.replace(lab, "<lab>").rstrip())
    if done.returncode != 0:
        last = [line for line in done.stderr.splitlines() if line.strip()][-1:] or ["no message"]
        print(PROGRAM, "stopped with exit status", done.returncode, "-", last[0].replace(lab, "<lab>"))
        sys.exit(done.returncode)

PySpark 4.2.0, local mode on 2 cores
spark.sql.adaptive.enabled = true (not set by this program)
spark.sql.adaptive.skewJoin.enabled = true (not set by this program)

orders = 1000000, cities = 41
city 0 holds 800000 orders (80.0 % of all); each other city holds 5000

hashed on order_id into 8 partitions (a key with 1,000,000 values)
  rows per partition: [124954, 124686, 124925, 125210, 125164, 125249, 125087, 124725]
  largest = 125249, mean = 125000.0, largest / mean = 1.00
  the largest partition holds 12.5 % of the rows; 4 of 8 partitions hold less than the mean

hashed on city_id into 8 partitions (a key with 41 values, one of them hot)
  rows per partition: [25000, 30000, 15000, 40000, 10000, 30000, 30000, 820000]
  largest = 820000, mean = 125000.0, largest / mean = 6.56
  the largest partition holds 82.0 % of the rows; 7 of 8 partitions hold less than the mean

"""PySpark 4.2.0, local mode: measure skew as a number, from the rows each partition holds.

One million generated orders over 41 cities, one of which takes 80 % of them, are spread over 8 partitions twice:
once by a hash of the order id, once by a hash of the city id. For each spread the program prints the rows of every
partition, the largest against the mean, and the share of all rows the largest partition holds. Every number is a
row count of a fixed dataset, so two runs print the same lines.
"""

from pyspark.sql import SparkSession
from pyspark.sql.functions import spark_partition_id

spark = (SparkSession.builder
         # learn-lint-allow inclusive: PySpark 4.2.0 names the builder method that sets the cluster URL master()
         .master("local[2]")
         .appName("skew-as-a-number")
         .config("spark.ui.enabled", "false")
         .config("spark.driver.bindAddress", "127.0.0.1")
         .config("spark.driver.host", "127.0.0.1")
         .config("spark.driver.memory", "1g")
         .config("spark.sql.shuffle.partitions", "8")
         .getOrCreate())
spark.sparkContext.setLogLevel("ERROR")

print(f"PySpark {spark.version}, local mode on {spark.sparkContext.defaultParallelism} cores")
for key in ("spark.sql.adaptive.enabled", "spark.sql.adaptive.skewJoin.enabled"):
    print(f"{key} = {spark.conf.get(key)} (not set by this program)")

# Every fifth order goes to one of cities 1 to 40 in turn; the other four in five go to city 0.
orders = (spark.range(0, 1000000)
          .selectExpr("id as order_id",
                      "case when id % 5 = 0 then cast(1 + (id div 5) % 40 as int) else 0 end as city_id",
                      "cast(100 + pmod(hash(id), 900) as int) as amount"))

per_city = {row["city_id"]: row["count"] for row in orders.groupBy("city_id").count().collect()}
others = sorted({count for city, count in per_city.items() if city != 0})
print(f"\norders = {sum(per_city.values())}, cities = {len(per_city)}")
print(f"city 0 holds {per_city[0]} orders ({100 * per_city[0] / sum(per_city.values()):.1f} % of all); "
      f"each other city holds {', '.join(str(c) for c in others)}")


def rows_per_partition(df):
    """The row count of every partition of df, in partition order, with 0 for a partition that holds no rows."""
    held = {row["part"]: row["count"]
            for row in df.withColumn("part", spark_partition_id()).groupBy("part").count().collect()}
    return [held.get(part, 0) for part in range(df.rdd.getNumPartitions())]


def report(label, df):
    """Print one spread: the rows of each partition, the largest against the mean, and its share of all rows."""
    rows = rows_per_partition(df)
    mean = sum(rows) / len(rows)
    print(f"\n{label}")
    print(f"  rows per partition: {rows}")
    print(f"  largest = {max(rows)}, mean = {mean:.1f}, largest / mean = {max(rows) / mean:.2f}")
    print(f"  the largest partition holds {100 * max(rows) / sum(rows):.1f} % of the rows; "
          f"{sum(1 for r in rows if r < mean)} of {len(rows)} partitions hold less than the mean")


report("hashed on order_id into 8 partitions (a key with 1,000,000 values)", orders.repartition(8, "order_id"))
report("hashed on city_id into 8 partitions (a key with 41 values, one of them hot)", orders.repartition(8, "city_id"))

spark.stop()

Both spreads put 1,000,000 rows into 8 partitions, a mean of 125,000 each, so the largest partition of each spread says how even it is:

Hashed on the order id, the largest partition is 249 rows above the mean. Hashed on the city id, one holds 82.0 % of the rows, city 0’s 800,000 plus four small cities, so its task does the work of 6.56 average ones on any cluster. The run also printed spark.sql.adaptive.enabled and spark.sql.adaptive.skewJoin.enabled as true, though the program sets neither. What each shape needs next:

• For a skewed sort-merge or shuffled hash join, AQE’s skew-join handlingskew-join handling: Adaptive query execution splitting an oversized partition of a sort-merge or shuffled hash join into several tasks at runtime, switched by spark.sql.adaptive.skewJoin.enabled, which is on by default. splits the hot partition into several tasks by itself.
• Where it cannot, saltingsalting: Adding a computed suffix to a hot key so that its rows spread over several partitions, then combining the parts in a second step that removes the suffix. adds a computed suffix to the hot key so its rows spread out, and a second aggregation removes it.
• For spills, more shuffle partitions or narrower rows shrink each task’s share before anyone adds memory.
• Cachingcaching: Keeping a DataFrame's computed rows in executor memory or on local disk, so that later actions that read the DataFrame use those rows instead of computing it again. keeps a DataFrame that several actions reuse, and unpersist() releases it afterwards.
• Sizing works the partition count out from the data size, then the cores and the executor memoryexecutor memory: The heap memory of each Spark executor process, set by spark.executor.memory and shared by the tasks that run in that executor..

Practise

Six questions on the three shapes and the skew ratio, each choice explained once you answer.

Skew: what adaptive execution splits, and the salting that follows when it cannot Pro: opens with a Pro, Premium or Ultimate pass

Spills and caching, decided from counts Pro: opens with a Pro, Premium or Ultimate pass

Sizing partitions, cores and memory, worked from the data size Pro: opens with a Pro, Premium or Ultimate pass

Interview questions

Warm-up: what is data skew in Spark, and how do you see it before you change anything?

Data skew is an uneven spread of rows over the partitions of a stage, nearly always because one key value is far more common than the rest. Spark sends every row of a key to the partition its hash picks, so a hot key fills one partition and its task becomes the straggler the stage waits for. You see it by counting rows per partition with spark_partition_id() and comparing the largest with the mean, or on the stage page as the one task with a far larger shuffle read. A ratio near 1 is an even stage; the hot city here gave 6.56.

Three questions carry the same diagnosis into cases with numbers:

Key takeaways

• A slow stage usually has one of three shapes: one task far above the rest, every task too large for its memory, or thousands of near-empty tasks.
• Counting rows per partition turns skew into a number: the hot city put 820,000 of 1,000,000 rows into one of 8 partitions, 6.56 times the mean.
• Adaptive execution (AQE) and its skew-join handling are on by default in PySpark 4.2.0, and the skew handling is documented for sort-merge and shuffled hash joins.