Data Engineering

PySpark Performance Tuning: 7 Techniques for Faster Data Pipelines

At scale, a slow PySpark job is not just inconvenient — it translates directly into compute cost and delayed insights. These seven techniques, applied in the right order, can cut job runtimes by 50–80% without touching your business logic.

Why PySpark Performance Matters

Apache Spark is a distributed compute engine — it parallelises work across a cluster. But parallelism amplifies both good and bad patterns. An unnecessary shuffle on a 10-node cluster moves 10× more data than the same mistake on a single machine. Understanding where Spark spends time — I/O, shuffle, serialisation, or CPU — determines which technique to apply first.

Start every optimisation session with the Spark UI. Look at stage durations, shuffle read/write sizes, and task skew before guessing at fixes.

Technique 1: Partitioning

Partitioning controls how data is distributed across executor cores. Too few partitions under-utilise the cluster; too many creates scheduling overhead.

# Too few partitions — most cores idle
df.repartition(4).count()

# Rule of thumb: 2–4 partitions per CPU core in the cluster
df.repartition(200).count()

# Partition by a column for join or groupBy performance
df.repartition(200, "customer_id").write.parquet("/output/")

# coalesce() reduces partitions WITHOUT a full shuffle (use before write)
df.coalesce(10).write.parquet("/output/")

Use repartition() when increasing partition count or when partitioning by a column. Use coalesce() only when reducing partitions before a final write — it avoids the full shuffle that repartition() triggers.

Technique 2: Caching & Persistence

If a DataFrame is used more than once in a pipeline — for example, as both a filter target and a join input — cache it. Without caching, Spark recomputes it from scratch on each action.

from pyspark import StorageLevel

# Default cache: MEMORY_AND_DISK
df.cache()

# Explicit storage level
df.persist(StorageLevel.MEMORY_AND_DISK_SER)

# Always unpersist when done — releases cluster memory
df.unpersist()

# Check if a DataFrame is cached
print(df.is_cached)

MEMORY_AND_DISK_SER serialises data before storing, using less memory at the cost of CPU. For very large DataFrames that don't fit in memory, DISK_ONLY is safer than letting Spark evict and recompute partitions unpredictably.

Technique 3: Broadcast Joins

A standard join between two large DataFrames triggers a shuffle — both sides are redistributed across the cluster so matching keys land on the same node. For joins where one side is small (<10 MB by default, configurable up to ~200 MB), a broadcast join eliminates the shuffle entirely by sending the small DataFrame to every executor.

from pyspark.sql.functions import broadcast

# Without broadcast: both tables shuffle (slow)
result = large_orders.join(customers, "customer_id")

# With broadcast: customers sent to all nodes (fast)
result = large_orders.join(broadcast(customers), "customer_id")

# Tune the auto-broadcast threshold (bytes)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", 209715200)  # 200 MB

In Azure Databricks, Photon accelerator further speeds up broadcast joins. Watch the Spark UI SQL plan — a BroadcastHashJoin node confirms the broadcast is working.

Technique 4: Avoiding Unnecessary Shuffles

Shuffles are the most expensive operations in Spark — they write data to disk and transfer it across the network. Common shuffle triggers to avoid or minimise:

• groupByKey() — shuffles all values before aggregation. Replace with reduceByKey() or DataFrame groupBy().agg().
• distinct() — triggers a shuffle. Use dropDuplicates(["key_col"]) on a subset of columns when possible.
• Multiple join() calls on the same key — repartition once by the join key before the chain of joins.

# Bad: groupByKey shuffles all values (high memory pressure)
rdd.groupByKey().mapValues(sum)

# Good: reduceByKey combines values locally before shuffle
rdd.reduceByKey(lambda a, b: a + b)

# Good: DataFrame API avoids RDD entirely
df.groupBy("customer_id").agg({"amount": "sum"})

Technique 5: Predicate Pushdown

Spark can push filter conditions down to the data source so that only the rows you need are read from storage. This is automatic with Parquet and Delta Lake — but you must filter before joins and aggregations for Spark to optimise correctly.

# Bad: filter after join (all data loaded first)
result = orders.join(products, "product_id") \
               .filter(col("order_date") >= "2026-01-01")

# Good: filter before join (Spark pushes filter to Parquet reader)
recent_orders = orders.filter(col("order_date") >= "2026-01-01")
result = recent_orders.join(products, "product_id")

# Verify pushdown in the query plan
result.explain(True)  # Look for "PushedFilters" in the plan

With Delta Lake, filtering on a partition column (e.g., order_date if it is the partition column) skips entire file directories — called partition pruning, and is dramatically faster than row-level filtering.

Technique 6: Columnar Formats (Parquet & Delta Lake)

Never store intermediate data as CSV. Parquet and Delta Lake use columnar storage, which means:

• Only the columns referenced in a query are read from disk.
• Each column is compressed independently (often 10–20× smaller than CSV).
• Min/max statistics per row group enable file skipping.

# Write as Parquet with Snappy compression (default in Spark)
df.write.format("parquet") \
    .option("compression", "snappy") \
    .save("/output/parquet/")

# Delta Lake adds ACID, Z-ORDER optimisation and VACUUM
df.write.format("delta").save("/output/delta/")

# Z-ORDER on high-cardinality filter columns speeds point queries
spark.sql("""
    OPTIMIZE delta.`/output/delta/`
    ZORDER BY (customer_id, order_date)
""")

Technique 7: Adaptive Query Execution (AQE)

Introduced in Spark 3.0 and mature in Spark 3.4+, Adaptive Query Execution re-optimises query plans at runtime based on actual statistics collected during execution. It handles three common problems automatically:

• Coalescing shuffle partitions — reduces 200 post-shuffle partitions to a smaller number based on actual data size.
• Converting sort-merge joins to broadcast joins — if one side turns out small after filtering.
• Skew join optimisation — splits skewed partitions automatically to even out task durations.

# Enable AQE (on by default in Databricks Runtime 7.3+)
spark.conf.set("spark.sql.adaptive.enabled", True)

# Tune the target size for coalesced shuffle partitions
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128mb")

# Enable skew join handling
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", True)
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", 5)

With AQE enabled, Spark can cut shuffle partition counts from 200 to 20 on a small dataset — reducing task scheduling overhead and memory pressure simultaneously. Enable it in every production pipeline.

Summary: Optimisation Priority Order

• First: Check the Spark UI. Identify whether the bottleneck is I/O, shuffle, or skew.
• If I/O is slow: Switch to Parquet/Delta, enable predicate pushdown, use partition pruning.
• If shuffle is large: Use broadcast joins, reduce shuffles via better aggregation patterns.
• If tasks are skewed: Enable AQE skew join or repartition by the skewed key with a salt.
• If recomputing DataFrames: Cache the right checkpoints.
• Always: Enable AQE and use columnar formats — zero-cost gains.