Python UDF versus pandas UDF versus Arrow-optimized UDF?

Summary

• A Python UDF calls your function once per row. An Arrow-optimized Python UDF is the same row-by-row udf that uses Arrow for "efficient batch data transfer and serialization". It was introduced in Spark 3.5.0 and is on by default since Spark 4.2.0.
• A pandas UDF (@pandas_udf) receives pandas.Series batches. An Arrow UDF (@arrow_udf, new in Spark 4.1.0) receives pyarrow.Array batches. Don't confuse the Arrow UDF with an Arrow-optimized Python UDF.
• Pick by what your function should receive: single values, pandas objects or PyArrow arrays. Iterator variants suit logic that needs one-time setup, such as loading a model.
• All Arrow-based UDFs need PyArrow on every node, and pandas UDFs also need pandas.

Python UDF versus pandas UDF versus Arrow-optimized UDF?

"Arrow-optimized UDF" can mean two different things in Spark. In the Spark 4.2.0 docs, an Arrow-optimized Python UDF is a regular udf that uses Arrow to (de)serialize data. Since Spark 4.1.0 there is also a separate Arrow UDF, arrow_udf, which works on pyarrow.Array batches. The table covers all four.

Kind	Declare with	Function receives	Role of Arrow	Availability
Pickled Python UDF	@udf with useArrow=False	One value per call	None, pickle-based serialization	Default before Spark 4.2.0
Arrow-optimized Python UDF	@udf with useArrow=True	One value per call	Batch data transfer and serialization	Introduced in 3.5.0, default since 4.2.0
pandas UDF	@pandas_udf	A pandas.Series per batch	Transfers data, pandas does the work	Since 2.3.0, type hints since 3.0
Arrow UDF	@arrow_udf	A pyarrow.Array per batch	Transfers data, PyArrow does the work	New in 4.1.0

The same arithmetic in each style:

import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
from pyspark.sql.functions import arrow_udf, pandas_udf, udf
from pyspark.sql.types import IntegerType

@udf(returnType=IntegerType(), useArrow=True)  # one value per call
def calc_row(a, b):
    return a + 10 * b

@pandas_udf(returnType=IntegerType())  # one pandas.Series per batch
def calc_pandas(a: pd.Series, b: pd.Series) -> pd.Series:
    return a + 10 * b

@arrow_udf(returnType=IntegerType())  # one pyarrow.Array per batch
def calc_arrow(a: pa.Array, b: pa.Array) -> pa.Array:
    return pc.add(a, pc.multiply(b, 10))

What changes with an Arrow-optimized Python UDF?

You write the function the same way, because the guide says its type hints are specified as for default, pickled Python UDFs. The guide describes Arrow Python UDFs as "executed row-by-row, utilizing Arrow for efficient batch data transfer and serialization". Turn one on per function with useArrow=True, or for the whole session with spark.sql.execution.pythonUDF.arrow.enabled, which takes effect only when useArrow isn't set. Since Spark 4.2, Arrow optimization is on by default for regular Python UDFs, and setting that configuration to false restores the pickled path. The guide says Arrow Python UDFs provide "a more coherent type coercion mechanism" when a UDF returns values that don't match its declared return type, so re-test results when you upgrade.

When do pandas UDFs and Arrow UDFs fit?

A pandas UDF is executed "using Arrow to transfer data and Pandas to work with the data, which allows vectorized operations". An Arrow UDF does the same with PyArrow: Spark transfers and works with the data using Arrow, "which allows pyarrow.Array operations". For both, a Series-to-Series or Arrays-to-Arrays function should return output of the same length as its input, and the input is one internal batch, not the whole column. Data partitions are converted into Arrow record batches, and spark.sql.execution.arrow.maxRecordsPerBatch caps the rows per batch (default 10,000).
The Iterator variants (Iterator[pd.Series] or Iterator[pa.Array]) suit work that needs one-time setup. The docs show an expensive initialization before the loop over batches, and Databricks gives loading a machine learning model file as the example.
Since Spark 4.1.0, @udf can also define a vectorized function when its type hints match a pandas UDF or Arrow UDF pattern and useArrow isn't set.

How Databricks approaches this

Databricks documents the same APIs. Its udf reference shows the Arrow-optimized Python UDF as the default since Spark 4.2, with useArrow=False to use pickle-based serialization. Its pandas UDF page says Spark runs one "by splitting the data into batches of rows, calling the function for each batch, and then concatenating the results".

FAQs

Do UDFs short-circuit conditions?

No. The API docs say UDFs "do not support conditional expressions or short circuiting in boolean expressions" and that all of it is executed. If a function can fail on special rows, put the condition inside the function.

What do I need installed?

Arrow-based UDFs need PyArrow. With pip, pip install pyspark[sql] brings it in as an extra dependency. Otherwise PyArrow must be available on all cluster nodes. The Spark 4.2.0 docs list minimum supported versions of pandas 2.2.0 and PyArrow 18.0.0.

Are there table-function equivalents?

Yes. udtf (new in Spark 3.5.0) defines a table function whose handler yields rows, and it accepts useArrow. arrow_udtf (new in Spark 4.1.0) is the PyArrow-native version, where eval receives PyArrow RecordBatches or Arrays.