A job spends most of its time in a Python UDF. How do you make it faster?

Replace the UDF with built-in Spark functions if you can. If you cannot, turn it into a vectorized pandas UDF that works on whole batches. A plain Python UDF is slow because every row is sent from the JVM to a Python process, handled one at a time by your function, and sent back. Built-in functions run inside the JVM and skip that trip.

1. Confirm the UDF is the cost

In the Spark UI's SQL tab, the query plan shows a separate Python evaluation step for the UDF. Find the stage that runs it and compare its time with the rest of the job. A quick test is to run the job once with the UDF column replaced by a constant.

2. Rewrite it with built-in functions

Most UDFs in real code do things Spark already has: string cleanup, parsing dates, mapping codes, simple maths.

# Before: a Python UDF
@udf("string")
def clean_phone(s):
    return re.sub(r"[^0-9]", "", s) if s else None

# After: the same with a built-in, no Python involved
df = df.withColumn("phone", F.regexp_replace("phone", r"[^0-9]", ""))

Other common swaps: F.when chains for if/else logic, F.to_date with a format, a join to a small mapping table instead of a dictionary lookup, and higher-order functions (F.transform, F.filter, F.aggregate) for arrays.

3. If the logic really needs Python, vectorize it

A pandas UDF receives a batch of rows as a pandas Series, transferred with Arrow, so your code can use vectorized operations:

@pandas_udf("double")
def score(amount: pd.Series, days: pd.Series) -> pd.Series:
    return np.log1p(amount) / (days + 1)

This is the right choice for numeric work with NumPy, or for calling a library (a tokenizer, a model) on a batch at a time.

4. If you must keep a row-at-a-time UDF

Use Arrow for the transfer: @udf(..., useArrow=True), or the session setting spark.sql.execution.pythonUDF.arrow.enabled. From Spark 4.2 this is on by default for regular Python UDFs. The function still runs once per row, so it helps less than a pandas UDF.

Other things to check

• Calling the UDF on more rows than needed. Filter first, then apply it.
• Expensive setup inside the function. Loading a model or compiling a pattern on every call. Do it once per batch (pandas UDF) or once per worker.

What interviewers want: why Python UDFs are slow (row by row, data moved between the JVM and Python), the order of fixes (built-ins, then pandas UDFs, then Arrow), and an example of a real rewrite.