Shuffles: Beginner

A narrow transformation computes every output row from rows that sit in a single input partition. The task holding partition 2 reads its rows and applies the call, then writes partition 2 of the result; no row arrives from any other task and none leaves. This is how filter and select work, and withColumn and drop as well. union is narrow too, along with every column function you wrap inside any of them.

Take an orders table of 8 rows that Spark reads as 4 partitions of 2 rows each, and a job that keeps rows where status equals Completed. The task on the first partition sees orders 101 and 102 and keeps 101, which finishes its work. The task on the second sees 103 and 104 and keeps 103. Each task decides the fate of its rows by reading a column value already in front of it, so there is nothing to request from another executor and nothing to wait for.

Each of the 4 tasks reads its own 2 rows and hands the survivors straight to the driver. No arrow crosses from one task to another, and the result still has 4 partitions: 3 of them now hold 1 row each and the fourth holds 2. filter and select can shrink a partition or change its columns, as can withColumn, and none of them moves the partition count or the executor holding each partition. A union places both inputs' partitions side by side, unchanged.

The partition count after a filter equals the count before it, so keeping 1% of the rows leaves you with the same 200 partitions, most of them nearly empty. The Partitioning lessons deal with what to do about that.

Because a narrow call needs nothing from outside its partition, Spark runs a chain of them inside the same task. In the program below, each task takes a row, tests status, tests profit, computes profit_x2, drops the columns that select leaves out, and moves to the next row. The row stays in that executor's memory from the first call to the last. The code has 4 transformations, and they still add up to a single pass over the data.

Run the chain and read the 5 rows it prints. Then change the withColumn to compute F.col("profit") / 2 and run again. The values you see change, while the number of rows and the number of partitions stay the same, because every change you made is decided row by row inside the task that read the row.

completed = (orders
    .filter(F.col("status") == "Completed")
    .filter(F.col("profit").isNotNull())
    .withColumn("profit_x2", F.col("profit") * 2)
    .select("order_id", "status", "profit", "profit_x2"))

completed.show(5)

Spark reads an orders table of 7 rows as 2 partitions. The first holds orders 1 to 4, where 1 and 3 are US orders and 2 is EU, with 4 from APAC. The second holds orders 5 to 7. Order 5 is from EU and order 6 from LATAM; order 7 is another US order. Ask for the list of order ids per region and the difficulty shows itself: US appears in both partitions, and so does EU. No task can write the US row of the result from the rows it holds, and the same is true for EU.

Read the argument of the call to find the key. groupBy("region") gathers by region, dropDuplicates(["user_id"]) gathers by user, distinct() gathers by the whole row, and orderBy("profit") gathers by ranges of profit so that one task holds the smallest values and another the largest. Each of these has the same shape underneath, an answer that depends on every row sharing a key, and that shared need is what makes a call wide.

Run the aggregation below. Spark has to bring every row for a region into one task before it can print a line for that region, so the rows move between the filter and the count. The filter itself runs where the rows were read, in the same task, before anything travels.

by_region = (orders
    .filter(F.col("region").isNotNull())
    .groupBy("region")
    .count())

by_region.show()

The destination is computed, never chosen by hand. Spark hashes the key and takes the remainder when divided by the number of output partitions. Hash the string US and you get the same number on every executor, so every US row, whatever partition it started in, is sent to the same destination. Different keys can share a destination, and that causes no trouble: the receiving task sorts them into separate groups when they arrive. Everything else in a shuffle depends on each key having exactly one destination.

The diagram has 3 writing tasks and 3 reading tasks. Each writer splits its 3 rows into piles, one pile per reader that has rows coming, and each reader collects its pile from every writer. Watch the US rows: 4 of them, spread across several input partitions, converge on a single reading task, which then counts 4. The reader that receives JP counts a single row, and it still waited for every writer to finish, because any writer might have been holding another JP.

The number of destinations comes from a setting named spark.sql.shuffle.partitions, whose default is 200. It is independent of how many partitions were read and of how many distinct keys exist. A table read as 8 partitions and grouped by region with 4 distinct regions is split by every writer into 200 destinations, and at most 4 of them receive rows. By default Spark then merges small and empty destinations, so far fewer reading tasks run, and the Adaptive Query Execution lessons cover how. Either way, the input partition count stops mattering once a shuffle happens, because the reading side has a count of its own.

by_country = (user_events
    .filter(F.col("country").isNotNull())
    .groupBy("country")
    .count())

by_country.collect()

Say a job reads 50GB of order rows. A filter over those rows touches 50GB once, in the tasks that read them, and produces nothing on disk. An orderBy("order_date") over the same rows writes about 50GB of shuffle files, then reads about 50GB over the network, so the rows are handled 3 times. A dropDuplicates on the full row does the same when nearly every row is unique, because nearly every row survives to the shuffle. Sort the table and then deduplicate it and the rows cross twice: 100GB written, 100GB read, before the job has produced its first row of output.

When the aggregate is sum or count, Spark always adds up each task's rows per key before sending, and avg gets the same treatment. So a groupBy moves at most a single row per key per partition. The Shuffles: Intermediate lesson covers how.

Since a shuffle charges per row, the cheapest shuffle is the one with the fewest rows. A filter placed before a wide call drops rows in the task that read them, at CPU cost only. With 1,000,000 orders of which 250,000 are Completed, filtering before an orderBy means 250,000 rows cross instead of 1,000,000. Before a groupBy, it cuts the rows every writer hashes and adds up before sending. Any condition you can test on a raw row belongs before the wide call, and a condition on the aggregated result, such as a total above 1000, can only come after.

Run the program, then move the status filter after the agg at the end of the chain and run it again. It fails as soon as the line defining totals runs, before any action reads data, with an error that status cannot be resolved, since grouping by region left no status column. A condition on a raw column can only go before the shuffle, which is also where it costs least.

totals = (orders
    .filter(F.col("region").isNotNull())
    .filter(F.col("status") == "Completed")
    .groupBy("region")
    .agg(F.sum("profit").alias("total_profit"),
         F.count("order_id").alias("orders"))
    .filter(F.col("orders") >= 1))

totals.show()

Before you run the program below, count its shuffles. There are 7 calls between the table and the action, and 2 of them are wide. Run it with collect() so the sort is a true global sort. Then delete the orderBy line and run again: the same rows come back, 1 shuffle fewer, in whatever order the readers returned them.

report = (orders
    .filter(F.col("status") == "Completed")
    .filter(F.col("region").isNotNull())
    .withColumn("profit_cents", F.col("profit") * 100)
    .groupBy("region")
    .agg(F.sum("profit_cents").alias("total"))
    .filter(F.col("total") > 0)
    .orderBy("total"))

report.collect()

distinct() reads like a tidy-up and runs as a groupBy on every column: to know a row is unique, Spark has to see all its copies in one task. dropDuplicates(["user_id"]) does the same with a chosen key. And orderBy followed by collect() or a file write sorts across the whole table, so every row is sent to the partition that owns its range of values. Each of these costs the same shuffle as a groupBy on its key, whatever its name suggests.

orderBy followed by show() or limit(n) runs as a different plan: each partition keeps its top n rows and the driver merges those few, with no range shuffle. The range shuffle is paid when the full sorted result is needed.