Conversation
Comet implements Spark's round-robin shuffle as hash partitioning over every
column of every row. On a wide nested schema that dominates the shuffle write:
`create_murmur3_hashes` recurses into every struct child per row, and the
resulting row-level scatter forces `interleave_record_batch` to walk every
column and child again on flush. It is also not round robin. Placement is a
pure function of a row's contents, so a column of one repeated value lands
entirely on one reducer where Spark's round robin spreads it.
Adds `RoundRobinStrategy::RowGroups`, which places rows the way Spark does: the
row at task-global ordinal i goes to
`(mapPartitionId + i / groupRows) % numPartitions`. The counter is over rows,
not batches, and carries across batch boundaries, so placement does not depend
on how the reader frames its input. That matters because no Spark contract
covers framing: `DETERMINATE` promises the same rows in the same order and says
nothing about chunking, so an operator that spills can reframe under different
memory pressure while still honouring it. Keying on a row ordinal reduces the
residual assumption to exactly the one Spark's own round robin makes.
Positional placement is used only where that assumption is established, in two
independent places that both have to hold:
* `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused
into the writer, which the RDD graph cannot see because the subtree
collapses into one `CometNativeShuffleInputRDD`. Short allowlist rather than
a denylist: a native scan under nothing but projections and filters.
* `CometNativeShuffleInputRDD.getOutputDeterministicLevel` applies Spark's own
`isOrderSensitive` rule to everything below that RDD, reporting
INDETERMINATE over a non-determinate parent so the DAGScheduler rolls the
stage back instead of re-running one task into consumed output.
Anything else keeps `HashAll`, which is safe to re-execute whatever its input
does. Both are behind
`spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled`,
default false.
The index representation gains a second shape: positional placement records
`(batch, start, len)` runs instead of one `(batch, row)` pair per row, which is
smaller against the spill reservation and lets the flush copy whole ranges.
`RunIterator` builds a chunk by slicing and concatenating runs, and passes a run
covering an entire buffered batch straight through without copying. A schema
containing `Utf8View` or `BinaryView` falls back to `HashAll`, because those are
the one family whose buffers the IPC writer does not truncate for a slice.
github-actions
Bot
added
enhancement
Capping `groupRows` at `batch_size` silently ignored what the user asked for. A longer group is meaningful and works as written: it sends several consecutive input batches to the same output partition, which is a legitimate way to trade balance for fewer, larger shuffle blocks.
|
Ran the partitioning benchmarks on this branch so the two approaches can be compared on numbers, Method
This is the Results
Speedup over
The #5449 column is from a separate binary, so I normalised everything to the Why this is faster than #5449, not just safer
So the run representation looks worth having independently of which placement key wins. On
|
Resolve the positional group size into the partitioning itself instead of a parallel field, give each round-robin strategy its own match arm, and buffer runs without moving the scratch vector in and out. Skip the unused partition_starts scratch under positional placement, pre-size the run iterator's chunk scratch, and compute the driver-side positional decision once per exchange so the RDD and the writer read the same value.
A ShuffleWriterExec publishes its partition offsets through a OnceLock, so re-executing one exec across criterion iterations fails on the second run with "partition offsets were already published". Build a fresh exec per iteration with iter_batched, keeping construction out of the timing.
|
Thanks for running these, and for spotting the bench breakage — that one had been sitting on I had put the fix up as #6104, but on reflection it makes no sense as a separate PR when this is the The finding I had not expected is that most of the speedup is the run representation rather than the
edit: I'll add the benchmark here |
The end-to-end shuffle benches are dominated by IPC encoding and the file write, so they cannot separate two placement strategies. Add a group that stops before both: one arm times placement and per-partition index buffering alone, the other adds the flush through a writer that discards its batches, which is where a run-indexed partitioner diverges from a row-indexed one. benches/ compiles as its own crate, so reaching the partitioners needs a pub seam. Rather than export MultiPartitionShuffleRepartitioner and the PartitionWriter trait, add one opaque handle in a doc(hidden) module and leave the rest crate-private. The nested fixture grows a per-row fill. The existing one repeats a single value down every leaf, so every row hashes alike and a hash strategy would place the whole input on one output partition, never performing the scatter that the gather exists to undo. The encoding benches keep the constant fill and their current numbers.
|
The partitioning bench is in tree now, and running it past the flush changes the conclusion enough It lives in M3 Max, intervals under 1% almost throughout:
Your suspicion that the flush is where more of the win is turns out to be half right, and the What I did not expect is Same caveat as yours on the host, though: this is a laptop, and the nested |
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
@andygrove, keying on a row ordinal instead of a batch ordinal is the right call, and positional_placement_survives_reframing pins the property that makes it right. Feeding the same 1000 rows through five framings and asserting identical output is the test I would have asked for. The two-gate design (plan allowlist plus the isOrderSensitive rule on the RDD) also reads cleanly, and the new suite is registered in both PR workflows.
On your question about defaulting groupRows to the batch size (#6095 (comment)), I'd keep the small default. The benchmark runs a single map task, so it can't show the cost. The per-task bound doesn't compose across map tasks: each task walks ceil(rows / groupRows) consecutive partitions from its start, so when a task emits fewer groups than there are output partitions, the stage-wide spread grows with groupRows times the number of map tasks. I simulated the placement formula with the start at the map partition id, as this PR does:
| map tasks | rows per task | partitions | groupRows |
min rows per reducer | max rows per reducer | empty reducers |
|---|---|---|---|---|---|---|
| 50 | 1,000,000 | 200 | 8192 | 0 | 409,600 | 28 |
| 50 | 1,000,000 | 200 | 64 (auto) | 249,600 | 251,200 | 0 |
| 10 | 500,000 | 200 | 8192 | 0 | 81,920 | 129 |
| 10 | 500,000 | 200 | 64 (auto) | 24,960 | 25,600 | 0 |
The mean in the first row is 250,000, so a batch-sized group turns repartition(200) over 50 files of a million rows into a skewed stage with 28 idle reducers. That's the workload positional placement is supposed to fix. The auto default avoids it because a task wraps around the partitions once per batch. If a large group is worth offering for the zero-copy path, it can stay an explicit setting. The config doc should then say that balance depends on each task emitting many more groups than there are partitions.
.ai/skills/review-comet-shuffle-pr/SKILL.md still says Comet's round robin is hash-based on purpose, and that "a PR that implements true round robin to fix skew breaks that" (lines 91-97). Once this lands, a reviewer following that skill will flag the positional path as a bug. Could you update that checklist item here to describe the positional path and the conditions that gate it?
| // `ShuffleWriterExec::execute` cannot supply it. See | ||
| // `RoundRobinStrategy::RowGroups` for why it has to be this value. | ||
| RoundRobinStrategy::RowGroups { | ||
| start_partition: self.partition.max(0) as usize, |
There was a problem hiding this comment.
Using the map partition id as the start keeps retries reproducible, but consecutive tasks start on consecutive partitions. That's the correlation Spark fixed in SPARK-21782, which is why Spark scrambles the start with new XORShiftRandom(partitionId).nextInt(numPartitions) (Spark). Comet's JVM path does the same (CometShuffleExchangeExec.scala). With adjacent starts, every task's run of partitions overlaps its neighbours', and the partitions past numMapTasks + groupsPerTask get nothing. Ten map tasks of 5,000 rows into 200 partitions at the auto group of 64 leave 112 reducers empty with adjacent starts and none with scrambled starts.
What do you think about computing the start in CometNativeShuffleWriter.buildUnifiedPlan, which already runs per task with context in scope, and passing it in the proto? new XORShiftRandom(context.partitionId()).nextInt(numPartitions) is still a pure function of the map partition, so retries stay safe. The planner would also stop depending on self.partition, which removes the jni_api partition-0 caveat in the comment above. Spark increments before its first use, so a start of nextInt(numPartitions) + 1 would make groupRows = 1 place rows the way Spark does for the same row order. A native test that runs several map tasks with a small number of groups each and asserts the stage-wide spread would cover this.
There was a problem hiding this comment.
You're right, and I had taken "distinct is enough" on faith rather than simulating it. Done the way
you suggested: the start is computed in buildUnifiedPlan as
XORShiftRandom(context.partitionId()).nextInt(numPartitions) + 1 and passed in the proto, so the
planner no longer reads self.partition and the jni_api partition-0 caveat goes with it. I took
the + 1 as well — Spark increments before its first use, so with it groupRows = 1 places rows
exactly where Spark's round robin would for the same row order, which makes the "this is Spark's
round robin at a coarser granularity" claim literally true rather than nearly true.
Your ten-tasks-of-5,000-into-200-at-64 figure reproduces exactly: 112 empty with adjacent starts,
zero with scrambled. CometNativePositionalRoundRobinSuite now asserts both, running the placement
formula over the real positionalStartPartition, and keeping the adjacent case in the test so the
hazard stays visible rather than becoming folklore. On the native side the new test asserts that a
task's groups walk consecutive partitions from its start, which is the property that makes the
stage-wide spread a function of how the starts are chosen rather than of the data — the thing that
makes your argument load-bearing on this side of the boundary.
| /// `start_partition` must be the Spark map partition id. It has to be distinct across mappers, | ||
| /// or every task starts at partition 0 and a task emitting fewer groups than there are output | ||
| /// partitions leaves the tail empty stage-wide; and it has to be a pure function of the map | ||
| /// partition, or a re-executed task does not reproduce its own placement. Spark seeds | ||
| /// `XORShiftRandom(partitionId)` for the same two reasons. | ||
| /// |
There was a problem hiding this comment.
The doc comment says a distinct start per mapper prevents an empty tail, and that Spark seeds XORShiftRandom(partitionId) for the same reason. Distinct isn't enough. Adjacent starts still leave the tail empty whenever the number of map tasks plus the groups per task is less than the partition count (see the comment on planner.rs). Spark's reason for the random seed is to decorrelate the starts. Could this say that, whichever start function you end up with? The same text is in native_shuffle.md lines 351-356.
There was a problem hiding this comment.
Fixed, in the Rust doc and in native_shuffle.md. Both now say the starts have to be decorrelated
rather than distinct, and give your overlap argument as the reason instead of the empty-tail one.
| /// Smallest automatically chosen group. A multiple of 8 so that a run starts on a byte | ||
| /// boundary of a validity bitmap, which keeps the per-run copy a memcpy rather than a | ||
| /// bit-shift for every column. | ||
| const MIN_AUTO_GROUP_ROWS: usize = 64; |
There was a problem hiding this comment.
The multiple-of-8 rationale only holds when every input batch starts on a group boundary. row_seq counts rows across batches, so after a CometFilterExec under the exchange, which the allowlist admits, a batch starts at an arbitrary ordinal and its first run ends at an arbitrary offset. Every run after it in that batch starts off a byte boundary. The same happens with any explicit groupRows that isn't a multiple of 8. Is the 64-row floor there to cap the run count instead? If so, could the comment say that?
There was a problem hiding this comment.
Capping the run count is the real reason, and the alignment claim was wrong for exactly the reason
you give — I had it in my head that a batch starts on a group boundary, which is only true with no
filter and a group that divides the batch size. Reworded to say the floor is there so that a
partition count far larger than the batch size cannot round batch_size / num_partitions down
towards a handful of rows and turn the flush back into the per-row gather the strategy exists to
avoid. I kept a sentence saying alignment is explicitly not guaranteed, since 64 looks like an
alignment number and someone will assume it back otherwise.
| "Rows per contiguous group under positional round robin. Imbalance between any two " + | ||
| "output partitions is bounded by this many rows however the reader frames its " + | ||
| "batches, so smaller groups balance better while larger groups produce fewer, longer " + |
There was a problem hiding this comment.
"Imbalance between any two output partitions is bounded by this many rows" is true for the rows one map task writes, but a reducer sees the sum over all map tasks. With 50 map tasks and groupRows = 8192, the simulation in the review body gives one reducer 409,600 rows and another 0. This is user-facing documentation, so could it say "within one map task", and say that stage-wide balance needs each task to emit many more groups than there are output partitions? The same sentence is in comet_partitioning.rs lines 56-58 and native_shuffle.md lines 358-359.
There was a problem hiding this comment.
Changed in all three. The config doc now says the bound holds within one map task, that a reducer
sees the sum over all of them, and that the stage is only evenly balanced when each task emits many
more groups than there are output partitions.
Starting each map task at its own partition id made consecutive tasks start on consecutive partitions. A task walks ceil(rows / groupRows) consecutive partitions from its start, so adjacent starts overlap and the partitions past numMapTasks + groupsPerTask get nothing: ten tasks of 5,000 rows into 200 partitions at the auto group of 64 leave 112 reducers empty. That is the correlation SPARK-21782 fixed. Compute the start the way Spark does, XORShiftRandom(partitionId) .nextInt(numPartitions) + 1, in CometNativeShuffleWriter where the TaskContext is in scope, and pass it in the proto. Still a pure function of the map partition, so a retried task reproduces its own placement, and the + 1 matches Spark's pre-increment so groupRows = 1 places rows exactly where Spark's round robin would. The planner no longer reads self.partition, which removes the jni_api partition-0 caveat.
Three wording fixes from review, none of them behavioural. The groupRows bound is per map task. A reducer sees the sum over all of them, which only evens out when each task emits many more groups than there are output partitions, so say that in the config doc rather than promising a stage-wide bound. MIN_AUTO_GROUP_ROWS is a cap on how finely a batch is cut, not an alignment guarantee. A run only starts on a byte boundary when the batch starts on a group boundary, and row_seq counts across batches, so after a filter every run in a batch is offset. The start has to be decorrelated across mappers, not merely distinct, which is what XORShiftRandom is for. Also update the round robin item in the shuffle review skill, which still said a positional round robin was a bug by construction.
|
Thanks — all four are addressed, in two commits: the scrambled start, and the wording corrections. I'm keeping the auto default, on your argument. The benchmark runs a single map task, so the 204x One correction to your table, though, because the scrambled start moves it. Every empty-reducer
So the case against a batch-sized group is no longer "28 idle reducers" but "1.57x on that shape,
|
3 participants
Which issue does this PR close?
Alternative to #5449, for #5397. Not intended to land alongside it — this is the same
optimization with a different placement key, put up so the two can be compared directly.
Rationale for this change
Comet implements round robin as hash partitioning over every column of every row. On the wide
nested schema that motivated #5449 that dominates the shuffle write:
create_murmur3_hashesrecurses into every struct child per row, and the row-level scatter then forces
interleave_record_batchto walk every column and child again on flush.It is also not round robin. Placement is a pure function of a row's contents, so a column of one
repeated value lands entirely on one reducer where Spark's round robin spreads it evenly. There is
a test in this PR that pins both halves of that.
#5449 replaces the hash with a batch counter. My concern with that, laid out in
#5449 (comment), is that batch framing
is the one property no Spark contract covers.
DeterministicLevel.DETERMINATEpromises the samerows in the same order and says nothing about how a downstream operator chunks them, so an operator
that spills can reframe under different memory pressure while still honouring it. That gap is what
failOnRetryin #5449 exists to paper over, andfailOnRetryand theINDETERMINATEdeclarationcancel each other out: every rollback the declaration asks for goes through a stage resubmission,
and
DAGScheduler.submitMissingTasksbumps the stage attempt on every submission, so everyrolled-back task hits the throw.
So this PR keys on a row ordinal instead. The residual assumption becomes exactly the one
Spark's own round robin makes — that the map task replays rows in the same order — which is the
assumption the determinism level already describes and the DAGScheduler already knows how to act
on.
failOnRetryis then unnecessary and there is no equivalent here: Spark's normal faulttolerance is preserved.
What changes are included in this PR?
RoundRobinStrategyreplaces the baremax_hash_columnsargument ofCometPartitioning::RoundRobin.HashAllis the existing behaviour and stays the default.RowGroups { start_partition, group_rows }places the row at task-global ordinaliat(start_partition + i / group_rows) % num_partitions. The counter runs over rows and carriesacross batch boundaries, so a group one input batch leaves part-way through is finished by the
next and placement is independent of framing.
start_partitionis the map partition id scrambledthe way Spark scrambles it,
XORShiftRandom(mapPartitionId).nextInt(numPartitions) + 1, computedper task in
CometNativeShuffleWriter.buildUnifiedPlanand passed in the proto. Distinct startsare not enough: a task walks
ceil(rows / group_rows)consecutive partitions from its start, soadjacent starts overlap and leave the tail of the partition space empty, which is the correlation
SPARK-21782 fixed. The
+ 1matches Spark's pre-increment, sogroup_rows = 1places rows exactlywhere Spark's round robin would.
group_rowsdefaults toclamp(batch_size / num_partitions, 64, batch_size), resolved on the driver.Two independent gates decide where it is used, and both must hold:
CometShuffleExchangeExec.replaysRowsInOrderwalks the native subtree fused into the writer.The RDD graph cannot see it, because the whole subtree collapses into one
CometNativeShuffleInputRDDwhose dependencies are its leaves. Deliberately a short allowlistrather than a denylist of known-bad operators: a native scan under nothing but projections and
filters. Operators that spill are the interesting exclusion. Anything else keeps
HashAll.CometNativeShuffleInputRDD.getOutputDeterministicLevelapplies Spark's ownisOrderSensitiverule to everything below that RDD, reporting
INDETERMINATEover a non-determinate parent. Thispart is the same as feat: RR partition with vectorized distribution #5449's, and it is the right mechanism.
On the flush side,
PartitionIndicesgains a second shape. Positional placement records(batch, start, len)runs rather than one(batch, row)pair per row, which is smaller againstthe spill reservation and lets the copy move whole ranges instead of gathering row by row.
RunIteratorbuilds each output chunk by slicing and concatenating runs, and hands a run coveringan entire buffered batch straight through with no copy. A repartitioner picks one shape for its
lifetime from the partitioning it was built with, so the scatter paths are byte-for-byte unchanged.
One schema-level restriction: positional placement is the only path that hands a sliced array to
the IPC writer. arrow-rs truncates a slice's buffers for every type reached here —
ArrayData::slicepushes the slice into struct children,
get_or_truncate_bufferhandles numeric and temporal,reencode_offsetsthe byte arrays,get_list_array_bufferslists and maps,bit_slicebooleans —except
Utf8View/BinaryView, where it truncates the views buffer but serializes every shared databuffer in full.
create_repartitionerfalls back toHashAllfor a schema containing one.One commit is unrelated to the placement key, folded in from #6104, now closed in favour of
this. Every end-to-end bench in
native/shuffle/benches/shuffle_writer.rspanics during criterionwarmup, the pre-existing hash and range ones on
mainincluded, withpartition offsets were already published. Each builds oneShuffleWriterExecoutsideb.iter()and re-executes it, but the writer publishes its offsets through a
OnceLockthat errors on asecond
set.pr_benchmark_check.ymlonly runscargo check --benches, so CI never saw it. Abench_end_to_endhelper now builds a fresh exec per iteration withb.iter_batched, withconstruction counted as setup and left out of the timing. It is here rather than on its own branch
because it blocks measuring this PR at all.
Another adds the partitioning benchmark @comphead had to carry locally.
benches/compilesas its own crate, so reaching the partitioners means exporting something; rather than make
MultiPartitionShuffleRepartitionerand thePartitionWritertraitpub, there is one opaquehandle in a
#[doc(hidden)] bench_supportmodule that places batches and then gathers them backout through a writer that discards, and everything it is built from stays crate-private. The nested
fixture grows a per-row fill, because the existing one repeats a single value down every leaf: every
row hashes alike, so a hash strategy puts the whole input on one output partition and never performs
the scatter that the gather exists to undo. The encoding benches keep the constant fill and their
current numbers.
How are these changes tested?
Native, in
comet_partitioning.rsandmulti_partition.rs:positional_placement_is_independent_of_batch_framingandpositional_placement_survives_reframingfeed the same 1000 rows through five different batchframings (one batch, 1000 single-row batches, ragged, and so on) and assert an identical
partitioning each time. This is the property the whole design rests on and the one the batch
counter cannot offer.
positional_placement_is_a_partition_of_the_inputasserts every row written exactly once, eachpartition in input order.
positional_placement_is_balanced_within_one_groupasserts the spread stays within one groupunder deliberately ragged framing.
positional_placement_walks_consecutive_partitions_from_its_startasserts a task of five groupstouches exactly five consecutive partitions from its start, which is what makes the stage-wide
spread a function of how the starts are chosen rather than of the data, and pins the overlap that
adjacent starts produce.
positional_placement_begins_at_the_start_partition,positional_placement_records_runs_not_rows(four runs rather than 256 row entries, and no row scratch allocated), plus run-iterator tests for
chunking, the zero-copy whole-batch path, and the slice-not-copy path.
view_types_are_detected_at_any_depthfor the fallback.JVM, in
CometNativePositionalRoundRobinSuite(new, registered in both PR workflows) andCometNativeShuffleInputRDDSuite:exchange does not, the config gates it, a hash repartition never takes it;
checkSparkAnswerover a repartition, including with a group larger than a batch so thezero-copy path runs;
HashAllcollapses them all onto one — the skew described above;copyForLocalShuffle;with the scrambled start, against 112 empty if the start were the bare map partition id, which
the test keeps asserting so the hazard stays visible;
million rows into 200 partitions comes out at 1.00x max/min at the auto group and 1.57x at a
batch-sized one, which is why the auto default stays.
For the bench fix,
cargo bench -p datafusion-comet-shuffle --bench shuffle_writer -- --quick "end to end"now runs all twelve end-to-end benches to completion. On the commit before it, thesame command panics on the first one.
cargo bench -p datafusion-comet-shuffle --bench shuffle_writer -- shuffle_partitioning, on anM3 Max. 8 batches of 8192 rows into 50 output partitions, so
AUTO_GROUP_ROWSresolves to 163.plainis the flat four-column schema the file already used;nestedis 40 struct columns over athree-field leaf, 120 leaf arrays, filled per row. Intervals were under 1% almost throughout.
HashAllHashAll{1}RowGroups(auto)RowGroups(8192)placeplace+gatherplaceplace+gatherPlacement alone is 108x cheaper on the flat schema and 125x on the nested one. That number is not
the one to quote, though, and adding the gather is what shows why: subtracting the rows gives
18.58 ms of gather for
HashAllagainst 8.13 ms forRowGroups(auto), only 2.3x, because bothstill copy all 65536 rows and the run shape only changes the chunk they copy in. So at the default
group size the honest end-to-end figure is 4.5x on nested, not 125x.
RowGroups(8192)is the case worth arguing about. A group as long as a batch makes every run covera whole buffered batch,
RunIteratorhands it through with no copy at all, and the gather drops to38 µs — 204x over
HashAllrather than 4.5x. TheAUTO_GROUP_ROWSdefault is therefore buyingbalance at roughly a 46x gather cost on this shape.
Resolved in favour of keeping the auto default, on @mbutrovich's argument in
#6095 (review): the benchmark
runs a single map task, so the 204x is the per-task gather cost with the composition problem
invisible. The per-task bound does not compose — each task walks
ceil(rows / groupRows)consecutive partitions from its start, so a task emitting fewer groups than there are output
partitions cannot cover the space even once, and the stage-wide spread grows with
groupRows.Simulated over 50 map tasks of a million rows into 200 partitions with the scrambled start, a
batch-sized group gives 1.57x max/min against 1.00x for the auto group, and 6.00x against 1.01x at
10 tasks of 500,000. A
repartition(200)that hands one reducer six times another's rows is notworth a gather that is already 4.5x better than
HashAll. The large group stays available as anexplicit setting, and the config doc now says what it gives up.
For comparison, @comphead measured
placeon both branches against a common baseline in#6095 (comment) — 57.9x flat and 62.2x
nested here against 11.3-22.2x and 47.2-48.8x for #5449, most of it from the run representation
rather than the placement key. The
placecolumn above agrees with theirs in shape; the nestedHashAllrow is slower here (18.25 ms against 12.72 ms) because the per-row fill gives murmur3varying strings to walk instead of a repeated
"x".