PySpark UDFs: custom functions on columns

A PySpark UDF wraps a Python function so it can run on DataFrame columns. Define it with `udf(fn, returnType)` or the `@udf` decorator. UDFs are flexible but slow — they serialize data to Python row-by-row and bypass Spark's optimizer, so always prefer built-in `pyspark.sql.functions` when one exists.

Defining a UDF

from pyspark.sql.functions import udf; from pyspark.sql.types import StringType; upper_udf = udf(lambda s: s.upper(), StringType()); df.withColumn("up", upper_udf(col("name"))).

Prefer built-ins / Pandas UDFs

Built-in functions run in the JVM and are far faster. If you must use Python logic on large data, a vectorized Pandas UDF (applyInPandas / @pandas_udf) is much faster than a plain row-at-a-time UDF.

Why a Python UDF is slow, and what to use instead

Built-in Spark functions run inside the JVM and are visible to the Catalyst optimiser. A Python UDF is not: every row must be serialised out of the JVM into a Python worker, processed, and serialised back. That round trip, plus the fact that the optimiser treats the UDF as an opaque black box and cannot push filters through it, is why UDFs are often the slowest part of a job. Before writing one, check whether the logic can be expressed with existing functions — when/otherwise, regexp_extract, date arithmetic and higher-order array functions such as transform and filter cover far more cases than people expect.

Pandas UDFs, null safety and return types

When you genuinely need Python, prefer a pandas UDF (a vectorised UDF). It uses Apache Arrow to move data in batches rather than row by row, which is typically several times faster than a plain UDF. Whichever kind you use, declare the return type accurately — a mismatch produces nulls rather than an error, which is easy to miss. Handle nulls explicitly inside the function, because Spark will happily pass None in and a Python exception on one row fails the whole task. Keep the function pure and deterministic; a UDF that depends on external state or the current time behaves unpredictably when tasks are retried.

Example (PySpark)

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame([("alice",), ("bob",)], ["name"])

upper_udf = udf(lambda s: s.upper(), StringType())
df.withColumn("name_upper", upper_udf(col("name"))).show()

Wraps a Python upper() in a UDF (a built-in upper() would be faster in practice).

Run this example in the free online PySpark compiler

Frequently asked questions

What is a UDF in PySpark?

A user-defined function that wraps custom Python logic so it can be applied to DataFrame columns via udf(fn, returnType) or the @udf decorator.

Why are PySpark UDFs slow?

They serialize data to Python and run row-by-row, bypassing Spark's Catalyst optimizer and Tungsten execution. Built-in functions and vectorized Pandas UDFs are much faster.

When should I use a UDF vs a built-in function?

Prefer built-in pyspark.sql.functions whenever possible; use a UDF only for logic that has no built-in equivalent, and consider a Pandas UDF for performance.

Open the free PySpark compiler · Data Engineering challenges · Data Engineering jobs