feat: RR partition with vectorized distribution

.github/workflows/pr_build_linux.yml

@@ -507,6 +507,7 @@ jobs:
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometNativeRoundRobinRetrySuite
               org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite

.github/workflows/pr_build_macos.yml

@@ -155,6 +155,7 @@ jobs:
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometNativeRoundRobinRetrySuite
               org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite

docs/source/contributor-guide/native_shuffle.md

@@ -49,7 +49,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition
    columnar output. Row-based Spark operators require JVM shuffle.
 
 3. **Supported partitioning type**: Native shuffle supports:
-
    - `HashPartitioning`
    - `RangePartitioning`
    - `SinglePartition`
@@ -152,7 +151,6 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet-
 1. **Plan construction**: `CometNativeShuffleWriter` builds a protobuf operator tree with a
    `ShuffleWriter` operator at the root and `childNativeOp` as its child. `childNativeOp` takes
    one of two shapes:
-
    - The child plan's `nativeOp` directly, when `CometShuffleExchangeExec`'s child is a
      `CometNativeExec` subtree. The upstream operators run inside the same `CometExecIterator`
      as the writer, with no JVM-to-native batch boundary between them.
@@ -165,15 +163,13 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet-
 2. **Native execution**: A single `CometExecIterator` per partition runs the unified plan.
 
 3. **Partitioning**: `ShuffleWriterExec` receives batches and routes to the appropriate partitioner:
-
    - `MultiPartitionShuffleRepartitioner`: For hash/range/round-robin partitioning
    - `SinglePartitionShufflePartitioner`: For single partition (simpler path)
 
 4. **Buffering and spilling**: The partitioner buffers rows per partition. When memory pressure
    exceeds the threshold, partitions spill to temporary files.
 
 5. **Encoding**: `ShuffleBlockWriter` encodes each partition's data as compressed Arrow IPC:
-
    - Writes compression type header
    - Writes field count header
    - Writes compressed IPC stream
@@ -203,7 +199,6 @@ read time. See [Direct Read](#direct-read-shufflescan) below for how the choice
 1. `CometBlockStoreShuffleReader` fetches shuffle blocks via `ShuffleBlockFetcherIterator`.
 
 2. For each block, `NativeBatchDecoderIterator`:
-
    - Reads the 8-byte compressed length header
    - Reads the 8-byte field count header
    - Reads the compressed IPC data
@@ -305,7 +300,11 @@ batch is written as a single block that may exceed the batch size.
 
 ### Round Robin Partitioning
 
-Comet implements round robin partitioning using hash-based assignment for determinism:
+`CometPartitioning::RoundRobin` carries a `RoundRobinStrategy` that decides how rows reach output
+partitions. The default is `HashAll`; `WholeBatch` is opt-in through
+`spark.comet.shuffle.native.partitioning.roundrobin.batchGranular`.
+
+#### `HashAll`: hash-based assignment (default)
 
 1. Computes a Murmur3 hash of columns (using seed 42)
 2. Assigns partitions directly using the hash: `partition_id = hash % num_partitions`
@@ -321,6 +320,71 @@ which Arrow's layout does not reproduce, unsorted output can land in different p
 Spark's. Sorted output is identical. That difference is why
 `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`.
 
+`spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` caps how many leading columns
+are hashed. `0`, the default, hashes all of them.
+
+#### `WholeBatch`: positional assignment
+
+Hashing every column of every row dominates the shuffle write on wide nested schemas, because
+`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.
+`WholeBatch` assigns each incoming `RecordBatch` whole to `counter % num_partitions`.
+`partition_row_indices` is left unmaterialized and the flush clones the source batch instead of
+interleaving it, which removes both the per-row hash and the per-row `interleave_record_batch`
+gather. Those two are what dominate the shuffle write on a wide nested schema.
+
+Two caveats on the cost, because the path is not free per batch. It still writes and rescans
+`partition_starts` across all `num_partitions` slots, and it still appends one `(batch, row)` pair
+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. The clone fast path
+also only fires when a flush chunk lines up exactly with one buffered batch, so a partition holding
+several batches shorter than `batch_size` falls back to interleaving.
+
+The counter starts at the Spark map partition id, carried on the strategy as
+`RoundRobinStrategy::WholeBatch { start_partition }` and filled in by
+`PhysicalPlanner::create_partitioning` from the planner's partition. It cannot come from
+`ShuffleWriterExec::execute`, because `jni_api` runs every native root plan with partition 0 (one
+Comet execution per Spark task). The map partition id is the right seed on both counts: it is
+distinct across mappers, so a task emitting fewer batches than there are output partitions does not
+leave the tail empty stage-wide, and it is 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.
+
+Distribution is even at batch granularity rather than row granularity: fewer batches than
+partitions leaves partitions empty, and unequal batch sizes give unequal partitions.
+
+#### Retry safety under `WholeBatch`
+
+Positional assignment is not a function of the rows, so it is only reproducible when the upstream
+operator replays the same batches in the same order and with the same framing. Re-executing one map
+task against differently framed input writes 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 faces the same problem with its own round robin and answers it
+in two ways, both of which Comet's native path mirrors:
+
+- **Declaring the risk.** Spark wraps a round-robin repartition in a `MapPartitionsRDD` with
+  `isOrderSensitive = true`, which reports `INDETERMINATE` whenever its parent is `UNORDERED`. The
+  `DAGScheduler` then rolls the whole stage back rather than re-running a single task, and aborts
+  the job outright when a result stage has already consumed output. The native path has no
+  `MapPartitionsRDD` to carry the flag, so `CometNativeShuffleInputRDD.getOutputDeterministicLevel`
+  applies the same rule directly. A determinate parent such as a plain scan stays determinate and
+  keeps cheap per-task retry; anything below another exchange is unordered, because reduce tasks
+  see shuffle blocks in arrival order, and goes indeterminate.
+
+- **Refusing the retry.** Rollback is only as sound as the parent's determinism level being an
+  accurate description of what the upstream operator replays. While the strategy is opt-in and
+  unproven, `spark.comet.shuffle.native.partitioning.roundrobin.batchGranular.failOnRetry`
+  (default `true`) makes `CometNativeShuffleWriter` refuse to run at all 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. Spark has no API for failing an application from inside a task, so this throws an
+  ordinary exception; because the condition only gets more true on each attempt, the task set
+  aborts after `spark.task.maxFailures` and takes the job with it. Setting the config to `false`
+  leaves the indeterminate declaration above as the only defence.
+
+Neither applies to `HashAll`, whose output is a pure function of the rows it sees.
+
 ## Memory Management
 
 Native shuffle uses DataFusion's memory management with spilling support:
@@ -370,21 +434,23 @@ independently compressed, allowing parallel decompression during reads.
 
 ## Configuration
 
-| Config                                                              | Default | Description                                                   |
-| ------------------------------------------------------------------- | ------- | ------------------------------------------------------------- |
-| `spark.comet.shuffle.enabled`                                       | `true`  | Enable Comet shuffle                                          |
-| `spark.comet.shuffle.mode`                                          | `auto`  | Shuffle mode: `native`, `jvm`, or `auto`                      |
-| `spark.comet.shuffle.directRead.enabled`                            | `true`  | Decode shuffle blocks in native code, bypassing Arrow FFI     |
-| `spark.comet.shuffle.compression.codec`                             | `lz4`   | Compression codec                                             |
-| `spark.comet.shuffle.compression.zstd.level`                        | `1`     | Zstd compression level                                        |
-| `spark.comet.shuffle.native.writeBufferSize`                        | `1MB`   | Write buffer size                                             |
-| `spark.comet.shuffle.native.maxBufferBytes`                         | `0`     | Fixed spill threshold. `0` disables it, leaving pool pressure |
-| `spark.comet.shuffle.native.partitioning.hash.enabled`              | `true`  | Allow `HashPartitioning` on the native path                   |
-| `spark.comet.shuffle.native.partitioning.hash.nested.enabled`       | `false` | Allow struct and array hash keys, and map keys on Spark 4.0+  |
-| `spark.comet.shuffle.native.partitioning.range.enabled`             | `true`  | Allow `RangePartitioning` on the native path                  |
-| `spark.comet.shuffle.native.partitioning.roundrobin.enabled`        | `false` | Allow `RoundRobinPartitioning` on the native path             |
-| `spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` | `0`     | Columns to hash for round robin. `0` hashes all of them       |
-| `spark.comet.shuffle.jvm.batchSize`                                 | `8192`  | Target rows per batch                                         |
+| Config                                                                         | Default | Description                                                                |
+| ------------------------------------------------------------------------------ | ------- | -------------------------------------------------------------------------- |
+| `spark.comet.shuffle.enabled`                                                  | `true`  | Enable Comet shuffle                                                       |
+| `spark.comet.shuffle.mode`                                                     | `auto`  | Shuffle mode: `native`, `jvm`, or `auto`                                   |
+| `spark.comet.shuffle.directRead.enabled`                                       | `true`  | Decode shuffle blocks in native code, bypassing Arrow FFI                  |
+| `spark.comet.shuffle.compression.codec`                                        | `lz4`   | Compression codec                                                          |
+| `spark.comet.shuffle.compression.zstd.level`                                   | `1`     | Zstd compression level                                                     |
+| `spark.comet.shuffle.native.writeBufferSize`                                   | `1MB`   | Write buffer size                                                          |
+| `spark.comet.shuffle.native.maxBufferBytes`                                    | `0`     | Fixed spill threshold. `0` disables it, leaving pool pressure              |
+| `spark.comet.shuffle.native.partitioning.hash.enabled`                         | `true`  | Allow `HashPartitioning` on the native path                                |
+| `spark.comet.shuffle.native.partitioning.hash.nested.enabled`                  | `false` | Allow struct and array hash keys, and map keys on Spark 4.0+               |
+| `spark.comet.shuffle.native.partitioning.range.enabled`                        | `true`  | Allow `RangePartitioning` on the native path                               |
+| `spark.comet.shuffle.native.partitioning.roundrobin.enabled`                   | `false` | Allow `RoundRobinPartitioning` on the native path                          |
+| `spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns`            | `0`     | Columns to hash for round robin. `0` hashes all of them                    |
+| `spark.comet.shuffle.native.partitioning.roundrobin.batchGranular`             | `false` | Assign each batch whole to one partition instead of hashing rows           |
+| `spark.comet.shuffle.native.partitioning.roundrobin.batchGranular.failOnRetry` | `true`  | Fail a retried positional round-robin map task instead of rewriting output |
+| `spark.comet.shuffle.jvm.batchSize`                                            | `8192`  | Target rows per batch                                                      |
 
 ## Comparison with JVM Shuffle
 

native/core/src/execution/planner.rs

@@ -98,7 +98,7 @@ use datafusion_comet_spark_expr::{
 use iceberg::expr::Bind;
 
 use crate::execution::operators::ExecutionError::GeneralError;
-use crate::execution::shuffle::{CometPartitioning, CompressionCodec};
+use crate::execution::shuffle::{CometPartitioning, CompressionCodec, RoundRobinStrategy};
 use crate::execution::spark_plan::SparkPlan;
 use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url;
 use crate::parquet::parquet_support::prepare_object_store_with_configs;
@@ -3652,15 +3652,22 @@ impl PhysicalPlanner {
             }
             PartitioningStruct::SinglePartition(_) => Ok(CometPartitioning::SinglePartition),
             PartitioningStruct::RoundRobinPartition(rr_partition) => {
-                // Treat negative max_hash_columns as 0 (no limit)
-                let max_hash_columns = if rr_partition.max_hash_columns <= 0 {
-                    0
+                let strategy = if rr_partition.batch_granular {
+                    // The Spark map partition id, not the DataFusion one: the root plan is
+                    // always executed with partition 0 (one Comet execution per Spark task),
+                    // so `ShuffleWriterExec::execute` cannot supply it. See
+                    // `RoundRobinStrategy::WholeBatch` for why it has to be this value.
+                    RoundRobinStrategy::WholeBatch {
+                        start_partition: self.partition.max(0) as usize,
+                    }
                 } else {
-                    rr_partition.max_hash_columns as usize
+                    // Treat negative max_hash_columns as 0 (no limit).
+                    let max_hash_columns = rr_partition.max_hash_columns.max(0) as usize;
+                    RoundRobinStrategy::HashAll { max_hash_columns }
                 };
                 Ok(CometPartitioning::RoundRobin(
                     rr_partition.num_partitions as usize,
-                    max_hash_columns,
+                    strategy,
                 ))
             }
         }

native/proto/src/proto/partitioning.proto

@@ -57,4 +57,8 @@ message RoundRobinPartition {
   int32 num_partitions = 1;
   // Maximum number of columns to hash. 0 means no limit (hash all columns).
   int32 max_hash_columns = 2;
+  // When true, assign each incoming RecordBatch as a whole to one output partition,
+  // skipping per-row hashing entirely. Retry-safe when the upstream operator
+  // emits deterministic batches under retry.
+  bool batch_granular = 3;
 }