Conversation
|
@andygrove @sunchao need your brain on the below: Although the performance is great, we need to think if this PR addresses task retry correctly, namely: Spark requires round-robin shuffle to be deterministic under task retry. If a retried map Spark's round-robin assignment is positional, not content-based var position = new XORShiftRandom(partitionId).nextInt(numPartitions)
(_: InternalRow) => { position += 1; position } // then HashPartitioner does the modA counter seeded per map task, bumped per row. Placement depends purely on a row's index
Why determinism is requiredA map task can be retried after downstream reducers have already consumed the original How this PR provides determinism
The difference is where the order guarantee comes from:
So the PR provides determinism by assumption, not by construction. It takes Spark's What it does not have is Spark's enforcement. Spark does not trust its upstream, it sorts. |
|
I agree with the retry concern, and I think the performance improvement makes this worth pursuing. One extra detail is that the batch boundaries need to stay the same too. Even with identical row order, Another option is to count rows across incoming batches and send each fixed-size group to the next partition. That would make the assignment independent of how the reader divides rows into batches. We would still need repeatable row order. There’s also a separate issue with the starting partition: the counter receives zero for every task, rather than the Spark input partition ID. If each task produces one batch, everything goes to partition 0. Passing the actual input partition ID through should fix that. Before enabling this more broadly, maybe we should add a test that fails and retries a task after some output has already been consumed, then checks for missing or duplicated rows. I’d also check partition sizes and whole-job runtime, so we know the faster writes aren’t offset by uneven work downstream. |
|
Thanks @sunchao that looks really promising and pretty popular operation, I'm thinking how to apply defensive mechanisms from your suggestions and some other ideas as well |
|
One thing worth pulling in before we go further on enumerating safe upstreams: Spark already
The thin RDD isn't a What I like about this over an upstream allowlist is that it doesn't require us to be right about |
|
Picking up the fixed-size row group idea from @sunchao's comment, because I think it interacts The hash isn't the only per-cell cost in the current round-robin path. Because Which is why I'd be careful with the row-group variant as stated. That fast path fires only when Would it work to change On sizing, I think the balance question then answers itself analytically rather than needing to be Two things I'd want to check before committing to it. First, Second, there's an ownership property we'd be giving up. With whole-batch assignment each buffered Separately, since |
|
we would need to fallback if |
|
Another thing to restrict issues is make this approach enabled by config, false by default and make the job fail if task/executor failed on RR partitioning, the Spark already support such mechanisms. So we prioritize correctness over stability |
|
So the main idea is to fallback if RR recomputation happens due to task retry or executor failure, preserving correctness over stability and investigate alt options like in https://github.com/vecbricks/varka |
`RoundRobinStrategy::WholeBatch` assigns each Arrow batch to an output partition from a per-task counter, so a row's destination is a function of the order and the framing of the batches the upstream operator produced rather than of the row itself. Re-executing a map task can write a different partitioning of the same rows, and once any consumer has fetched the output that attempt replaces, the reduce side silently gets some rows twice and others not at all. Spark answers this for its own round robin in two ways, neither of which the native path had. `spark.sql.execution.sortBeforeRepartition` makes placement content-determined, but the native path never sees it because the sort lives in `prepareJVMShuffleDependency`. Failing that, a round-robin repartition is wrapped in a `MapPartitionsRDD` with `isOrderSensitive`, which reports INDETERMINATE when its parent is UNORDERED, so the DAGScheduler rolls the whole stage back instead of re-running one task and aborts the job when a result stage has already consumed output. The native path has no `MapPartitionsRDD` to carry that flag, so apply the rule directly in `CometNativeShuffleInputRDD.getOutputDeterministicLevel`. The dependency's RDD is this one, so `Stage.isIndeterminate` reads it. A determinate parent such as a plain scan stays determinate and keeps cheap per-task retry. Rollback is only as sound as the parent's determinism level describing what the upstream operator actually replays, so while the strategy is opt-in and unproven, `CometNativeShuffleWriter` also refuses to run when `TaskContext` reports a task attempt after the first or a re-submitted stage attempt. Both counters are needed: `TaskSetManager.executorLost` re-enqueues a dead executor's map tasks inside the current task set with a fresh task attempt, while a fetch failure against that executor's output resubmits the stage instead. Gated on the new `spark.comet.shuffle.native.partitioning.roundrobin.batchGranular.failOnRetry` (default true), which can be turned off to leave the indeterminate declaration as the only defence. Spark has no API for failing an application from inside a task, and the two exceptions it declines to retry are matched by class name, so this throws an ordinary exception. The condition only gets more true on each attempt and fires before any work, so the task set aborts after `spark.task.maxFailures` and takes the job with it. Note this is stricter than the data requires: a task re-run inside a still-running map stage, and a speculative attempt, are both safe and both refused. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
7dc2368 to
3aba03d
Compare
sunchao
left a comment
There was a problem hiding this comment.
Correctness
This PR targets the cost of Comet's existing round-robin implementation, which hashes row contents across the selected columns and gathers rows again when writing partitions. The opt-in WholeBatch strategy sends each nonempty input batch slice to (mapPartitionId + batchSequence) % outputPartitions. The planner now supplies the Spark map partition ID; native execution's fixed partition 0 is not used as the seed. The counter survives spills, empty batches do not advance it, and oversized input batches are sliced at the configured native batch size. The shared clone fast path checks that a chunk selects every row of one source batch in natural order.
I compared this with the maintained Spark 3.5 and 4.0 round-robin and scheduler implementations. Spark assigns rows by position, with a mapper-seeded starting position, and uses local sorting or order-sensitivity metadata to handle replay. Comet's batch assignment deliberately gives different partition contents and only batch-level balance. Row values, nulls and nested values are preserved; the existing hash strategy remains the default. Single-output-partition and empty-schema routing keep their existing implementations.
The current revision addresses the earlier mapper-ID and unguarded retry concerns. Both native-child and columnar-input dependency construction carry the positional determinism declaration, including the local-shuffle sibling. With failOnRetry=true (the default), the writer rejects any nonzero task or stage attempt before creating shuffle output. These counters cover different scheduler paths: requeued tasks within a task set and resubmitted stages. This deliberately sacrifices fault tolerance when the experimental mode is enabled; it does not automatically fall back to hashing. Turning the guard off still relies on the upstream determinism declaration and does not establish stable batch framing for every upstream plan. I found no additional verified P1/P2 runtime correctness defect within this scope.
Validation and limits
Reviewed head 6f595108c7db6b702d9208537ac7906aa739efd1 against base 5ca149928f7743bfe7a96feadea5e0f9bed1412f. The shuffle CI job passed 506 Scala tests, including the new retry suite and input-RDD declarations; the Rust job passed 1,597 tests with five skipped, including all three WholeBatch tests. Those jobs actually checked out merge c9895e6f932e7f67b6909db2ef26e2be7f55f867, with parents 09b44ad6fa17f58f4bbccf958c5cf02d790f7334 and the reviewed head. All 26 authored changed files are byte-identical between that merge and the head; the entire trees and assigned base are different.
Local validation was source inspection and 24 source-contract assertions, with no local Spark/JNI execution, full suite or benchmark run. The maintained Spark 3.4 and 4.1 source branches were unavailable; Spark 4.1 CI is execution evidence, not a substitute for that missing source comparison. The new tests cover guard decisions, declarations and replay of an identical supplied batch sequence. They do not demonstrate recovery from a fetch failure after a reducer has consumed output.
Performance
[P2] Build a fresh writer for each RoundRobin benchmark iteration
The new end-to-end benchmark creates one ShuffleWriterExec outside b.iter. Every iteration shares its one-shot partition-offset slot, so the second successful write returns partition offsets were already published and the benchmark panics at collect(stream).unwrap(). The current PR description claims this benchmark uses a fresh exec through iter_batched, but that fix is absent from the reviewed code. Use a per-iteration setup closure and verify the focused RoundRobin benchmark completes warmup and sampling. This finding concerns the newly added RoundRobin block; equivalent pre-existing benchmark failures can be fixed separately.
The runtime optimization removes row-content hashing and can avoid column gathering for an exactly aligned whole-batch chunk. It still fills and scans the partition-start array, appends and accounts for one index pair per row, scans the candidate chunk and clones Arrow references. Multiple short batches in one chunk still use interleave. Pinned-buffer accounting and spill release remain in place; this is not a claim that native resident memory disappears.
The author reports nested partitioning-only times of 12.77–12.84 ms for HashAll and 263–270 µs for WholeBatch. Those are author measurements, not independently reproduced here. The author also explicitly qualifies the noisy disk-writing measurements. The checked-in new fixture covers the nested shape; the body also reports a plain-schema comparison whose corresponding new RoundRobin fixture is absent at this head. Preserve the measurement revision and runnable harness before treating those figures as evidence for this revision. Batch balance and downstream query time remain workload-dependent, particularly with few or uneven batches; the existing stage screenshots do not establish whole-query performance.
Design
The strategy enum keeps the existing content-hash behavior and experimental positional behavior explicit. Carrying the real Spark partition ID through the planner fixes the right boundary, and placing the retry guard before output creation makes its safety/availability tradeoff easy to inspect. The RDD declaration also reaches both input paths and survives the existing local-shuffle copy.
Whole-batch placement is a reasonable bounded experiment for avoiding nested-column hashing and copying. Stable logical row groups would reduce dependence on incoming batch framing, but would require a different representation and would not automatically retain the same clone opportunity. That alternative was already discussed on this PR. A focused failure-after-consumption integration test would strengthen any future claim that disabling the retry guard is safe; the current unit tests should not be described as proving that recovery behavior.
Abstraction & complexity
RoundRobinStrategy expresses a real choice in placement and retry behavior. The optional row-index slice lets the existing spill path represent the identity mapping without a second writer implementation, and the generic whole-source-batch check safely allows other partitioning modes to use the clone opportunity. The benchmark-facing exports widen visibility but add no separate runtime execution path. I found no additional actionable abstraction issue; the concrete change needed is to honor the existing writer lifecycle in the new benchmark.
| b.iter(|| { | ||
| let task_ctx = ctx.task_ctx(); | ||
| let stream = exec.execute(0, task_ctx).unwrap(); | ||
| rt.block_on(collect(stream)).unwrap(); |
There was a problem hiding this comment.
Performance
[P2] Build a fresh writer for each RoundRobin benchmark iteration
The new end-to-end RoundRobin benchmark constructs exec before b.iter, so both strategies reuse one Arc<PartitionOffsets> across iterations. ShuffleWriterExec::execute clones that destination, and LocalPartitionWriter::finish_all unconditionally publishes into its OnceLock. After the first successful iteration, the next collect(stream).unwrap() therefore panics with partition offsets were already published. This prevents the documented cargo bench -p datafusion-comet-shuffle --bench shuffle_writer -- RoundRobin command from producing the new comparison, despite the PR description saying this case now uses iter_batched. Move creation of the RoundRobin exec into a per-iteration setup closure, then run the focused benchmark through warmup and sampling. The finding is limited to this new benchmark; the older benchmarks' equivalent problem can remain a separate fix.
|
Coming back to this after reading the two new defences together, and I think they cancel each
The deeper reason the flag is there is that Which is the argument I'd now make for the row-counter variant from my earlier comment, and it's a One thing the override can't see, separately from all that. On the two things I said I wanted to check before committing to the run-based index, both turn out On sizing Last thing, and it belongs in its own issue rather than here: |


Which issue does this PR close?
Experiment for #5397
Motivation
Comet implements Spark's round-robin shuffle (
df.repartition(n)) as hash partitioning overevery column of every row.
create_murmur3_hashesrecurses into everyStructchild per row,and the row-level scatter forces
interleave_record_batchto walk every column and nestedchild again on flush. On a wide nested schema (the motivating workload had ~194 nested
columns) this dominates the shuffle write.
Round-robin has no content-based placement requirement. The only reason to hash row content
is determinism under retry, and when the upstream operator emits the same batches in the same
order under retry (Comet's Parquet scan and other order-preserving operators), batch identity
is already a deterministic key.
Changes
CometPartitioning::RoundRobin(usize, RoundRobinStrategy)replacesRoundRobin(usize, usize).Existing behaviour becomes
RoundRobinStrategy::HashAll { max_hash_columns }, still default.RoundRobinStrategy::WholeBatch { start_partition }assigns eachRecordBatchwhole toone partition via
round_robin_batch_seq.start_partitionis the Spark map partition id,filled in by
PhysicalPlanner::create_partitioningfrom the planner's partition, becauseShuffleWriterExec::executecannot supply it:jni_apiruns every native root plan withpartition 0, one Comet execution per Spark task. It has to be the map partition id on two
counts. Distinct across mappers, or every task starts at partition 0 and a task emitting
fewer batches than there are output partitions leaves the tail empty stage-wide. And a pure
function of the partition, so a re-executed task reproduces its placement. Spark's own round
robin seeds
XORShiftRandom(partitionId)for the same two reasons.partition_startsisbuilt directly and
partition_row_indicesis left unmaterialized;buffer_partitioned_batch_may_spilltakesOption<&[u32]>and treatsNoneas an identitymapping.
PartitionedBatchIterator::nextclones the source batch instead of interleaving when a chunkcovers one whole source batch in order. Generic, so
HashAllbenefits opportunistically.spark.comet.shuffle.native.partitioning.roundrobin.batchGranular(defaultfalse),plumbed through
CometNativeShuffleWriterand thebatch_granularprotobuf field intoPhysicalPlanner::create_partitioning.partition_ids,partition_row_indices, ~64 KB/task) no longer allocated forSinglePartitionorWholeBatch.Trade-offs
Distribution is even at batch granularity, not row granularity: fewer batches than partitions
leaves partitions empty, and unequal batch sizes give unequal partitions. Retry safety holds
only under order-preserving upstreams. Hence default
false, documented on bothRoundRobinStrategyand theCometConfentry.HashAllsemantics are unchanged.Testing
test_round_robin_batch_granular_retry_deterministicasserts byte-identical data and indexfiles across two runs over a nested-struct schema, with no rows dropped.
test_round_robin_batch_granular_whole_batch_per_partitionasserts every IPC block holds awhole input batch.
test_round_robin_batch_granular_starts_at_map_partitionasserts thatmapper 3 starts at output partition 3 rather than 0, which is the assertion that fails without
the
start_partitionplumbing.Benches in
native/shuffle/benches/shuffle_writer.rscompare both strategies on two schemashapes, end-to-end and in a partitioning-only microbench. For a real parquet dataset use the
existing
shuffle_benchbinary, which now accepts--partitioning round-robin-whole-batchalongsideround-robin.MultiPartitionShuffleRepartitioner,ShufflePartitioner,PartitionWriter,LocalPartitionWriter, andShufflePartitionerMetricsarepuband re-exported#[doc(hidden)]for the benches.Flow
flowchart TB subgraph BEFORE["Before: RoundRobin(n, max_hash_columns)"] direction TB A1["input RecordBatch"] --> A2["create_murmur3_hashes<br/>per row, recurses into every Struct child"] A2 --> A3["partition_ids[i] = pmod(hash, n)<br/>+ partition_row_indices"] A3 --> A4["rows scattered across all n partitions"] A4 --> A5["interleave_record_batch<br/>walks every column + nested child"] A5 --> A6["ShuffleBlockWriter"] end subgraph AFTER["After: RoundRobin(n, WholeBatch)"] direction TB B1["input RecordBatch"] --> B2["target = round_robin_batch_seq % n"] B2 --> B3["partition_starts direct<br/>partition_row_indices = None"] B3 --> B4{"chunk = whole source batch in order?"} B4 -- yes --> B5["clone source batch, no interleave"] B4 -- no --> B6["interleave_record_batch"] B5 --> B7["ShuffleBlockWriter"] B6 --> B7 end BEFORE ~~~ AFTERRemaining per-batch work: one modulo, two
fillcalls overpartition_starts(
num_partitions + 1slots, rescanned bybuffer_partitioned_batch_may_spill), one(batch, row)pair appended per row to the target partition's index list, which is chargedagainst the spill reservation and scanned again by
PartitionedBatchIteratorto recognise thewhole-batch case, and a
RecordBatchclone (Arcbumps, not a data copy). So per-row work isreduced, not eliminated.
The clone fast path only fires when a flush chunk lines up exactly with one buffered batch.
Input batches of exactly
batch_sizerows align and take it. A partition holding severalshorter batches does not, and falls back to
interleave_record_batch.Benchmark findings
cargo bench -p datafusion-comet-shuffle --bench shuffle_writer -- RoundRobin, 8 batches of8192 rows into 50 output partitions,
CompressionCodec::None. Two schema shapes, row and batchcounts held equal so they are comparable:
decimal128), ~34 B/row
struct<a: int64, b: utf8, c: float64>,~840 B/row. This is the shape the optimization targets.
Partitioning only (
insert_batch, so partition assignment plus index buffering, no flush,no IPC encode, no disk). Two runs, tight intervals:
HashAllWholeBatchHashAllis reproducible to well under 1% on both shapes.WholeBatchon the plain schemavaries (35-70 µs) because the absolute time is small enough to be dominated by allocation
bookkeeping, hence the wide range on that row. The nested figure is stable.
This is the honest measurement of what the patch changes, and the nested-vs-plain gap is the
interesting part: skipping
create_murmur3_hashesmatters ~4x more on the nested schema,which is exactly the recursion-into-every-struct-child cost the change is aimed at.
End-to-end (full
ShuffleWriterExec, including IPC encode and the write to disk). Threeruns, ratio per run:
HashAll(ms)WholeBatch(ms)I do not trust these end-to-end numbers and they should be reproduced on a quiet host before
anyone quotes them. The
HashAllnested figure drifts monotonically downward across runs(55.8, 44.0, 38.6) with 15-19% of samples flagged as outliers, which is warming rather than
noise. The cause is that this bench writes uncompressed output to
/tmp: ~55 MB per iterationon the nested schema, so ~5.5 GB per 100-sample bench function, and the machine I ran on was at
99% disk. The end-to-end path is I/O-bound here, so it measures the disk more than the patch.
The direction is consistently in
WholeBatch's favour, the magnitude is not established.Takeaway: the partitioning step gets an order of magnitude faster and that part is solid. How
much of that survives to end-to-end shuffle-write time depends on how much of the total is IPC
encode and I/O, which this microbenchmark cannot answer on a full disk. The cluster numbers in
the comment below are the better evidence for the end-to-end claim.
Note on the bench harness
Every end-to-end bench in
native/shuffle/benches/shuffle_writer.rspanicked during criterionwarmup before this, including the pre-existing hash and range ones on
main:One
ShuffleWriterExecwas built outsideb.iter()and re-executed per iteration, butPartitionOffsetsis aOnceLockthat errors on a secondset. CI never caught it becausepr_benchmark_check.ymlonly runscargo check --benches, never executes them. The round-robinbenches here now build a fresh exec per iteration via
b.iter_batched. The pre-existingbenches for the other partitionings are still broken the same way and need a separate fix.