PySpark UDFs, Arrow & the Python boundary

Where Python runs

Your PySpark program runs in a Python process on the driver, but DataFrame calls — filter, groupBy, even F.upper — only build a plan that the JVM executes (see RDDs, DataFrames & Spark SQL). When built-in expressions run on executors, no Python is involved.

A Python UDF is different: its function must run in Python beside the data. Executors hand that work to Python worker processes, reused across tasks by default (spark.python.worker.reuse=true). In a local run, a UDF returning os.getpid() reported worker PIDs 31531 and 31532, while the driver's Python PID was 31455 — the function never ran in the driver's interpreter.

For an ordinary Python UDF, the BatchEvalPython operator projects the UDF's input columns, groups rows into batches, pickles them, and writes them over a socket to the worker. The worker unpickles each batch, calls your function once per row, pickles the results, and sends them back. The JVM unpickles each result and pairs it with the original row, which it held in a queue meanwhile. Every row pays for conversion into Python objects and back, plus a Python function call.

Check the sources: EvalPythonExec · BatchEvalPythonExec · Python worker settings

Why the optimizer stops at a UDF

Catalyst cannot read Python code, so a UDF is an opaque expression with a declared return type. ExtractPythonUDFs moves Python UDFs out of projections and filters into dedicated operators, and the plan makes the boundary visible. Upper-casing a name column four ways, followed by filter(id > 1):

The trailing number is the Python evaluation type: 101 is an Arrow-batched Python UDF and 200 a scalar pandas UDF. Two consequences follow.

• Codegen splits around the Python operator; operators below and above still generate code, in separate stages.
• Predicates on the UDF's result cannot be pushed. On a Parquet dataset, filter(amount > 400) pushed GreaterThan(amount,400) into the scan; the same test written as a Python UDF left PushedFilters: [] and a Filter above BatchEvalPython, so every row was read first. Both returned 594 rows. Predicates that do not touch the UDF still move below it: id > 1 ran before the UDF in all four plans.

Spark treats a UDF as deterministic unless you call asNondeterministic(), and the API notes warn that it may be invoked more or fewer times than it appears in the query. Selecting a UDF result and then filtering on that column pushed the filter below the projection by copying the expression, so the plan evaluated the UDF in two BatchEvalPython operators; marked non-deterministic, it was evaluated once. See Catalyst, code generation & AQE for how to read these plans.

Check the sources: ExtractPythonUDFs · Python eval types · udf notes

Arrow: batches across the boundary, pandas inside it

Apache Arrow is a columnar in-memory format. Spark uses it to move column batches between the JVM and Python instead of pickled rows. What your function receives depends on the API:

• Arrow-optimized Python UDF — @udf(..., useArrow=True), or for every eligible UDF spark.sql.execution.pythonUDF.arrow.enabled (default false; the setting dates from Spark 3.4.0). Transfer uses Arrow, but the function is still called row by row: the 3.5.7 worker loops over each batch and calls it per row, and the guide calls these UDFs "executed row-by-row." Arrow's return-type coercion can differ from the pickled path, which the source gives as the reason the option is off by default.
• pandas UDF — @pandas_udf with Python type hints. Series to Series (output length must equal input); Iterator of Series to Iterator of Series, useful when execution needs expensive state initialized once; Iterator of multiple Series; and Series to scalar, an aggregate for groupBy().agg() or unbounded windows that loads each group into memory with no partial aggregation (planned as AggregateInPandas after an Exchange and Sort). The function receives whole batches, so vectorized pandas or NumPy code avoids a Python call per row. Type hints are preferred; the guide says PandasUDFType will be deprecated.
• pandas function APIs take and return pandas DataFrames. mapInPandas maps an iterator of DataFrames per partition and may return any number of rows; it planned as MapInPandas over the scan, with no shuffle. groupBy().applyInPandas hands each group to your function as one DataFrame and requires a full shuffle: Exchange, Sort, FlatMapGroupsInPandas. cogroup().applyInPandas does the same for two grouped DataFrames.

Batch size. spark.sql.execution.arrow.maxRecordsPerBatch (default 10,000) caps rows per Arrow batch. A pandas UDF over one 25,000-row partition saw batches of 10,000, 10,000, and 5,000 rows; at 1,000, every batch had 1,000 rows. The cap counts rows, not bytes, so wide rows may need a lower value. It does not apply to applyInPandas or cogroup: each whole group is one batch, and a skewed group can exhaust the worker.

Check the sources: Arrow and pandas guide · worker.py · pandas_udf · mapInPandas · applyInPandas

Memory for Python workers

Python workers are native processes outside the executor's JVM heap, so spark.executor.memory does not size them. By default spark.executor.pyspark.memory is unset: Spark does not limit Python's memory, and Python workers share spark.executor.memoryOverhead with other non-JVM memory in the container. On YARN and Kubernetes, a container that exceeds its total can be killed while the heap looks healthy.

Setting spark.executor.pyspark.memory does two things. On YARN and Kubernetes it adds that amount to the executor's resource request. And PythonRunner divides it by the executor's cores and passes the result to each worker, which applies it as an address-space limit through Python's resource module. The configuration reference notes that the limit is not supported on Windows and not enforced on macOS, so it cannot be tested on a Mac laptop.

Two neighbors are easy to confuse. spark.python.worker.memory (512m) only sets when a Python worker spills to disk during aggregation. And Arrow batches take JVM memory while they are built — the guide warns this "can temporarily lead to high memory usage in the JVM." toPandas() collects the whole result to the driver whether or not Arrow is enabled. Executor memory explains how the container total adds up.

Check the sources: PySpark memory settings · PythonRunner · worker.py limits · Arrow batch memory · toPandas

Choose the least opaque tool

1. Built-in functions and SQL expressions. Catalyst can optimize them, push eligible predicates, and generate code. Check the built-in list first — string, date, JSON, regex, and conditional logic are mostly covered.
2. Higher-order functions such as transform, filter, exists, and aggregate for array and map columns. They stay in the JVM with no Python round trip, although in the observed plan the Project using transform fell out of whole-stage codegen.
3. pandas UDFs or pandas function APIs when the logic genuinely needs Python libraries, written with vectorized operations. Series.apply inside a pandas UDF is still a Python call per row.
4. A row-at-a-time Python UDF as the last resort. Consider useArrow=True after checking that its type coercion gives the results you expect.

An illustration, not a benchmark. On one laptop (local[2], 2,000,000 cached strings, best of three runs of summing the length of upper-cased values), the built-in took 0.15 s, the pandas UDF 0.38 s, the Arrow Python UDF 0.60 s, and the plain Python UDF 1.22 s. Those ratios come from one trivial function on one machine. The gap changes with the function's own cost, data types, batch size, row width, and cluster layout, so there is no general "10–100×" multiplier. Compare plans first, then time a representative sample of your own job.

Check the sources: Built-in functions · Arrow and pandas guide · udf

Try it on a small dataset

Run this in a PySpark notebook with an existing spark session. Create a local session if needed. The example was run for review; this website does not execute Spark.

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

people = spark.createDataFrame(
    [(1, "ada"), (2, "grace"), (3, None)], "id INT, name STRING"
)

@F.udf("string")
def upper_py(s):
    return s.upper() if s is not None else None

@F.udf("string", useArrow=True)
def upper_arrow(s):
    return s.upper() if s is not None else None

@F.pandas_udf("string")
def upper_pandas(s: pd.Series) -> pd.Series:
    return s.str.upper()

for label, column in [
    ("built-in", F.upper("name")),
    ("Python UDF", upper_py("name")),
    ("Arrow Python UDF", upper_arrow("name")),
    ("pandas UDF", upper_pandas("name")),
]:
    print(label)
    people.select("id", column.alias("u")).filter(F.col("id") > 1).explain()

Ran on local PySpark 3.5.7 (local[2], JDK 17, pandas and PyArrow installed) on 2026-09-15; output above is observed.

A teammate wraps a string cleanup in @udf, finds it slow, and adds useArrow=True. What changed, and what did not?

Say your answer out loud before opening the notes.