Skip to main content
PySpark intermediate Lesson 7 of 10

UDFs and Pandas UDFs in PySpark

Measure what a Python UDF costs against a built-in function, then rewrite it as a pandas UDF and watch most of that cost disappear.

A UDF lets you run arbitrary Python on your data. It is also the most common reason a Spark job is ten times slower than it should be. Both facts are worth understanding precisely.

The baseline

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StringType, DoubleType
import time

spark = SparkSession.builder.appName("udfs").master("local[*]").getOrCreate()

df = spark.range(5_000_000).select(
    F.col("id"),
    (F.rand() * 1000).alias("amount"),
    F.concat(F.lit("user_"), F.col("id") % 10000).alias("username"),
).cache()
df.count()   # materialise the cache
$ python bench_setup.py
5000000

Built-in first

start = time.perf_counter()
df.withColumn("upper", F.upper("username")).count()
print(f"built-in upper:  {time.perf_counter() - start:.2f}s")
built-in upper:  0.38s

The same thing as a UDF

@F.udf(returnType=StringType())
def upper_udf(s):
    return s.upper() if s else None

start = time.perf_counter()
df.withColumn("upper", upper_udf("username")).count()
print(f"python udf:      {time.perf_counter() - start:.2f}s")
python udf:      9.72s

Twenty-five times slower for identical output. Every one of five million rows was serialised from the JVM, sent to a Python process, uppercased, and sent back.

The pandas UDF

import pandas as pd

@F.pandas_udf(StringType())
def upper_pandas(s: pd.Series) -> pd.Series:
    return s.str.upper()

start = time.perf_counter()
df.withColumn("upper", upper_pandas("username")).count()
print(f"pandas udf:      {time.perf_counter() - start:.2f}s")
pandas udf:      1.51s

Six times faster than the plain UDF. The function now receives a whole pd.Series at a time, transferred as an Arrow record batch, and str.upper() is a vectorised C loop rather than a Python one.

Summarising all three:

built-in upper:  0.38s   (1.0x)
pandas udf:      1.51s   (4.0x slower)
python udf:      9.72s  (25.6x slower)

The order is always this. Reach for a built-in; if there is none, reach for a pandas UDF; use a plain UDF last.

Why the optimiser cannot help

df.withColumn("upper", upper_udf("username")).filter(F.col("amount") > 900).explain()
== Physical Plan ==
*(2) Filter (isnotnull(amount#2) AND (amount#2 > 900.0))
+- *(2) Project [id#0L, amount#2, username#3, pythonUDF0#25 AS upper#20]
   +- BatchEvalPython [upper_udf(username#3)#19], [pythonUDF0#25]
      +- ...

BatchEvalPython runs below the Filter. Spark computed the UDF for all five million rows, then discarded 90% of them. Catalyst cannot push the filter below the UDF because it has no idea what the UDF does.

Reorder it yourself:

start = time.perf_counter()
df.filter(F.col("amount") > 900).withColumn("upper", upper_udf("username")).count()
print(f"filter first: {time.perf_counter() - start:.2f}s")
filter first: 1.14s

Nine times faster, same result. Filter before a UDF, never after. This single reordering recovers most of what a UDF costs.

Null handling

A UDF receives Python None for null and will raise if you do not expect it:

with_nulls = spark.createDataFrame([("alice",), (None,), ("carol",)], ["name"])

@F.udf(returnType=StringType())
def naive(s):
    return s.upper()

with_nulls.withColumn("up", naive("name")).show()
org.apache.spark.api.python.PythonException:
AttributeError: 'NoneType' object has no attribute 'upper'

Built-ins handle null for you; UDFs do not. Guard explicitly:

@F.udf(returnType=StringType())
def safe(s):
    return s.upper() if s is not None else None

with_nulls.withColumn("up", safe("name")).show()
+-----+-----+
| name|   up|
+-----+-----+
|alice|ALICE|
| NULL| NULL|
|carol|CAROL|
+-----+-----+

Wrong return type fails quietly

@F.udf(returnType=DoubleType())
def bad_type(x):
    return f"value-{x}"          # returns a string, declared double

df.limit(3).withColumn("v", bad_type("id")).show()
+---+------------------+---------+----+
| id|            amount| username|   v|
+---+------------------+---------+----+
|  0| 412.331...       |   user_0|NULL|
|  1| 887.204...       |   user_1|NULL|
|  2|  91.556...       |   user_2|NULL|
+---+------------------+---------+----+

Nulls, no error. Spark trusts your declared type and discards anything that does not match. A column of unexplained nulls after a UDF is nearly always a return-type mismatch — check it before looking anywhere else.

Pandas UDF variants

Series to Series — the scalar case shown above, one column in, one column out.

Iterator of Series — for expensive per-partition setup, like loading a model once:

from typing import Iterator

@F.pandas_udf(DoubleType())
def scored(batches: Iterator[pd.Series]) -> Iterator[pd.Series]:
    model = load_model()          # runs once per partition, not per batch
    for batch in batches:
        yield batch * model.coefficient

Grouped map — a whole group as a DataFrame, useful for per-group modelling:

def normalise(pdf: pd.DataFrame) -> pd.DataFrame:
    pdf["z"] = (pdf["amount"] - pdf["amount"].mean()) / pdf["amount"].std()
    return pdf

result = (
    df.limit(1000)
      .withColumn("grp", F.col("id") % 3)
      .groupBy("grp")
      .applyInPandas(normalise, schema="id long, amount double, username string, grp long, z double")
)
result.select("grp", "amount", "z").show(5)
+---+------------------+-------------------+
|grp|            amount|                  z|
+---+------------------+-------------------+
|  0|412.33127716044655|-0.2937847320441283|
|  0| 91.55663290023102|-1.4021559206733611|
|  0| 887.2047186413429| 1.3486722841255197|
|  1|233.90112884562913|-0.9124556207784412|
|  1|764.3319218301235| 0.9231044877213622|
+---+------------------+-------------------+

applyInPandas collects each group into a single pandas DataFrame in one executor’s memory, so a group larger than that executor’s heap will fail. It is the right tool for thousands of modest groups, the wrong one for a handful of enormous ones.

Check for a built-in first

Things people write UDFs for that already exist:

df.limit(2).select(
    F.regexp_replace("username", r"\d+", "#").alias("masked"),
    F.split("username", "_").getItem(1).alias("suffix"),
    F.when(F.col("amount") > 500, "high").otherwise("low").alias("band"),
    F.date_format(F.current_date(), "yyyy-MM").alias("period"),
    F.sha2(F.col("username").cast("string"), 256).substr(1, 12).alias("hash"),
).show(truncate=False)
+--------+------+----+-------+------------+
|masked  |suffix|band|period |hash        |
+--------+------+----+-------+------------+
|user_#  |0     |low |2026-02|8f2a1c943e7b|
|user_#  |1     |high|2026-02|3c91e5b7d248|
+--------+------+----+-------+------------+

Regex, splitting, conditionals, dates, hashing, JSON parsing, array and map operations — all built in, all running in the JVM at full speed. pyspark.sql.functions has several hundred entries; read it before writing a UDF.

Practice

1. Write a UDF that computes string length, and compare it to F.length.
F.length:    0.31s
python udf:  8.94s
pandas udf:  1.22s

Same 25x gap. The lesson generalises: the cost is the serialisation boundary, not the complexity of your function. A trivial UDF is just as expensive per row as a complex one.

2. Apply a UDF then filter, and filter then apply. Measure both.
udf then filter: 9.81s
filter then udf: 1.14s

Nearly 9x, because the second version runs the UDF on a tenth of the rows. Catalyst normally does this reordering for you — a UDF is where you have to do the optimiser’s job yourself.

3. Declare IntegerType and return a Python float from a UDF.
+---+----+
| id|   v|
+---+----+
|  0|NULL|
+---+----+

Null again. Spark does not coerce, and does not warn. When debugging unexpected nulls after a UDF, print type() of what your function actually returns — the mismatch is usually numpy.int64 versus Python int, which is easy to miss.

4. Time a pandas UDF with and without Arrow enabled.
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", False)
arrow enabled:  1.51s
arrow disabled: 8.63s

Without Arrow the pandas UDF falls back to row-by-row transfer and performs like a plain UDF. Arrow is on by default in Spark 3, but some environments disable it — if a pandas UDF is not faster than a plain one, this is the first setting to check.

Next: what Spark is actually doing between your code and the cluster.

Frequently Asked Questions

Why are Python UDFs slow?
Spark runs on the JVM, so every row has to be serialised, sent to a Python worker process, deserialised, processed, then sent back. That round trip happens per row and is invisible to Catalyst, which also cannot optimise through the UDF. Built-in functions run entirely in the JVM with no serialisation at all.
What makes a pandas UDF faster?
It moves batches of rows using Apache Arrow's columnar format instead of one row at a time, and your function operates on a whole pandas Series rather than a scalar. That amortises the serialisation cost across thousands of rows and lets you use vectorised NumPy and pandas operations.
Can Catalyst optimise through a UDF?
No. A UDF is an opaque black box to the optimiser, so it cannot push filters through it, prune columns based on it, or reorder around it. Filtering before a UDF instead of after is therefore something you have to do yourself.
When is a plain Python UDF acceptable?
When no built-in equivalent exists, the data volume is small, or the logic calls a Python library with no Spark counterpart. Always check pyspark.sql.functions first — it is large, and most things people write UDFs for are already there.