Optimization: replacing slow Python UDFs on Databricks - built-in functions, pandas UDFs and Arrow

#Why plain Python UDFs are slow

• Each row is serialized from the JVM to a Python worker and back.
• The optimizer can't see inside the UDF, so it can't push filters or reorder work.
• Photon can't run the stage, so it falls back to the slower engine.

#Option 1: built-in functions (best)

from pyspark.sql import functions as F

# Before: Python UDF
@F.udf("string")
def clean_phone(p):
    return "".join(c for c in p if c.isdigit()) if p else None

# After: built-in, Photon-friendly
df = df.withColumn("phone_clean", F.regexp_replace("phone", r"[^0-9]", ""))

Most string, date, JSON (from_json, get_json_object), array (transform, filter) and conditional (when) logic can be expressed with built-ins. SQL higher-order functions cover many loops.

#Option 2: pandas UDFs (vectorized)

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

@pandas_udf("double")
def score(amount: pd.Series, days: pd.Series) -> pd.Series:
    return amount * (0.95 ** days)

df = df.withColumn("decayed", score("amount", "days_since"))

Data moves in Arrow batches, which is far cheaper than row by row.

#Option 3: Arrow-optimized Python UDFs

For simple scalar UDFs you don't want to rewrite, Arrow optimization can speed up the transfer:

@F.udf("string", useArrow=True)
def legacy_logic(x):
    ...

Or set spark.sql.execution.pythonUDF.arrow.enabled=true. Check support on your runtime version.

#Avoid driver-side traps

#Measure

• Spark UI: stages with BatchEvalPython / ArrowEvalPython show UDF cost.
• Compare run time and Photon coverage before and after.

Tip

Search the codebase for @udf and udf(; list each with the job that uses it and the job's runtime. Rewriting the UDFs in the three slowest jobs usually gives most of the benefit.

#Checklist

• UDFs inventoried by job
• Built-in rewrites done where possible
• pandas or Arrow UDFs for the rest
• Driver-side pandas/collect removed from production jobs
• Run time and Photon coverage compared