Skip to content

feat: RR partition with vectorized distribution - #5449

Open
comphead wants to merge 4 commits into
apache:mainfrom
comphead:shuffle_writes_vectorized
Open

comphead wants to merge 4 commits into
apache:mainfrom
comphead:shuffle_writes_vectorized

Conversation

@comphead

@comphead comphead commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Experiment for #5397

Motivation

Comet implements Spark's round-robin shuffle (df.repartition(n)) as hash partitioning over
every column of every row. create_murmur3_hashes recurses into every Struct child per row,
and the row-level scatter forces interleave_record_batch to walk every column and nested
child 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) replaces RoundRobin(usize, usize).
    Existing behaviour becomes RoundRobinStrategy::HashAll { max_hash_columns }, still default.
  • New RoundRobinStrategy::WholeBatch { start_partition } assigns each RecordBatch whole to
    one partition via round_robin_batch_seq. start_partition is the Spark map partition id,
    filled in by PhysicalPlanner::create_partitioning from the planner's partition, because
    ShuffleWriterExec::execute cannot supply it: jni_api runs every native root plan with
    partition 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_starts is
    built directly and partition_row_indices is left unmaterialized;
    buffer_partitioned_batch_may_spill takes Option<&[u32]> and treats None as an identity
    mapping.
  • PartitionedBatchIterator::next clones the source batch instead of interleaving when a chunk
    covers one whole source batch in order. Generic, so HashAll benefits opportunistically.
  • Opt-in via spark.comet.shuffle.native.partitioning.roundrobin.batchGranular (default false),
    plumbed through CometNativeShuffleWriter and the batch_granular protobuf field into
    PhysicalPlanner::create_partitioning.
  • Row scratch (partition_ids, partition_row_indices, ~64 KB/task) no longer allocated for
    SinglePartition or WholeBatch.

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 both
RoundRobinStrategy and the CometConf entry. HashAll semantics are unchanged.

Testing

test_round_robin_batch_granular_retry_deterministic asserts byte-identical data and index
files across two runs over a nested-struct schema, with no rows dropped.
test_round_robin_batch_granular_whole_batch_per_partition asserts every IPC block holds a
whole input batch. test_round_robin_batch_granular_starts_at_map_partition asserts that
mapper 3 starts at output partition 3 rather than 0, which is the assertion that fails without
the start_partition plumbing.

Benches in native/shuffle/benches/shuffle_writer.rs compare both strategies on two schema
shapes, end-to-end and in a partitioning-only microbench. For a real parquet dataset use the
existing shuffle_bench binary, which now accepts
--partitioning round-robin-whole-batch alongside round-robin.
MultiPartitionShuffleRepartitioner, ShufflePartitioner, PartitionWriter,
LocalPartitionWriter, and ShufflePartitionerMetrics are pub and 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 ~~~ AFTER
Loading

Remaining per-batch work: one modulo, two fill calls over partition_starts
(num_partitions + 1 slots, rescanned by buffer_partitioned_batch_may_spill), one
(batch, row) pair appended per row to the target partition's index list, which is charged
against the spill reservation and scanned again by PartitionedBatchIterator to recognise the
whole-batch case, and a RecordBatch clone (Arc bumps, not a data copy). So per-row work is
reduced, not eliminated.

The clone fast path only fires when a flush chunk lines up exactly with one buffered batch.
Input batches of exactly batch_size rows align and take it. A partition holding several
shorter batches does not, and falls back to interleave_record_batch.

Benchmark findings

cargo bench -p datafusion-comet-shuffle --bench shuffle_writer -- RoundRobin, 8 batches of
8192 rows into 50 output partitions, CompressionCodec::None. Two schema shapes, row and batch
counts held equal so they are comparable:

  • plain: the flat 4-column schema the rest of that bench file uses (int32, utf8, date32,
    decimal128), ~34 B/row
  • nested: 40 struct columns of depth 2, leaf 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:

schema HashAll WholeBatch speedup
plain 785 µs / 782 µs 70 µs / 35 µs 11x - 22x
nested 12.77 ms / 12.84 ms 270 µs / 263 µs 47x - 49x

HashAll is reproducible to well under 1% on both shapes. WholeBatch on the plain schema
varies (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_hashes matters ~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). Three
runs, ratio per run:

schema HashAll (ms) WholeBatch (ms) ratio
plain 3.21 / 3.40 / 3.17 1.87 / 2.10 / 2.68 1.72x / 1.62x / 1.18x
nested 55.8 / 44.0 / 38.6 33.0 / 38.5 / 37.4 1.69x / 1.14x / 1.03x

I do not trust these end-to-end numbers and they should be reproduced on a quiet host before
anyone quotes them.
The HashAll nested 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 iteration
on 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.rs panicked during criterion
warmup before this, including the pre-existing hash and range ones on main:

Execution("Error in shuffle write: Execution error: shuffle write error:
           partition offsets were already published")

One ShuffleWriterExec was built outside b.iter() and re-executed per iteration, but
PartitionOffsets is a OnceLock that errors on a second set. CI never caught it because
pr_benchmark_check.yml only runs cargo check --benches, never executes them. The round-robin
benches here now build a fresh exec per iteration via b.iter_batched. The pre-existing
benches for the other partitionings are still broken the same way and need a separate fix.

@comphead
comphead marked this pull request as draft August 24, 2026 15:35
@comphead

Copy link
Copy Markdown
Contributor Author

Spark

Image

Comet

Image

Comet with vectorized shuffle

Image

@comphead

Copy link
Copy Markdown
Contributor Author

@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
task places rows on different reducer partitions than the original attempt, and downstream
reducers have already consumed the original output, the job silently loses or duplicates
rows. This is SPARK-23207.

Spark's round-robin assignment is positional, not content-based
(CometShuffleExchangeExec.scala:969-975, copied from ShuffleExchangeExec):

var position = new XORShiftRandom(partitionId).nextInt(numPartitions)
(_: InternalRow) => { position += 1; position }   // then HashPartitioner does the mod

A counter seeded per map task, bumped per row. Placement depends purely on a row's index
in the input iterator. Nothing about the row itself is inspected.

sortBeforeRepartition=true therefore wraps the input in a local sort by binary UnsafeRow
(:999-1035, UnsafeExternalRowSorter with RecordBinaryComparator) before the counter
runs. The sort canonicalizes arrival order, which makes index a function of content, which
makes placement deterministic.

Why determinism is required

A map task can be retried after downstream reducers have already consumed the original
attempt's output (fetch failure, executor loss, speculation). If attempt 2 assigns rows to
different reducers than attempt 1, reducers that already fetched keep attempt-1 rows while
reducers fetching after the retry get attempt-2 rows. The union is not the input: rows are
lost and duplicated, silently, with no failure. That is SPARK-23207, and it is why the sort
is on by default despite costing a full external sort per map task.

How this PR provides determinism

RoundRobinStrategy::WholeBatch is structurally the same design as Spark's, one level
coarser: round_robin_batch_seq % n per batch instead of position % n per row. Both are
positional. Both are deterministic exactly when input order is.

The difference is where the order guarantee comes from:

Spark HashAll (Comet default) WholeBatch (this PR)
Placement key row index pmod(murmur3(row), n) batch index
Order-sensitive yes no yes
Determinism from forced local sort row content assumed order-preserving upstream
Cost external sort/task hash every column of every row one modulo per batch

So the PR provides determinism by assumption, not by construction. It takes Spark's
positional design and drops the sort that makes it safe, betting that Comet's Parquet scan
already emits identical batches in identical order on retry, which for a plain scan it does.
The bet is expressed as an opt-in flag defaulting to false, prose on RoundRobinStrategy
and the CometConf entry, and
test_round_robin_batch_granular_retry_deterministic asserting byte-identical output across
two runs.

What it does not have is Spark's enforcement. Spark does not trust its upstream, it sorts.
WholeBatch trusts its upstream and nothing checks that the trust is warranted:
CometNativeShuffleWriter.scala:332 sets the flag regardless of what the child is, and the
support check at CometShuffleExchangeExec.scala:515 tests only whether round-robin is
enabled. Note also that Comet's native shuffle path ignores sortBeforeRepartition
entirely, since the sort lives in prepareJVMShuffleDependency and not in
prepareNativeShuffleDependency (:764). For HashAll that is correct. For WholeBatch
it removes the one mechanism Spark relies on.

@sunchao

sunchao commented Aug 24, 2026

Copy link
Copy Markdown
Member

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, [a,b] [c,d] and [a] [b,c,d] can send rows to different partitions. Could we initially enable this only for cases where we’ve verified that retries produce the same batches in the same order, and use the existing hashing approach elsewhere?

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.

@comphead

Copy link
Copy Markdown
Contributor Author

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

@andygrove

Copy link
Copy Markdown
Member

One thing worth pulling in before we go further on enumerating safe upstreams: Spark already
has a declarative mechanism for exactly this, and our JVM path already uses it. In
prepareShuffleDependency we set isOrderSensitive = isRoundRobin && !SQLConf.get.sortBeforeRepartition
(CometShuffleExchangeExec.scala:1046) and thread it into both mapPartitionsWithIndexInternal
calls. That flag is what lets Spark run safely with sortBeforeRepartition=false at all —
MapPartitionsRDD.getOutputDeterministicLevel returns INDETERMINATE when the map function is
order-sensitive and the parent is UNORDERED, and RDD.getOutputDeterministicLevel marks every
reduce-side RDD UNORDERED because, in Spark's own words, "the arrival order of these shuffle
blocks are totally random." So a round-robin repartition downstream of another exchange declares
itself indeterminate, and the DAGScheduler rolls the whole stage back rather than re-running a
single task into a partially-consumed output.

prepareNativeShuffleDependency (:771) never sets that bit. That's correct today, because
HashAll places rows by content and is genuinely deterministic regardless of input order. With
WholeBatch it becomes load-bearing and we'd be the only round-robin path in either engine that
is positional and neither sorts nor declares itself order-sensitive.

The thin RDD isn't a MapPartitionsRDD, so we can't just pass the flag, but I don't think we need
to: CometNativeShuffleInputRDD declares OneToOneDependency on each leaf input RDD, so the
determinism level already propagates up from the real parents, and getOutputDeterministicLevel
is protected on RDD while the class lives in an org.apache.spark package. Could we override
it there to return INDETERMINATE when the parent level is UNORDERED and batch-granular
round-robin is enabled?

What I like about this over an upstream allowlist is that it doesn't require us to be right about
which operators preserve order and framing. A plain scan keeps a DETERMINATE parent, stays
determinate, and still gets cheap per-task retry. Anything downstream of an exchange goes
indeterminate on its own and we get either a correct rollback or a loud job abort instead of
silently dropping and duplicating rows. It also composes with the sortBeforeRepartition question
rather than replacing it, so we could keep the flag defaulting to false and still not be relying
on the default to stay safe.

@andygrove

Copy link
Copy Markdown
Member

Picking up the fixed-size row group idea from @sunchao's comment, because I think it interacts
with the flush path in a way that isn't obvious and that decides whether it's worth doing.

The hash isn't the only per-cell cost in the current round-robin path. Because HashAll scatters
adjacent rows across all N partitions, partition_indices ends up as a list of (batch, row)
pairs and the flush goes through interleave_record_batch
(partitioners/partitioned_batch_iterator.rs:111), a per-row gather that re-walks every column
and every nested child a second time. A good share of what WholeBatch buys us comes from the
new fast path skipping that, not from skipping the hash.

Which is why I'd be careful with the row-group variant as stated. That fast path fires only when
a flush chunk covers one whole source batch in natural order. With fixed-size groups of B rows a
group is a strict sub-range of a source batch, and an 8192-row flush chunk gets stitched from
~batch_size / B groups spanning several source batches, so the condition never holds and we
fall straight back to per-row interleave. We'd keep the hash savings and hand back the interleave
savings.

Would it work to change partition_indices from Vec<Vec<(u32, u32)>> to runs of
(batch_idx, start, len) and build the output with slice + concat_batches, keeping the
zero-copy return this PR adds for when a single run already fills the chunk? Then the group size
becomes a knob rather than a cliff: at B = num_rows it degenerates to what this PR does today,
at B = 1 it's Spark's per-row positional assignment, and in between the copy is a memcpy per run
per buffer instead of a gather per row. It also drops index memory from 8 bytes per row to 12
bytes per run, which is not nothing when we're buffering wide batches.

On sizing, I think the balance question then answers itself analytically rather than needing to be
measured. With a global row counter the imbalance between any two partitions is bounded by B rows
regardless of how the reader frames batches, so something like
B = clamp(batch_size / num_partitions, 64, batch_size) gives 64-row runs at the default shuffle
partition count and a worst case of 64 rows of skew. That seems more defensible than reasoning
about how many batches a task happens to produce.

Two things I'd want to check before committing to it. First, WholeBatch today never produces a
sliced batch — the fast path clones the source batch whole — so nothing here exercises arrow-rs
IPC serialization of a sliced nested array. If a slice writes the parent's full buffers rather
than just its window, shuffle write volume regresses and the win evaporates. Could we add a test
asserting the bytes written for a sliced batch are within noise of the same rows copied into a
fresh batch?

Second, there's an ownership property we'd be giving up. With whole-batch assignment each buffered
batch is pinned by exactly one partition, so flushing that partition can release it. With runs,
every partition holds a slice of every batch again and nothing can be released until spill or
finish. That's neutral against HashAll today, and count_new_buffers dedups by buffer address
so the accounting stays honest either way, but it is a real regression against WholeBatch and
worth being deliberate about rather than discovering later.

Separately, since HashAll { max_hash_columns } is being reshaped here anyway — is
roundrobin.maxHashColumns worth keeping? Hashing the first N columns quietly collapses to N
distinct partitions when the leading columns are low cardinality, so someone whose leading column
is a date or a partition key gets severe skew with nothing in the plan to indicate why.

@andygrove andygrove added enhancement New feature or request performance area:shuffle Shuffle (JVM and native) labels Sep 6, 2026
@comphead

comphead commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

we would need to fallback if spark.sql.execution.sortBeforeRepartition is off

@comphead

comphead commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

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

@comphead

Copy link
Copy Markdown
Contributor Author

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>
@comphead
comphead force-pushed the shuffle_writes_vectorized branch from 7dc2368 to 3aba03d Compare September 21, 2026 16:07
@comphead
comphead marked this pull request as ready for review September 21, 2026 21:06

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +204 to +207
b.iter(|| {
let task_ctx = ctx.task_ctx();
let stream = exec.execute(0, task_ctx).unwrap();
rt.block_on(collect(stream)).unwrap();

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@andygrove

Copy link
Copy Markdown
Member

Coming back to this after reading the two new defences together, and I think they cancel each
other out at the defaults, which is the thing to settle before the placement question.

failIfRetryingPositionalRoundRobin throws when attemptNumber > 0 or stageAttemptNumber > 0.
But every path the INDETERMINATE declaration exists to enable goes through a stage resubmission,
and DAGScheduler.submitMissingTasks calls stage.makeNewStageAttempt() on every submission, so
every rolled-back task carries stageAttemptNumber >= 1 and hits the throw before it does any
work. The rollback that CometNativeShuffleInputRDD.getOutputDeterministicLevel asks for can
never complete while failOnRetry is at its default, so the fallback sentence in the config doc
describes a path you can only reach by turning the flag off. The operational side is worse than
the design side: an executor lost mid-stage re-enqueues its map tasks with a fresh
attemptNumber, an executor lost after the stage completes produces a fetch failure and a
resubmit, both throw, both exhaust spark.task.maxFailures, and the job dies. Speculation adds a
spurious failure on every long-tail task. I don't think we can ask anyone to enable this on the
kind of cluster where the 47x is worth having.

The deeper reason the flag is there is that round_robin_batch_seq keys placement on batch
framing, and framing is the one property no Spark contract covers. DETERMINATE promises the same
rows in the same order and says nothing about how a downstream operator chunks them, so a spilling
operator between the scan and the writer can reframe under different memory pressure while still
satisfying the contract. failOnRetry is covering exactly that uncovered gap.

Which is the argument I'd now make for the row-counter variant from my earlier comment, and it's a
better one than the performance argument: partition = (seed + rows_seen / B) % n closes the gap,
so the residual assumption becomes exactly the one Spark's own round robin makes. That's the
assumption the determinism level already describes and the DAGScheduler already knows how to act
on, so rollback becomes load-bearing instead of decorative and failOnRetry can go away entirely.
Skew also stops depending on how the reader happened to frame things — a task emitting 8 batches
into 200 output partitions touches 8 reducers today, which the trade-offs section concedes.

One thing the override can't see, separately from all that. CometNativeShuffleInputRDD's
dependencies are ctx.inputs (CometShuffleExchangeExec.scala:126), the leaf RDDs of the inlined
native subtree. Everything between those leaves and the writer runs inside the same
CometExecIterator and never appears in the RDD graph, so a spilling native aggregate under a
Parquet scan still reports DETERMINATE. Spark is blind in the same way for MapPartitionsRDD,
which is why sortBeforeRepartition defaults to true — the determinism level is Spark's backstop
for people who turn the sort off, not its primary defence. So what we have restores parity with
sortBeforeRepartition=false, not with the default. Could we close that with a plan-level check on
the native subtree instead? With the row counter the predicate shrinks from "preserves order and
framing" to "preserves order", which is short enough to be checkable — scan, project, filter, not
aggregate or join or anything that spills — and much more defensible than the allowlist I was
skeptical of in August. And when the check fails I'd rather fall back to HashAll at plan time
than commit to positional placement and then fail the job from inside a task, which is the
opposite of how we handle every other unsupported case.

On the two things I said I wanted to check before committing to the run-based index, both turn out
to be answerable from arrow-rs rather than from a benchmark. ArrayData::slice propagates the
slice into struct children (arrow-data-59.3.0/src/data.rs:630) and the IPC writer truncates per
type — get_or_truncate_buffer for numeric and temporal, reencode_offsets for the byte arrays,
get_list_array_buffers for lists and maps, bit_slice for booleans — so a sliced nested struct
writes only its window and the write-amplification worry doesn't apply. The exception is
Utf8View/BinaryView, where the views buffer is truncated but every variadic data buffer is
written in full, which we already know about at rss_partition_writer.rs:744. So the test worth
writing is narrow and specific to view columns. The wrinkle is that arrow::compute::concat
short-circuits at one input with array.slice(0, array.len())
(arrow-select-59.3.0/src/concat.rs:506), so a chunk that is a single strict sub-range comes back
as a slice rather than a fresh batch and lands on that path. Worth deciding deliberately rather
than letting it fall out.

On sizing B, I'd now measure the default rather than derive it from the formula I suggested.
clamp(batch_size / num_partitions, 64, batch_size) gives 64 at the default partition count, which
is ~128 runs per 8192-row batch and, on a 194-column nested schema, tens of thousands of small
bulk copies per batch. Still far better than a per-row gather, but it also eliminates the zero-copy
path completely, where B = batch_size keeps it whenever framing is aligned and degenerates to
exactly what this PR does today. The two ends behave differently enough that I wouldn't guess.

Last thing, and it belongs in its own issue rather than here: HashAll isn't a good fallback
either. pmod(murmur3(whole row), n) sends identical rows to the same partition, so
repartition(200) over low-cardinality data collapses onto a handful of reducers where Spark's
round robin spreads it evenly, and maxHashColumns makes that strictly worse. That's an
independent reason to get positional placement right rather than treating the hash path as the
safe harbour.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle Shuffle (JVM and native) enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants