`pandas_udf` and Arrow-based vectorization

Description

Lesson 4 forced an @udf's classic path to see, in the plan, the BatchEvalPython node: row-by-row communication between the JVM and a separate Python process, with cloudpickle carrying the cost. This lesson resolves that exact same problem — margin_category, the same logic, the same result — with a different mechanism: pandas_udf. The difference isn't cosmetic: instead of receiving one value at a time, your function receives a complete batch of rows as a pandas.Series, transported with Apache Arrow, and returns another complete batch. Same result, radically different mechanism.

Connection to the module. This lesson is the direct answer to lesson 4: there you saw the problem (row-by-row communication, opaque to Catalyst); here you see the solution the official PySpark documentation recommends for "performance bottlenecks" with large data. Lesson 6 is going to put both versions side by side, over the same data, so the comparison is impossible to ignore.

An analogy: a single shipping manifest, for the whole batch

Pick back up lesson 4's supplier overseas. Instead of mailing a handwritten letter for every package, Kiosko changes strategy: it groups forty packages into a single shipment, and sends a single manifest — a sheet with one row per package, all together, in a standardized format the supplier can process in one pass. The supplier doesn't open forty different envelopes: it receives the complete manifest, checks it start to finish in a single operation, and returns a single document with all forty classifications. The final work — every package classified as high or low margin — is identical to lesson 4's. What changes is that the cost of "opening and closing the envelope" gets paid once per shipment, not once per package. That's, precisely, what pandas_udf does: it groups a partition's rows into batches, transports them with Arrow as a single unit, and your Python function processes the entire batch at once.

Worked example: margin_category, with pandas_udf

Step 1 — The same data from lesson 4

# pandas_udf_margin_category.py
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, pandas_udf
from pyspark.sql.types import (
    StructType, StructField, StringType, IntegerType, DoubleType, TimestampType,
)
import pandas as pd

spark = SparkSession.builder.appName("kiosko-spark").master("local[*]").getOrCreate()

orders_schema = StructType([
    StructField("order_id", StringType(), False),
    StructField("store_id", StringType(), False),
    StructField("product_id", StringType(), False),
    StructField("quantity", IntegerType(), False),
    StructField("unit_price", DoubleType(), False),
    StructField("order_ts", TimestampType(), False),
])
dim_product_schema = StructType([
    StructField("product_id", StringType(), False),
    StructField("product_name", StringType(), False),
    StructField("category", StringType(), False),
    StructField("unit_cost", DoubleType(), False),
])

orders_df = spark.read.csv("orders_2026-08-*.csv", schema=orders_schema, header=True, enforceSchema=False)
dim_product_df = spark.read.csv("dim_product.csv", schema=dim_product_schema, header=True, enforceSchema=False)

priced_df = orders_df.join(dim_product_df, "product_id").select(
    "order_id", "product_id", "product_name", "quantity", "unit_price", "unit_cost"
)

Important note, to leave no ambiguity. import pandas as pd in this lesson does not contradict the pandas ban you already saw in python-for-data-engineering-guide (where DuckDB and Polars replace pandas as the working engine). That ban is about pandas as a DataFrame engine — loading a complete file into memory with pd.read_csv() and working over it — something this guide never does either. pandas.Series here is only the data type Spark uses to represent a batch of one column inside a pandas_udf's boundary — pd.read_csv(), pd.DataFrame(), or any other pandas engine never get used outside this function's signature. It's a PySpark API that uses pandas.Series internally, not a return to pandas as a working tool.

Step 2 — The pandas_udf: the same logic, a different signature

@pandas_udf(StringType())
def margin_category(unit_price: pd.Series, unit_cost: pd.Series) -> pd.Series:
    return ((unit_price - unit_cost) > 1.0).map({True: "high", False: "low"})

classified_df = priced_df.withColumn(
    "margin_category", margin_category(col("unit_price"), col("unit_cost"))
)

Compare this signature against lesson 4's: def margin_category(unit_price: float, unit_cost: float) -> str received a scalar value; def margin_category(unit_price: pd.Series, unit_cost: pd.Series) -> pd.Series receives a complete column — a batch of values, not one. unit_price - unit_cost is no longer a subtraction between two numbers: it's a vectorized pandas operation over two complete Series, evaluated all at once with the same internal mechanism pandas uses to process any arithmetic operation over a column. .map({True: "high", False: "low"}) converts the boolean result into the same two text labels, also in a single operation over the complete batch, not row by row with an if.

Step 3 — Verify the result: identical to lesson 4's

classified_df.select("product_id", "product_name", "unit_price", "unit_cost", "margin_category").distinct().orderBy("product_id").show(truncate=False)

counts = {r["margin_category"]: r["n"] for r in classified_df.groupBy("margin_category").count().withColumnRenamed("count", "n").collect()}
print(f"count by margin_category (40 rows) = {counts}")
assert counts == {"low": 33, "high": 7}
print("Verification: same result as lesson 4's regular @udf -> OK")

What to expect (executed in this run):

+----------+---------------------+----------+---------+---------------+
|product_id|product_name         |unit_price|unit_cost|margin_category|
+----------+---------------------+----------+---------+---------------+
|P001      |Bottled Water 600ml  |0.55      |0.4      |low            |
|P002      |Energy Bar           |1.2       |0.6      |low            |
|P003      |Instant Coffee Sachet|0.75      |0.35     |low            |
|P004      |Phone Charger Cable  |4.5       |2.1      |high           |
+----------+---------------------+----------+---------+---------------+

count by margin_category (40 rows) = {'low': 33, 'high': 7}
Verification: same result as lesson 4's regular @udf -> OK

Identical, row for row, to lesson 4's result — {"low": 33, "high": 7}, only P004 in "high". This is the central point that needs backing with evidence, not just intuition: changing the transport and execution mechanism never changes the result, when the business logic is the same.

Step 4 — The execution plan: ArrowEvalPython, with a vectorized signature

classified_df.explain()

What to expect (executed in this run):

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Project [order_id#0, product_id#2, product_name#7, quantity#3, unit_price#4, unit_cost#9, pythonUDF0#58 AS margin_category#13]
   +- ArrowEvalPython [margin_category(unit_price#4, unit_cost#9)#12], [pythonUDF0#58], 200
      +- Project [order_id#0, product_id#2, quantity#3, unit_price#4, product_name#7, unit_cost#9]
         +- BroadcastHashJoin [product_id#2], [product_id#6], Inner, BuildRight, false, false
            :- Filter isnotnull(product_id#2)
            :  +- FileScan csv [order_id#0,product_id#2,quantity#3,unit_price#4] ...
            +- BroadcastExchange HashedRelationBroadcastMode(List(input[0, string, false]),false), [plan_id=588]
               +- Filter isnotnull(product_id#6)
                  +- FileScan csv [product_id#6,product_name#7,unit_cost#9] ...

The node is ArrowEvalPython — the same node name you also saw with lesson 4's @udf when Arrow was enabled by default. This is an important piece not to overlook: the node's name in the plan alone doesn't distinguish an @udf with Arrow enabled from a pandas_udf — both transport data with Arrow, and that's why both produce the same kind of node. The real difference — the one that actually matters for performance at scale — lives in the function's signature, not the plan: an @udf (even with Arrow enabled) still calls your Python function once per row; a pandas_udf calls your function once per batch, with complete pandas.Series as the argument. The plan tells you Arrow is involved; your own function's signature is what confirms whether real vectorization is happening.

Diagram: a complete batch's journey, not a single row's

flowchart LR
    A["Executor (JVM)\ncomplete partition:\nN rows of unit_price, unit_cost"] -->|"Arrow\nserializes the COMPLETE column,\nnot row by row"| B["Executor's\nPython process"]
    B -->|"deserializes as\npandas.Series (complete batch)"| C["Runs the function ONCE\nover the whole batch\n(unit_price - unit_cost) vectorized"]
    C -->|"result: pandas.Series\nwith N 'high'/'low' values"| D["Arrow\nserializes the complete Series"]
    D -->|"back to the JVM,\na single trip per batch"| A

Going deeper: what "operating in batches" means, with the official quote

The official PySpark documentation defines pandas_udfs like this: "Pandas UDFs (a.k.a. Vectorized UDFs) are UDFs that take/return pandas Series or DataFrame serialized/deserialized by Apache Arrow and operate block by block". Three pieces of that definition, each verified in this lesson's worked example: take/return pandas.Series (step 2's signature, pd.Series -> pd.Series, instead of lesson 4's float -> str); serialized by Apache Arrow (the same in-memory columnar format you've already seen mentioned in this module's earlier lessons, for efficiently transferring data between the JVM and Python); and operate block by block, not row by row — each block's size is controlled by spark.sql.execution.arrow.maxRecordsPerBatch, with an officially documented default value of 10,000 rows per batch.

It's worth summarizing, in a single table, the three paths you've already seen in this module:

The middle row is the most surprising, and it's worth making it clear: Spark 4.2.0 already speeds up any @udf's transport by default, but that doesn't vectorize your logic — only a pandas_udf, with its Series-based signature, makes your own Python code operate over the complete batch at once, instead of receiving the same old scalar value, one at a time, just arriving faster.

Common mistakes

Writing a pandas_udf's body with a row-by-row for, without taking advantage of vectorization. What happens: someone, used to a regular @udf's scalar if/else, writes a pandas_udf's body like this: return pd.Series(["high" if (p - c) > 1.0 else "low" for p, c in zip(unit_price, unit_cost)]) — a plain Python for, iterating over the Series row by row. Why it happens: the signature changed (pd.Series instead of float), but the habit of thinking "one row at a time" didn't change with it. How to spot it: if your function's body has a for or a list comprehension iterating over the arguments, you're wasting the vectorization — the result is still correct, but you gave up pandas_udf's real benefit in the first place. How to fix it: use vectorized pandas operations — unit_price - unit_cost, comparisons (> 1.0), .map(), .where() — the way this lesson's example does; these operations are already implemented in C/Cython inside pandas and NumPy, and process the complete batch with no explicit Python loop.

Forgetting the function's return type (-> pd.Series) in the signature, even though the decorator already declares Spark's StringType(). What happens: someone writes def margin_category(unit_price: pd.Series, unit_cost: pd.Series): with no -> pd.Series annotation, and the code still works — Python doesn't require type annotations — but loses clarity about what kind of Python object the function returns. Why it happens: @pandas_udf(StringType()) already declares Spark's column type, so annotating the Python type too seems redundant. How to spot it: if your function returns a value that isn't a pandas.Series — a plain Python list, for example — some Spark versions throw an error at runtime, not at definition time. How to fix it: always keep both annotations — pd.Series in the Python signature, and Spark's type in the decorator (StringType()) — as executable documentation that your function honors the vectorized contract, this entire guide's same "declare, never infer" discipline.

Assuming any pandas operation inside a pandas_udf is automatically fast, regardless of what it does. What happens: someone writes complex logic inside a pandas_udf — applying .apply() with a lambda function over the Series, for example, instead of a native vectorized operation — and assumes that just being inside a pandas_udf automatically inherits the full performance advantage. Why it happens: pandas_udf gives data-transport vectorization (Arrow) unconditionally, but the logic's vectorization depends on which pandas operations you use inside. How to spot it: Series.apply(lambda x: ...) is, in itself, a form of row-by-row iteration inside pandas — faster than a plain Python for for some implementation reasons, but far from the real vectorization of operations like Series - Series or direct comparisons. How to fix it: inside a pandas_udf, always prefer native vectorized pandas/NumPy operations (direct arithmetic between Series, comparisons, .where(), .map() with a dictionary as in this lesson) over .apply() with a custom function — pandas_udf's payoff depends on two layers working together: batched transport (Arrow) and vectorized execution (pandas/NumPy).

Exercises

Exercise 1 — Rewrite margin_category with .where() instead of .map(), and confirm the same result. Use pandas.Series.where() to produce the same classification, and verify the count is still {"low": 33, "high": 7}.

@pandas_udf(StringType())
def margin_category_where(unit_price: pd.Series, unit_cost: pd.Series) -> pd.Series:
    is_high = (unit_price - unit_cost) > 1.0
    return pd.Series(["low"] * len(is_high)).where(~is_high, "high")

where_df = priced_df.withColumn("margin_category", margin_category_where(col("unit_price"), col("unit_cost")))
counts_where = {r["margin_category"]: r["n"] for r in where_df.groupBy("margin_category").count().withColumnRenamed("count", "n").collect()}
print(f"count (with .where()) = {counts_where}")
assert counts_where == {"low": 33, "high": 7}
print("Verification: .where() and .map() give the same result -> OK")

count (with .where()) = {'low': 33, 'high': 7}
Verification: .where() and .map() give the same result -> OK

Exercise 2 — Confirm, with type(), that the argument your function receives inside a pandas_udf really is a pandas.Series, not a Spark column. Temporarily modify margin_category to print type(unit_price) inside the function body, and confirm the real type.

@pandas_udf(StringType())
def margin_category_debug(unit_price: pd.Series, unit_cost: pd.Series) -> pd.Series:
    print(f"DEBUG (inside the executor): type(unit_price) = {type(unit_price)}", flush=True)
    return ((unit_price - unit_cost) > 1.0).map({True: "high", False: "low"})

debug_df = priced_df.withColumn("margin_category", margin_category_debug(col("unit_price"), col("unit_cost")))
debug_df.count()

DEBUG (inside the executor): type(unit_price) = <class 'pandas.core.series.Series'>

Exercise 3 — Explain, without code, why the ArrowEvalPython node name alone isn't enough to know whether an operation is vectorized. In 2-3 sentences, using this lesson's "going deeper" table, explain what additional evidence is needed to confirm it.

Summary and next step

In this lesson you rewrote margin_category as a pandas_udf, with a signature taking and returning pandas.Series instead of scalar values, and confirmed, with assert, the result is identical to lesson 4's @udf: {"low": 33, "high": 7}, only P004 in "high". You read the plan (ArrowEvalPython) and learned this whole module's subtlest piece: that same node name also shows up for a regular @udf with Arrow enabled, so real vectorization can only be confirmed by checking your function's signature, not the plan.

Before moving on you should be able to: write a pandas_udf's signature from memory (pd.Series -> pd.Series, with the @pandas_udf(type) decorator); explain the difference between "Arrow speeds up the transport" and "the function is vectorized"; and explain why import pandas as pd within this lesson doesn't contradict other guides in the ecosystem's ban on pandas as an engine.

Lesson 6 puts both versions — lesson 4's @udf and this lesson's @pandas_udf — side by side, over the same data, with the complete mechanism contrast documented in one place.

Resources

• PySpark — "Unleashing UDFs & UDTFs" (the official quote: "Pandas UDFs... are UDFs that take/return pandas Series or DataFrame serialized/deserialized by Apache Arrow and operate block by block", and @pandas_udf's exact syntax). spark.apache.org/docs/latest/api/python/user_guide/udfandudtf.html.
• PySpark — Apache Arrow in PySpark (spark.sql.execution.arrow.maxRecordsPerBatch, default value 10,000, and the complete distinction between scalar UDFs and pandas_udf). spark.apache.org/docs/latest/api/python/tutorial/sql/arrow_pandas.html.
• python-for-data-engineering-guide's DESIGN doc (the ban on pandas as a working engine, and why this lesson doesn't contradict it). src/guides/python-for-data-engineering-guide/DISENO.md
• This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — the explicit clarification that pandas_udf is a PySpark API allowed in this guide, distinct from pandas as a DataFrame engine.