Why plain Python UDFs are slow: `cloudpickle` and the executor

Description

Since module 3, dim_product.csv has carried a column that's never been used: unit_cost. That module's lesson 2 said so explicitly, in its common-mistakes section: "unit_cost doesn't show up again in this guide until module 7, when you compute margin (unit_price - unit_cost) to classify margin_category". This is that lesson. You're going to write your first Python UDF in this entire guide — a function no native Spark function resolves directly — and you're going to see, with the real execution plan, exactly why Spark's official documentation warns this path "encounters performance bottlenecks."

Connection to the module. This lesson isn't just "how to write an @udf" — it's, deliberately, the problem half of a two-part problem. Lesson 5 is going to resolve this exact same business logic with pandas_udf, and the comparison between the two only makes sense once you first understand, with evidence, what makes this version slow.

An analogy: a handwritten letter, for every package

Imagine Kiosko subcontracts a supplier overseas to check every package of merchandise and decide whether it's high or low margin. Instead of sending a single document with the criterion and letting the supplier apply it to every package in a shipment, Kiosko mails a handwritten letter, different for each individual package: "Package 1: price 0.55, cost 0.40 — classify it." The supplier has to open the envelope, read the letter, do the calculation, write the answer in another letter, and seal it back up — and repeat that whole process, envelope by envelope, for every single package in the shipment, with no way to process them together. That's, precisely, what a plain Python @udf does in its classic form: every row of the DataFrame turns into an individual Python object, gets serialized, gets sent to the executor's Python process, gets processed one row at a time, and the result comes back, also row by row.

Worked example: margin_category, with a regular @udf

Step 1 — The data: orders joined with dim_product, to have unit_price and unit_cost together

# regular_udf_margin_category.py
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, udf
from pyspark.sql.types import (
    StructType, StructField, StringType, IntegerType, DoubleType, TimestampType,
)

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"
)

This is the first time, in this entire guide, unit_price (from orders) and unit_cost (from dim_product) live in the same DataFrame, ready for the subtraction no earlier lesson needed.

Step 2 — The @udf: your own Python function, wrapped for Spark

@udf(returnType=StringType())
def margin_category(unit_price: float, unit_cost: float) -> str:
    return "high" if (unit_price - unit_cost) > 1.0 else "low"

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

@udf(returnType=StringType()) is the exact syntax the official PySpark guide documents: the decorator wraps a normal Python function — nothing different from any function you've already written — and explicitly declares to Spark what type it returns (StringType()), because Spark can't infer it from the Python code the way it does with a native expression like col("quantity") * col("unit_price").

Step 3 — Verify the result, over Kiosko's real week

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}

high_products = sorted({r["product_id"] for r in classified_df.filter(col("margin_category") == "high").select("product_id").distinct().collect()})
print(f"products in 'high' = {high_products}")
assert high_products == ["P004"]
print("Verification: only P004 falls in high (4.50 - 2.10 = 2.40 > 1.0) -> OK")

What to expect. Running python3 regular_udf_margin_category.py, the output is exactly this (executed in this run, PySpark 4.2.0):

+----------+---------------------+----------+---------+---------------+
|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}
products in 'high' = ['P004']
Verification: only P004 falls in high (4.50 - 2.10 = 2.40 > 1.0) -> OK

Check it by hand, with Kiosko's four products: P001 (0.55 - 0.40 = 0.15), P002 (1.20 - 0.60 = 0.60), P003 (0.75 - 0.35 = 0.40) — all three below 1.0, all three "low". Only P004 (4.50 - 2.10 = 2.40) clears the threshold, and falls in "high". Since P004 shows up in 7 of the real week's 40 orders (the same figure you already used in module 1 to count products), the final count is {"low": 33, "high": 7} — thirty-three rows from the other three products, seven from P004.

Step 4 — The execution plan: forcing the classic path

By default, since Spark 4.2, a plain @udf already uses Arrow to transport the data (faster than traditional pickle), even though — as you're going to confirm in lesson 5 — the function still runs one row at a time. To see the classic mechanism the official UDF documentation describes — the one that really uses cloudpickle row by row, with no batched transport at all — disable that transport explicitly before defining the UDF:

# classic_udf_plan.py -- the same margin_category, forcing the no-Arrow path
spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "false")

@udf(returnType=StringType())
def margin_category_classic(unit_price: float, unit_cost: float) -> str:
    return "high" if (unit_price - unit_cost) > 1.0 else "low"

classic_df = priced_df.withColumn(
    "margin_category", margin_category_classic(col("unit_price"), col("unit_cost"))
)
classic_df.explain()

What to expect (executed in this run, with spark.sql.execution.pythonUDF.arrow.enabled set to false):

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Project [order_id#0, product_id#2, unit_price#4, unit_cost#9, pythonUDF0#14 AS margin_category#13]
   +- BatchEvalPython [margin_category_classic(unit_price#4, unit_cost#9)#12], [pythonUDF0#14]
      +- Project [order_id#0, product_id#2, unit_price#4, 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,unit_price#4] ...
            +- BroadcastExchange HashedRelationBroadcastMode(List(input[0, string, false]),false), [plan_id=36]
               +- Filter isnotnull(product_id#6)
                  +- FileScan csv [product_id#6,unit_cost#9] ...

There's the node the official documentation describes: BatchEvalPython. Despite the name — "Batch" — this node doesn't process your data in vectorized batches: it's the internal name of the Spark operator that manages row-by-row communication with the executor's Python process, using pickle/cloudpickle to serialize each value. The rest of the plan is exactly the one you already know from module 5 (BroadcastHashJoin) — the UDF gets inserted as an additional step, opaque to Catalyst, after the join has already been resolved.

Diagram: a row's journey, from the JVM to the Python process and back

flowchart LR
    A["Executor (JVM)\nrow: unit_price=4.50, unit_cost=2.10"] -->|"cloudpickle\nserializes the row"| B["Executor's\nPython process"]
    B -->|"deserializes,\ncalls margin_category(4.50, 2.10)"| C["Runs ONE row\nat a time"]
    C -->|"result: 'high'"| D["cloudpickle\nserializes the result"]
    D -->|"back to the JVM"| A
    A -->|"repeats, row by row,\nfor all 40 rows"| A

Going deeper: what cloudpickle serializes, and why that's the bottleneck

It's worth distinguishing two things cloudpickle serializes, because they're different and both matter. First, it serializes the function itself — margin_category's code, along with any variable it captures from its environment (a closure) — to send it, once, from the driver to every executor. This cost is fixed, paid once per executor, and doesn't grow with the row count. Second — and this is the one the official documentation flags as the real bottleneck — on the classic path (Arrow disabled, the one you forced in step 4), cloudpickle/pickle also serialize every data row, individually, in each direction: from the JVM to the Python process, and back. With forty rows, that cost is insignificant. With ten million — the fact_orders_at_scale this module already wrote in lesson 2 — that same row-by-row pattern repeats ten million times, each one crossing the boundary between two separate processes (the executor's JVM and a separate Python process, communicating over a socket).

The official quote from the PySpark UDF guide says it plainly: "Scalar Python UDFs rely on cloudpickle for serialization and deserialization, and encounter performance bottlenecks, particularly when dealing with large data inputs and outputs". That same documentation also confirms that even the "Arrow-optimized" version Spark 4.2.0 uses by default — the one you didn't force disabled in this lesson — still runs the function row by row; the only thing Arrow speeds up there is transporting the data, not the way your Python function processes it one at a time. Lesson 5 shows the only one of the three mechanisms that genuinely changes that part.

Common mistakes

Forgetting to declare returnType in the @udf, and getting an unexpected type. What happens: someone writes @udf with no arguments, trusting Spark to infer the return type from the Python annotations (-> str). Why it happens: Python does allow type annotations, and it's reasonable to expect Spark to use them automatically. How to spot it: with no explicit returnType, PySpark defaults to StringType() for any UDF with no recognized annotation — a numeric value returned by a function with no declared returnType can end up converted to text or to null, with no visible error. How to fix it: always declare returnType explicitly, as in this lesson's @udf(returnType=StringType()) — the same "never infer when you can declare" discipline this guide has held since module 1 with StructType.

Confusing BatchEvalPython with real batch processing. What happens: someone reads the name BatchEvalPython in step 4's plan, and assumes Spark is already processing data in vectorized batches, with no need for pandas_udf. Why it happens: the word "Batch" in the operator's name is misleading — it describes how Spark internally organizes communication between processes, not how your Python function receives the data. How to spot it: check your function's signature: def margin_category(unit_price: float, unit_cost: float) -> str receives a float, a scalar value, not a collection — if your function receives a single value at a time, it's running row by row, no matter the plan node's name. How to fix it: lesson 5 shows a vectorized function's real signature (pandas.Series in and out) — that signature, not the node's name in the plan, is the reliable sign your function processes data in a batch.

Thinking a Python UDF is slower because Python "is a slow language." What happens: someone attributes an @udf's cost to Python being, in general, slower than the JVM, without connecting the explanation to the real data-transport mechanism. Why it happens: it's an intuitive simplification, and not entirely false, but incomplete. How to spot it: if your explanation of "why it's slow" never mentions serialization between processes, you're missing half the picture — the cost isn't just that Python interprets code more slowly than the JVM, it's that every individual row has to cross the boundary between two separate processes (JVM and Python), a round trip that doesn't exist for native Spark functions, which run entirely inside the JVM. How to fix it: the complete explanation, the one you're going to be able to recite after lesson 5, has two parts: (1) inter-process communication, row by row, with cloudpickle on the classic path; and (2) the lack of vectorization, even on the path with Arrow enabled by default.

Exercises

Exercise 1 — Rewrite margin_category with a different threshold, and confirm the new result by hand. Change the threshold from 1.0 to 0.5, and before running anything, calculate by hand which products should fall in "high" with the new threshold (recall the margins: P001=0.15, P002=0.60, P003=0.40, P004=2.40).

@udf(returnType=StringType())
def margin_category_05(unit_price: float, unit_cost: float) -> str:
    return "high" if (unit_price - unit_cost) > 0.5 else "low"

classified_05_df = priced_df.withColumn("margin_category", margin_category_05(col("unit_price"), col("unit_cost")))
high_products_05 = sorted({r["product_id"] for r in classified_05_df.filter(col("margin_category") == "high").select("product_id").distinct().collect()})
print(f"products in 'high' with threshold 0.5 = {high_products_05}")
assert high_products_05 == ["P002", "P004"]

products in 'high' with threshold 0.5 = ['P002', 'P004']

Exercise 2 — Force the classic path (BatchEvalPython) yourself, and confirm the result is identical to the path with Arrow enabled. Run margin_category twice: once with spark.sql.execution.pythonUDF.arrow.enabled set to false (set before the @udf), once with the default value (true). Confirm both give the same count {"low": 33, "high": 7}.

# Important: the setting gets resolved at the MOMENT the @udf is defined, not when it's called later.
spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "false")

@udf(returnType=StringType())
def margin_category_no_arrow(unit_price: float, unit_cost: float) -> str:
    return "high" if (unit_price - unit_cost) > 1.0 else "low"

no_arrow_df = priced_df.withColumn("margin_category", margin_category_no_arrow(col("unit_price"), col("unit_cost")))
counts_no_arrow = {r["margin_category"]: r["n"] for r in no_arrow_df.groupBy("margin_category").count().withColumnRenamed("count", "n").collect()}
print(f"count (arrow.enabled=false) = {counts_no_arrow}")
assert counts_no_arrow == {"low": 33, "high": 7}
print("Verification: the result does NOT change based on the transport mechanism -> OK")

count (arrow.enabled=false) = {'low': 33, 'high': 7}
Verification: the result does NOT change based on the transport mechanism -> OK

Exercise 3 — Explain, without code, why .explain() can't show margin_category's function content at any level of detail. In 2-3 sentences, comparing against what you already know about RDDs (module 2) and Catalyst (module 6), explain why the BatchEvalPython/ArrowEvalPython node shows up as an opaque box in any plan.

Summary and next step

In this lesson you wrote your first Python UDF in this guide — margin_category, finally using the unit_cost column dim_product loaded since module 3 — and confirmed the result by hand: only P004 (margin 2.40) falls in "high", with {"low": 33, "high": 7} over Kiosko's real week. You forced an @udf's classic path (spark.sql.execution.pythonUDF.arrow.enabled=false) to see, in the plan, the BatchEvalPython node the official documentation describes: row-by-row communication between the JVM and a separate Python process, with cloudpickle along the way. You also confirmed that even the path with Arrow enabled by default in Spark 4.2.0 still runs the function one row at a time — Arrow there speeds up the transport, not the execution.

Before moving on you should be able to: write an @udf with an explicit returnType from memory; explain what cloudpickle serializes, and why that cost grows with the row count; and explain why BatchEvalPython isn't batch processing, despite the name.

Lesson 5 introduces the mechanism that genuinely changes execution, not just transport: pandas_udf, where your function receives a complete batch of rows as a pandas.Series, not one row at a time.

Resources

• PySpark — "Unleashing UDFs & UDTFs" (the exact @udf syntax, and the official quote: "Scalar Python UDFs rely on cloudpickle for serialization and deserialization, and encounter performance bottlenecks"). spark.apache.org/docs/latest/api/python/user_guide/udfandudtf.html.
• PySpark — Apache Arrow in PySpark (the spark.sql.execution.pythonUDF.arrow.enabled setting, and the distinction between Arrow-optimized scalar UDFs — row by row, sped-up transport — and pandas_udf — vectorized). spark.apache.org/docs/latest/api/python/tutorial/sql/arrow_pandas.html.
• This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — margin_category's exact specification: "high" if unit_price - unit_cost > 1.0, verified by hand against Kiosko's real unit_cost values.