The Word Count Shuffle Trap

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

Domain

Spark

Difficulty

easy

Seniority

L5

Asked in

Apple

Problem

Ever since the input tripled, the nightly tally of calls per response status that an on-call engineer owns has died with 48 GB of shuffle write and three executors out of memory, because it collects every single call for a status code onto one machine before counting them. She needs the same per-status call totals from api_calls, produced in a way that stays small no matter how many calls pile onto the busiest status code.

groupByKey works. Your cluster disagrees.

Tables

api_calls

Columns: call_id, endpoint, method, status, latency, user_id, call_time, err_msg

Expected output

Worked solution and explanation

This is word count with status codes standing in for words, and the whole question is where the adding happens. Everyone gets the right totals. What separates candidates is whether the sum runs before the shuffle or after it. If you gather first and add later, every single (status, 1) pair crosses the network, and the busiest key's entire value list piles into one executor's heap. With most traffic answering 200, that one key carries most of the job. This is how a job that ran fine for months starts killing executors in the quarter the data triples.

Why it only broke now

groupByKey has no combiner. It ships raw values and materializes each key's full Iterable on the reduce side, and only then does mapValues(sum) run. On a small input that list fits and nobody notices. At 2.1B records, the hottest key alone is over a hundred million boxed integers in a 16 GB heap, and the executors spend 85% of their time in garbage collection without freeing enough memory.

calls = spark.table("api_calls")

# every (status, 1) pair is shuffled; each status code's full list lands on one executor
status_counts = (
    calls.rdd
    .map(lambda row: (row.status, 1))
    .groupByKey()
    .mapValues(sum)
)

The fix

Step 1: Name the operation as associative

A count is a sum of ones, and sums can be split: add within each partition, then add the partials. That property is what makes it safe to combine early. A median has no such property, which is why a median is the legitimate reason to gather values.

Step 2: Let the engine combine on the map side

On the RDD API that means reduceByKey(lambda a, b: a + b). On DataFrames, groupBy("status").agg(F.count("*")) already plans a partial HashAggregate per partition, then the Exchange, then a final merge. Either way, at most one partial per key per partition crosses the wire.

Step 3: Check the plan, not the output

Run explain() and look for two HashAggregate nodes around the Exchange. If the shuffle appears before any aggregation, you have rebuilt the incident.

calls = spark.table("api_calls")

status_counts = (
    calls
    .groupBy("status")
    .agg(F.count("*").alias("call_count"))
)

Gather then add

groupByKey() then mapValues(sum): the shuffle carries every raw record, and the reduce side holds each key's whole list in memory.

Add then gather

reduceByKey or DataFrame groupBy().agg(): each partition emits one partial per key, and the reduce side adds a few hundred integers per key.

Common follow-up questions

• The 200 key still dominates one reduce task after the fix. What now? (Salting a skewed key and merging the salted partials.)
• How would you compute each status code's average `latency` without gathering values? (Carrying a `(sum, count)` pair through `aggregateByKey`.)
• What changes if you need the exact median `latency` per status code? (Non-associative aggregates and when gathering is justified.)