Data Engineering · Batch Processing
Adaptive Partitioning and Runtime Repartitioning
How a shuffle writer reports per-partition sizes at the stage barrier, and how coalescing and splitting use that report to correct a partition count chosen before the data was read.
The partition count is fixed before the first record is read, because it is part of the function that assigns records to partitions. That makes it a guess, informed by table statistics that may be stale, missing, or simply wrong about the day’s data. Adaptive partitioning is the family of techniques that correct the guess using measurements taken from the shuffle itself, at the one moment where a correction is cheap: the stage barrier.
Why a static count is wrong at both ends
A fixed default partition count is unrelated to input size, so it is too large for small jobs and too small for large ones. Splitting one gigabyte across 200 partitions yields about five megabytes per task, and a task reading five megabytes spends far longer starting up, warming its runtime, and reporting than it does processing — with fixed per-task costs on the order of hundreds of milliseconds, the fixed part dominates. Splitting ten terabytes across the same 200 partitions yields fifty gigabytes per task, which exceeds any per-task memory budget, forces repeated disk passes in the reduce-side merge, and leaves the cluster’s slots mostly idle.
Empty and near-empty partitions waste work at the other extreme. If the stage’s data has fewer distinct keys than the partition count, some partitions receive nothing, yet their tasks are still scheduled and started; a job with 1,000 reduce tasks over 50 distinct keys spends 950 slots producing empty output.
Planning-time statistics would fix this if they existed and were current. They frequently do not: a table may have no statistics because it was never analyzed, statistics may predate the previous day’s load, and a query may produce its input from an upstream stage whose size is not knowable until that stage runs. The plan is a guess about data volume, and the shuffle stage is where the truth shows up.
What the shuffle writer knows and when it knows it
The measurements needed for a better decision are produced as a side effect of the shuffle. When a map task finishes, its output file is laid out as one segment per partition, and the writer naturally knows each segment’s length and record count. Reporting those counts costs one small record per map task; the coordinator aggregates them across the map stage into a per-partition size distribution.
The crucial property is timing. A reduce stage cannot begin until every map task has finished, because a reduce task cannot produce a correct result without all of a key’s records — the barrier described in the shuffle data path. The full per-partition size distribution is therefore available at exactly the point where reduce tasks would have been created anyway, so acting on it adds no extra pass over the data and no extra barrier. Adaptive decisions made after a shuffle are almost free; decisions that would require an extra materialization are not.
A measured distribution makes the failure visible in one line: 1,000 map tasks feeding 200 partitions, of which 190 partitions hold about 2 MB each and 10 hold about 5 GB each. The defect is not the partition count — that is a property of the keys — but the assignment of work to tasks, and it can be corrected in two directions.
Coalescing small partitions
Coalescing combines several adjacent partitions into one reduce task. The task then processes a group of partitions, merging each as it goes, exactly as a normal reduce task would do for one partition and then for the next. Total bytes read and merged are unchanged; the saving is the elimination of per-task fixed costs that would otherwise be paid many times for tiny payloads.
The arithmetic makes the case. A task with a startup cost in the low hundreds of milliseconds processing about two megabytes at a typical disk read rate finishes its read in a few tens of milliseconds, so overhead is several times the useful work. Coalescing 190 such partitions into, say, ten tasks of about 38 MB each amortizes the same startup cost over nineteen times the payload. The limit is the opposite direction: coalescing until a task’s input exceeds its merge memory reintroduces disk passes, so the grouping is driven by a target minimum partition size rather than by a fixed grouping factor.
Coalescing has a second effect that is easy to miss. The stage’s output is typically one file per partition, so reducing the partition count reduces the number of files the next stage reads as its input splits, which relieves the small-file pressure described in input splits and task granularity. A pipeline whose stages each write 50,000 tiny files will accumulate them until the scan cost exceeds the processing cost.
Splitting an oversized partition
Splitting does not change the partitioning function or the assignment of keys. A partition’s data is spread across one segment per map task, so the partition can be handed to several tasks by giving each a disjoint subset of those segments. Every byte is still read exactly once by exactly one task, no network traffic is added, and no key moves between partitions.
What makes this legitimate is the reduce function’s algebraic structure. A sum, a count, a maximum, or a union of sets computed over disjoint subsets of a key’s records can be combined afterward into the correct total, which is the same associativity precondition that lets a combiner run map-side. Splitting is therefore safe for decomposable aggregations and unsound for operations that need the whole group — an exact median, a top-$k$ that is not re-mergeable, or the construction of a single ordered result per key.
Splitting also does not create more partitions. The stage still has the same number of output partitions and, if the write path reconciles the sub-task results into the partition’s file, the same file count; several partial outputs per partition must be combined, which costs a merge pass. And splitting spreads work without fixing skew inside a key: if one key accounts for most of the partition’s bytes, the sub-task holding that key’s segments still does a disproportionate share, because no partitioning of segments can divide a single key group. Skew of that shape needs the salting approach described in hot-key skew mitigation.
Re-planning downstream of the same measurement
Per-partition sizes are one input to a larger re-planning decision. Once a stage’s real output size is known, the choice of algorithm for the next operator can be revisited: a sort-merge join whose build side turned out to be a few megabytes can be replaced with a broadcast, removing the second shuffle entirely — the same threshold logic that motivates co-partitioned layouts when the size is known in advance and is stable. Conversely, a plan that assumed a small broadcast side can be corrected to a shuffled join when the measurement shows otherwise, which is a correctness-adjacent repair: the broadcast threshold being wrong is a memory failure, not merely a slow plan.
Because the plan is re-decided at each stage barrier, the effects propagate: a stage’s partition count determines the next stage’s input partition count, so a coalesced stage hands a smaller partition count downstream and the following shuffle’s join decisions are made against that. Chains of shuffled stages therefore improve monotonically in estimate quality while accumulating variability. The same job on the same data can end up with a different number of tasks, a different plan shape, and a different number of output files on different runs, which matters for anything downstream that depends on layout rather than content.
Where adaptivity does not pay is worth stating plainly. Short jobs spend a larger share of their runtime in re-planning and in the extra bookkeeping of publishing and reading shuffle statistics. I/O-bound scans gain nothing from a better task distribution when the bottleneck is the read path. And workloads with stable, well-known data sizes are usually served better by a fixed plan derived from maintained statistics, whose partition count can be reviewed and tested rather than varying per run.
References
- Performance Tuning — Apache Spark documentation
- Spark SQL: Relational Data Processing in Spark — Armbrust, Xin, Lian, Huai, Liu, Bradley, Meng, Kaftan, Franklin, Ghodsi and Zaharia, SIGMOD 2015
- Making Sense of Performance in Data Analytics Frameworks — Ousterhout, Rasti, Ratnasamy, Zeldovich and Shenker, NSDI 2015