Partitions: repartition, coalesce and Skew

Right-sized work units

A partition is the unit of parallelism: one task per partition. Too few partitions leave cores idle and risk out-of-memory; too many create tiny tasks and scheduling overhead. A common target is partitions of roughly 100 to 200 MB and at least a few times the number of cores. repartition(n) does a full shuffle to exactly n balanced partitions (and can increase the count); coalesce(n) only merges existing partitions without a full shuffle and can only reduce. After a shuffle, the count comes from spark.sql.shuffle.partitions (default 200; AQE can coalesce it). Data skew means one key has far more rows than others, so one task runs long while the rest wait; fixes include AQE skew handling, salting the hot key, or broadcasting the small side.

Counting partitions, run

I ran this on Apache Spark 4.0.0 (PySpark, local mode, in the official Docker image). With local[2] the small DataFrame starts with 2 partitions. repartition(8) produces 8 and coalesce(1) merges down to 1.

print("default", orders.rdd.getNumPartitions(), "repartition(8)", orders.repartition(8).rdd.getNumPartitions(), "coalesce(1)", orders.repartition(8).coalesce(1).rdd.getNumPartitions())

Output:

default 2 repartition(8) 8 coalesce(1) 1

Salting a hot key, run

I ran this on Apache Spark 4.0.0 (PySpark, local mode, in the official Docker image). Key hot has 8 rows and cold has 1. Adding a random salt (0 to 2) splits hot into up to three sub-groups that are summed first, then combined: the totals stay correct (28 and 1) while the work is spread.

skew = spark.createDataFrame([("hot", i) for i in range(8)] + [("cold", 1)], ["k", "v"])
salted = skew.withColumn("salt", F.floor(F.rand(seed=1) * 3)).groupBy("k", "salt").agg(F.sum("v").alias("s")).groupBy("k").agg(F.sum("s").alias("total")).orderBy("k")
print([tuple(r) for r in salted.collect()])

Output:

[('cold', 1), ('hot', 28)]

Quick check: Which call reduces partitions without a full shuffle?

• coalesce(n)
• repartition(n)
• groupBy()
• distinct()