feat: positional round robin shuffle keyed on a row ordinal

.ai/skills/review-comet-shuffle-pr/SKILL.md

@@ -88,13 +88,31 @@ Partitioning is where shuffle silently produces wrong answers rather than failin
 - [ ] **Hash partitioning uses Murmur3 with seed 42** and `partition_id = hash % num_partitions`,
       matching Spark. Any change to the hash, the seed, or the modulo changes which rows land in
       which partition, which breaks a join between a Comet-shuffled side and a Spark-shuffled side.
-- [ ] **Round robin is hash-based on purpose, and off by default.** Comet assigns partitions from a
-      Murmur3 hash rather than cycling row by row, because determinism across task retries is
-      required for correctness under fault tolerance. A PR that implements "true" round robin to fix
-      skew breaks that. Two costs are accepted: low-cardinality data distributes unevenly, and
-      unsorted rows land in different partitions than Spark's `UnsafeRow`-sorted assignment would
-      put them, which is why `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults
-      to `false`. Sorted output is identical either way.
+- [ ] **Round robin defaults to a content hash, and is off by default.** Comet's default
+      `RoundRobinStrategy::HashAll` assigns partitions from a Murmur3 hash rather than cycling row
+      by row, because placement has to be reproducible across task retries. Two costs are accepted:
+      low-cardinality data distributes unevenly, and unsorted rows land in different partitions
+      than Spark's `UnsafeRow`-sorted assignment would put them, which is why
+      `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`. Sorted
+      output is identical either way.
+- [ ] **Positional round robin is allowed, but only where retry reproducibility is established.**
+      `RoundRobinStrategy::RowGroups` places the row at task-global ordinal `i` at
+      `(startPartition + i / groupRows) % numPartitions`, which is Spark's own round robin at a
+      coarser granularity. It is not a bug, but it is only correct behind two gates, and a PR that
+      widens either one needs an argument. `CometShuffleExchangeExec.replaysRowsInOrder` is an
+      allowlist over the native subtree fused into the writer, which the RDD graph cannot see: a
+      native scan under nothing but projections and filters, with anything that spills staying out
+      because it reorders between attempts. `CometNativeShuffleInputRDD.getOutputDeterministicLevel`
+      applies Spark's `isOrderSensitive` rule to everything below that RDD.
+- [ ] **Positional placement keys on a row ordinal, and the start is scrambled.** The counter is
+      over rows and carries across batch boundaries. A PR that keys on a batch ordinal instead is
+      relying on framing, which no Spark contract covers: `DETERMINATE` promises the same rows in
+      the same order and says nothing about how an operator chunks them. `startPartition` must be
+      decorrelated across mappers, not merely distinct — each task walks consecutive partitions
+      from its start, so adjacent starts overlap and leave the tail of the partition space empty
+      (SPARK-21782). And `groupRows` bounds imbalance within one map task only; a reducer sees the
+      sum over all of them, which only evens out when each task emits many more groups than there
+      are partitions.
 - [ ] **Range partitioning bounds come from the driver.** Spark's `RangePartitioner` samples and
       computes boundaries, they are serialized into the native plan, and native does a binary
       search over comparable-row-format keys. A change to the comparison or the row encoding must

.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.CometNativePositionalRoundRobinSuite
               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.CometNativePositionalRoundRobinSuite
               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

@@ -305,7 +305,12 @@ 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`; `RowGroups` is opt-in through
+`spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled`, and is used only where the
+planner can establish that the map task replays its rows in the same order.
+
+#### `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`
@@ -314,13 +319,97 @@ This approach guarantees determinism across retries, which is critical for fault
 However, unlike true round robin which cycles through partitions row-by-row, hash-based
 assignment only provides even distribution when the data has sufficient variation in the
 hashed columns. Data with low cardinality or identical values may result in skewed partition
-sizes.
+sizes: because placement is a pure function of a row's contents, a column of one repeated value
+lands entirely on one reducer.
 
 Because Spark assigns round robin partitions by sorting rows on their binary `UnsafeRow` form,
 which Arrow's layout does not reproduce, unsorted output can land in different partitions than
 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.
+
+#### `RowGroups`: 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.
+`RowGroups` places rows the way Spark's own round robin does: the row at task-global ordinal `i`
+goes to `(startPartition + i / groupRows) % numPartitions`. That removes the per-row hash, and it
+replaces the per-row gather with a bulk copy per contiguous run, because adjacent rows now stay
+together. It also spreads duplicate rows evenly, which `HashAll` cannot.
+
+The counter is over **rows**, not batches, and it carries across batch boundaries: a group that one
+input batch leaves part-way through is finished by the next. That is deliberate. Spark's
+determinism contract, `DeterministicLevel`, describes row _order_ and says nothing about how a
+downstream operator frames rows into batches, so an operator that spills can reframe under
+different memory pressure while still honouring `DETERMINATE`. Keying on a row ordinal means the
+strategy depends only on the property Spark actually publishes.
+
+`start_partition` is the output partition a map task's first group goes to. It is computed per task
+on the JVM, in `CometNativeShuffleWriter.buildUnifiedPlan` where the Spark map partition id is in
+scope, and passed down in the proto. It has to be _decorrelated_ across mappers, not merely
+distinct: a task walks `ceil(rows / groupRows)` consecutive partitions from its start, so if
+consecutive tasks start on consecutive partitions their runs all overlap and the partitions past
+`numMapTasks + groupsPerTask` get nothing — ten map tasks of 5,000 rows into 200 partitions at a
+group of 64 would leave 112 reducers empty. It also has to be a pure function of the map partition,
+or a re-executed task does not reproduce its own placement. Spark satisfies both by scrambling the
+map partition id through `XORShiftRandom`
+([SPARK-21782](https://issues.apache.org/jira/browse/SPARK-21782)), and
+`CometShuffleExchangeExec.positionalStartPartition` does the same. Spark increments its counter
+before the first row uses it, so the start is `nextInt(numPartitions) + 1`, which makes
+`groupRows = 1` place rows exactly where Spark's round robin would for the same row order.
+
+`groupRows` trades balance against copying. Within one map task, imbalance between any two output
+partitions is bounded by `groupRows` rows however the reader frames its batches, so small groups
+balance better; large groups produce fewer, longer runs to copy, and a group as large as the batch
+size lets a whole input batch pass through to one partition untouched. That bound does not compose
+across map tasks — a reducer sees the sum over all of them, which is only even when each task emits
+many more groups than there are output partitions. `0`, the default, derives it as
+`clamp(batch_size / num_partitions, 64, batch_size)`, which keeps a task wrapping around the output
+partitions roughly once per batch. The 64-row floor caps how finely a batch is cut: with far more
+partitions than rows in a batch, `batch_size / num_partitions` rounds down towards one row and the
+flush degenerates into the per-row gather positional placement exists to avoid.
+
+Internally, `MultiPartitionShuffleRepartitioner` records `(batch, start, len)` runs rather than
+one `(batch, row)` pair per row, so the index list charged against the spill reservation is
+smaller, and `RunIterator` builds each output chunk by slicing and concatenating runs. A run that
+covers an entire buffered batch and already fills a chunk is passed through without copying.
+
+#### Retry safety under `RowGroups`
+
+Positional assignment is not a function of the rows, so it is reproducible only when the map task
+replays the same rows in the same order. Re-executing one map task against differently ordered
+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-23207](https://issues.apache.org/jira/browse/SPARK-23207)). Spark faces the same problem
+with its own round robin. Comet establishes the condition in two places, both of which must hold:
+
+- **In the plan.** `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused
+  into the writer, which the RDD graph cannot see because the whole subtree collapses into one
+  `CometNativeShuffleInputRDD`. It is a short allowlist, not a denylist: a native scan under
+  nothing but projections and filters. Operators that spill are the interesting exclusion, since
+  an aggregate or sort under memory pressure emits output in an order that depends on how many
+  times it spilled, which differs between attempts. Anything else keeps `HashAll`, which is safe
+  to re-execute whatever its input does.
+
+- **In the RDD graph.** 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. 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.
+
+Neither applies to `HashAll`, whose output is a pure function of the rows it sees.
+
+One schema-level restriction is applied in `create_repartitioner`: positional placement is the only
+strategy that hands a sliced array to the IPC writer, and while the writer truncates a slice's
+buffers for every other type, for `Utf8View` and `BinaryView` it serializes every shared data
+buffer in full. A schema containing a view type anywhere therefore falls back to `HashAll`.
+
 ## Memory Management
 
 Native shuffle uses DataFusion's memory management with spilling support:
@@ -370,21 +459,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.positional.enabled`   | `false` | Place round-robin rows by position where the plan allows it                   |
+| `spark.comet.shuffle.native.partitioning.roundrobin.positional.groupRows` | `0`     | Rows per positional group. `0` derives it from batch size and partition count |
+| `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,25 @@ 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.positional {
+                    RoundRobinStrategy::RowGroups {
+                        // Computed per task on the JVM, where the Spark map partition id is in
+                        // scope, and scrambled the way Spark scrambles it. See
+                        // `CometShuffleExchangeExec.positionalStartPartition`.
+                        start_partition: rr_partition.positional_start_partition.max(0) as usize,
+                        // Negative or zero means "derive it from the batch size and partition
+                        // count", which the repartitioner does once it knows both.
+                        group_rows: rr_partition.positional_group_rows.max(0) as usize,
+                    }
                 } else {
-                    rr_partition.max_hash_columns as usize
+                    // Treat negative max_hash_columns as 0 (no limit).
+                    RoundRobinStrategy::HashAll {
+                        max_hash_columns: rr_partition.max_hash_columns.max(0) as usize,
+                    }
                 };
                 Ok(CometPartitioning::RoundRobin(
                     rr_partition.num_partitions as usize,
-                    max_hash_columns,
+                    strategy,
                 ))
             }
         }