Python UDFs and pandas UDFs in PySpark: create, register, invoke

1.Creating and invoking a scalar Python UDF

A scalar Python UDF starts as an ordinary Python function. Spark gives you two ways to make it callable. spark.udf.register(name, func, returnType) registers it under a name that SQL can use. pyspark.sql.functions.udf(func, returnType) wraps it as a column function for the DataFrame API. A third form is the @udf("long") decorator, which does the same job as udf() with annotation syntax. All three take the same optional return type, and that argument is where candidates slip.

from pyspark.sql.types import LongType
def squared_typed(s):
  return s * s
spark.udf.register("squaredWithPython", squared_typed, LongType())

After registration, you call the function by name in SQL, for example select id, squaredWithPython(id) as id_squared from test. For the DataFrame API you don't need a name. You use the wrapped object directly as a column expression inside select or withColumn, as below. The return type can be given as a type object (LongType()) or as a DDL string ("long").

from pyspark.sql.functions import udf

@udf("long")
def squared_udf(s):
  return s * s
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

Return types aren't limited to scalars. The Databricks examples declare returnType=StructType([...]), ArrayType(...) and MapType(StringType(), ...), and return a Python dict or list that matches the declared shape.

Sources1

2.Null handling and evaluation order

The Databricks UDF page warns about subexpression evaluation-order caveats in Spark SQL. In practice, a WHERE s is not null sitting next to a UDF call is not a reliable guard for that UDF. The documentation gives two safe patterns. The first is to make the UDF itself null-aware. The second is to put the null check in an IF or CASE WHEN expression and call the UDF only inside the branch where the value is known to be non-null.

spark.udf.register("strlen_nullsafe", lambda s: len(s) if not s is None else -1, "int")
spark.sql("select s from test1 where s is not null and strlen_nullsafe(s) > 1") # ok
spark.sql("select s from test1 where if(s is not null, strlen(s), null) > 1")   # ok

Sources1

3.Vectorized pandas UDFs: Series to Series

A scalar Python UDF handles one row per call. A pandas UDF (also called a vectorized UDF) uses Apache Arrow to move data and pandas to process it. Spark splits the column into batches of rows, calls your function once per batch, and concatenates the results. You declare one with pandas_udf, either as a function or as a decorator, and the Python type hints tell Spark which kind it is. For a Series-to-Series UDF, the function must return a Series of the same length as its input. It plugs into select and withColumn just like a scalar UDF.

# Declare the function and create the UDF
def multiply_func(a: pd.Series, b: pd.Series) -> pd.Series:
    return a * b

multiply =  ? (multiply_func, returnType=LongType())

Sources2

4.Iterator pandas UDFs: one-time setup inside a UDF

An iterator pandas UDF receives an iterator of batches and yields an iterator of batches. The type hint is Iterator[pd.Series] -> Iterator[pd.Series], or, for several input columns, Iterator[Tuple[pandas.Series, ...]] -> Iterator[pandas.Series]. Across all batches, the total output must still be as long as the total input. The reason to use this shape is that code placed before the loop runs once and then serves every batch, such as loading a model file. Wrap the loop in try/finally so resources are released at the end.

y = 1  # value captured by the UDF closure

@pandas_udf("long")
def plus_y(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
    try:
        for x in batch_iter:
            yield x + y
    finally:
        pass  # release resources here, if any

df.select(plus_y(col("x"))).show()

The 'state' here is setup state that lives inside one UDF execution. Stateful streaming operators work differently: they keep state information for each grouping key in a state store, and that is a separate mechanism with its own API.

Sources23

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.A Python UDF that returns numbers produces a numeric column even without a declared return type.Why is that wrong?

   Spark does not infer the type. With no return type, a Python UDF's result is StringType, so pass LongType() or "long" to get numbers back.

   Covered in Creating and invoking a scalar Python UDF

2. 2.Because an iterator pandas UDF yields batches freely, it can return fewer rows than it received, for example to filter.Why is that wrong?

   Across the whole iterator, the output must be exactly as long as the input. Iterator UDFs are a column transformation, not a filter.

   Covered in Iterator pandas UDFs: one-time setup inside a UDF