repartition vs coalesce

Concepts covered: Shuffle Optimization

Both repartition(n) and coalesce(n) change how many partitions a DataFrame has, and they get there by opposite routes. repartition is a wide transformation: the plan gains an Exchange, every row passes through shuffle files, and a new stage starts. coalesce is narrow: each output partition reads several existing partitions where they already sit, with no Exchange and no new stage.

repartition: a full shuffle into even pieces

Below, 9 orders arrive in 2 input partitions and repartition(4) spreads them over 4. Each task of the first stage writes its rows to shuffle files, then the 4 tasks of the second stage each fetch a share. With no column given, Spark deals rows out round-robin, so the 4 outputs end up close to equal in size.

orders.repartition(4).collect()

The physical plan shows this as Exchange RoundRobinPartitioning(4). The price is that of any shuffle: every row is serialized, written to local disk and fetched over the network by the next stage. In return, repartition can raise the count as well as lower it, and the pieces come out balanced however lopsided the inputs were.

coalesce: merging partitions in place

Here the same 9 orders start in 4 partitions and coalesce(2) brings them to 2. No row touches a shuffle file. Each of the 2 tasks reads 2 of the original partitions back to back, inside the same stage that read them from storage.

orders.coalesce(2).collect()

The plan shows a Coalesce 2 node with no Exchange between it and the scan. That has 2 consequences. coalesce only lowers the count: coalesce(10) on a 4-partition DataFrame leaves 4. And each merged partition is the sum of the inputs it absorbed, so lopsided inputs stay lopsided.

Repartitioning by a column

repartition("region") hashes on region, so every row with the same region lands in the same partition. Without a count it uses spark.sql.shuffle.partitions, and repartition(3, "region") sets the count directly. The plan shows Exchange hashpartitioning(region, 3). A groupBy("region") right after needs exactly that arrangement, so the planner adds no second Exchange. The run below has 1 shuffle, and the aggregation happens inside the tasks that received the rows.

orders.repartition(3, "region").groupBy("region").count().collect()

When coalesce shrinks the stage above it

Since coalesce adds no stage, it shrinks the stage it lands in, and that stage includes every narrow step above it. Take a read of 400 partitions, a filter that keeps 2% of rows, and coalesce(1) so the output is a single file. The whole stage now runs as 1 task. That task reads all 400 input partitions one after another and filters every row while the rest of the cluster sits idle.

logs = spark.read.parquet("/lake/raw_logs")
errors = logs.filter(F.col("level") == "ERROR")

errors.coalesce(1).write.parquet("/lake/errors_a")

errors.repartition(1).write.parquet("/lake/errors_b")

repartition(1) pays for a shuffle and still wins here. The read and filter keep all 400 tasks, and only the 2% of rows that survive cross the Exchange to the single writing task. The rule I use: coalesce when the work sharing its stage is light, repartition when that work is heavy, when the count has to go up, or when the partitions need evening out.

Do

• Use coalesce to cut the count when the work sharing its stage is light
• Use repartition(n) to raise the count or even out lopsided partitions
• Use repartition(n, col) when the next step groups or writes by that column
• Check explain() for an Exchange or a Coalesce node before trusting a count

Avoid

• Don't coalesce(1) right after a heavy read or filter over hundreds of partitions
• Don't expect coalesce(n) to raise the count above what the DataFrame has
• Don't repartition by one key right before a groupBy on another key

The read below produces 400 partitions. How many tasks run the filter?

events = spark.read.parquet("/lake/events")
br = events.filter(F.col("country") == "BR")
br.coalesce(10).write.parquet("/lake/br_events")

• All 400, because the filter comes before coalesce in the code
• Only 10, because coalesce shares the read's stage and sets its task count
• Only 10, after a shuffle moves the 400 read partitions into 10
• A single task, because any write after coalesce runs as 1 task

The light case looks like this: a small table, a filter, and an output that needs only a couple of partitions. The merged stage reads few rows, so running it on 2 tasks costs little, and coalesce skips the shuffle that repartition would add. Check both counts with df.rdd.getNumPartitions(), which reports how many partitions the DataFrame has. Then look for Coalesce 2 directly above the filter in explain(), with no Exchange between them.

About this lesson

This section is part of the Partitioning: Intermediate 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.