TopolearnA careless Python UDF is the classic reason a PySpark job runs ten times slower than it should.
Built-in functions run inside the JVM, are optimized by Catalyst, and are code-generated. A Python UDF runs in separate Python worker processes: rows are serialized from the JVM to Python and back, and Catalyst can't see inside it. So first look for a built-in in pyspark.sql.functions or a SQL expression; there almost always is one.
When you do need Python:
- Row-at-a-time UDF:
@F.udf(returnType="string")on a normal function, called once per row. Classic execution pickles every row, which is slow. Arrow-optimized execution (useArrow=True, orspark.sql.execution.pythonUDF.arrow.enabled, on by default in Spark 4.2) moves data in Arrow batches and is much faster, but the function is still called row by row. - pandas UDF (vectorized):
@F.pandas_udf("double")on a function that takes and returnspd.Series. Data arrives in Arrow batches of thousands of rows, so you can use vectorized pandas and NumPy code. This is usually the best choice. - Python UDTFs (
@udtf) return multiple rows per input.
Rules: always declare the return type, handle nulls (None or NaN) inside the function, keep functions pure, and import heavy libraries inside the function or at module level, never pickling large objects in the closure. Register a UDF with spark.udf.register("name", f) to call it from SQL.
Example
from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.master("local[2]").getOrCreate()
import pandas as pd
df = spark.createDataFrame([("Ana Lopez",), ("ben ODA",), (None,)], "name STRING")
@F.udf(returnType="string", useArrow=True) # Arrow-optimized Python UDF
def initials(s):
return None if s is None else "".join(w[0].upper() for w in s.split())
@F.pandas_udf("int") # vectorized: whole Series at once
def name_len(s: pd.Series) -> pd.Series:
return s.str.len().astype("Int32")
out = df.select("name", initials("name").alias("ini"), name_len("name").alias("n"))
print([tuple(r) for r in out.collect()])
# [('Ana Lopez', 'AL', 9), ('ben ODA', 'BO', 7), (None, None, None)]
# The built-in way: no Python at all, and fully optimized
print(df.select(F.length("name").alias("n")).first().n) # 9Remember
- Prefer built-in functions; UDFs are opaque and cross the JVM-Python boundary
- pandas UDFs process Arrow batches and are far faster than row UDFs
- Declare return types and handle nulls inside UDFs
Watch out for
- Writing a Python UDF for something like lowercasing or date parsing that a built-in already does.