Back to Advanced Data Engineering
Q8Question 8 of 24

Pandas UDF (vectorized UDF)

Key Concepts
  • A Pandas UDF processes whole columns as pandas Series via Apache Arrow batches, not one row at a time.
  • Arrow is a columnar in-memory format both the JVM and Python understand, amortizing the JVM-to-Python boundary cost across many rows.
  • Pandas UDFs are typically 10-100x faster than plain row-at-a-time Python UDFs.
  • Series-to-Series maps columns; Series-to-scalar aggregates in groupBy; iterator-of-Series initializes resources once per partition.
  • Enable with spark.sql.execution.arrow.pyspark.enabled=true and tune batch size with maxRecordsPerBatch.
  • Pandas UDFs still cannot be optimized by Catalyst, so a built-in expression is always faster when one exists.

Definition

A Pandas UDF processes a whole column at once via Arrow batches as a pandas Series, turning per-row overhead into per-batch overhead.
  • Pandas UDF – also called a vectorized UDF, it operates on batches of data as pandas Series instead of one row at a time.
  • Apache Arrow – a columnar in-memory format used to transfer data between the JVM and Python without per-value conversion, amortizing the boundary cost.
  • Speedup – 10-100x faster than a plain Python UDF because the expensive JVM-to-Python crossing happens once per batch, not once per row.
  • Flavors – Series-to-Series for column mapping, Series-to-scalar for aggregation, iterator-of-Series for expensive per-partition setup.

A plain UDF mails the National Bank's transactions to a Python analyst one letter at a time, paying postage on each. A Pandas UDF ships a whole box of transactions at once; the analyst processes the entire box with a single sweep and ships the box back. Same work, almost no postage per item, so it finishes far faster.

Understanding the Concepts

What is the problem with plain Python UDFs?

A plain Python UDF serializes and deserializes every single row across the JVM-to-Python boundary. With millions of rows, the overhead dwarfs the actual computation.

PYTHON
from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import DoubleType

spark = SparkSession.builder.appName("UDFDemo").getOrCreate()

txn = spark.read.parquet("/lake/silver/transactions")

# Plain UDF: one row at a time, slow
@F.udf(DoubleType())
def add_fee_slow(amount):
    return float(amount) * 0.001

result = txn.withColumn("fee", add_fee_slow(F.col("amount")))
# Each of 100 million rows crosses the JVM-Python boundary individually

What is a Pandas UDF and how does it fix this?

A Pandas UDF operates on batches of rows as pandas Series. Apache Arrow transfers data in columnar batches, so the boundary crossing happens once per batch instead of once per row.

PYTHON
import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType

@pandas_udf(DoubleType())
def add_fee_fast(amount: pd.Series) -> pd.Series:
    return amount * 0.001

result = txn.withColumn("fee", add_fee_fast(F.col("amount")))
# 100 million rows processed in batches of ~10,000 each
# Only ~10,000 boundary crossings instead of 100 million

The speedup is typically 10-100x over the plain UDF.

How do you enable Apache Arrow?

Arrow must be enabled for the fast transfer path to work.

PYTHON
# Enable Arrow for Pandas UDFs
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# Tune the batch size (default: 10000 rows per batch)
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "10000")

# Verify it is on
print(spark.conf.get("spark.sql.execution.arrow.pyspark.enabled"))  # true

If Arrow is disabled or an unsupported type appears, Spark falls back to the slower non-Arrow path.

What is a Series-to-Series Pandas UDF?

This is the most common type. It maps one or more input columns to one output column, processing each batch as a pandas Series.

PYTHON
@pandas_udf(DoubleType())
def risk_score(amount: pd.Series, tx_count: pd.Series) -> pd.Series:
    """Custom risk scoring formula using NumPy"""
    import numpy as np
    base = np.log1p(amount) * 0.3
    frequency_factor = np.where(tx_count > 100, 1.5, 1.0)
    return base * frequency_factor

result = txn.withColumn(
    "risk",
    risk_score(F.col("amount"), F.col("tx_count"))
)
result.select("account_id", "amount", "tx_count", "risk").show(5)
TEXT
+----------+--------+--------+------------------+
|account_id|  amount|tx_count|              risk|
+----------+--------+--------+------------------+
|  ACC-1001|50000.00|     120|4.8901234567890125|
|  ACC-1002| 3000.00|      15|2.4012345678901234|
+----------+--------+--------+------------------+

What is a Series-to-scalar Pandas UDF?

This type acts as a custom aggregate inside groupBy().agg(). It takes a Series and returns a scalar.

PYTHON
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType

@pandas_udf(DoubleType())
def weighted_avg(amount: pd.Series, weight: pd.Series) -> float:
    return (amount * weight).sum() / weight.sum()

result = txn.groupBy("branch_code").agg(
    weighted_avg(F.col("amount"), F.col("tx_count")).alias("weighted_avg_amount")
)
result.show(5)

What is an iterator-of-Series Pandas UDF?

The iterator form lets you initialize an expensive resource once per partition, then process many batches with it.

PYTHON
from typing import Iterator

@pandas_udf(DoubleType())
def score_with_model(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
    # Load the model ONCE per partition, not once per batch
    import joblib
    model = joblib.load("/models/fraud_scorer.pkl")

    for amount_batch in batch_iter:
        scores = model.predict(amount_batch.values.reshape(-1, 1))
        yield pd.Series(scores)

result = txn.withColumn("fraud_score", score_with_model(F.col("amount")))

This avoids loading the model thousands of times when there are many batches.

When should you use a built-in instead?

Pandas UDFs cannot be optimized by Catalyst. A built-in Spark function runs entirely on the JVM with no Python boundary.

PYTHON
# Built-in: fastest, Catalyst-optimized, no Python
result = txn.withColumn("fee", F.col("amount") * 0.001)

# Pandas UDF: 10-100x faster than plain UDF, but still slower than built-in
result = txn.withColumn("fee", add_fee_fast(F.col("amount")))

# Plain UDF: slowest, row-at-a-time
result = txn.withColumn("fee", add_fee_slow(F.col("amount")))

Use Pandas UDFs only when you need genuine Python or a library (NumPy, scikit-learn) that has no SQL equivalent.

What should a senior engineer know?

  • Always check if a built-in function exists before writing a UDF.
  • Use Kryo serialization alongside Arrow for best overall performance.
  • The iterator form is essential for ML inference where model loading is expensive.
  • Test UDFs on sample data first to catch type mismatches early.
  • In Spark 3.x, mapInPandas and applyInPandas extend the pattern to whole DataFrames.
PYTHON
# applyInPandas: process each group as a pandas DataFrame
def normalize_amounts(pdf: pd.DataFrame) -> pd.DataFrame:
    pdf["normalized"] = (pdf["amount"] - pdf["amount"].mean()) / pdf["amount"].std()
    return pdf

result = txn.groupBy("branch_code").applyInPandas(
    normalize_amounts,
    schema=txn.schema.add("normalized", DoubleType())
)

What are common interview follow-up questions?

  • Does Arrow work with all Spark types? Most types are supported, but complex nested types may fall back to non-Arrow.
  • Can Pandas UDFs be used in streaming? Yes, they work in structured streaming the same way as in batch.
  • What happens if Arrow is not installed? Spark raises an error or falls back to non-Arrow conversion, depending on the configuration.

Visual Overview

DIAGRAM
  Plain UDF vs Pandas UDF Data Transfer
  ──────────────────────────────────────

  Plain UDF (row-at-a-time):
  JVM ──row1──▶ Python ──result1──▶ JVM
  JVM ──row2──▶ Python ──result2──▶ JVM
  JVM ──row3──▶ Python ──result3──▶ JVM
  ...millions of crossings...

  Pandas UDF (vectorized via Arrow):
  JVM ══[batch of 10K rows]══▶ Python
       pd.Series operation
  JVM ◀══[batch of 10K results]══ Python
  ...only thousands of crossings...

  Speedup: 10-100x

Key Takeaway

A Pandas UDF, or vectorized UDF, processes data in batches rather than one row at a time. Apache Arrow transfers a column of values as a pandas Series to the Python worker, your function runs vectorized operations over the whole Series, and the result comes back via Arrow. Because the expensive JVM-to-Python boundary crossing happens once per batch instead of once per row, a Pandas UDF is typically 10-100x faster than a plain Python UDF.

There are three flavors depending on the shape of the work. Series-to-Series maps input columns to an output column and is the default for custom math. Series-to-scalar acts as a custom aggregate inside groupBy().agg(). Iterator-of-Series lets you load an expensive resource like a model once per partition and process many batches with it. Enable the fast path by setting spark.sql.execution.arrow.pyspark.enabled to true.

At the National Bank of India, a custom risk-scoring formula that used NumPy ran far too slowly as a plain UDF on millions of transactions, so rewriting it as a Series-to-Series Pandas UDF finished in a fraction of the time. The iterator form was used when a later step loaded a heavy fraud-scoring model once per partition. The important caveat is that Pandas UDFs still cannot be optimized by Catalyst, so a pure built-in expression is always faster when one exists.

Remember one sentence: a Pandas UDF ships data in Arrow batches as pandas Series, turning millions of per-row boundary crossings into thousands of per-batch crossings for a 10-100x speedup over plain Python UDFs.

---

Share
Use ← → arrow keys to navigate