The Shuffle: Spark's Performance Killer

In the chapter "Transformations, Actions & Lazy Evaluation", a wide step needed a shuffle. A shuffle moves rows between partitions, so rows that belong together end up in one place.

A shuffle is often the most costly step in a Spark job. In an interview, you are often asked to find it in a slow job and make it smaller.

This chapter has four parts.

1. How a shuffle works, step by step, on the shop data.
2. Which steps cause a shuffle, and why a shuffle is slow.
3. How many pieces a shuffle makes, and the setting for it.
4. How to make a shuffle smaller, and how to read its numbers in the Spark UI.

The job is total revenue for each customer, over all orders. We do not filter by status this time. That makes customer 3 appear in two partitions, so the shuffle has real work to do.

We use the model from the architecture chapter: 3 partitions of 2 orders. Each partition has one task. Look at the job and the three partitions.

revenue = orders.groupBy("customer_id").sum("amount")

Customer 3 is now in Task 1 and in Task 2. To get one total for customer 3, those two rows must meet in one place.

Spark chooses the place with a hash. A hash turns the customer_id into a number. Spark divides that number by N, the number of new partitions, and keeps the remainder.

Real Spark uses its own hash. We use a simple model instead: customer_id % 3. The symbol % gives the remainder after dividing by 3.

Now comes the shuffle write.

Each stage 0 task sorts its partial rows by new partition. It writes them to one file on the local disk of its machine.

A small index file says where each part starts. Look at the parts.

Stage 0 is now finished, and stage 1 starts.

This is the shuffle read. Each stage 1 task fetches its part from every stage 0 file, over the network.

# Stage 1: three tasks, one for each partition after the shuffle.
revenue.show()

Every row for one customer reached one task. That was the goal of the shuffle.

A shuffle happens when rows from many partitions must meet. The groupBy above is one case. Here are the common ones, with the reason for each.

A narrow step, like filter or withColumn, never causes a shuffle. Each task works on its own rows.

To find a shuffle, call explain() and look for the Exchange line. Each Exchange is one shuffle.

In the plan below, the 200 is the number of partitions after the shuffle. Earlier plans left the 200 out to keep them short.

== Physical Plan (simplified) ==
HashAggregate(keys=[customer_id], functions=[sum(amount)])
+- Exchange hashpartitioning(customer_id, 200)
   +- HashAggregate(keys=[customer_id], functions=[partial_sum(amount)])
      +- FileScan parquet [customer_id, amount]

Count the Exchange lines. Two groupBy steps on different keys give two Exchanges, so two shuffles.

A narrow step works in memory, inside one task. A shuffle uses the disk and the network as well. There are four costs.

1. Serialize and write: Spark turns each row into bytes (this is serialization) and saves them to local disk.
2. Network: the new tasks fetch those files from other machines. The network is slower than memory.
3. Read and rebuild: each new task reads the bytes and turns them back into rows.
4. A stage boundary: stage 1 cannot start until every task of stage 0 has finished writing.

Spark keeps the files on disk for a reason. If a stage 1 task fails, it can fetch the files again. The stage 0 tasks do not need to run again.

Sometimes there is a fifth cost.

If a task gets more data than fits in memory, Spark writes the extra to disk while it works. This is called a spill. A spill adds more disk work.

The setting spark.sql.shuffle.partitions says how many partitions a shuffle makes. The default is 200. In our model we used 3, so we could draw it.

Without help, 200 partitions would give our 5 partial rows 200 tiny tasks. Most would get no rows at all. That is a waste.

Spark has a feature for this, called AQE (Adaptive Query Execution). It is on by default since Spark 3.2. After the shuffle, AQE can see how small the partitions are, and it can merge small ones into fewer tasks.

You can still set the number yourself. Look at one example.

spark.conf.set("spark.sql.shuffle.partitions", "3")

orders.groupBy("customer_id").sum("amount").explain()
# Exchange hashpartitioning(customer_id, 3)

The chapter "Partitioning & Parallelism" explains how to choose the number. The chapter "Data Skew & Salting" explains what happens when one customer has most of the rows.

You often cannot remove a shuffle. You can often make it smaller. A smaller shuffle writes less to disk and sends less over the network.

The first way is to filter early. Compare the job from the start of this chapter with the same job on paid orders only.

paid = orders.filter(F.col("status") == "paid")
paid.groupBy("customer_id").sum("amount")

Spark often moves a filter down by itself. Writing the filter early still makes your plan easy to read. Do not expect it when the filter uses the result of a UDF (a Python function you give to Spark).

Here are the other ways. Each one removes rows or removes a whole shuffle.

• Select only the columns you need. Rows with fewer columns are cheaper to write and send.
• Use built-in aggregates like sum and count. Spark adds up inside each task first, as in step 1. A Python function that gets all rows of a group cannot do this. Then every row crosses the shuffle.
• Join with a small table by broadcast. Spark sends the small table to every machine, so the big table does not move. The chapter "Joins in Spark" teaches this.
• Skip steps you do not need. Do not use repartition or distinct without a reason. Do one orderBy at the end, not in the middle.

The Spark UI shows how much a shuffle moved. Open a job, then click one stage to see its numbers.

A stage that ends in an Exchange shows a Shuffle Write number. A stage that starts after it shows a Shuffle Read number.

Look at what each number means for our model job.

The UI says "records". A record is a row. The UI also shows the size in bytes.

A real run has larger numbers that will not match this table.

Three checks help you find the problem.

• Shuffle Write is very large compared with your input. Filter or select earlier.
• There is spill. The partitions are too big for memory.
• A few tasks read much more than the rest. This is skew, and the next chapter covers it.

Four points to keep.

• A shuffle moves rows so rows with the same key meet. Each task writes files to local disk, and the next stage reads them over the network.
• groupBy, join, distinct, orderBy, repartition and window partitionBy cause shuffles. An Exchange line in explain() marks each one.
• A shuffle costs disk, network, serialization and a stage boundary. Too much data in one task makes it spill.
• Make a shuffle smaller: filter early, select fewer columns, broadcast small tables, and skip needless steps.

Answer these without looking back.

1. In our model, which two partial rows for customer 3 meet in the shuffle, and what is the final total?
2. Why is the shuffle write kept on disk and not only in memory?
3. Name three steps that cause a shuffle and one that does not.
4. The default for spark.sql.shuffle.partitions is 200. Why can that be a poor choice for a tiny job?