Partitioning in Spark: repartition(), coalesce() and Adaptive Query Execution

The partition count decides whether a Spark job crawls, crashes or runs well. Too few partitions, or badly skewed ones, and executors run out of heap. Too many, and the job drowns in tiny tasks and writes thousands of tiny files. repartition() and coalesce() are the two tools for changing it, and they look interchangeable but aren't: one shuffles, the other can quietly collapse the parallelism of everything that runs before it.

The mechanics haven't changed, but how often you need them has. Since Spark 3.2, Adaptive Query Execution (AQE) is on by default and does most of this tuning at runtime, using real statistics instead of your guess. This article explains what each call really does, where coalesce goes wrong, what AQE does for you, and the few settings still worth touching.

What the two calls actually do

The PySpark docs for DataFrame.coalesce are refreshingly direct. It creates a narrow dependency: going from 1,000 partitions to 100 doesn't shuffle, each new partition just claims 10 of the old ones. Ask for more partitions than you have and nothing happens. You stay at the current count.

repartition is the opposite trade. It always shuffles, and the result is hash partitioned. You can pass a number, one or more columns, or both. If you pass only columns, the default number of partitions is used, which for DataFrames means spark.sql.shuffle.partitions.

df = spark.read.parquet("s3://bucket/events/")      # say, 2,000 input partitions

df.coalesce(100)                    # no shuffle; 100 partitions, uneven sizes possible
df.repartition(100)                 # full shuffle; 100 roughly even partitions
df.repartition("country")           # hash by country; count = spark.sql.shuffle.partitions
df.repartition(50, "country")       # hash by country into exactly 50

print(df.coalesce(100).rdd.getNumPartitions())     # 100
print(df.coalesce(5000).rdd.getNumPartitions())    # still 2000, coalesce never grows

The coalesce(1) trap

This is the part people miss, and the docs call it out explicitly. Because coalesce is narrow, Spark doesn't add a stage boundary. It pushes the reduced partition count up into the stage that produces the data. A drastic coalesce(1) can therefore make your whole upstream computation run as a single task on a single executor.

# Looks harmless. The join, the filter and the UDF now run in ONE task.
(orders.join(customers, "customer_id")
       .filter("amount > 100")
       .withColumn("score", score_udf("amount"))
       .coalesce(1)
       .write.mode("overwrite").csv("s3://bucket/report/"))

If you really need one output file, repartition(1) is usually the better choice. It adds a shuffle, but the expensive upstream work stays parallel and only the final write is single-threaded. Better still, ask whether you need one file at all. A downstream system that can read a directory saves you the trouble.

Adaptive Query Execution: tuning at runtime

Before AQE was switched on by default in Spark 3.2, the shuffle partition count was effectively fixed at plan time. You set spark.sql.shuffle.partitions (default 200) once, every join and aggregation used that number, and you tuned it per job by trial and error. Too few partitions and executors ran out of heap; too many and you drowned in tiny tasks and tiny files.

AQE changes this. The performance tuning guide says AQE uses runtime statistics to choose the most efficient plan, and that it has been enabled by default since Apache Spark 3.2.0. The umbrella switch is spark.sql.adaptive.enabled. Three features matter most here.

1. Coalescing shuffle partitions

With spark.sql.adaptive.coalescePartitions.enabled (default true), Spark looks at the actual shuffle output and merges contiguous small partitions. The target size is spark.sql.adaptive.advisoryPartitionSizeInBytes, default 64MB.

One detail almost nobody reads: spark.sql.adaptive.coalescePartitions.parallelismFirst defaults to true. In that mode Spark ignores the 64MB target and only respects coalescePartitions.minPartitionSize (default 1MB), to maximize parallelism and avoid regressions for people turning AQE on. The docs recommend setting it to false on a busy cluster so you don't end up with many small tasks. On shared clusters, I recommend changing this one first.

The role of spark.sql.shuffle.partitions has changed too. With AQE on, it's effectively the initial partition count before coalescing (coalescePartitions.initialPartitionNum falls back to it). So the modern move is to set it generously high and let AQE merge down, not to hunt for the perfect number.

2. Skew join handling

spark.sql.adaptive.skewJoin.enabled (default true) splits, and replicates if needed, skewed partitions in sort-merge joins. A partition counts as skewed when it's larger than skewedPartitionFactor (default 5.0) times the median partition size and larger than skewedPartitionThresholdInBytes (default 256MB). That used to mean salting keys by hand. Now it's usually automatic, and it removes a common cause of "Java heap space" errors in joins.

3. The REBALANCE hint

It's the modern replacement for using coalesce just to get fewer output files. According to the hints reference, REBALANCE tries to make every output partition a reasonable size, not too small and not too big. It's best-effort, splits skewed partitions, and is meant exactly for writing a query result to a table. It's ignored if AQE is not enabled.

from pyspark.sql import SparkSession

spark = (SparkSession.builder
         .appName("daily-events")
         # AQE is on by default since 3.2; being explicit documents intent
         .config("spark.sql.adaptive.enabled", "true")
         # on a shared cluster, honour the 64MB target instead of max parallelism
         .config("spark.sql.adaptive.coalescePartitions.parallelismFirst", "false")
         # start high, let AQE merge down
         .config("spark.sql.shuffle.partitions", "800")
         .getOrCreate())

events = spark.read.parquet("s3://bucket/raw/events/")

(events.filter("event_date = '2026-09-22'")
       .hint("rebalance", "country")      # same as /*+ REBALANCE(country) */
       .write.mode("overwrite")
       .partitionBy("country")
       .parquet("s3://bucket/curated/events/"))

The same thing in SQL:

INSERT OVERWRITE TABLE curated.events
SELECT /*+ REBALANCE(country) */ *
FROM raw.events
WHERE event_date = '2026-09-22';

When manual repartitioning is still right

AQE isn't magic. It only works at shuffle boundaries, and it can't read your mind. Explicit calls are still the right tool when:

• I need a specific layout for the write. repartition("country") before partitionBy("country") keeps each output folder from getting a file from every task.
• The input itself is badly partitioned. A handful of huge gzip files gives you a handful of tasks, and AQE can't fix a scan. A repartition(n) right after the read spreads the work.
• I'm handing data to something non-Spark, like a JDBC sink with a connection limit. coalesce(n) is fine there as long as the upstream work stays parallel (check the Spark UI).
• Iterative algorithms that reuse a DataFrame many times, where partitioning once and caching pays off.

One subtlety: an explicit number is treated as a promise. AQE won't coalesce a repartition(400) behind your back. If you want AQE's help, leave the number out or use REBALANCE.

Sidebar: Spark 4.0 and ANSI mode

If you're upgrading older jobs, also note this one. The migration guide says that since Spark 4.0, spark.sql.ansi.enabled is on by default. Casts that used to quietly return NULL, like CAST('abc' AS INT), and integer overflow that used to wrap around now raise errors. That's a good thing, but it breaks pipelines that relied on silent nulls. Use try_cast and the other try_* functions where bad input is expected. Turning ANSI mode back off should be a short-term bridge, not a fix.

from pyspark.sql import functions as F

clean = raw.withColumn("qty", F.expr("try_cast(qty_str AS INT)"))   # NULL instead of error

Checklist

• Confirm AQE is on (spark.sql.adaptive.enabled), especially on older or vendor-tuned configs.
• On busy shared clusters, set coalescePartitions.parallelismFirst=false.
• Treat spark.sql.shuffle.partitions as a starting point: set it high and let AQE merge.
• Never coalesce(1) after expensive work. Use repartition(1), or question the single file.
• For output file sizing, try REBALANCE before hand-tuning counts.
• Remove old "magic number" repartitions and compare in the Spark UI. Many turn out to do nothing.
• On Spark 4.x, audit casts and arithmetic for ANSI errors.

repartition() and coalesce() are still worth understanding, because they explain what the knobs do. With AQE, most of the time you shouldn't be the one turning them.