5 PySpark Performance Anti-Patterns on Databricks (And How to Fix Them)

Databricks makes scaling big data processing feel almost effortless. But under the hood, writing PySpark without understanding how Apache Spark and Delta Lake execute your code can quietly ruin performance, spike DBU costs, and cause out-of-memory (OOM) crashes.

Here are 5 common PySpark anti-patterns seen in production, along with exact code fixes to keep your pipelines running fast and lean.

Anti-Pattern 1: The Python UDF Trap

The Problem

Python User-Defined Functions (UDFs) are convenient when built-in PySpark functions feel restrictive. However, Python UDFs force Spark to serialize data out of the JVM into a Python worker process and then back into the JVM. This eliminates Spark’s Catalyst Optimizer from analyzing your execution plan, dramatically slowing down transformations.

Python

# BAD: Using Python UDF for basic conditional logic
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
def categorize_amount(amt):
if amt is None:
return "UNKNOWN"
elif amt > 1000:
return "HIGH"
return "LOW"
categorize_udf = udf(categorize_amount, StringType())
df = df.withColumn("tier", categorize_udf(df["amount"]))

The Fix

Use native PySpark column expressions (when, otherwise). Native functions execute directly within the JVM, enabling memory optimization and vectorized evaluation.

Python

# GOOD: Using native PySpark functions
from pyspark.sql.functions import col, when
df = df.withColumn(
"tier",
when(col("amount").isNull(), "UNKNOWN")
.when(col("amount") > 1000, "HIGH")
.otherwise("LOW")
)

Pro-Tip: If complex logic forces you to write custom code, use Pandas UDFs (Vectorized UDFs) via Apache Arrow, which processes entire Arrow record batches at once rather than row-by-row.

Anti-Pattern 2: Overusing .collect() or Calling .toPandas() on Large Datasets

The Problem

Running .collect() or .toPandas() pulls every partition of a distributed DataFrame back to the Driver Node. On massive datasets, this chokes driver memory, causing driver OOM crashes and completely invalidating Spark’s distributed architecture.

Python

# BAD: Pulling millions of rows to the Driver
pandas_df = df.toPandas()
high_value = pandas_df[pandas_df["amount"] > 1000]

The Fix

Keep computation distributed on the workers. Filter or aggregate using Spark expressions first, or use display() / .take(n) if inspecting sample data in Databricks notebooks.

Python

# GOOD: Filter on distributed workers, or use PySpark Pandas API for large sets
from pyspark.sql.functions import col
high_value_df = df.filter(col("amount") > 1000)
# If you strictly need Pandas API syntax on distributed data:
import pyspark.pandas as ps
ps_df = df.pandas_api()

Anti-Pattern 3: Ignoring Data Skew on Join Keys

The Problem

When performing joins or aggregations, Spark reshuffles rows across nodes based on the join key. If 80% of your events belong to a single null key or generic tenant_id, one worker node ends up handling most of the load while other nodes sit idle. This causes tasks to hang at 99% completion forever.

Python

# BAD: Joining tables with heavy skew on customer_id
joined_df = orders_df.join(customers_df, on="customer_id", how="inner")

The Fix

Leverage Databricks Adaptive Query Execution (AQE), or manually apply Salting to spread skewed keys across multiple tasks.

Python

# GOOD: Enable AQE Skew Join handling in Databricks (Enabled by default in modern DBR)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
# MANUAL FIX (Salting): Append a random integer to split skewed keys
from pyspark.sql.functions import concat, lit, floor, rand
# Add a salt key (0 to 3) to the left table
orders_salted = orders_df.withColumn("salt", floor(rand() * 4)) \
.withColumn("salted_id", concat(col("customer_id"), lit("_"), col("salt")))
# Replicate right table 4 times with matching salts
customers_salted = customers_df.crossJoin(
spark.range(4).withColumnRenamed("id", "salt")
).withColumn("salted_id", concat(col("customer_id"), lit("_"), col("salt")))
# Join on the salted key to balance load across executors
joined_df = orders_salted.join(customers_salted, on="salted_id", how="inner")

Anti-Pattern 4: Over-Partitioning Small Datasets

The Problem

Partitioning a table by small granular keys (like transaction_date or hour) creates millions of tiny files on storage. Each file request incurs HTTP/storage metadata overhead, leading to the “Small File Problem” and crippling query performance.

Python

# BAD: Writing 1 GB of total data partitioned by date and hour
df.write.format("delta") \
.partitionBy("year", "month", "day", "hour") \
.saveAsTable("my_analytics_table")

The Fix

Replace rigid traditional table partitioning with Liquid Clustering (CLUSTER BY). Liquid Clustering dynamically adjusts data layout without the risks of over-partitioning, and allows you to redefine clustering keys at any time.

Python

# GOOD: Using Liquid Clustering instead of rigid partitioning
# SQL approach
%sql
CREATE TABLE my_analytics_table
CLUSTER BY (transaction_date, customer_id)
AS SELECT * FROM temp_view;
# Python DataFrame approach (Databricks Runtime 14.2+)
df.writeTo("my_analytics_table") \
.using("delta") \
.clusterBy("transaction_date", "customer_id") \
.create()

Anti-Pattern 5: Chain-Calling .withColumn() in a Loop

The Problem

Every call to .withColumn() creates a new internal projection step in Spark’s logical execution plan. In loops with dozens of columns, this generates a massive logical plan that takes minutes just for Catalyst to analyze, leading to driver memory pressure before the job even starts.

Python

# BAD: Modifying 20 columns in a sequential loop
columns_to_clean = ["col1", "col2", "col3", "col4", "col5"]
for col_name in columns_to_clean:
df = df.withColumn(col_name, trim(col(col_name)))

The Fix

Apply all column transformations simultaneously using select() or .withColumns() with a dictionary.

Python

# GOOD: Using select with list comprehension to build plan in a single pass
from pyspark.sql.functions import trim, col
clean_exprs = [
trim(col(c)).alias(c) if c in columns_to_clean else col(c)
for c in df.columns
]
df = df.select(*clean_exprs)
# ALTERNATIVE: PySpark 3.3+ withColumns dict syntax
transformations = {c: trim(col(c)) for c in columns_to_clean}
df = df.withColumns(transformations)

Summary Checklist

Anti-PatternOperational ImpactRecommended Fix
Python UDFsSerialization overhead, breaks CatalystBuilt-in PySpark functions (when/otherwise) or Pandas UDFs
.toPandas() / .collect()Driver Out-Of-Memory (OOM) crashesPerform filtering on workers or use pyspark.pandas
Data SkewStraggler tasks hanging at 99%AQE Skew Join settings or Salting
Over-PartitioningSmall file problem, slow readsDelta Lake Liquid Clustering (CLUSTER BY)
Looping .withColumn()Explosive plan analysis timesSingle select() expression list or .withColumns()

Leave a Reply

Discover more from Srinimf

Subscribe now to keep reading and get access to the full archive.

Continue reading