Session-scoped Python and pandas UDFs in PySpark

1.When a UDF is the right tool

A user-defined function lets you reuse custom logic that Spark does not provide out of the box. Databricks is clear that this has a cost. Built-in Apache Spark functions are optimized for distributed processing and perform better at scale, so use a UDF only for logic that is hard to express with them. The recommended uses are ad hoc queries, manual data cleaning, exploratory analysis and small to medium-sized datasets. Typical jobs are encryption, decryption, hashing, JSON parsing and validation. For large datasets and for anything that runs regularly or continuously, such as ETL jobs and streaming, Databricks recommends native Spark methods instead.

Every UDF type is defined by two things: how many rows go in and come out, and where the function lives. A UC-governed function is registered in Unity Catalog and shared under its permissions. A session-scoped function exists only in the notebook or SparkSession that defined it. This page covers the session-scoped Python and pandas UDFs you write in PySpark.

Sources1

2.Scalar Python UDFs: three ways to declare them

A scalar Python UDF runs one row at a time. You can declare it in three ways. Use spark.udf.register to make it callable by name from SQL, for example from a %sql cell against a temp view. Use udf(func, returnType) to wrap a function for the DataFrame API. Or put the @udf("long") decorator on the function. In every case the return type is optional, and if you leave it out you get StringType. That is easy to miss when a numeric function quietly returns strings.

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

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")))

from pyspark.sql.functions import udf
from pyspark.sql.types import LongType
squared_udf = udf(squared,  ? )
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

Sources2

3.Null handling and evaluation order

Spark SQL, including the DataFrame API, does not fix the order in which an operator's inputs are evaluated. Boolean AND and OR do not short-circuit from left to right, and the query optimizer can reorder WHERE and HAVING clauses. A null check written to the left of a UDF call therefore protects nothing. There are two safe patterns. Either make the UDF itself null-aware, or put the call inside an IF or CASE WHEN branch that only runs when the value is not 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

Sources2

4.pandas UDFs: vectorized batches over Arrow

A pandas UDF, also called a vectorized UDF, removes the row-at-a-time overhead. Spark uses Apache Arrow to move data in batches, and your function works on those batches with pandas. You declare one with pandas_udf, as a decorator or a wrapper, and the Python type hints tell Spark which of four forms you are writing. The simplest form is Series to Series. Spark splits the data into batches, calls your function once per batch, and concatenates the results. The output Series must be the same length as the input.

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

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

The iterator forms take an iterator of batches and yield an iterator of batches, and the total output length must still match the input. They exist so you can do expensive setup, such as loading a model, once before the loop instead of once per batch. Use try/finally or a context manager to release resources at the end. Iterator of Series takes one column. Iterator of multiple Series takes a tuple of Series per batch, one for each input column. Series to scalar is the odd one out: it works like a Spark aggregate function and can be used with select, withColumn, groupBy.agg and Window. It does not support partial aggregation, so all the data for each group is loaded into memory.

@pandas_udf("long")
def plus_one(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for x in batch_iter:
        yield x + 1

Sources3

Exam traps

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

1. 1.A Python UDF's return type is inferred from what the function returns.Why is that wrong?

   Nothing is inferred. If you don't pass a return type, the UDF returns StringType.

   Covered in Scalar Python UDFs: three ways to declare them

2. 2.Putting s IS NOT NULL AND before a UDF call protects the UDF from NULLs.Why is that wrong?

   Spark SQL fixes no evaluation order, and AND does not short-circuit. Do the null check inside the UDF or in an IF/CASE WHEN branch.

   Covered in Null handling and evaluation order

3. 3.A Series to scalar pandas UDF aggregates incrementally like a built-in sum.Why is that wrong?

   It does not partially aggregate. Each group's data is loaded into memory in full.

   Covered in pandas UDFs: vectorized batches over Arrow