Module 7: Parquet At Scale And Python Udfs

Rewriting `margin_category` as a `pandas_udf`

Next →

Description

Lessons 4 and 5 built margin_category twice, separately: first as a classic @udf, then as an @pandas_udf. This lesson puts both in the same script, over the same DataFrame, and confirms with assert — not anyone's word — that both produce exactly the same result, row by row. Then, it takes the vectorized version and applies it over fact_orders_at_scale, the ten million rows this module wrote partitioned in lesson 2, closing the complete argument for why this is the right way to write a UDF that reuses at any scale.

Connection to the module. This lesson is the synthesis of lessons 4 and 5: it introduces no new mechanism, it integrates the two you already know into a single direct comparison, and it extends the correctness proof — already confirmed over forty rows — up to ten million.

An analogy: the same shipment, two ways to process it, one final result

Pick back up, for the last time in this module, the supplier overseas. Imagine that, to confirm both ways of working — letter per package, or a single manifest — give the same result, Kiosko sends the same forty-package shipment twice: once with forty individual letters, once with a single manifest. At the end of the day, it compares both paths' forty classifications, package by package — and they match, all forty. That comparison isn't a decorative exercise: it's proof that changing the communication mechanism with the supplier never changes the classification criterion itself. Once that's confirmed, the remaining question is purely about scale: what happens when the shipment isn't forty packages, but ten million? This lesson answers that question with evidence, over this module's real fact_orders_at_scale.

Worked example: both versions, side by side

Step 1 — Define both versions in the same script

# compare_udf_vs_pandas_udf.py
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, udf, 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", "unit_price", "unit_cost")


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


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

Two functions, same business name (margin_category), same condition (unit_price - unit_cost > 1.0), two different decorators, two different signatures.

Step 2 — Apply both, and compare row by row

both_df = (
    priced_df
    .withColumn("margin_category_udf", margin_category_udf(col("unit_price"), col("unit_cost")))
    .withColumn("margin_category_pandas", margin_category_pandas(col("unit_price"), col("unit_cost")))
)

mismatches = both_df.filter(col("margin_category_udf") != col("margin_category_pandas")).count()
print(f"rows where @udf and @pandas_udf differ = {mismatches}")
assert mismatches == 0
print("Verification: @udf and @pandas_udf produce EXACTLY the same result, row by row -> OK")

both_df.select("product_id", "unit_price", "unit_cost", "margin_category_udf", "margin_category_pandas").distinct().orderBy("product_id").show(truncate=False)

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

rows where @udf and @pandas_udf differ = 0
Verification: @udf and @pandas_udf produce EXACTLY the same result, row by row -> OK
+----------+----------+---------+-------------------+----------------------+
|product_id|unit_price|unit_cost|margin_category_udf|margin_category_pandas|
+----------+----------+---------+-------------------+----------------------+
|P001      |0.55      |0.4      |low                |low                   |
|P002      |1.2       |0.6      |low                |low                   |
|P003      |0.75      |0.35     |low                |low                   |
|P004      |4.5       |2.1      |high               |high                  |
+----------+----------+---------+-------------------+----------------------+

Zero discrepancies, over Kiosko's real forty-row week. This assert — mismatches == 0 — is direct evidence of something lessons 4 and 5 already foreshadowed, now confirmed in the same script, without depending on comparing two separate runs: the transport and execution mechanism (cloudpickle row by row versus Arrow in batches) is completely independent of the business result.

Step 3 — The vectorized version, applied to fact_orders_at_scale (10,000,000 rows)

# pandas_udf_at_scale.py
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import col, pandas_udf
from pyspark.sql.types import StringType
import pandas as pd

spark = SparkSession.builder.appName("kiosko-spark").master("local[*]").config("spark.driver.memory", "3g").getOrCreate()

fact_at_scale_df = spark.read.parquet("fact_orders_at_scale.parquet")
print(f"fact_at_scale_df.count() = {fact_at_scale_df.count()}")


@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 = fact_at_scale_df.withColumn(
    "margin_category", margin_category(col("unit_price"), col("unit_cost"))
)

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 (10,000,000 rows) = {counts}")
assert counts == {"low": 8_250_000, "high": 1_750_000}

total = classified_df.agg(F.round(F.sum("revenue"), 2).alias("total")).collect()[0][0]
print(f"total_revenue after adding margin_category = {total}")
assert total == 26_537_500.00

print("assert OK: 1,750,000 rows in high (P004), 8,250,000 in low, revenue unchanged -> OK")

spark.stop()

What to expect (executed in this run, over lesson 2's partitioned fact_orders_at_scale.parquet):

fact_at_scale_df.count() = 10000000
count by margin_category (10,000,000 rows) = {'low': 8250000, 'high': 1750000}
total_revenue after adding margin_category = 26537500.0
assert OK: 1,750,000 rows in high (P004), 8,250,000 in low, revenue unchanged -> OK

Check it with module 4's same arithmetic: 7 P004 rows per franchise (the figure you already confirmed in lesson 4) × 250,000 franchises = 1,750,000 rows in "high"; the other 33 rows per franchise (P001, P002, P003) × 250,000 = 8,250,000 in "low". And the total revenue — 26,537,500.00, this guide's same anchor figure since module 4 — doesn't change at all: margin_category is a new derived column, touching revenue nowhere in the calculation.

Diagram: the same proof, two scales

flowchart TD
    A["priced_df\n40 rows (real week)"] --> B["@udf margin_category_udf\nBatchEvalPython / ArrowEvalPython\nrow by row"]
    A --> C["@pandas_udf margin_category_pandas\nArrowEvalPython\ncomplete batch"]
    B --> D{"mismatches == 0?"}
    C --> D
    D -->|"yes -- confirmed"| E["Same result,\ndifferent mechanism"]
    E --> F["fact_orders_at_scale.parquet\n10,000,000 rows"]
    F --> G["Only @pandas_udf applied at scale\n(the vectorized version)"]
    G --> H["high=1,750,000, low=8,250,000\nrevenue=26,537,500.00 (unchanged)"]

Going deeper: why this lesson doesn't run the classic @udf over the 10,000,000 rows

Notice something deliberate: step 3 of this lesson applies only the pandas_udf version over fact_orders_at_scale, not the classic @udf version. It isn't an oversight — it's this guide's same discipline since module 1: never measure with a stopwatch to justify a performance decision. Running the classic @udf over ten million rows and comparing wall-clock time against the pandas_udf would produce a number of seconds — but that number would depend on the machine, the system load, the exact moment of the run, none of which is evidence reproducible across different environments.

The evidence that does back the decision is already complete, and it's all mechanism, not clock: lesson 4 confirmed, with .explain(), that the classic path serializes every individual row with cloudpickle/pickle, crossing the boundary between the JVM and a separate Python process, once per row. Lesson 5 confirmed, with the same tool, that pandas_udf transports complete batches with Arrow, and that your Python function processes them in a single vectorized pass. This lesson confirmed, with assert over zero discrepancies, that both paths produce an identical result. With those three pieces of evidence — never a stopwatch — the argument is complete: at ten million rows, the number of inter-process boundary crossings the row-by-row path requires grows linearly with the data volume, while the batched path divides it by the batch size (spark.sql.execution.arrow.maxRecordsPerBatch, 10,000 rows by default) — a structural difference in the mechanism, documented by PySpark's own official UDF guide, not a laptop measurement.

Common mistakes

Comparing two string-type columns with == when one of the two has null values, and getting false negatives. What happens: someone rewrites this lesson's comparison with their own data where unit_price or unit_cost can be null, and filter(col("a") != col("b")) doesn't catch the expected discrepancy. Why it happens: in SQL (and in Spark), any comparison against null — including != — returns null, not true or false, so a row with null in either column simply disappears from the filter() result, with no warning. How to spot it: if your data has null values and your mismatches count comes out suspiciously low, first check how many rows have null in unit_price or unit_cost with .filter(col("unit_price").isNull() | col("unit_cost").isNull()).count(). How to fix it: in Kiosko, neither column has null — both declared False (not nullable) in their respective StructType since module 3 — so this lesson doesn't need it, but on your own data, use <=> (Spark's null-safe equality operator) instead of != if null values are possible.

Running the classic @udf over the ten million rows "just to see what happens," with no measurement goal. What happens: someone, driven by curiosity after reading this lesson's "going deeper" section, runs the classic @udf over the complete fact_orders_at_scale anyway. Why it happens: the curiosity of "seeing the number" is understandable, even after understanding why this guide doesn't use it as evidence. How to spot it: if you end up writing down how many seconds that run took on your laptop, and using it as an argument in any future conversation about performance, you fell exactly into the pattern this guide has warned against avoiding since module 1. How to fix it: there's nothing wrong with running the experiment out of personal curiosity — but the argument backing this lesson, and one you should be able to repeat yourself, is never "it took X seconds": it's "it crosses the process boundary once per row, versus once per batch of 10,000," a mechanical difference, not a single run's measurement on a single machine.

Assuming pandas_udf always wins, regardless of data size. What happens: someone generalizes this lesson's conclusion — pandas_udf is the right way at scale — toward any UDF, regardless of whether the DataFrame has ten rows or ten million. Why it happens: once you understand the mechanism, it's tempting to treat it as a universal rule with no conditions. How to spot it: over Kiosko's forty real-week rows, the mechanism difference between @udf and @pandas_udf is real — confirmed in this lesson — but practically irrelevant: forty inter-process boundary crossings, with or without batching, isn't a volume that justifies worrying about the transport mechanism. How to fix it: the criterion isn't "pandas_udf always," it's this guide's same criterion since module 1: the larger the volume, the larger the number of boundary crossings the row-by-row path accumulates, and therefore the larger batched processing's structural advantage — the decision depends on your data's real scale, not a fixed rule.

Exercises

Exercise 1 — Add a third column, computed with Spark's native expression (F.when()), and confirm it also matches the other two. Without using any UDF, write margin_category with F.when(col("unit_price") - col("unit_cost") > 1.0, "high").otherwise("low"), and confirm it matches margin_category_udf and margin_category_pandas.

See solution 
from pyspark.sql import functions as F

native_df = both_df.withColumn(
    "margin_category_native",
    F.when(col("unit_price") - col("unit_cost") > 1.0, "high").otherwise("low"),
)
three_way_mismatches = native_df.filter(
    (col("margin_category_udf") != col("margin_category_native"))
    | (col("margin_category_pandas") != col("margin_category_native"))
).count()
print(f"discrepancies among the 3 versions = {three_way_mismatches}")
assert three_way_mismatches == 0
print("Verification: the native F.when() also matches both UDFs -> OK")

Expected output:

discrepancies among the 3 versions = 0
Verification: the native F.when() also matches both UDFs -> OK

This exercise leaves an open question on purpose, for you to answer with what you already know: if the native F.when() also resolves margin_category correctly, and is additionally an expression Catalyst can actually optimize (unlike any UDF), when would it make sense to use a UDF instead of a native expression? The answer — that a UDF only makes sense when the logic is too complex to express with pyspark.sql.functions's native functions — is why this guide chose margin_category as a simple pedagogical example: in a real case, F.when() would be the right choice, and that honesty is part of the lesson.

Exercise 2 — Verify, over the complete fact_orders_at_scale, that margin_category's breakdown per store keeps the same proportion as over the real week. Group by store_id and margin_category, and confirm the proportion of "high" rows per store is the same across the 10,000,000 rows as across the original 40.

See solution 
by_store_margin = {
    (r["store_id"], r["margin_category"]): r["n"]
    for r in classified_df.groupBy("store_id", "margin_category").count().withColumnRenamed("count", "n").collect()
}
print(f"breakdown by store and margin = {by_store_margin}")
# P004 shows up at all 3 stores within the real week: 1 time in S01, 3 times in S02, 3 times in S03
assert by_store_margin[("S01", "high")] == 1 * 250_000 == 250_000
assert by_store_margin[("S02", "high")] == 3 * 250_000 == 750_000
assert by_store_margin[("S03", "high")] == 3 * 250_000 == 750_000
print("Verification: 250,000 + 750,000 + 750,000 = 1,750,000, the same figure from step 3 -> OK")

Expected output (relevant excerpt):

breakdown by store and margin = {..., ('S01', 'high'): 250000, ('S02', 'high'): 750000, ('S03', 'high'): 750000, ...}
Verification: 250,000 + 750,000 + 750,000 = 1,750,000, the same figure from step 3 -> OK

250,000 + 750,000 + 750,000 = 1,750,000 — the same total figure you already confirmed in step 3, now broken down by store, and consistent with the real week's exact proportion replicated across every franchise.

Exercise 3 — Explain, without code, why step 2's assert mismatches == 0 is stronger proof than only comparing lessons 4 and 5's aggregated counts ({"low": 33, "high": 7}) separately. In 2-3 sentences, explain what kind of error could go unnoticed if you only compare totals, and why comparing row by row avoids it.

See solution

Comparing only the aggregated counts ({"low": 33, "high": 7} in both versions) could, in theory, hide an error where both versions misclassify the same number of rows, but different rows — for example, if a bug made one version classify a P002 row as "high" and another P004 row as "low", the total count would balance out even though the result wasn't correct. Comparing row by row, as this lesson's assert mismatches == 0 does — checking every individual order_id has the same classification in both versions — eliminates that possibility: the test only passes if absolutely every row matches, not if the totals happen to add up. This is the same verification discipline this guide used since module 3, comparing Spark's fact_orders against the four earlier engines not just on the total (106.15), but on the complete per-store breakdown.

Summary and next step

In this lesson you confirmed, with direct evidence in a single script, that @udf and @pandas_udf produce exactly the same result — zero discrepancies over Kiosko's real forty-row week — and applied the vectorized version over fact_orders_at_scale's ten million rows, verifying with assert the exact breakdown: 1,750,000 rows in "high", 8,250,000 in "low", with total revenue (26,537,500.00) untouched. You also confirmed, in the "going deeper" section, why this lesson doesn't run the classic @udf at that same scale: the complete argument is already backed by documented mechanism, not a stopwatch.

Before moving on you should be able to: write both versions of margin_category (@udf and @pandas_udf) from memory; explain why comparing results row by row is stronger than only comparing aggregated totals; and recite the complete argument for why pandas_udf scales better, without using the word "seconds" anywhere in the explanation.

With the UDF problem fully resolved — problem, solution, comparison — lesson 7 shifts register entirely: it names, in a single paragraph and with no line of streaming code, where this same DataFrame API extends to when data stops arriving in batches and starts arriving continuously.

Resources

← Previous`pandas_udf` and Arrow-based vectorizationNext →Naming Structured Streaming, without building it