Spotting Shuffles in Your Code

You can count a program's shuffles without running it. Read the chain from the first DataFrame to the action, and at each call ask the question from the first section: can a task answer this for its own rows alone? When the answer is no, there is a shuffle at that line. Narrow calls add none, however many you stack, and a wide call adds exactly 1.

The wide calls, as they appear in code

• groupBy(...) followed by agg, count, sum or avg: 1 shuffle, keyed on the groupBy columns
• distinct(): 1 shuffle, keyed on every column
• dropDuplicates([...]): 1 shuffle, keyed on the listed columns
• orderBy(...) or sort(...) ending in collect() or a write: 1 shuffle that places rows by range
• join(...) of 2 large tables: 1 shuffle of each side, keyed on the join column

Walk a chain line by line

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.

The chain has 2 shuffles with 3 stretches of narrow work around them. The first shuffle moves the filtered raw rows. The second moves the aggregated rows, a single row per region, so it carries almost nothing. Where a shuffle sits in the chain matters as much as how many there are, because the rows alive at that line are the rows that cross.

Calls that stay narrow despite appearances

• union: stacks 2 DataFrames partition by partition, and no row changes partition
• filter after a groupBy: runs on the aggregated rows inside the reading task
• withColumn with F.when: a case expression evaluated per row, with no key involved
• select with arithmetic or string functions: each row rewritten where it sits
• show() after narrow calls: the driver pulls 20 rows, and nothing moves between executors

Calls that shuffle without saying so

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.

How many shuffles does the program below run?

(page_views
   .filter(F.col("device") == "mobile")
   .select("user_id", "page_url")
   .distinct()
   .withColumn("is_home", F.col("page_url") == "/")
   .filter(F.col("is_home"))
   .collect())

• 3, one for each call that follows select, since each of them sees a changed row shape
• 1, at distinct(), since every other call is tested or computed per row
• 2, one at distinct() and one at the final filter, which has to check rows across partitions
• 0, because distinct() on 2 columns can be finished inside each partition on its own

Do

• Count the wide calls before you run: each one is a shuffle
• Put every raw-column filter before the first groupBy
• Name the key of each shuffle: region, user_id, all columns
• Treat distinct and dropDuplicates as a groupBy in disguise

Avoid

• Assume a short chain is cheap: a single orderBy can move the whole table
• Add a second groupBy when one agg can compute both numbers
• Trust the method name; ask what the task needs to answer
• Call collect() on a shuffled result larger than the driver can hold

About this lesson

This section is part of the Shuffles: Beginner lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.