Pandas UDF variants: iterator and Series to scalar

1.Iterator of Series to Iterator of Series

A Series to Series UDF does all its work inside one call per batch. If that work needs expensive setup, such as loading a machine learning model, the setup runs again for every batch. The iterator variant avoids this. Spark calls the function once and gives it an iterator of batches. The function yields one output batch for each input batch, so you can do the setup once before the loop. Spark still picks the variant from the type hint, here Iterator[pd.Series] -> Iterator[pd.Series].

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

The length rule is relaxed here, but it still applies. Each yielded batch doesn't have to match the batch it came from, but the total output across all batches must be as long as the total input. To make sure resources acquired before the loop are released, the docs recommend a try/finally block or a context manager:

@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

Sources1

2.Iterator of multiple Series: several input columns

To pass several columns to an iterator UDF, declare the input as an iterator of tuples: Iterator[Tuple[pd.Series, ...]] -> Iterator[pd.Series]. Each element of the iterator is a tuple with one Series per column, all covering the same batch, and you unpack the tuple in the for loop. Everything else is the same as the single-column iterator: you can set up state once, and the total output length must match the total input length.

@pandas_udf("long")
def multiply_two_cols(
        iterator: Iterator[Tuple[pd.Series, pd.Series]]) -> Iterator[pd.Series]:
    for a, b in iterator:
        yield a * b

df.select(multiply_two_cols("x", "x")).show()

Sources1

3.Series to scalar: pandas UDFs as aggregates

Every variant so far returns as many rows as it receives. A Series to scalar UDF returns one value, so it works like a Spark aggregate function. Its type hint is pd.Series, ... -> Any, and in practice you replace Any with a specific scalar type such as float. The returnType must be a primitive data type, and the function can return a Python primitive like int or float or a NumPy scalar like numpy.float64. You can use it in select (aggregating the whole DataFrame), in groupBy().agg(), and over a Window.

@pandas_udf("double")
def mean_udf(v: pd.Series) -> float:
    return v.mean()

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
df.groupby("id").agg(mean_udf(df['v'])).show()

w = Window.partitionBy('id').orderBy('v').rowsBetween(-1, 0)
df.withColumn('mean_v', mean_udf("v").over(w)).show()

There is a cost. Built-in aggregates can partially aggregate data before the shuffle, but a Series to scalar pandas UDF can't. All of a group's data is loaded into memory before the function runs, so one very large group can exhaust a worker's memory. Newer Spark releases add an Iterator of Series to scalar form, Iterator[pd.Series] -> float, which Databricks describes as memory-efficient grouped aggregation. It receives a group's data batch by batch and keeps running totals instead of holding the whole group in memory.

@pandas_udf("double")
def pandas_mean_iter(it: Iterator[pd.Series]) -> float:
    sum_val = 0.0
    cnt = 0
    for v in it:
        sum_val += v.sum()
        cnt += len(v)
    return sum_val / cnt

Sources12

4.Choosing the variant from the type hint

All of these variants are created the same way, with pandas_udf and a returnType. They differ only in their Python type hints, and Spark uses those hints to decide how to call the function. The table lists the four variants in the Databricks pandas UDF guide by signature and by where you invoke them.

Sources1

Exam traps

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

1. 1.An Iterator[pd.Series] -> Iterator[pd.Series] UDF can take several columns, one per element of the iterator.Why is that wrong?

   That signature takes a single Spark column. For several columns, the hint must be an iterator of tuples of Series.

   Covered in Iterator of multiple Series: several input columns

2. 2.A Series to scalar pandas UDF partially aggregates like built-in aggregates, so skewed groups are safe.Why is that wrong?

   It doesn't support partial aggregation. All of a group's data is loaded into memory before the function runs.

   Covered in Series to scalar: pandas UDFs as aggregates

Practise it for real

Create a pandas UDF by explicit wrapping, test it locally, invoke it on a Spark DataFrame, then compare it with a decorator-based aggregate.

1. 1.Define multiply_func(a: pd.Series, b: pd.Series) -> pd.Series that returns a * b, and wrap it with multiply = pandas_udf(multiply_func, returnType=LongType()).

   Why: Explicit wrapping keeps the plain function available for local testing and gives you a separate UDF object for Spark.

   You should see: multiply is a pandas UDF, and multiply_func is still an ordinary Python function.

2. 2.Call multiply_func(x, x) on x = pd.Series([1, 2, 3]).

   Why: The function must work on local pandas data, so you can check its logic without Spark.

   You should see: A pandas Series containing 1, 4, 9 with dtype int64.

3. 3.Create df = spark.createDataFrame(pd.DataFrame(x, columns=["x"])) and run df.select(multiply(col("x"), col("x"))).show().

   Why: This calls the vectorized UDF on Spark, which runs the function once per Arrow batch.

   You should see: A column named multiply_func(x, x) with values 1, 4, 9.

4. 4.Define @pandas_udf("double") def mean_udf(v: pd.Series) -> float: return v.mean(), then run it with df.groupby("id").agg(mean_udf(df['v'])) on the five-row (id, v) DataFrame from the Series to scalar section.

   Why: The float return hint turns the same pandas_udf call into an aggregate.

   You should see: One row per id: 1.5 for id 1 and 6.0 for id 2.

5. 5.Register a Series to Series UDF with spark.udf.register("add_one", add_one) and run spark.sql("SELECT add_one(id) FROM range(3)").collect().

   Why: Registration lets you call the same pandas UDF from SQL.

   You should see: Rows with add_one(id) equal to 1, 2 and 3.