feat: RR partition with vectorized distribution

[comphead]

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 assigns each RecordBatch whole to one partition via
  round_robin_batch_seq, seeded with the input partition id so concurrent mappers do not all
  target partition 0 first. 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. Benches in native/shuffle/benches/shuffle_writer.rs compare both
strategies end-to-end on a 40-column nested schema, in a partitioning-only microbench, and
over a real parquet clickstream dataset. 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 resize calls, and a RecordBatch clone (Arc
bumps, not a data copy).

[comphead]

Spark

Comet

Comet with vectorized shuffle

[comphead]

@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]

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.