PySpark ML performance: partitions, data skew, caching, and the small-data trap

LEARN · DISTRIBUTED ML WITH PYSPARK

PySpark Performance problems are often blamed on cluster size, executor memory, or CPU capacity. Those factors matter, but many slow jobs fail for a simpler reason: the data is distributed badly.

A Spark application can waste minutes or hours because of:

Too many small partitions creating scheduling overhead.
Too few partitions limiting parallelism.
DataSkew sending most work to one unlucky task.
Caching the wrong datasets.
Using Spark for data that should run on a single machine.

This guide focuses on the practical mechanics behind PySpark Performance tuning:

Understanding partitions and shuffles.
Detecting and fixing DataSkew.
Using caching correctly.
Letting Adaptive Query Execution (AQE) help.
Knowing when Spark is the wrong engine.

The examples use current Apache Spark 4.0 concepts and standard PySpark DataFrame APIs. Always validate behavior against the exact Spark distribution and managed runtime you deploy.

The Spark execution model: partitions become tasks

Spark does not process a DataFrame as one continuous object. It divides data into partitions.

A simplified execution flow:

Input data
    |
    v
Partitions
    |
    v
Tasks on executors
    |
    v
Shuffle stages
    |
    v
Final result

Each partition normally becomes a task. The number and size of partitions directly influence performance.

Too few partitions can cause:

Idle executors.
Large tasks that run for a long time.
Memory pressure.
Poor cluster utilization.

Too many partitions can cause:

Excessive scheduling overhead.
Large amounts of task metadata.
Thousands of tiny tasks doing little work.

There is no universal partition count. A streaming pipeline, a large warehouse join, and a machine-learning feature pipeline have different needs.

Measure first.

print(df.rdd.getNumPartitions())

The physical plan can reveal expensive operations:

df.explain("formatted")

Look for:

Exchange operations, which usually indicate a shuffle.
Large joins that redistribute data.
Aggregations that require repartitioning.

Shuffles are expensive because Spark must:

Serialize data.
Transfer data across executors.
Sort or hash records.
Spill to disk if memory is insufficient.

Common shuffle triggers include:

join
groupBy
distinct
orderBy
repartition

Repartition versus coalesce

repartition() changes data distribution and performs a shuffle.

Use it when you need a new partition layout.

balanced_df = df.repartition(200, "customer_id")

This can help before a large join or aggregation when the current layout is inefficient.

coalesce() reduces partitions without a full shuffle.

A common use case is reducing output files after processing:

result_df = processed_df.coalesce(50)

(
    result_df.write
    .mode("overwrite")
    .parquet("/data/output")
)

Avoid accidental single-threaded execution:

df.repartition(1)

Writing one output file can be reasonable for a tiny result. For large datasets, it creates a single bottleneck.

DataSkew: when one key breaks a cluster

DataSkew occurs when values are distributed unevenly.

Imagine a transaction table:

customer_id	rows
CUSTOMER_001	900 million
CUSTOMER_002	1 million
CUSTOMER_003	1 million

Spark hash-partitions identical keys together. The partition containing CUSTOMER_001 becomes enormous while other executors finish quickly.

Typical symptoms:

Most tasks finish quickly.
A few tasks run much longer.
Executors appear healthy.
The Spark UI shows extreme task-duration differences.

The cluster is not necessarily too small. The workload may simply have a bad data distribution.

Find dominant keys before increasing hardware:

from pyspark.sql.functions import col

(
    df.groupBy("customer_id")
    .count()
    .orderBy(col("count").desc())
    .show(20)
)

Common causes:

A default value such as "UNKNOWN".
A single popular customer or product.
Bot-generated events.
Missing values mapped to one constant.

A runnable skew example

The following example intentionally creates one hot key.

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

hot_rows = [
    ("HOT_KEY", i)
    for i in range(900000)
]

normal_rows = [
    (f"KEY_{i}", i)
    for i in range(100000)
]

data = hot_rows + normal_rows

df = spark.createDataFrame(
    data,
    ["key", "value"]
)

print("Partitions:", df.rdd.getNumPartitions())

result = (
    df.groupBy("key")
    .count()
)

result.show()

In a production pipeline, the hot key usually comes from real business data rather than generated test records.

Fixing skew with salting

Salting spreads a hot key across multiple intermediate keys.

The pattern:

Add a salt value.
Group using the original key plus salt.
Perform partial aggregation.
Combine the partial results.

Example:

from pyspark.sql.functions import floor, rand, count, sum

salted = df.withColumn(
    "salt",
    floor(rand() * 20)
)

partial = (
    salted.groupBy("key", "salt")
    .agg(count("*").alias("partial_count"))
)

fixed = (
    partial.groupBy("key")
    .agg(sum("partial_count").alias("rows"))
)

fixed.show()

Salting is useful for:

Large aggregations with hot keys.
Large joins with uneven keys.
Historical datasets with known distribution problems.

It is not a free optimization. It creates extra intermediate work, so apply it where skew actually exists.

Why repartitioning does not automatically solve skew

A common assumption is:

df.repartition("key")

will fix uneven data.

It will not.

If one key owns most rows, repartitioning by that key simply sends the same problem to another partition layout.

Better strategies:

Salt extreme hot keys.
Process exceptional keys separately.
Choose a better partition key.
Use AQE where applicable.
Fix upstream data modeling problems.

The best partition strategy depends on the operation that follows.

Caching: store the right thing at the right time

Spark transformations are lazy. Creating a DataFrame does not execute the work immediately.

Caching helps when:

A transformed DataFrame is reused.
The transformation is expensive.
The cached data fits comfortably in memory.

Example:

from pyspark.sql.functions import col

cleaned = (
    raw_df
    .filter(col("status") == "active")
    .select("id", "category", "amount")
)

cleaned.cache()

cleaned.count()

summary = (
    cleaned.groupBy("category")
    .sum("amount")
)

details = cleaned.filter(
    col("amount") > 100
)

The first action materializes the cache.

Remove it when finished:

cleaned.unpersist()

Common caching mistakes:

Caching every intermediate DataFrame.
Caching data used only once.
Filling executor memory with data that is cheaper to recompute.
Ignoring cache eviction.

A bad cache strategy can slow down the pipeline.

Adaptive Query Execution: Spark learns during runtime

Adaptive Query Execution (AQE) allows Spark SQL to adjust parts of a query plan after seeing runtime statistics.

In Apache Spark 4.0, AQE remains a core SQL optimization feature.

AQE can help with:

Coalescing shuffle partitions.
Improving join strategies.
Handling some skewed joins.

Check your runtime configuration:

print(
    spark.conf.get(
        "spark.sql.adaptive.enabled"
    )
)

print(
    spark.conf.get(
        "spark.sql.shuffle.partitions"
    )
)

Enable AQE explicitly if your environment does not already do so:

spark.conf.set(
    "spark.sql.adaptive.enabled",
    "true"
)

spark.conf.set(
    "spark.sql.adaptive.skewJoin.enabled",
    "true"
)

AQE is valuable, but it is not magic.

It does not replace:

Data modeling.
Partition analysis.
Hot-key investigation.
Choosing the correct processing engine.

The small-data trap: sometimes Spark is the wrong tool

Spark is designed for distributed computation. Distribution creates overhead.

For data that fits comfortably on one machine, Spark can add unnecessary costs:

Cluster startup time.
Serialization.
Network communication.
Operational complexity.

Ask:

Does the dataset fit in RAM?
Is Spark startup longer than the actual computation?
Are there only a few transformations?
Do you need distributed execution?

A practical rule:

Workload	Often a better fit
Data fitting comfortably on one machine	pandas, Polars, DuckDB
Large distributed joins	Spark
Repeated production pipelines	Spark
Interactive local analysis	Local analytical tools
Very large-scale processing	Spark or another distributed engine

More machines do not automatically mean better performance.

Cherry on the cake: a Spark vulnerability caused by event logs

Performance engineers spend a lot of time analyzing Spark jobs, but the surrounding infrastructure also matters.

A surprising example is CVE-2025-54920, a Spark History Server vulnerability involving unsafe Jackson deserialization of event log data. The issue affected Apache Spark versions before 3.5.7 and 4.0.1, and upgrading to Spark 3.5.7 or 4.0.1+ is the recommended fix.

The lesson is unusual: the same event logs engineers use to debug performance can also become a security boundary.

Good operational habits:

Keep Spark versions patched.
Protect event log storage.
Restrict access to History Server infrastructure.
Treat operational data as production input.

A fast cluster is useful only when it is also safe to operate.

A practical PySpark Performance workflow

When a PySpark job is slow:

Open the Spark UI.
Find the slowest stage.
Look for task-duration outliers.
Inspect partition counts.
Check key frequency for joins and aggregations.
Remove unnecessary shuffles.
Apply salting only for real DataSkew.
Cache only reused expensive datasets.
Confirm AQE behavior.
Verify Spark is the right tool.

Performance tuning is not about memorizing a magic partition number. It is about matching execution plans, data distribution, and workload size.

Final checklist

Before shipping a PySpark pipeline:

Measure partitions instead of guessing.
Look for DataSkew before large joins.
Use salting for genuine hot-key problems.
Avoid unnecessary repartition() calls.
Use coalesce() carefully for output sizing.
Cache only valuable reused data.
Validate AQE settings.
Keep Spark components updated.
Question whether Spark is necessary.

Pick one slow PySpark job this week. Open the Spark UI, identify the bottleneck stage, measure the data distribution, make one targeted change, and record the before-and-after runtime. Build performance knowledge from evidence, not from random configuration changes.