UDFs & the Pandas API on Spark

A UDF is a user-defined function: your own Python function, used on a column. You write one when no built-in function does the job.

The built-in functions in F run inside Spark, in the JVM, the Java program that runs Spark. A Python UDF runs in a separate Python program on each executor.

So every row must travel from the JVM to Python and back. And Spark cannot see inside your function, so it cannot improve it.

Use a UDF only when F has no function for the job. This chapter shows when, and how to make it less slow.

The @udf decorator turns a Python function into a column function. You must give the type it returns:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

@udf(returnType=StringType())
def size_label(amount):
    if amount is None:
        return "unknown"
    return "large" if amount >= 50 else "small"

orders.select("order_id", "amount", size_label("amount").alias("size")).show()

The same result needs no Python at all. when and otherwise (chapter 8) run inside Spark:

orders.select(
    "order_id",
    F.when(F.col("amount").isNull(), "unknown")
     .when(F.col("amount") >= 50, "large")
     .otherwise("small").alias("size"),
).show()

explain() on the UDF shows a step called ArrowEvalPython. In Spark 4.2, rows travel to Python in batches, in the Arrow format, by default. But your function still runs once for each row.

The declared type matters. In our run, a UDF declared as IntegerType that returned text stopped the job:

@udf(returnType=IntegerType())
def as_text(amount):
    return str(amount)

orders.select("order_id", as_text("amount")).show()

pandas is a Python library for tables. A pandas Series is one column of values. A pandas UDF gets a whole batch of a column as a Series, not one value at a time.

So Python runs once for each batch, and pandas does the work on the whole column:

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

@pandas_udf(DoubleType())
def to_cents(amount: pd.Series) -> pd.Series:
    return amount * 100

orders.select("order_id", "amount", to_cents("amount").alias("cents")).show()

pandas UDFs need the pandas and pyarrow libraries on every executor.

Choose in this order:

1. A built-in function in F, or SQL. It runs inside Spark, and the optimizer can see it.
2. A pandas UDF, when you need Python. It works on batches.
3. A plain Python UDF, only for logic that pandas cannot do on a whole column.

The pandas API on Spark lets you write pandas-style code that runs on Spark. It lives in pyspark.pandas. pandas_api() turns a Spark DataFrame into one:

psdf = orders.pandas_api()

print(psdf.groupby("customer_id")["amount"].sum().sort_index())
print(psdf[psdf["status"] == "paid"]["amount"].mean())

It helps when a team already knows pandas, or when pandas code must grow past one machine.

Some pandas habits are slow here. Row order is not kept unless you sort, and some operations move all data to one place. For new pipelines, the DataFrame API is the clearer choice.