Project: Kiosko's partitioned Parquet and a vectorized UDF

Description

This project closes the module by pulling its previous seven lessons together into a single script: you rebuild fact_orders_at_scale with module 5's broadcast joins, write it partitioned by store_id (lesson 2), confirm both partition pruning and predicate pushdown in the same plan (lesson 3), and apply margin_category as a pandas_udf over the complete ten million rows (lessons 4 through 6). All of it, verified with assert against the numbers known since modules 1 and 4: 26,537,500.00 in total revenue, and now also 1,750,000 rows classified as "high" margin.

Connection to the module. This project introduces no new concept — it's the final integration of lessons 2 through 7, applied in a single pipeline, instead of isolated one by one.

An analogy: the organized warehouse, and the manifest that classifies every box

Pick back up this module's two complete analogies. This project is the full day's work: first the warehouse gets organized into three labeled aisles (lesson 2's partitioned write); then, a forklift walks straight into the right aisle without opening the other two, and checks the labels of the boxes it does need without unpacking every one (lesson 3's partition pruning and predicate pushdown); and finally, a single shipping manifest — not forty thousand handwritten letters — classifies every one of the ten million boxes as high or low margin, in a single vectorized pass (lessons 5 and 6's pandas_udf).

The material: everything this module built, in a single flow

You need kiosko_orders_at_scale.csv (module 4), dim_store.csv, and dim_product.csv (module 3), in the same folder where you're going to run this script.

The verified reference solution

Part 1 — Rebuild fact_orders_at_scale, just like in modules 5 and 6

# kiosko_partitioned_parquet_and_vectorized_udf.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 (
    StructType, StructField, StringType, IntegerType, DoubleType, TimestampType,
)
import pandas as pd

print("=== Partitioned Parquet and a vectorized UDF: module 7's final delivery ===\n")

print("Part 1 -- SparkSession, fact_orders_at_scale rebuilt (M5/M6)")
spark = SparkSession.builder.appName("kiosko-spark").master("local[*]").config("spark.driver.memory", "3g").getOrCreate()

scale_schema = StructType([
    StructField("order_id", StringType(), False),
    StructField("franchise_id", IntegerType(), 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_store_schema = StructType([
    StructField("store_id", StringType(), False),
    StructField("store_name", StringType(), False),
    StructField("city", StringType(), 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_at_scale_df = spark.read.csv("kiosko_orders_at_scale.csv", schema=scale_schema, header=True, enforceSchema=False)
dim_store_df = spark.read.csv("dim_store.csv", schema=dim_store_schema, header=True, enforceSchema=False)
dim_product_df = spark.read.csv("dim_product.csv", schema=dim_product_schema, header=True, enforceSchema=False)

fact_orders_at_scale_df = (
    orders_at_scale_df
    .join(dim_store_df, "store_id")
    .join(dim_product_df, "product_id")
    .withColumn("revenue", F.round(F.col("quantity") * F.col("unit_price"), 2))
)
print(f"fact_orders_at_scale_df.count() = {fact_orders_at_scale_df.count()}\n")

Part 2 — Partitioned write by store_id (lesson 2)

print("Part 2 -- partitioned write by store_id")
fact_orders_at_scale_df.write.mode("overwrite").partitionBy("store_id").parquet("fact_orders_at_scale_project.parquet")
print("Written: fact_orders_at_scale_project.parquet\n")

Part 3 — Partition pruning vs predicate pushdown, in the same script (lesson 3)

print("Part 3 -- partition pruning (store_id) vs predicate pushdown (product_id)")
fact_at_scale_df = spark.read.parquet("fact_orders_at_scale_project.parquet")

s01_only_df = fact_at_scale_df.filter(col("store_id") == "S01")
print(f"s01_only_df.count() = {s01_only_df.count()}")
assert s01_only_df.count() == 4_000_000

p004_only_df = fact_at_scale_df.filter(col("product_id") == "P004")
print(f"p004_only_df.count() = {p004_only_df.count()}")
assert p004_only_df.count() == 1_750_000
print("assert OK: partition pruning (S01) and predicate pushdown (P004) verified\n")

Part 4 — pandas_udf over the complete ten million rows (lessons 4 through 6)

print("Part 4 -- pandas_udf margin_category, applied over 10,000,000 rows")

@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 = {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 = {total}")
assert total == 26_537_500.00

by_store = {
    r["store_id"]: r["total_revenue"]
    for r in classified_df.groupBy("store_id").agg(F.round(F.sum("revenue"), 2).alias("total_revenue")).collect()
}
print(f"breakdown by store = {by_store}")
assert by_store == {"S01": 9_575_000.00, "S02": 9_700_000.00, "S03": 7_262_500.00}
print("assert OK: 26,537,500.00 total, exact per-store breakdown, margin_category correct\n")

What to expect. Running python3 kiosko_partitioned_parquet_and_vectorized_udf.py, Parts 1 through 4 produce exactly this (executed in this run, PySpark 4.2.0, end to end in a single SparkSession):

=== Partitioned Parquet and a vectorized UDF: module 7's final delivery ===

Part 1 -- SparkSession, fact_orders_at_scale rebuilt (M5/M6)
fact_orders_at_scale_df.count() = 10000000

Part 2 -- partitioned write by store_id
Written: fact_orders_at_scale_project.parquet

Part 3 -- partition pruning (store_id) vs predicate pushdown (product_id)
s01_only_df.count() = 4000000
p004_only_df.count() = 1750000
assert OK: partition pruning (S01) and predicate pushdown (P004) verified

Part 4 -- pandas_udf margin_category, applied over 10,000,000 rows
count by margin_category = {'low': 8250000, 'high': 1750000}
total_revenue = 26537500.0
breakdown by store = {'S01': 9575000.0, 'S02': 9700000.0, 'S03': 7262500.0}
assert OK: 26,537,500.00 total, exact per-store breakdown, margin_category correct

Notice something worth verifying by hand, with nothing to run: p004_only_df.count() (filtering by product_id == "P004", regardless of store) gives 1,750,000 — exactly the same figure as counts["high"] in Part 4. That isn't a coincidence — it's proof, via two independent paths over the same data, that margin_category == "high" and product_id == "P004" are, in this dataset, the same set of rows: you already confirmed it by hand in lesson 4 (only P004 clears the 1.0 margin threshold), and here you confirm it again, at full scale, with two distinct queries arriving at the same number.

Part 5 — Closing out the module

print("Part 5 -- module 7's final checklist")
print("=== spark.stop() -- module 7 closed ===")
spark.stop()

Diagram: the four parts, closing out the complete module

flowchart TD
    A["Part 1: fact_orders_at_scale rebuilt (M5/M6)\n10,000,000 rows, 13 columns"] --> B
    B["Part 2 (L2): partitioned write\nby store_id -- 3 folders, 36 files"] --> C
    C["Part 3 (L3): partition pruning (S01=4M rows)\nand predicate pushdown (P004=1.75M rows)"] --> D
    D["Part 4 (L4-L6): pandas_udf margin_category\nover 10,000,000 rows -- high=1.75M, low=8.25M"] --> E["Module 7 closed:\npartitioned Parquet and vectorized UDF verified"]

Closing out this module's checklist, piece by piece

Seven lessons resolved out of seven, plus this closing project. With this, you have the complete criterion for writing Parquet at scale with the right organization, reading with evidence exactly what Spark prunes in every query, and knowing — with documented mechanism, never a stopwatch — when a pandas_udf replaces a plain Python UDF.

Common mistakes

Writing this project's partitioned Parquet over the same folder as lesson 2's, without noticing. What happens: someone runs this project in the same working folder where they already ran lesson 2, and fact_orders_at_scale.parquet (from lesson 2) gets confused with fact_orders_at_scale_project.parquet (from this project) — two different folders, with identical data. Why it happens: both scripts rebuild exactly the same DataFrame from the same source, so it's easy not to notice this project uses a deliberately different file name. How to spot it: if your script fails looking for a folder that doesn't exist, or if du -sh shows double the expected disk space, check you're using each script's correct name. How to fix it: this project uses fact_orders_at_scale_project.parquet, with the _project suffix, precisely to avoid overwriting lesson 2's artifact — if you'd rather reuse the same file, you can point both scripts at fact_orders_at_scale.parquet with no problem at all, since the content is identical.

Reading p004_only_df.count() == counts["high"] as a coincidence without verifying it with your own evidence. What happens: someone reads the number in this project's "What to expect," assumes it's always going to be that way for any dataset, and generalizes the observation beyond Kiosko. Why it happens: seeing two identical numbers in the same script invites generalizing the relationship without checking the exact condition producing it. How to spot it: the equality p004_only_df.count() == counts["high"] is specific to this dataset, because P004 is the only product whose margin clears 1.0 — in a dataset with a different threshold, or different prices, "the set of rows with product_id == 'P004'" and "the set of rows with margin_category == 'high'" could be completely different. How to fix it: treat this coincidence as a valid cross-check for this specific dataset — the same double-verification discipline this guide has used since module 3 — not as a general rule about the relationship between products and margins.

Running Part 4 without having run Part 2 first, and getting a file-not-found error. What happens: someone copies only Part 4 of this project — the pandas_udf — and runs it against fact_orders_at_scale_project.parquet without having written that Parquet first. Why it happens: it's tempting to copy only the part that matters at the moment, without running the complete script end to end. How to spot it: spark.read.parquet("fact_orders_at_scale_project.parquet") fails with a path-not-found error if Part 2 didn't run before in the same work session (even from an earlier run of the same script, since the Parquet stays persisted on disk). How to fix it: this project is designed to run end to end, Part 1 through Part 5, in a single execution — if you need to repeat only one part, first confirm the earlier parts already left their artifacts on disk (fact_orders_at_scale_project.parquet must exist before trying to read it).

Exercises

Exercise 1 — Add a Part 6: write classified_df (with margin_category already computed) as a new Parquet, partitioned by margin_category instead of by store_id. Use .write.mode("overwrite").partitionBy("margin_category").parquet(...), and confirm the result has exactly two folders (margin_category=high, margin_category=low).

classified_df.write.mode("overwrite").partitionBy("margin_category").parquet("fact_orders_by_margin.parquet")

import os
partitions = sorted(os.listdir("fact_orders_by_margin.parquet"))
print(f"partitions = {[p for p in partitions if p.startswith('margin_category')]}")

reread_high = spark.read.parquet("fact_orders_by_margin.parquet").filter(col("margin_category") == "high")
print(f"reread_high.count() = {reread_high.count()}")
assert reread_high.count() == 1_750_000
print("Verification: partitioned by margin_category, 2 folders, high=1,750,000 -> OK")

partitions = ['margin_category=high', 'margin_category=low']
reread_high.count() = 1750000
Verification: partitioned by margin_category, 2 folders, high=1,750,000 -> OK

Exercise 2 — Confirm, with .explain(), that filtering fact_orders_by_margin.parquet (from exercise 1) by margin_category == "high" produces partition pruning, just like store_id did in lesson 3. Verify PartitionFilters shows up full for this new Parquet.

high_only = spark.read.parquet("fact_orders_by_margin.parquet").filter(col("margin_category") == "high")
high_only.explain()

PartitionFilters: [isnotnull(margin_category#N), (margin_category#N = high)], PushedFilters: []

Exercise 3 — Explain, without code, why this project verified the result with four distinct asserts (Part 3 and Part 4), instead of trusting that "if there was no error, it's fine." In 2-3 sentences, connecting to the same discipline you already saw in modules 3 and 6's projects, explain why the absence of errors is never sufficient evidence of correctness.

Summary and next step: the end of module 7

With this mini-project you close out the complete module 7. You wrote fact_orders_at_scale partitioned by store_id, confirmed with .explain() the difference between partition pruning (store_id, free at the folder level) and real predicate pushdown (product_id, using Parquet statistics), and applied pandas_udf over the complete ten million rows, verifying margin_category with four independent asserts: 4,000,000 rows for S01, 1,750,000 rows for P004 (identical to the 1,750,000 "high"-margin rows), and total revenue intact at 26,537,500.00.

You took this module's central step: you learned to organize Parquet on disk with a partitioning criterion, to read with evidence exactly what Spark skips in every query, and to write your own business logic over a DataFrame the way that genuinely scales — never with a stopwatch as evidence, always with documented mechanism and the execution plan.

Where you're headed. Module 8 — this guide's capstone — assembles Kiosko's complete distributed pipeline, from generating the at-scale dataset to partitioned Parquet with pandas_udf, and closes with the question that opened this guide in module 1: when does Kiosko — or any real pipeline — genuinely need Spark, and when doesn't it?

Resources

• PySpark — "Unleashing UDFs & UDTFs" (the complete UDF reference, the foundation for lessons 4 through 6 of this module, integrated in Part 4 of this project). spark.apache.org/docs/latest/api/python/user_guide/udfandudtf.html.
• Apache Spark — SQL Data Sources: Parquet (partition discovery and spark.sql.parquet.filterPushdown, the foundation for lessons 2 and 3, integrated in Parts 2 and 3 of this project). spark.apache.org/docs/latest/sql-data-sources-parquet.html.
• This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — module 7's full objective and its place in the eight-module plan.