How do I avoid Python UDF overhead?

Summary

• Use a built-in Spark function when one exists. A Python UDF must move data out of the JVM to the Python interpreter.
• When the logic needs Python, use a pandas UDF. Spark splits columns into batches and calls your function once per batch instead of once per row.
• Since Spark 4.2.0, regular Python UDFs use Arrow for serialization by default, but they still process data row by row.
• spark.sql.execution.arrow.maxRecordsPerBatch sets the Arrow batch size (default 10,000 records per batch). Adjust it for wide tables to avoid out-of-memory errors.
• On Databricks, Photon's documented limitations include UDFs, which is one more reason to prefer built-in functions.

How do I avoid Python UDF overhead?

Replace the UDF with a built-in function whenever one exists, so Spark evaluates the logic without a trip to the Python interpreter. When the logic needs Python, use a pandas UDF so Spark calls your function once per batch of rows instead of once per row.

Why do Python UDFs add overhead?

Databricks' UDF guidance explains that Python UDFs and pandas UDFs tend to be slower than Scala UDFs because they must serialize data and move it out of the JVM to the Python interpreter. Two costs follow from that:

1. Serialization. Before Spark 4.2, regular Python UDFs were pickled by default and Arrow was opt-in. Spark 4.2.0 enables Arrow-optimized Python UDFs by default (SPARK-54555).
2. One call per row. An Arrow-optimized Python UDF still processes data row by row with regular Python types. Arrow makes the data transfer more efficient but does not batch your function calls.

UDFs also have an evaluation quirk. The udf reference says they do not support conditional expressions or short circuiting in boolean expressions, so a guard in the query may not stop the function from running on the rows the guard was meant to exclude. If the function can fail on special rows, put the check inside the function.

Which option should I use?

Option	When to use it	Trade-off
Built-in function	The logic maps to a function in Spark's function library	Limited to the functions Spark provides
Pandas UDF	You need pandas or Python code on whole columns	Your function receives pandas objects, and large batches use more memory
Iterator pandas UDF	The function must initialize state, such as loading a model file	You write a loop over batches and yield results
Python UDF (Arrow-optimized from Spark 4.2)	The logic is row-wise and cannot be vectorized	Still one Python call per row, with no short circuiting

How do I replace a Python UDF with a built-in function?

from pyspark.sql.functions import col, udf, upper

@udf(returnType="string")
def to_upper(s):
    if s is not None:
        return s.upper()

df.select(to_upper(col("name")))  # Python UDF: runs in Python
df.select(upper(col("name")))     # built-in: runs in Spark

Check the PySpark function reference first. Its string functions include upper, substring and regexp_replace. Its conditional functions include coalesce and when. It also lists date_format, round, and the collection functions transform and filter, which work on array columns.

How do I write a pandas UDF?

A pandas UDF is executed by Spark using Arrow to transfer data and pandas to work with the data. Spark splits columns into batches, calls the function for each batch, and concatenates the results. For a Series to Series UDF, the output should always be the same length as the input.

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

@pandas_udf("string")
def to_upper(s: pd.Series) -> pd.Series:
    return s.str.upper()

@pandas_udf("long")
def plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    # initialize state here, such as loading a model file
    for s in iterator:
        yield s + 1

How do I use an Arrow-optimized Python UDF?

Arrow Python UDFs are executed row by row and use Arrow for efficient batch data transfer and serialization. Since Spark 4.2, regular Python UDFs are Arrow-optimized by default. To opt out, set useArrow=False on one UDF, or set spark.sql.execution.pythonUDF.arrow.enabled to false for the session. On Spark 4.1, opt in with useArrow=True or by setting that configuration to true. The configuration takes effect only when useArrow is unset.

from pyspark.sql.functions import udf

@udf(returnType="int", useArrow=True)  # An Arrow Python UDF
def arrow_slen(s):
    return len(s)

Type hints are specified the same way as for pickled Python UDFs. Check one behavior difference when you upgrade: the Spark Arrow guide says pickled UDFs rely on None as a fallback for type mismatches, which can cause ambiguity and data loss, and that Arrow Python UDFs standardize type coercion instead.

How do I control batch size and memory?

Spark converts data partitions into Arrow record batches, which can temporarily lead to high memory usage in the JVM. spark.sql.execution.arrow.maxRecordsPerBatch sets the maximum number of rows per batch, and the default is 10,000 records per batch. Adjust it when your tables have many columns. In iterator-based pandas UDFs, the length of each series is the length of one internal batch, so the same setting applies.

How does Databricks approach Python UDF performance?

Databricks documents the same Spark APIs and adds guidance on choosing among them.

• The UDF overview says built-in Apache Spark functions are optimized for distributed processing and offer better performance at scale. It lists built-in functions and SQL UDFs as the most efficient options.
• Photon accelerates SQL workloads, DataFrame API calls, ETL pipelines and stateless streaming workloads. Its limitations section states that Photon does not support UDFs.
• On serverless compute, PySpark UDFs have a memory limit of 1GB per PySpark UDF.

FAQs

Are Python UDFs still row by row on Spark 4.2?

Yes. Arrow-optimized Python UDFs use Arrow to move data but still process it row by row. To process batches, write a pandas UDF.

When is a UDF acceptable?

Databricks recommends UDFs for ad hoc queries, manual data cleaning, exploratory data analysis, and operations on small to medium-sized datasets. For large datasets and for workloads that run regularly or continuously, it recommends Apache Spark methods.