Shuffles: Intermediate

The rest of this lesson follows that job over a year of plays read as 400 partitions on a cluster of 20 executors. Stage 0 runs 400 map tasks. Each reads its split, applies the filter in the same pass, and for every surviving row works out which of the 200 reduce partitions owns its (listener_id, artist_id) pair. The 200 is the default of spark.sql.shuffle.partitions, the setting that fixes how many partitions a DataFrame shuffle produces.

For a groupBy, the owner comes from a hash of the grouping columns, taken modulo the number of reduce partitions. Every map task applies the same function, so a given listener and artist pair goes to the same reducer whichever task happened to read it. That agreement lets reducer 17 collect every row for its keys without asking anyone where they went. A later section of this lesson opens up the hash and the other rule Spark uses.

The files go under spark.local.dir, which is /tmp unless set, on the machine running the executor; cluster managers often substitute directories of their own. A job writing 500GB of shuffle across 20 executors therefore needs about 25GB of free local disk on each, and when one machine fills up, its map tasks fail with No space left on device while the cluster's main storage sits half empty. The files also outlive the tasks that wrote them, because the next stage still has to fetch them.

Fetches go out in parallel batches, so a reducer pulls from several executors at once while it processes blocks that have already arrived. For the groupBy, processing means merging: the reducer feeds each incoming row into a hash map keyed by (listener_id, artist_id), adding counts for rows that share a key. Once the last block is merged, every key's count is final, because no other reducer holds rows for those keys. The finished rows then flow into the narrow work after the shuffle in the same stage, here the parquet write.

Each byte stage 0 writes, stage 1 reads, so the reduce stage's Shuffle Read total matches the map stage's Shuffle Write total. The split differs. Writes were divided by input partition, and reads are divided by key, each reducer taking whatever share of the data its keys own. In this job that share stays close across reducers. As the next section shows, each map task sends no more than 1 row for each (listener_id, artist_id) pair, so a listener who streams all day adds at most 400 rows per artist, and the hash spreads their artists over many reducers. Read sizes split apart when raw rows cross the shuffle, as after repartition("listener_id"): every play of a heavy listener lands on 1 reducer, that reducer reads far more than the median, and the stage waits on it.

The plan sets the reduce task count to spark.sql.shuffle.partitions, whatever the data holds. With adaptive execution turned off (spark.sql.adaptive.enabled set to false), a groupBy on a column with 3 distinct values launches all 200 reduce tasks. Each key's rows land in a single partition, so at least 197 tasks find every block empty and finish with no output, more if 2 of the keys hash to the same partition. Each empty task still costs the driver a scheduling round trip. With adaptive execution on, the default since Spark 3.2, Spark measures the map output once the map stage ends and merges the small and empty partitions into a few tasks, often 1 for 3 keys. The Adaptive Query Execution lessons cover how.

The cheapest shuffle byte is the one never written. For count, sum, min, max, avg and most other built-in aggregates, Spark computes the result in 2 passes. A partial pass runs inside each map task before the shuffle write and collapses that task's rows to 1 row per key it saw, carrying a running result. A final pass runs in the reducer after the read and merges each key's partial rows. Every DataFrame groupBy using these aggregates gets the partial pass without being asked.

Both outputs give the reducer what it needs: it adds US 4 to the partial US counts from every other map task and reaches the same total. The partial version writes 2 rows where the raw version writes 6. For avg, the partial row carries a sum and a count, since an average of averages would be wrong.

The same groupBy code cuts the shuffle by a factor of about 14,000 for country, 50 for artist, and almost nothing for listener: a task holding 2.5 million rows in arrival order sees nearly every listener once. Layout matters as much as the key. Had the files been written so each partition held a contiguous range of listeners, each task would see all of a listener's plays together, around 10 on average, and the partial pass would cut the shuffle about 10 times.

Run as written, exact takes 2 shuffles and 3 stages: the first shuffle groups by (artist_id, listener_id) so duplicate pairs meet and drop, and the second groups the surviving pairs by artist_id to count them. approx takes 1 shuffle and 2 stages, and each map task writes 1 sketch per artist it saw, the same size whatever the listener count.

The cell below counts rows and distinct values for 3 columns of the practice transactions table. Run it, then divide rows by each key count. The quotient is the most a partial pass could shrink a groupBy on that column if 1 task saw every row; how the rows are split across tasks decides how close a real job gets.

from pyspark.sql import functions as F

tx = spark.table("transactions")

tx.agg(
    F.count("*").alias("rows"),
    F.countDistinct("product_id").alias("product_keys"),
    F.countDistinct("user_id").alias("user_keys"),
    F.countDistinct("transaction_id").alias("transaction_keys"),
).show()

Every map task has to send each row to exactly 1 reducer, using a rule every other map task also applies. For shuffles keyed on columns, Spark has 2 rules. Hash partitioning sends a row to partition pmod(hash(key), n), which hashes the key columns with Murmur3, divides the hash by the partition count n, and keeps the non-negative remainder. Range partitioning cuts the key's sorted values into n contiguous intervals and sends each row to the interval its key falls in. A third rule, round-robin, deals rows out in turn with no key at all; repartition(n) uses it.

A groupBy reducer only needs every row for its keys, and a hash delivers that. A global sort needs more. Each reducer sorts its own rows, and reading partition 0, then 1, then 2 has to yield 1 ordered result, which holds only if every key in partition 0 sorts before every key in partition 1. A hash sends adjacent values like 99 and 100 to unrelated partitions, so orderBy uses ranges.

To place the boundaries, Spark first samples the data: around 100 keys per output partition by default, the value of spark.sql.execution.rangeExchange.sampleSizePerPartition, drawn from every input partition. It sorts the sample and picks n minus 1 cut points that split it evenly. That sampling is a job of its own, the kind you saw run ahead of an action's main job, and it computes the input once for the sample and again for the shuffle.

The animation sorts 9 orders by amount into 3 reduce partitions, with cut points drawn from its own sample. Watch which amounts reach each reducer: low ones at the first, high ones at the last, and each reducer sorts only its own share. Collected in partition order, the result runs 15, 30, 45, 60, 70, 80, 95, 120, 140.

Because the rule ends in a remainder, the same key lands in unrelated partitions under different counts: US might sit in partition 41 of 200 and partition 7 of 50. Two outputs hash-partitioned on the same column line up partition by partition only when they used the same n. Range layouts line up even less often, since their boundaries come from a sample of 1 run's data.

Neither rule splits a key: every row with country = 'US' lands in 1 partition under hash and under range alike.

A groupBy("region") needs every row of a region in 1 partition. After repartition("region") that already holds: the repartition hashed every row by region. Before adding a shuffle, Spark compares what each operation needs with how its input is partitioned, and here they match, so the aggregation runs inside the repartition's reduce tasks with no second shuffle.

Spark's test is whether the existing layout puts all rows that share the new grouping key in 1 partition. Hashing on a subset of the grouping columns passes it, hashing on extra columns fails it, and round-robin always fails. When the layout passes, its partition count carries forward, so repartition(50, "region") sets the task count of everything after it up to the next shuffle.

The first collect() runs 400 map tasks that read plays and write partial sums, then 200 reduce tasks. The second runs the same query, finds all 400 map outputs registered, skips stage 0, and runs only the 200 reducers, so plays is never read a second time.

The reuse belongs to 1 query. by_artist.filter(F.col("ms") > 0).collect() is a new query that plans its own shuffle, so both stages run again even though the rows feeding the shuffle are identical. The files are also temporary: Spark deletes them some time after the driver drops its last reference to the query that wrote them, and an executor that dies usually takes its files with it. Reuse you control on purpose comes from caching, covered in a later lesson.

The last question joins the partial pass from section 2 to the layout reuse in this one. A repartition has no aggregate to apply, so its map tasks write every row they receive, as they read it. Whatever runs after it, partial pass included, runs in the repartition's reduce tasks, on rows that have already crossed the network.