Spark Pandas UDFs, in depth: type-hint dispatch, struct columns, the batch-boundary trap, iterator setup, grouped aggregates and cogroup

The contract: batches in, batches out

Every pandas UDF has the same underlying contract. Spark splits each partition of the input into Arrow record batches of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, 10,000 by default. Each batch becomes pandas objects, your function runs, and the result is converted back to Arrow and checked against the declared return type.

Since Spark 3.0, the kind of UDF is chosen from Python type hints rather than from an explicit function-type constant. The older PandasUDFType style still appears in old code but is deprecated. Four hint shapes are recognised:

One rule changes the object type: when an input or output column is a StructType, that column appears as a pd.DataFrame rather than a pd.Series. This is how you return several values from one UDF without running it twice:

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

@pandas_udf("lat double, lon double")
def parse_point(raw: pd.Series) -> pd.DataFrame:
    parts = raw.str.split(",", expand=True)
    return pd.DataFrame({"lat": parts[0].astype(float), "lon": parts[1].astype(float)})

df = df.withColumn("pt", parse_point("raw_point")).select("*", "pt.lat", "pt.lon")

Annotate every pandas UDF fully, inputs and return, so the shape Spark infers is the one you meant.

Batch boundaries are not semantics

The most damaging pandas UDF bug produces no error. A Series-to-Series function sees one batch, not the column. Batch boundaries depend on partitioning, on the batch size setting and on how many rows earlier stages produced, so anything your function computes across the batch is a statistic over an arbitrary slice of the data.

Here is the classic version, a z-score written as if the Series were the whole column:

@pandas_udf("double")
def zscore_wrong(x: pd.Series) -> pd.Series:
    return (x - x.mean()) / x.std()          # mean and std of THIS BATCH only

Run it on the values 1 to 20,000 in one partition with the default batch size and Spark makes two batches. The first sees a mean of 5,000.5, the second 15,000.5, while the true column mean is 10,000.5. The value 10,000 is the last row of the first batch, so it gets a z-score of about +1.73 when the right answer is close to zero. The same pure-pandas computation shows it directly:

import numpy as np, pandas as pd

x = pd.Series(np.arange(1, 20001, dtype="float64"))
batches = [x.iloc[:10000], x.iloc[10000:]]
wrong = pd.concat([(b - b.mean()) / b.std() for b in batches])
right = (x - x.mean()) / x.std()
print(round(wrong.iloc[9999], 2), round(right.iloc[9999], 2))   # 1.73 -0.0, i.e. zero

Change the partitioning or the batch size and the answers change too, so the bug cannot be reproduced reliably, which is the worst property a data bug can have. The fix is to compute global statistics in Spark, where they are correct, and pass them in as columns or constants:

from pyspark.sql import functions as F

stats = df.agg(F.mean("x").alias("mu"), F.stddev("x").alias("sigma")).first()

@pandas_udf("double")
def zscore(x: pd.Series) -> pd.Series:
    return (x - stats.mu) / stats.sigma       # constants captured in the closure

The same reasoning applies to ranking, de-duplication, forward fill, rolling windows and anything that looks at neighbouring rows. Rows in a batch are not sorted unless you made them so, and a batch need not contain the whole partition you might expect. If the logic needs a whole group, use a grouped API, which gives you the group by contract.

Iterator UDFs: setup once per task

The iterator forms exist for one reason: work that should happen once per task rather than once per batch. Loading a model, opening a connection pool or compiling a regular expression table inside a Series-to-Series function repeats that work for every 10,000 rows. With an iterator UDF it happens once and is reused for every batch the task processes:

from typing import Iterator, Tuple

@pandas_udf("double")
def churn_score(batches: Iterator[Tuple[pd.Series, pd.Series]]) -> Iterator[pd.Series]:
    model = load_model("/dbfs/models/churn-v7.joblib")   # your loader, once per task
    for tenure, spend in batches:
        feats = pd.DataFrame({"tenure": tenure, "spend": spend})
        yield pd.Series(model.predict_proba(feats)[:, 1])

The output contract is looser than for Series-to-Series: each yielded batch need not match its input batch, but the total number of rows yielded must equal the total number received. In practice, yield one result per input batch unless you have a reason to re-chunk. When the setup opens something that must be released, such as a database connection or an HTTP session, wrap the batch loop in try/finally and close it there; a task can be cancelled or fail mid-stream, and Python workers are reused, so a leaked connection outlives the task.

Keep the setup cheap relative to the task. A 2 GB model loaded by every task on a 16-core executor means 16 copies in Python memory at once, which is the usual cause of Python workers being killed for exceeding memory limits. Broadcast the model file, load a memory-mapped format, or reduce Python worker concurrency.

Series to scalar: custom aggregates

A Series-to-scalar UDF is a grouped aggregate. It works with groupBy().agg() and with window functions, which makes it the way to write an aggregate Spark does not have, such as a weighted median:

import numpy as np

@pandas_udf("double")
def weighted_median(v: pd.Series, w: pd.Series) -> float:
    order = np.argsort(v.to_numpy())
    v, w = v.to_numpy()[order], w.to_numpy()[order]
    cum = np.cumsum(w)
    return float(v[np.searchsorted(cum, cum[-1] / 2.0)])

df.groupBy("region").agg(weighted_median("price", "qty").alias("wmedian"))

Two limits come with it. First, there is no partial aggregation: Spark cannot combine per-partition partial results as it does for sum, so it shuffles every row of a group to one task and loads the whole group into memory. Second, as a window function it supports only unbounded windows, for example Window.partitionBy("region") with no frame or with an unbounded-preceding to unbounded-following frame. Rolling frames need a different approach, such as built-in window functions or a grouped map over sorted data.

Grouped map and cogroup

Grouped map and cogrouped map are not pandas UDFs in the decorator sense, but they run on the same machinery and are where most memory incidents happen. groupBy(...).applyInPandas(func, schema) calls your function once per group with the group as a DataFrame. groupBy(...).cogroup(other.groupBy(...)).applyInPandas(func, schema) calls it once per key with the matching groups from two DataFrames, which is the natural way to run an as-of join or a per-key reconciliation:

def asof_join(trades: pd.DataFrame, quotes: pd.DataFrame) -> pd.DataFrame:
    trades = trades.sort_values("ts")
    quotes = quotes.sort_values("ts")
    out = pd.merge_asof(trades, quotes[["ts", "bid", "ask"]], on="ts", direction="backward")
    return out[["symbol", "ts", "price", "bid", "ask"]]

joined = (trades.groupBy("symbol")
          .cogroup(quotes.groupBy("symbol"))
          .applyInPandas(asof_join, "symbol string, ts timestamp, price double, bid double, ask double"))

The documentation is explicit on the memory model: all data for a group is loaded into memory before the function runs, and maxRecordsPerBatch is not applied to groups. One skewed key therefore becomes one very large pandas DataFrame in one worker; data skew covers how to detect and split such keys.

The returned DataFrame is matched to the schema in one of two ways. If its column labels are strings, they are matched to schema fields by name; otherwise by position. Name matching is the safer choice, and selecting the output columns explicitly at the end of the function, as above, stops a stray helper column from causing a column-count mismatch error.

Type conversions that bite

Conversions between Spark types and pandas types are where correct-looking code returns subtly wrong data.

• Integers with nulls. pandas' default NumPy-backed integer dtype cannot hold missing values, so a nullable Spark int or bigint column with nulls typically arrives as float64 with NaN. Compare with isna(), and cast back deliberately before returning integers.
• Timestamps. Spark-to-pandas conversion produces nanosecond datetime64[ns] values converted to the session time zone and then made time-zone naive. Going back, values are treated as UTC microseconds and nanoseconds are truncated. Set spark.sql.session.timeZone explicitly so this does not depend on cluster defaults.
• Decimals. DecimalType columns arrive as Python Decimal objects in an object-dtype Series, which defeats vectorisation. Cast to double in Spark first if exactness is not required.
• Return type mismatch. Returning float64 for a declared long can fail or be cast depending on the values and on Arrow's safe-cast checks. Return exactly the dtype you declare.

Testing the logic without a cluster

Because the logic is plain pandas, most of it can be tested without Spark. Keep the body in an ordinary function, test that with pandas, and wrap it for Spark separately. pandas_udf accepts a function and a return type directly, so no decorator is needed:

def _clean_email(s: pd.Series) -> pd.Series:
    return s.str.strip().str.lower().where(s.str.contains("@", na=False))

clean_email = pandas_udf(_clean_email, "string")      # used in Spark

def test_clean_email():
    got = _clean_email(pd.Series([" A@X.COM ", "nope", None]))
    assert got.tolist()[0] == "a@x.com"
    assert got.isna().tolist()[1:] == [True, True]

Add two Spark-level tests on a local session: one that runs the UDF on a few hundred rows with maxRecordsPerBatch set to a small number such as 7, which flushes out any batch-local logic, and one with nulls in every input column. Both catch bugs that unit tests on a single Series cannot.

Failure modes and choosing an API

Before writing any UDF, check whether built-in functions can do the job. A native expression stays inside the JVM and the optimiser, while any Python UDF is opaque to it. When Python is needed, choose by shape: Series-to-Series for element-wise work, iterator for expensive setup, Series-to-scalar for custom aggregates, applyInPandas when the logic genuinely needs a whole group, and mapInPandas when the number of rows changes. Grouped work on very large keys is often better handled by pandas API on Spark or by plain DataFrame operations, which keep the work distributed.

What to do next

To put this into practice on a real job:

1. Search the codebase for @pandas_udf functions that call .mean(), .std(), .rank(), .shift() or .sort_values() on their inputs; each one is a candidate batch-boundary bug.
2. Re-run suspicious UDFs on a local session with spark.sql.execution.arrow.maxRecordsPerBatch set to 7 and compare the output with the default.
3. Move per-batch model loading into iterator UDFs, and close any connections they open in a try/finally.
4. For every applyInPandas job, measure the largest group by row count and bytes before it reaches production.
5. Set spark.sql.session.timeZone explicitly and add a test with nulls in every input column.
6. Read Spark memory management to size Python worker memory alongside JVM memory.