PySpark UDFs are slow: built-in functions, pandas UDFs and Arrow-optimized UDFs compared

A PySpark job spends most of its time in one stage, CPU is busy on the executors, and the Spark UI shows nothing unusual: no skew, no spill, small shuffles. The stage contains a Python UDF. Every row leaves the JVM, gets pickled, crosses a socket to a Python worker, runs through the interpreter and comes back. For simple logic, that serialization costs more than the work.

This guide shows how to find these UDFs, then replaces them in order of preference.

#Find the Python UDFs in the plan

Run explain on the DataFrame that feeds the slow write:

df_out.explain(mode="formatted")

Look for these nodes:

Two more costs that are easy to miss:

• No pushdown through a UDF. A filter(my_udf(col("x"))) cannot be pushed to Parquet or Delta, so Spark reads every row first. Check whether PushedFilters in the scan node is empty.
• No whole-stage codegen across it. Operators on either side of the Python node run in separate stages of evaluation, and Spark may evaluate the UDF more than once if the column is referenced several times.

#Option 1: replace the UDF with built-in functions

Most UDFs in production are string cleanup, date parsing, conditional mapping or JSON extraction, all of which exist as built-ins and run inside the JVM with codegen.

# Before: row-at-a-time UDF
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

@udf(StringType())
def normalise_phone(p):
    if p is None:
        return None
    digits = "".join(c for c in p if c.isdigit())
    return digits[-10:] if len(digits) >= 10 else None

# After: built-ins, no Python worker involved
from pyspark.sql import functions as F

digits = F.regexp_replace(F.col("phone"), r"[^0-9]", "")
df = df.withColumn(
    "phone_norm",
    F.when(F.length(digits) >= 10, F.substring(digits, -10, 10))
)

Useful replacements:

Tip

Test that the replacement returns the same results before switching. Run both versions on a sample and compare with exceptAll in both directions. Null handling and Unicode digits are where they usually differ.

#Option 2: Arrow-optimized Python UDFs (Spark 3.5+)

When the logic really needs Python but is still row-at-a-time, Spark 3.5 added Arrow-optimized UDFs. The function body is unchanged; data moves in Arrow batches instead of pickled rows.

@udf(returnType=StringType(), useArrow=True)
def classify(text):
    ...

Or for the session, which applies to UDFs that do not set useArrow explicitly:

spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "true")

The default for this setting has been changing between releases; check your version's docs or spark.conf.get rather than assuming. Type coercion differs slightly from pickled UDFs (for example how mismatched return types are handled), so run your tests after turning it on.

Do not confuse it with spark.sql.execution.arrow.pyspark.enabled, which controls toPandas() and createDataFrame(pandas_df), not UDFs.

#Option 3: pandas UDFs for vectorised logic

If the work can be expressed with pandas or NumPy over a whole column, a pandas UDF processes a batch at a time with vectorised code.

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def haversine_km(lat1: pd.Series, lon1: pd.Series,
                 lat2: pd.Series, lon2: pd.Series) -> pd.Series:
    import numpy as np
    r = 6371.0
    p1, p2 = np.radians(lat1), np.radians(lat2)
    dphi = p2 - p1
    dlmb = np.radians(lon2 - lon1)
    a = np.sin(dphi/2)**2 + np.cos(p1)*np.cos(p2)*np.sin(dlmb/2)**2
    return pd.Series(2 * r * np.arcsin(np.sqrt(a)))

df = df.withColumn("km", haversine_km("lat1", "lon1", "lat2", "lon2"))

For expensive set-up such as loading a model or opening a client, use the iterator form so it runs once per partition rather than once per batch:

from typing import Iterator

@pandas_udf("string")
def predict(batches: Iterator[pd.Series]) -> Iterator[pd.Series]:
    model = load_model("/dbfs/models/classifier.pkl")   # once per Python worker task
    for s in batches:
        yield pd.Series(model.predict(s.to_frame()))

Batch size is controlled by spark.sql.execution.arrow.maxRecordsPerBatch (default 10,000). Lower it for wide rows or memory-hungry models; raise it carefully for narrow numeric columns.

#Option 4: mapInPandas and applyInPandas

mapInPandas hands each partition to Python as a stream of DataFrames; use it when output rows do not map one-to-one to input rows.

groupBy(...).applyInPandas(fn, schema) is the one that causes outages. Each group is loaded fully into the Python worker's memory, regardless of batch size. One customer with 50 million rows fails the job.

# Check the largest groups before using applyInPandas
(df.groupBy("customer_id").count()
   .orderBy(F.desc("count")).show(10))

If the top group is large, pre-aggregate in Spark, split the key (for example by date) or rewrite the per-group logic with window functions.

#Memory: Python workers live outside the JVM heap

Python worker memory is not part of spark.executor.memory. On YARN and Kubernetes it comes out of the memory overhead, so heavy pandas UDFs produce container kills ("exceeding memory limits", exit code 137 or OOMKilled) rather than Java OOM errors.

spark-submit \
  --conf spark.executor.memory=8g \
  --conf spark.executor.memoryOverhead=4g \
  --conf spark.executor.pyspark.memory=3g \
  --conf spark.sql.execution.arrow.maxRecordsPerBatch=2000 \
  job.py

Platform-managed environments such as Databricks may size overhead for you; check your platform's documentation before overriding it.

Warning

On shared or Unity Catalog standard access mode clusters, Python UDFs run in a sandbox with extra restrictions and overhead. Check the access mode limitations for your runtime before moving UDF-heavy jobs onto shared compute.

#Measure the change

Compare before and after on the same input with the stage time in the Spark UI, or time it directly:

import time
t = time.time()
df_out.write.format("noop").mode("overwrite").save()
print(f"{time.time() - t:.1f}s")

The noop format runs the full plan without writing output, so you measure compute, not storage.

#Prevent it next time

• Add a review rule: a new Python UDF needs a short note on why built-ins do not work.
• Search the codebase for @udf and udf( periodically and track the count.
• Use explain in CI for key jobs and fail on unexpected BatchEvalPython nodes.
• Check group size distribution before any applyInPandas goes to production, and alert on it if keys can grow.