Project: Kiosko's optimized query plan

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 two broadcast joins, decide with a real criterion to cache it — because you're going to reuse it in three different queries, exactly lesson 5's criterion — confirm that savings by reading .explain(mode="formatted") over the final result, verify Adaptive Query Execution keeps coalescing partitions even over a cached DataFrame, and close by releasing memory with .unpersist() as soon as you're done. All of it, verified with assert against the total known since module 1: 26,537,500.00.

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

An analogy: the kitchen that serves three dishes from the same stew, and cleans up when closing

Pick back up this module's refrigerator. This project is the complete kitchen on a real workday: the stew gets cooked once (fact_orders_at_scale, with its two joins), it gets stored in the refrigerator because it's known, ahead of time, that three different dishes are going to get served from that same pot (the overall total, the per-store breakdown, the running total), it gets confirmed that every dish served after the first one never cooked anything from scratch again, and at the end of the shift, whatever's left over gets stored or the refrigerator gets cleaned out — food never stays cached indefinitely "just in case."

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, and decide to cache it with a real criterion

# kiosko_optimized_query_plan.py
from pyspark.sql import SparkSession, Window
from pyspark.sql import functions as F
from pyspark.sql.types import (
    StructType, StructField, StringType, IntegerType, DoubleType, TimestampType,
)

print("=== Kiosko's optimized query plan: module 6's final delivery ===\n")

print("Part 1 -- SparkSession, fact_orders_at_scale rebuilt (M5), caching decision")
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))
)
fact_orders_at_scale_df.createOrReplaceTempView("fact_orders_at_scale")

# Caching decision, with lesson 5's criterion explicitly justified:
# this DataFrame is going to get reused across THREE distinct queries (Part 2), not just once.
fact_orders_at_scale_df.cache()
print(f"isCached BEFORE any action: {spark.catalog.isCached('fact_orders_at_scale')}\n")

Part 2 — The three queries justifying the cache (lesson 5), verified against the known total

print("Part 2 (L5) -- three queries reusing the cached fact_orders_at_scale_df")

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

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"Query 2 -- per-store breakdown = {by_store}")
assert by_store == {"S01": 9_575_000.00, "S02": 9_700_000.00, "S03": 7_262_500.00}

store_window = Window.partitionBy("store_id").orderBy("order_ts", "franchise_id", "order_id")
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 running.groupBy("store_id").agg(F.max("running_total").alias("max_running")).collect()
}
print(f"Query 3 -- 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\n")
print(f"isCached AFTER the three queries: {spark.catalog.isCached('fact_orders_at_scale')}\n")

Part 3 — Reading the savings: .explain(mode="formatted") over the final result

print("Part 3 (L2-L4) -- explain(formatted) of Query 2, over the cached DataFrame")
fact_orders_at_scale_df.groupBy("store_id").agg(F.round(F.sum("revenue"), 2).alias("total_revenue")).explain(mode="formatted")

What to expect. Running python3 kiosko_optimized_query_plan.py, Parts 1 through 3 produce exactly this (executed in this run):

=== Kiosko's optimized query plan: module 6's final delivery ===

Part 1 -- SparkSession, fact_orders_at_scale rebuilt (M5), caching decision
isCached BEFORE any action: True

Part 2 (L5) -- three queries reusing the cached fact_orders_at_scale_df
Query 1 -- overall total = 26537500.0
Query 2 -- per-store breakdown = {'S02': 9700000.0, 'S01': 9575000.0, 'S03': 7262500.0}
Query 3 -- 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 AFTER the three queries: True

Part 3 (L2-L4) -- explain(formatted) of Query 2, over the cached DataFrame
== Physical Plan ==
AdaptiveSparkPlan (31)
+- HashAggregate (30)
   +- Exchange (29)
      +- HashAggregate (28)
         +- InMemoryTableScan (1)
               +- InMemoryRelation (2)
                     +- AdaptiveSparkPlan (27)
                        +- == Final Plan ==
                           ResultQueryStage (17)
                           +- * Project (16)
                              +- * BroadcastHashJoin Inner BuildRight (15)
                                 :- * Project (10)
                                 :  +- * BroadcastHashJoin Inner BuildRight (9)
                                 :     :- * Filter (4)
                                 :     :  +- Scan csv  (3)
                                 :     +- BroadcastQueryStage (8), Statistics(sizeInBytes=8.0 MiB, rowCount=3)
                                 :        +- BroadcastExchange (7)
                                 :           +- * Filter (6)
                                 :              +- Scan csv  (5)
                                 +- BroadcastQueryStage (14), Statistics(sizeInBytes=8.0 MiB, rowCount=4)
                                    +- BroadcastExchange (13)
                                       +- * Filter (12)
                                          +- Scan csv  (11)
                        +- == Initial Plan ==
                           Project (26)
                           +- BroadcastHashJoin Inner BuildRight (25)
                              ... (identical in shape to the Final Plan -- both joins were already
                                    BroadcastHashJoin from the initial plan, because dim_store and
                                    dim_product have a size known ahead of time, just as you confirmed
                                    in lesson 4's exercise 1)

(27) AdaptiveSparkPlan
Arguments: isFinalPlan=true
...
(31) AdaptiveSparkPlan
Arguments: isFinalPlan=false

This plan brings together, in a single .explain(), three complete lessons from this module. The outer part — HashAggregate, Exchange, HashAggregate — is this run's Query 2, with isFinalPlan=false because it itself hasn't run yet. Right below, InMemoryTableScan over an InMemoryRelation is lesson 5's evidence: this query never rereads any CSV, it reads directly from the fact_orders_at_scale already cached in Part 1. And inside that InMemoryRelation — something no earlier lesson has shown yet — lives the complete plan for how that cache got built the first time: module 5's two BroadcastHashJoins, with their own internal AdaptiveSparkPlan, marked isFinalPlan=true because that part really did already run, with its own Final Plan/Initial Plan from lesson 4. A single plan, three pieces of the module, all verifiable in the same text.

Part 4 — Confirming AQE still acts, even over a cached DataFrame

print("\nPart 4 (L4) -- AQE keeps coalescing partitions, even over already-cached data")
q_cached = fact_orders_at_scale_df.groupBy("store_id").agg(F.round(F.sum("revenue"), 2).alias("total_revenue"))
q_cached.collect()
print(f"post-shuffle partitions, over CACHED fact_orders_at_scale_df = {q_cached.rdd.getNumPartitions()}")
assert q_cached.rdd.getNumPartitions() == 1
print("Verification: AQE coalesces to 1 partition, the same behavior as lesson 4, regardless of caching -> OK")

What to expect (executed in this run):

Part 4 (L4) -- AQE keeps coalescing partitions, even over already-cached data
post-shuffle partitions, over CACHED fact_orders_at_scale_df = 1
Verification: AQE coalesces to 1 partition, the same behavior as lesson 4, regardless of caching -> OK

Part 5 — Closing with a real criterion: releasing memory that's no longer needed

print("\nPart 5 (L6-L7) -- retained memory report, and explicit release when done")
infos = spark.sparkContext._jsc.sc().getRDDStorageInfo()
print(f"Cached RDDs before closing: {len(infos)}")
for i in range(len(infos)):
    info = infos[i]
    print(f"    numCachedPartitions={info.numCachedPartitions()} memSize={info.memSize():,} diskSize={info.diskSize():,}")

fact_orders_at_scale_df.unpersist(blocking=True)
print(f"\nisCached AFTER unpersist(): {spark.catalog.isCached('fact_orders_at_scale')}")

print("\n=== spark.stop() -- module 6 closed, Catalyst/explain()/AQE/caching verified over 10M rows ===")
spark.stop()

What to expect (executed in this run):

Part 5 (L6-L7) -- retained memory report, and explicit release when done
Cached RDDs before closing: 1
    numCachedPartitions=12 memSize=595,211,368 diskSize=0

isCached AFTER unpersist(): False

=== spark.stop() -- module 6 closed, Catalyst/explain()/AQE/caching verified over 10M rows ===

This project's fact_orders_at_scale_df — with the two joins and thirteen columns, more columns than the smaller fact from lessons 5 through 7 — took up 595,211,368 bytes in memory: caching's real cost, paid once and amortized over three complete queries. Releasing it explicitly at the end, with .unpersist(blocking=True), is the same discipline lesson 7 recommended for any session that keeps running after this script.

Diagram: the five parts, closing out the complete module

flowchart TD
    A["Part 1: fact_orders_at_scale rebuilt (M5),\ncache() decided with real criteria -- 3 reuses planned"] --> B
    B["Part 2 (L5): 3 queries verified --\n26,537,500.00, per-store breakdown, running total"] --> C
    C["Part 3 (L2-L4): explain(formatted) --\nInMemoryTableScan + nested Final/Initial Plan"] --> D
    D["Part 4 (L4): AQE keeps coalescing\npartitions over cached data"] --> E
    E["Part 5 (L6-L7): 595 MB reported,\nreleased with unpersist() when closing"] --> F["Module 6 closed:\nCatalyst, explain(), AQE\nand caching verified"]

Closing out this module's checklist, piece by piece

Six modules resolved out of eight. With this project closed, you have the complete criterion for reading why Spark executes what it executes — at any of its five levels of detail — and for deciding, with measured evidence and never a stopwatch, when saving an intermediate result is worth it.

Common mistakes

Caching fact_orders_at_scale_df in this project without having counted, ahead of time, the three queries that justify it. What happens: someone copies Part 1 of this project, sees the .cache(), and replicates it in their own pipeline without having done the same explicit count — how many real queries are going to reuse this result. Why it happens: it's tempting to treat .cache() as a "standard" part of any Spark pipeline, instead of a decision depending on each case's real reuse. How to spot it: if your version of this project can't point to, with the same precision as Part 2 (three named queries, each with its own assert), how many times each cached DataFrame gets reused, you're missing the justification — not the code. How to fix it: before writing .cache() in any pipeline of your own, explicitly name, as Part 1 of this project did, how many later queries are going to reuse it; if the answer is "not sure," this module's lesson 7 already showed what happens if the answer turns out to be one.

Reading the Initial Plan identical to the Final Plan, inside Part 3's InMemoryRelation, as evidence AQE "did nothing" in that part of the pipeline. What happens: someone reviews Part 3's nested plan, sees the two BroadcastHashJoins' Initial Plan and Final Plan have the same shape, and concludes Adaptive Query Execution contributed nothing to building the cache. Why it happens: lesson 4 showed a case where Initial Plan and Final Plan really were different — SortMergeJoin switching to BroadcastHashJoin — so it's easy to expect the same contrast in any plan with AQE active. How to spot it: if your conclusion is "AQE did nothing here," check whether the initial plan already had the right strategy from the start — as lesson 4's exercise 1 confirmed, with dim_store_df. How to fix it: dim_store.csv and dim_product.csv have a size known ahead of time (real files, not shuffle results), so Catalyst already correctly chooses BroadcastHashJoin from the initial plan — AQE has nothing to fix there, and that's exactly expected, not a mechanism failure. The AdaptiveSparkPlan isFinalPlan=true node still confirms that part of the plan already ran, even if its shape didn't change from the initial plan.

Running Part 4 before Part 2, expecting the same result. What happens: someone reorders the script and runs the AQE comparison (Part 4) before Part 2's three queries have materialized the cache, and is surprised if the behavior doesn't exactly match what's documented. Why it happens: it seems reasonable that the order of the parts, within the same script, shouldn't matter. How to spot it: if you reordered this project and fact_orders_at_scale_df.storageLevel doesn't yet reflect data materialized in memory by the time you reach your version of Part 4, the AQE comparison is also going to trigger the work of materializing the cache for the first time — mixing two different costs into the same measurement. How to fix it: follow this project's order exactly as it is: Part 2's Query 1 materializes the cache first (lesson 5's cost), and only afterward does Part 4 measure AQE's behavior over data that's already cached — a clean measurement, with no two different costs mixed into a single figure.

Exercises

Exercise 1 — Add a Part 6: repeat lesson 6's experiment, comparing MEMORY_ONLY against .cache()'s default MEMORY_AND_DISK_DESER, over this fact_orders_at_scale_df with joins (not lessons 5-7's smaller fact). Use .persist(StorageLevel.MEMORY_ONLY) instead of .cache(), and compare the resulting memSize against the 595,211,368 bytes already measured in Part 5 of this project.

from pyspark import StorageLevel

fact_orders_at_scale_df.unpersist(blocking=True)
fact_orders_at_scale_df.persist(StorageLevel.MEMORY_ONLY)
fact_orders_at_scale_df.count()

infos = spark.sparkContext._jsc.sc().getRDDStorageInfo()
for i in range(len(infos)):
    info = infos[i]
    print(f"MEMORY_ONLY: memSize={info.memSize():,} diskSize={info.diskSize():,}")

MEMORY_ONLY: memSize=... diskSize=0

Exercise 2 — Confirm that skipping .cache() entirely in this project would make Query 2 and Query 3 reread the original CSV. Without running anything, using the evidence already measured in lesson 5 (602,074,847 bytes per complete CSV read), predict how many additional bytes would get read from disk in this project if the .cache() line got removed from Part 1.

Exercise 3 — Explain, without code, why this project verified the result with assert in Part 2, instead of just trusting that Part 3's .explain() "looked fine." In 2-3 sentences, justify why this module's evidence — execution plans, cached bytes, partition count — never replaces verifying the result's own correctness.

Summary and next step: the end of module 6

With this mini-project you close out the complete module 6. You rebuilt fact_orders_at_scale with module 5's two broadcast joins, decided to cache it with lesson 5's exact criterion — three real queries justifying it, not an assumption — and confirmed that savings by reading .explain(mode="formatted"): InMemoryTableScan replacing the repeated work, with the complete plan for how the cache got built — including its own Adaptive Query Execution Final Plan/Initial Plan — nested inside. You confirmed AQE keeps coalescing partitions even over already-cached data, and closed by releasing 595,211,368 bytes of memory with .unpersist(), the same discipline lesson 7 recommended.

You took this module's central step: you stopped treating AdaptiveSparkPlan as an unexplained word in every plan, and learned to read, with evidence — never a stopwatch — exactly what Catalyst revises, when AQE rewrites a plan, and when .cache() genuinely saves work.

Where you're headed. Module 7 — parquet-at-scale-and-python-udfs — takes this same fact_orders_at_scale and writes it as columnar Parquet partitioned by store_id, measures partition pruning and column pushdown on the read, and closes with the reason a plain Python UDF is slow — it serializes row by row with cloudpickle to every executor — versus a vectorized pandas_udf with Arrow.

Resources

• Apache Spark — SQL Performance Tuning (Catalyst, .explain(mode=...), Adaptive Query Execution, and caching settings — the complete reference for this entire module). spark.apache.org/docs/latest/sql-performance-tuning.html.
• Apache Spark — RDD Programming Guide, "Which Storage Level to Choose?" section (the memory/CPU trade-off between storage levels, the foundation for exercise 1 of this project). spark.apache.org/docs/latest/rdd-programming-guide.html#which-storage-level-to-choose.
• This guide's DESIGN doc (spark-and-distributed-processing-guide/DISENO.md) — module 6's full objective and its place in the eight-module plan.