Python UDFs vs Built-in Functions vs Pandas UDFs

What happens in a Python UDF

Spark runs in the JVM. For a regular Python UDF, each row is serialised, sent to a Python worker, processed and sent back. The optimizer cannot see inside the function, so it cannot push filters through it or prune columns.

Rewrite with built-ins

# UDF
@F.udf("string")
def band(amount):
    return "high" if amount > 1000 else "low"
 
# built-in: same result, optimised
df = df.withColumn("band", F.when(F.col("amount") > 1000, "high").otherwise("low"))

Common replacements: F.regexp_extract, F.split, F.coalesce, F.date_trunc, F.transform (for arrays) and F.from_json.

When you need Python: pandas UDF

import pandas as pd
 
@F.pandas_udf("double")
def score(amount: pd.Series, qty: pd.Series) -> pd.Series:
    return model.predict(pd.DataFrame({"amount": amount, "qty": qty}))

It processes whole Arrow batches, so serialisation cost is shared across thousands of rows.