Assembling the distributed pipeline, end to end

Description

This is this module's longest lesson, and rightly so: it's the complete Deliverable 1 from lesson 2's brief. In a single script, with a single SparkSession, you're going to read kiosko_orders_at_scale (ten million rows), join it against dim_store and dim_product with the default broadcast join, compute revenue, decide to cache the result because you're going to reuse it in three different queries, compute a running revenue total per store and a top-product-per-store-per-day ranking with window functions, write the result as Parquet partitioned by store_id, and close by applying margin_category with a vectorized pandas_udf over the complete ten million rows. Nine parts, a single run, zero shortcuts.

Connection to the module. This lesson introduces no new concept — you already built, tested, and verified every piece of this script in its own module (M3 for the join and revenue, M5 for the JOIN strategy and windows, M6 for caching, M7 for partitioned Parquet and pandas_udf). The only genuinely new thing here is the order: everything in a single SparkSession, with no interruptions, over the complete dataset.

An analogy: dress rehearsal, with the whole orchestra playing together

Every earlier module in this guide was, in a sense, a different orchestra section's rehearsal: the violins practiced their part alone (the broadcast join, module 5), the brass theirs (caching, module 6), the percussion theirs (partitioned Parquet and pandas_udf, module 7). Every section sounded good in its own rehearsal. But an orchestra is never truly tested until the dress rehearsal, where every section plays together, in the score's real order, with no pauses between movements. It's in that dress rehearsal — not in any of the partial ones — where the problems that only emerge from the interaction between parts show up: does the tempo the violins set match what the brass expects? Does the silence before the percussion enters last as long as it should? This lesson is that dress rehearsal: every piece that already sounded good on its own, played together, start to finish, over the complete ten-million-row score.

Worked example: the complete pipeline, in nine parts

Part 1 — The SparkSession

# kiosko_distributed_pipeline.py
from pyspark.sql import SparkSession, Window
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("=== Module 8, lesson 3: assembling the complete distributed pipeline ===\n")

print("Part 1 -- SparkSession")
spark = (
    SparkSession.builder
    .appName("kiosko-spark")
    .master("local[*]")
    .config("spark.driver.memory", "4g")
    .getOrCreate()
)
print(f"Spark version: {spark.version}\n")

spark.driver.memory goes up to 4g in this lesson, for the same reason you already saw in modules 6 and 7's mini-projects: this script keeps several views of the same ten-million-row DataFrame active at once (the cached result, the partitioned write, the reread from Parquet), and the driver needs enough headroom to coordinate all of it without running out of memory.

Part 2 — Reading with an explicit schema

print("Part 2 -- reading kiosko_orders_at_scale, dim_store, dim_product (explicit schema)")
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)

raw_count = orders_at_scale_df.count()
print(f"orders_at_scale_df.count() = {raw_count}")
assert raw_count == 10_000_000
print("Verification: 10,000,000 raw rows -> OK\n")

kiosko_orders_at_scale.csv is exactly the file generate_orders_at_scale(250_000) generated in module 4 — 250,000 synthetic franchises, 40 rows each, 10,000,000 rows total, not one row more or less.

Part 3 — The default broadcast join

print("Part 3 -- broadcast join (default) + revenue")
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(col("quantity") * col("unit_price"), 2))
)
print(f"spark.sql.autoBroadcastJoinThreshold = {spark.conf.get('spark.sql.autoBroadcastJoinThreshold')}")
print("Join's physical plan (note BroadcastHashJoin / BroadcastExchange, no shuffle Exchange):")
fact_orders_at_scale_df.explain()

Part 4 — Deciding to cache, with lesson 5 of module 6's criterion explicitly justified

print("Part 4 -- caching decision: this DataFrame is going to get reused across three distinct queries")
fact_orders_at_scale_df.createOrReplaceTempView("fact_orders_at_scale")
fact_orders_at_scale_df.cache()
print(f"isCached BEFORE any action: {spark.catalog.isCached('fact_orders_at_scale')}\n")

Part 5 — The three queries justifying the cache

print("Part 5 -- Query 1: overall total")
total = fact_orders_at_scale_df.agg(F.round(F.sum("revenue"), 2).alias("total")).collect()[0][0]
print(f"total_revenue = {total}")
assert total == 26_537_500.00

print("Part 5 -- Query 2: per-store breakdown")
by_store = {
    r["store_id"]: r["total_revenue"]
    for r in fact_orders_at_scale_df.groupBy("store_id").agg(F.round(F.sum("revenue"), 2).alias("total_revenue")).collect()
}
print(f"by_store = {by_store}")
assert by_store == {"S01": 9_575_000.00, "S02": 9_700_000.00, "S03": 7_262_500.00}

print("Part 5 -- Query 3: running revenue total per store (window function)")
store_window = Window.partitionBy("store_id").orderBy("order_ts", "franchise_id", "order_id")
with_running = fact_orders_at_scale_df.withColumn("running_total", F.round(F.sum("revenue").over(store_window), 2))
final_running = {
    r["store_id"]: r["max_running"]
    for r in with_running.groupBy("store_id").agg(F.max("running_total").alias("max_running")).collect()
}
print(f"final running_total per store = {final_running}")
assert final_running == {"S01": 9_575_000.0, "S02": 9_700_000.0, "S03": 7_262_500.0}
print("assert OK: all three queries match 26,537,500.00, isCached reused 3 times\n")

Part 6 — Top product per store per day

print("Part 6 -- top product per store per day (window function, row_number)")
fact_with_day = fact_orders_at_scale_df.withColumn("order_day", F.to_date("order_ts"))
revenue_by_product_day = (
    fact_with_day.groupBy("store_id", "order_day", "product_id")
    .agg(F.round(F.sum("revenue"), 2).alias("total_revenue"))
)
rank_window = Window.partitionBy("store_id", "order_day").orderBy(F.desc("total_revenue"), "product_id")
top_products = (
    revenue_by_product_day
    .withColumn("rank", F.row_number().over(rank_window))
    .filter(col("rank") == 1)
)
num_top = top_products.count()
print(f"num_top (rows with rank == 1) = {num_top}")
assert num_top == 20
print("Verification: 20 store x day combinations, each with its top product -> OK\n")

Part 7 — Partitioned write by store_id

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

Part 8 — Partition pruning and pandas_udf over the complete result

print("Part 8 -- filtered read (partition pruning) + pandas_udf margin_category")
fact_from_parquet = spark.read.parquet("fact_orders_at_scale_m8.parquet")

s01_only = fact_from_parquet.filter(col("store_id") == "S01")
s01_count = s01_only.count()
print(f"s01_only.count() = {s01_count}")
assert s01_count == 4_000_000

@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_from_parquet.withColumn(
    "margin_category", margin_category(col("unit_price"), col("unit_cost"))
)
margin_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 = {margin_counts}")
assert margin_counts == {"low": 8_250_000, "high": 1_750_000}
print("assert OK: partition pruning (S01=4,000,000) and margin_category (high=1,750,000) verified\n")

Part 9 — Releasing memory and closing

print("Part 9 -- releasing the cache and closing the session")
fact_orders_at_scale_df.unpersist(blocking=True)
print(f"isCached AFTER unpersist(): {spark.catalog.isCached('fact_orders_at_scale')}")
print("\n=== Complete distributed pipeline assembled and verified over 10,000,000 rows ===")
spark.stop()

What to expect. Running python3 kiosko_distributed_pipeline.py (executed in this run, PySpark 4.2.0, Java 17, local[*], end to end in a single SparkSession):

=== Module 8, lesson 3: assembling the complete distributed pipeline ===

Part 1 -- SparkSession
Spark version: 4.2.0

Part 2 -- reading kiosko_orders_at_scale, dim_store, dim_product (explicit schema)
orders_at_scale_df.count() = 10000000
Verification: 10,000,000 raw rows -> OK

Part 3 -- broadcast join (default) + revenue
spark.sql.autoBroadcastJoinThreshold = 10485760b
Join's physical plan (note BroadcastHashJoin / BroadcastExchange, no shuffle Exchange):
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Project [product_id#3, store_id#2, order_id#0, franchise_id#1, quantity#4, unit_price#5, order_ts#6, store_name#8, city#9, product_name#11, category#12, unit_cost#13, round((cast(quantity#4 as double) * unit_price#5), 2) AS revenue#28]
   +- BroadcastHashJoin [product_id#3], [product_id#10], Inner, BuildRight, false, false
      :- Project [store_id#2, order_id#0, franchise_id#1, product_id#3, quantity#4, unit_price#5, order_ts#6, store_name#8, city#9]
      :  +- BroadcastHashJoin [store_id#2], [store_id#7], Inner, BuildRight, false, false
      :     :- Filter (isnotnull(store_id#2) AND isnotnull(product_id#3))
      :     :  +- FileScan csv [order_id#0,franchise_id#1,store_id#2,product_id#3,quantity#4,unit_price#5,order_ts#6] Batched: false, DataFilters: [isnotnull(store_id#2), isnotnull(product_id#3)], Format: CSV, ...
      :     +- BroadcastExchange HashedRelationBroadcastMode(List(input[0, string, false]),false), [plan_id=75]
      :        +- Filter isnotnull(store_id#7)
      :           +- FileScan csv [store_id#7,store_name#8,city#9] Batched: false, DataFilters: [isnotnull(store_id#7)], Format: CSV, ...
      +- BroadcastExchange HashedRelationBroadcastMode(List(input[0, string, false]),false), [plan_id=79]
         +- Filter isnotnull(product_id#10)
            +- FileScan csv [product_id#10,product_name#11,category#12,unit_cost#13] Batched: false, DataFilters: [isnotnull(product_id#10)], Format: CSV, ...

Part 4 -- caching decision: this DataFrame is going to get reused across three distinct queries
isCached BEFORE any action: True

Part 5 -- Query 1: overall total
total_revenue = 26537500.0
Part 5 -- Query 2: per-store breakdown
by_store = {'S02': 9700000.0, 'S01': 9575000.0, 'S03': 7262500.0}
Part 5 -- Query 3: running revenue total per store (window function)
final running_total per store = {'S02': 9700000.0, 'S01': 9575000.0, 'S03': 7262500.0}
assert OK: all three queries match 26,537,500.00, isCached reused 3 times

Part 6 -- top product per store per day (window function, row_number)
num_top (rows with rank == 1) = 20
Verification: 20 store x day combinations, each with its top product -> OK

Part 7 -- partitioned write by store_id
Written: fact_orders_at_scale_m8.parquet

Part 8 -- filtered read (partition pruning) + pandas_udf margin_category
s01_only.count() = 4000000
count by margin_category = {'low': 8250000, 'high': 1750000}
assert OK: partition pruning (S01=4,000,000) and margin_category (high=1,750,000) verified

Part 9 -- releasing the cache and closing the session
isCached AFTER unpersist(): False

=== Complete distributed pipeline assembled and verified over 10,000,000 rows ===

Nine parts, not a single assert failed. Notice something you can only see by running the complete pipeline, not in any of its isolated modules: dim_store.csv and dim_product.csv get read twice in this script — once in Part 3 (the .join()), and implicitly again if you reran a query that didn't reuse the cached DataFrame — but the .join()'s own result, fact_orders_at_scale_df, gets read from the original CSV only once, thanks to Part 4's .cache(). Without that cache, each of Part 5's three queries would have reread the ten million rows and rebuilt both BroadcastHashJoins from scratch — the same savings module 6 already measured, now genuinely paying off inside a complete pipeline, not an isolated experiment.

Diagram: the nine parts, in order

flowchart TD
    A["Part 1-2: SparkSession + reading\n10,000,000 raw rows"] --> B
    B["Part 3: broadcast join + revenue\n(BroadcastHashJoin x2, no shuffle)"] --> C
    C["Part 4: cache() -- justified\nby 3 planned reuses"] --> D
    D["Part 5: 3 queries --\n26,537,500.00 verified 3 times"] --> E
    E["Part 6: top-product ranking\nper store and day -- 20 combinations"] --> F
    F["Part 7: partitioned write\nby store_id"] --> G
    G["Part 8: partition pruning (S01=4M)\n+ pandas_udf margin_category"] --> H["Part 9: unpersist() + spark.stop()\nComplete pipeline verified"]

Going deeper: why the nine parts' order isn't arbitrary

It's worth noting this script couldn't be freely reordered without changing its meaning. Part 4's cache has to happen before Part 5's three queries — caching after already running the queries would save nothing, because each query would have already paid its own read-and-join cost. Part 7's partitioned write has to happen before Part 8's filtered read — you can't read with partition pruning a Parquet that doesn't exist yet. And Part 9's unpersist() has to happen at the end, after every query needing the cache already used it, because releasing the memory ahead of time would force Spark to recompute the DataFrame from scratch on the next query that needed it.

This order isn't an aesthetic convention — it's, literally, the real dependency sequence among the nine parts, the same discipline you already saw in each earlier module's mini-project (M5, M6, M7), now applied to a pipeline integrating all three disciplines at once, instead of one per module.

Common mistakes

Running Parts 5 and 6 without having run Part 4 first, and not noticing the difference because the result stays correct. What happens: someone reorders the script, runs Part 5's three queries before Part 4's .cache(), and since the final result — 26,537,500.00 — stays identical, assumes the order didn't matter. Why it happens: the numeric result's correctness is independent of whether the DataFrame is cached or not — Spark always computes the same thing, cached or not, only how much work it repeats changes. How to spot it: if your version of this script gives the same 26,537,500.00 regardless of where you put the .cache(), that confirms caching was never necessary for correctness — but it is necessary to avoid rereading and rebuilding both BroadcastHashJoins three times, a real cost that shows up in no assert, only in the execution plan and in the Spark UI's job count. How to fix it: verify the cache's effect with the same evidence module 6 used — spark.catalog.isCached() before and after, and the job count in localhost:4040's Jobs tab — not just the final numeric result, which never changes with or without caching.

Skipping Part 9 (unpersist()), and leaving the DataFrame cached indefinitely if the script keeps running afterward. What happens: someone extends this script with additional steps after Part 8, and forgets fact_orders_at_scale_df keeps occupying memory in the driver, with no pending query needing it. Why it happens: while the script keeps running inside the same SparkSession, a cached DataFrame no longer in use generates no visible error — it just silently occupies memory. How to spot it: if your extended script starts failing with insufficient-memory errors in later steps with no apparent relation to fact_orders_at_scale_df, check whether you explicitly released any cache that already served its purpose. How to fix it: Part 9's discipline — releasing with .unpersist(blocking=True) as soon as a cached DataFrame is no longer going to get reused — is the same one module 6's lesson 7 already recommended, and it becomes more important, not less, when the script combines several stages like this one.

Assuming spark.driver.memory = "4g" is a magic number you always have to copy. What happens: someone copies config("spark.driver.memory", "4g") into any script of their own, without understanding why this specific pipeline needs it. Why it happens: seeing a configuration value work invites copying it without questioning it, especially when the script is long and looks "advanced." How to spot it: if your own script, much simpler than this nine-part pipeline, also includes "4g" with no concrete reason, you probably copied it without thinking. How to fix it: "4g" is appropriate here because this script keeps active, at some point, several views of the same ten-million-row dataset — the cache, the partitioned write, the reread — within a single session; a script that only does a single read and a simple .count(), like modules 1 through 3's, doesn't need that extra headroom. Adjust the driver's memory based on what your specific pipeline retains at once, not out of habit.

Exercises

Exercise 1 — Confirm Query 2's post-shuffle partition count is still the one AQE coalesced to in module 4. After Part 5, add a line capturing fact_orders_at_scale_df.groupBy("store_id").agg(F.sum("revenue")).rdd.getNumPartitions() and compare it against the value you already saw in module 4 (Adaptive Query Execution coalescing shuffle partitions).

q_check = fact_orders_at_scale_df.groupBy("store_id").agg(F.sum("revenue"))
q_check.collect()
print(f"post-shuffle partitions = {q_check.rdd.getNumPartitions()}")
assert q_check.rdd.getNumPartitions() == 1

post-shuffle partitions = 1

Exercise 2 — Extend Part 6 with a fourth assert: confirm S02's top product on 2026-08-08 is the same one you already verified in module 5. Using top_products, filter by store_id == "S02" and order_day == "2026-08-08", and print the resulting product_id and total_revenue.

sample = top_products.filter(
    (col("store_id") == "S02") & (col("order_day") == "2026-08-08")
).collect()[0]
print(f"S02 2026-08-08 -- top product = {sample['product_id']}, total_revenue = {sample['total_revenue']}")

S02 2026-08-08 -- top product = P001, total_revenue = 1750000.0

Exercise 3 — Explain, without code, why this pipeline uses .cache() on fact_orders_at_scale_df but never on orders_at_scale_df (the raw DataFrame, before the join). In 2-3 sentences, apply module 6's lesson 5 same criterion to justify why caching the raw DataFrame wouldn't make sense in this specific pipeline.

Summary and next step

In this lesson you assembled, for the first time in this guide, Kiosko's complete distributed pipeline in a single SparkSession: reading kiosko_orders_at_scale (ten million rows), a broadcast join against dim_store and dim_product, caching justified by three real reuses, running revenue total per store and top-product-per-store-per-day ranking with window functions, a partitioned write by store_id, and margin_category computed with pandas_udf over the complete ten million rows. Nine parts, all verified with assert, run end to end.

Before moving on you should be able to: name, from memory, the nine parts in order; explain why .cache() has to happen before the three queries justifying it; and explain why the partitioned write has to happen before Part 8's filtered read.

Lesson 4 takes this pipeline's result and puts it through the complete correctness test: the exact chain of numbers connecting 106.15 (the real forty-row week) to 26,537,500.00 (this lesson's ten million rows), verified with simple math anyone can reproduce without running a single line of Spark.

Resources

• Apache Spark — SQL Performance Tuning (Catalyst, joins, AQE, caching — the complete reference this pipeline integrates from modules 5 and 6). spark.apache.org/docs/latest/sql-performance-tuning.html.
• PySpark — pyspark.sql.Window (the partitionBy/orderBy syntax Part 5 and Part 6 of this script reuse from module 5). spark.apache.org/docs/latest/api/python/reference/pyspark.sql/window.html.
• PySpark — "Unleashing UDFs & UDTFs" (the @pandas_udf reference Part 8 of this script reuses from module 7). spark.apache.org/docs/latest/api/python/user_guide/udfandudtf.html.
• This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — the exact specification for the complete pipeline this lesson assembles. src/guides/spark-and-distributed-processing-guide/DISENO.md.