The Cache That Ate the Cluster

A medium Spark interview practice problem on DataDriven. Write the Spark code and run it against the tables, with instant feedback.

Domain

Spark

Difficulty

medium

Seniority

L5

Asked in

Microsoft

Problem

For each user in the transactions table, what is their total spend, how many transactions did they make, and what percentage of everyone's combined spend do they account for? Every one of those answers reads that same table, which in production runs to hundreds of millions of rows, so we want it pulled from storage only once and released after its last use, with the profiles coming back in any order.

You cached 200 GB and forgot to let go.

Expected output

Worked solution and explanation

This is a reuse-count question dressed up as a feature job. The share-of-total forces two plans over one source: a per-user groupBy can never see the all-users denominator, so transactions gets read twice no matter how you write it. The skill being probed is what you do with that fact. Anyone can compute a share of total. What separates candidates is counting the readers of each frame: transactions is read twice, so it earns a cache, every frame after it is read once, so caching it buys nothing, and the cache you take is a reservation you have to give back. Get that wrong at scale and you get the incident on the pager: five pinned frames, 218 GB on a 224 GB cluster, three executors dead.

Read the plan before the logs

In the incident, the step 5 plan gives it away: five nested InMemoryRelation nodes, from base_df at 24 features up to step4_df at 40. That is five copies of roughly the same 800M rows, each a little wider. The executor table confirms it: 55 GB used of 56 GB on the three dead nodes, and GC times in the hundreds of seconds.

Step 1: Count the readers of every frame

transactions feeds two independent plans: the per-user groupBy and the global agg that produces all_spend. Without a cache, both scan the source. per_user, overall and result are each consumed once, so a cache on any of them is pure cost.

Step 2: Pin the shared input, nothing else

Call cache() on transactions and let the downstream frames stay lazy. Joining the one-row overall onto per_user with crossJoin is a broadcast of a single row: no shuffle of the per-user side.

Step 3: Release it after the action that computes the result

Pair the cache() with an unpersist() placed after the action that computes the result, and read nothing from the pinned input after it. Here that action is result.collect(); in the incident fix it was the final write, which kept the footprint flat near 200 GB and executors near 28 GB used.

transactions = spark.table("transactions").cache()

per_user = transactions.groupBy("user_id").agg(
    F.sum("total_amount").alias("total_spend"),
    F.count("*").alias("purchase_count"),
).cache()  # read once: pinned for nothing

overall = transactions.agg(F.sum("total_amount").alias("all_spend")).cache()

features = per_user.crossJoin(overall).cache()  # pinned forever
features

# The one frame worth pinning: both readers below scan it.
transactions = spark.table("transactions").cache()

# Reader 1: per-user figures.
per_user = transactions.groupBy("user_id").agg(
    F.sum("total_amount").alias("total_spend"),
    F.count("*").alias("purchase_count"),
)

# Reader 2: the overall total every share is measured against.
overall = transactions.agg(F.sum("total_amount").alias("all_spend"))

# Nothing below is cached: each of these frames is read exactly once.
result = (
    per_user.crossJoin(overall)
    .withColumn("spend_share_pct", F.col("total_spend") * F.lit(100.0) / F.col("all_spend"))
    .select("user_id", "total_spend", "purchase_count", "spend_share_pct")
)

# Compute the result: this action fills the cache and both readers consume it
# (in a batch job it is the final write).
profiles = result.collect()

# Release the pinned input only after the result has been computed.
transactions.unpersist()

Cache every step

Five InMemoryRelation layers, about 218 GB pinned, GC above 40%, three executors lost, 142 minutes total.

Cache the shared input only

One InMemoryTableScan over the base, about 200 GB pinned, GC under 50 seconds per executor, all four alive, done in 66 minutes.

Common follow-up questions

• What if the base were 400 GB on the same 224 GB cluster? (Tests whether you drop the cache and reread from source, or use `DISK_ONLY`, instead of fighting the storage fraction.)
• The incident lineage grows by one join per step. When would you cut it? (Probes `checkpoint()` or a write-and-reread to bound plan size and retry cost.)
• Could you get `all_spend` without a second read of `transactions`? (Checks whether you see that an unpartitioned window funnels rows into one partition: tolerable over the small `per_user` side, a single-task stage over the base.)