PySpark Performance, in depth: the JVM-Python boundary, UDF cost, Arrow and vectorised UDFs, and a worked rewrite

Two processes, one boundary

On the driver, your Python program talks to a JVM through Py4J (or, with Spark Connect, through a gRPC client). When you write df.filter(F.col("x") > 3), Python only builds an expression tree; Catalyst optimises it and executors run generated JVM code. No row ever touches Python, which is why DataFrame code built from built-in functions performs the same in any language.

Apply a Python UDF and the plan gains a node that ships data to a Python worker process on each executor. For a classic row UDF that node is BatchEvalPython: the JVM pickles rows in small groups, the worker unpickles them, calls your function once per row, and pickles the results back. For vectorised UDFs the node is ArrowEvalPython and data moves as columnar Arrow batches. Everything that follows is about how often, and how expensively, a job crosses that line.

What a row UDF really costs

Consider a UDF that only lower-cases a string. Its cost per row is not the lower() call. It is converting a JVM row to bytes, moving them over a local socket, unpickling into Python objects, a Python function call with interpreter overhead, then the same in reverse. Each step is small, but each happens per row.

Worse, the UDF is opaque to the optimiser. Catalyst cannot push a filter through it, cannot fold it into whole-stage code generation, and knows nothing about its output. A UDF placed before a filter runs on rows the filter would have discarded.

So the first rule of PySpark performance has nothing to do with Python: if a built-in function can express it, use the built-in. pyspark.sql.functions covers strings, regular expressions, dates, arrays, maps, JSON and higher-order functions such as transform, filter and aggregate over arrays. Most UDFs in real code bases reimplement one of them.

Back-of-envelope: when the boundary matters

Turn the per-row cost into job time with one multiplication. Suppose, as an assumption to replace with your own measurement, that a trivial row UDF adds 2 microseconds per row for serialisation, transfer and the call. Over 2 billion rows that is 4,000 core-seconds, about 67 core-minutes. On 100 executor cores it adds roughly 40 seconds of wall-clock time per UDF per pass, and three chained UDFs triple it. The same arithmetic says a UDF on a 10,000-row lookup table costs nothing worth optimising, which is why the fix list should be ordered by rows touched, not by how ugly the code looks.

You can calibrate the assumption in a minute: run the same column expression as a built-in and as a UDF over a generated range, with the result written to a no-op sink so nothing else dominates.

import time
from pyspark.sql import functions as F

df = spark.range(50_000_000).withColumn("s", F.col("id").cast("string"))
slow = F.udf(lambda s: s.upper(), "string")   # on 4.2+ this runs over Arrow by default

for name, col in [("builtin", F.upper("s")), ("udf", slow("s"))]:
    t = time.time()
    df.select(col.alias("u")).write.format("noop").mode("overwrite").save()
    print(name, round(time.time() - t, 1), "s")

Worked rewrite

A typical enrichment step, written the way first versions often are:

from pyspark.sql import functions as F
from pyspark.sql.types import StringType, DoubleType

# Slow: three row-at-a-time Python UDFs for work Spark can do natively.
@F.udf(StringType())
def domain(email):
    return email.split("@")[1].lower() if email and "@" in email else None

@F.udf(DoubleType())
def eur(amount, rate):
    return float(amount) * float(rate) if amount is not None else None

@F.udf(StringType())
def bucket(x):
    return "high" if x is not None and x > 100 else "low"

out = (orders
       .withColumn("domain", domain("email"))
       .withColumn("eur", eur("amount", "fx_rate"))
       .withColumn("tier", bucket(F.col("eur"))))

Three UDFs means three trips through the Python worker per row; the eur value is computed in Python, sent back, then sent out again for bucket. The native version:

# Fast: the same logic as built-in expressions. No Python worker is involved.
out = (orders
       .withColumn("domain", F.when(F.col("email").contains("@"),
                                    F.lower(F.substring_index("email", "@", -1))))
       .withColumn("eur", F.col("amount").cast("double") * F.col("fx_rate"))
       .withColumn("tier", F.when(F.col("eur") > 100, "high").otherwise("low")))

out.explain()   # the plan should contain no BatchEvalPython / ArrowEvalPython node

The rewritten plan is generated JVM code with no Python node, and the saving grows with row count because the removed cost was per row. Measure it on your data rather than trusting any multiplier: run both versions on the same input, compare stage durations in the Spark UI, and confirm the plan change with explain().

When you need Python: vectorise

Some work genuinely needs Python: a scikit-learn model, a domain library, a numerical routine with no SQL equivalent. Use a UDF that receives a batch of rows as columns, so per-call overhead is paid once per batch and your code can use vectorised pandas, NumPy or PyArrow operations.

import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
from typing import Iterator
from pyspark.sql.functions import pandas_udf, arrow_udf   # arrow_udf: Spark 4.1+

# Series to Series: one call per Arrow batch (up to maxRecordsPerBatch rows).
@pandas_udf("double")
def zscore(x: pd.Series, mean: pd.Series, std: pd.Series) -> pd.Series:
    return (x - mean) / std

# Iterator of Series: expensive set-up runs once per task, not once per batch.
@pandas_udf("double")
def score(batches: Iterator[pd.Series]) -> Iterator[pd.Series]:
    model = load_model("/models/churn.bin")          # your loader; once per task
    for feats in batches:
        yield pd.Series(model.predict(feats.to_list()))

# Arrow-native: no pandas conversion at all.
@arrow_udf("double")
def discounted(price: pa.Array, pct: pa.Array) -> pa.Array:
    return pc.subtract(price, pc.multiply(price, pct))

# Whole groups as pandas DataFrames: each group must fit in one worker's memory.
def fit_group(pdf: pd.DataFrame) -> pd.DataFrame:
    return pd.DataFrame({"store": [pdf.store.iloc[0]], "trend": [slope(pdf)]})

trends = sales.groupBy("store").applyInPandas(fit_group, "store string, trend double")

applyInPandas deserves a warning: it shuffles every row of a group into one task, then materialises the group as one pandas DataFrame. One oversized key, a giant customer or a null key, runs one Python worker out of memory while every other task finishes. Check the group-size distribution first, and salt or pre-aggregate skewed keys.

Arrow settings and version-gated defaults

Arrow governs vectorised UDFs and the conversions toPandas() and createDataFrame(pandas_df). Its defaults changed recently, so state your Spark version before copying any tuning advice, including this article's.

spark = (SparkSession.builder
  .config("spark.sql.execution.arrow.pyspark.enabled", "true")      # toPandas / createDataFrame
  .config("spark.sql.execution.pythonUDF.arrow.enabled", "true")    # plain @udf over Arrow
  .config("spark.sql.execution.arrow.pyspark.fallback.enabled", "false")
  .config("spark.sql.execution.arrow.maxRecordsPerBatch", "10000")  # rows per Arrow batch
  .config("spark.executor.pyspark.memory", "2g")                    # cap Python worker memory
  .config("spark.executor.memoryOverhead", "3g")                    # non-heap room
  .getOrCreate())

# Spark 4.0+: profile UDFs per SparkSession, including over Spark Connect.
spark.conf.set("spark.sql.pyspark.udf.profiler", "perf")   # or "memory"
out.write.mode("overwrite").parquet("/tmp/probe")
spark.profile.show(type="perf")

Batch size trades memory for overhead: larger batches amortise calls, but each batch is fully materialised in the worker, so wide rows with large strings or arrays need a lower value. Fallback is convenient in notebooks and dangerous in production: when toPandas() or createDataFrame() meets something Arrow cannot convert, it quietly reverts to the slow row-based conversion. Disable it so the error is visible.

Types at the boundary changed too. From 4.2, nullable integer columns reach pandas UDFs as extension dtypes such as Int64 instead of float64, and 4.1 made Arrow conversion safe by default, so overflow or truncation raises an error. UDF output can change across an upgrade with no code change, so regression-test UDF-heavy jobs on every version bump.

Driver-side traps

• collect() or toPandas() on a large result. Everything funnels through one driver process and must fit in its memory. Write results to storage, or limit() first.
• Python loops over actions. A loop that filters and writes once per country launches one job per iteration and rescans the input each time. Express it as one grouped operation or a partitioned write.
• Growing plans. Adding a withColumn or union per loop iteration builds plans that take longer to analyse than to run. Use one select or withColumns, or checkpoint periodically.
• Large closures. A UDF that references a big dictionary or model pickles it with the function. Use a broadcast variable, or load it once inside an iterator UDF.

Python worker memory

Python workers are separate processes, so their memory is outside the JVM heap. On YARN and Kubernetes the container limit is roughly heap plus memory overhead plus, if set, spark.executor.pyspark.memory. If Python is not budgeted, a pandas UDF materialising a large batch pushes the container past its limit and the resource manager kills it, which shows up as lost executors rather than a Python traceback. The JVM side of the budget is covered in Spark memory management.

An executor runs a Python worker per concurrently running task, so a 5-core executor running an iterator UDF that loads a 1 GB model holds five copies. Budget for that or reduce concurrency for the stage. spark.python.worker.reuse, on by default, keeps workers alive between tasks so imports are not repeated.

Finding the expensive crossing

1. Run explain() and look for BatchEvalPython, ArrowEvalPython, MapInPandas and FlatMapGroupsInPandas. Each is a boundary crossing; note where filters sit relative to them.
2. In the Spark UI SQL tab, find the stage containing that node and compare its task time with the rest of the job.
3. On Spark 4.0 or later, set spark.sql.pyspark.udf.profiler to perf or memory for one run and read per-UDF results with spark.profile.show().
4. Rewrite the top offender natively if possible, vectorise it if not, and re-measure the same stage.

Failure modes and trade-offs

• Silent slow path: Arrow fallback in pandas conversion, or a plain @udf on a version where Arrow UDFs are off, runs the row path with no error. Check the plan and the effective config, not just the code.
• Container kills: Python memory outside the heap exceeds the container; fix with overhead budgeting and smaller batches.
• Dependency drift: executors with a different pandas or PyArrow from the driver produce type errors or subtly different results. Ship one environment, such as a packed virtual environment via --archives or a container image.
• Readability versus speed: a nested F.when chain is faster than a clear Python function but harder to review. Keep native expressions in small named helpers that return Columns.

What to do next

1. Grep your jobs for @udf, F.udf and spark.udf.register, and list what each does.
2. Replace every UDF that a built-in or higher-order array function can express; confirm with explain().
3. Convert remaining row UDFs to pandas or Arrow UDFs, using iterator variants for anything that loads a model.
4. Record your Spark version and set the Arrow flags explicitly, with fallback off in production.
5. Budget memory overhead and spark.executor.pyspark.memory for Python workers.
6. Remove driver-side loops over actions and any unbounded collect() or toPandas().
7. Profile one heavy job with the UDF profiler and fix its top offender first.