From 95c5ca31f6b04fe8935b835cf782a4d9717946d6 Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Mon, 24 Aug 2026 08:33:23 -0700
Subject: [PATCH 1/4] feat: RR partition with vectorized distribution

---
 native/core/src/execution/planner.rs          |  11 +-
 native/proto/src/proto/partitioning.proto     |   4 +
 native/shuffle/Cargo.toml                     |   3 +
 native/shuffle/benches/shuffle_writer.rs      | 220 +++++++++++++++--
 native/shuffle/src/bin/shuffle_bench.rs       |  11 +-
 native/shuffle/src/comet_partitioning.rs      |  21 +-
 native/shuffle/src/lib.rs                     |  11 +-
 native/shuffle/src/metrics.rs                 |   4 +-
 native/shuffle/src/partitioners/mod.rs        |   4 +-
 .../src/partitioners/multi_partition.rs       | 219 +++++++++++------
 .../partitioned_batch_iterator.rs             |  36 +++
 native/shuffle/src/partitioners/traits.rs     |   2 +-
 native/shuffle/src/rss_execution_tests.rs     |  23 +-
 native/shuffle/src/shuffle_writer.rs          | 232 +++++++++++++++++-
 .../writers/local/local_partition_writer.rs   |   4 +-
 native/shuffle/src/writers/mod.rs             |   2 +-
 .../shuffle/src/writers/partition_writer.rs   |   2 +-
 .../scala/org/apache/comet/CometConf.scala    |  12 +
 .../shuffle/CometNativeShuffleWriter.scala    |   2 +
 19 files changed, 695 insertions(+), 128 deletions(-)

diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 8f030da455b..22ca8c02d6d 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -3652,15 +3652,16 @@ 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 {
+                    crate::execution::shuffle::RoundRobinStrategy::WholeBatch
                 } 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;
+                    crate::execution::shuffle::RoundRobinStrategy::HashAll { max_hash_columns }
                 };
                 Ok(CometPartitioning::RoundRobin(
                     rr_partition.num_partitions as usize,
-                    max_hash_columns,
+                    strategy,
                 ))
             }
         }
diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto
index e70b8264f02..9aa6acc3224 100644
--- a/native/proto/src/proto/partitioning.proto
+++ b/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;
 }
diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml
index f0ed22ad730..6d6d85bd2c9 100644
--- a/native/shuffle/Cargo.toml
+++ b/native/shuffle/Cargo.toml
@@ -57,6 +57,9 @@ zstd = { version = "0.13.3", features = ["experimental"] }
 criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] }
 datafusion = { workspace = true, features = ["parquet_encryption", "sql"] }
 itertools = "0.15.0"
+# The shuffle_writer bench reads a real on-disk parquet dataset; the optional dependency above
+# is gated behind shuffle-bench, which benches do not activate.
+parquet = { workspace = true }
 tempfile = "3.26.0"
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 
diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs
index cc94375436e..0ac4bccb421 100644
--- a/native/shuffle/benches/shuffle_writer.rs
+++ b/native/shuffle/benches/shuffle_writer.rs
@@ -22,15 +22,18 @@ use arrow::row::{RowConverter, SortField};
 use criterion::{criterion_group, criterion_main, Criterion};
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::datasource::source::DataSourceExec;
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::physical_expr::expressions::{col, Column};
 use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
-use datafusion::physical_plan::metrics::Time;
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time};
 use datafusion::{
     physical_plan::{common::collect, ExecutionPlan},
     prelude::SessionContext,
 };
 use datafusion_comet_shuffle::{
-    CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext, ShuffleWriterExec,
+    CometPartitioning, CompressionCodec, LocalPartitionWriter, MultiPartitionShuffleRepartitioner,
+    PartitionOffsets, RoundRobinStrategy, ShuffleBlockWriter, ShuffleCodecContext,
+    ShufflePartitioner, ShufflePartitionerMetrics, ShuffleWriterExec,
 };
 use itertools::Itertools;
 use std::io::Cursor;
@@ -77,8 +80,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 let exec = create_shuffle_writer_exec(
                     compression_codec.clone(),
                     CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16),
-                    8192,
-                    10,
+                    create_batches(8192, 10),
                 );
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
@@ -130,8 +132,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 let exec = create_shuffle_writer_exec(
                     compression_codec.clone(),
                     partitioning.clone(),
-                    8192,
-                    10,
+                    create_batches(8192, 10),
                 );
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
@@ -161,8 +162,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 let exec = create_shuffle_writer_exec(
                     CompressionCodec::None,
                     CometPartitioning::SinglePartition,
-                    rows_per_batch,
-                    num_batches,
+                    create_batches(rows_per_batch, num_batches),
                 );
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
@@ -173,6 +173,137 @@ fn criterion_benchmark(c: &mut Criterion) {
             },
         );
     }
+
+    // RoundRobin on a wide nested schema: compares the row-level hash-all-columns strategy
+    // against the whole-batch strategy. The nested schema is where the win is largest:
+    // `create_murmur3_hashes` recurses into every Struct child per row on the row-level path,
+    // while the whole-batch path does one mod per batch and never inspects row contents.
+    // 40 top-level struct columns of shallow depth reflect the real-world shape where this
+    // optimization matters (event log workload was ~194 nested columns).
+    let num_partitions = 50usize;
+    let wide_batches: Vec<RecordBatch> = (0..8).map(|_| nested_schema_batch(8192, 40, 2)).collect();
+    let wide_schema = wide_batches[0].schema();
+    let round_robin_strategies = [
+        (
+            "hash_all_columns",
+            RoundRobinStrategy::HashAll {
+                max_hash_columns: 0,
+            },
+        ),
+        ("whole_batch", RoundRobinStrategy::WholeBatch),
+    ];
+    for (label, strategy) in &round_robin_strategies {
+        group.bench_function(
+            format!("shuffle_writer: RoundRobin nested schema (strategy={label})"),
+            |b| {
+                let ctx = SessionContext::new();
+                let rt = Runtime::new().unwrap();
+                let exec = create_shuffle_writer_exec(
+                    CompressionCodec::None,
+                    CometPartitioning::RoundRobin(num_partitions, strategy.clone()),
+                    wide_batches.clone(),
+                );
+                b.iter(|| {
+                    let task_ctx = ctx.task_ctx();
+                    let stream = exec.execute(0, task_ctx).unwrap();
+                    rt.block_on(collect(stream)).unwrap();
+                });
+            },
+        );
+    }
+
+    // Partitioning-only microbench. `insert_batch` runs partition assignment + index
+    // buffering but never flushes to disk, so the two RoundRobin strategies can be compared
+    // without the IPC encode + disk write cost that dominates the end-to-end bench above.
+    for (label, strategy) in &round_robin_strategies {
+        group.bench_function(
+            format!("partitioning_only: RoundRobin nested schema (strategy={label})"),
+            |b| {
+                let rt = Runtime::new().unwrap();
+                let batches = wide_batches.clone();
+                let runtime_env = Arc::new(
+                    RuntimeEnvBuilder::new()
+                        .with_memory_limit(1024 * 1024 * 1024, 1.0)
+                        .build()
+                        .unwrap(),
+                );
+                let dir = tempfile::tempdir().unwrap();
+                let data_path = dir.path().join("data.out").to_str().unwrap().to_string();
+                b.iter(|| {
+                    let block_writer =
+                        ShuffleBlockWriter::try_new(wide_schema.as_ref(), CompressionCodec::None)
+                            .unwrap();
+                    let writer = LocalPartitionWriter::try_new(
+                        data_path.clone(),
+                        Arc::new(PartitionOffsets::default()),
+                        block_writer,
+                        num_partitions,
+                        8192,
+                        1024 * 1024,
+                        Arc::clone(&runtime_env),
+                    )
+                    .unwrap();
+                    let metrics =
+                        ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+                    let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
+                        0,
+                        writer,
+                        CometPartitioning::RoundRobin(num_partitions, strategy.clone()),
+                        metrics,
+                        Arc::clone(&runtime_env),
+                        8192,
+                        false,
+                        None,
+                    )
+                    .unwrap();
+                    rt.block_on(async {
+                        for batch in &batches {
+                            repartitioner.insert_batch(batch.clone()).await.unwrap();
+                        }
+                    });
+                });
+            },
+        );
+    }
+
+    // End-to-end bench on a real on-disk parquet dataset at
+    // /tmp/clickstream_data_smoke (Spark-written, snappy-compressed, wide nested
+    // clickstream schema — the shape this optimization targets). Loaded once per
+    // process; each iteration streams the batches through `ShuffleWriterExec`
+    // for both RoundRobin strategies. Skipped if the dataset is absent.
+    let clickstream_batches = load_parquet_dir_batches("/tmp/clickstream_data_smoke");
+    if let Some(batches) = clickstream_batches {
+        let clickstream_num_partitions = 50usize;
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        let num_cols = batches[0].num_columns();
+        eprintln!(
+            "clickstream bench: {} batches, {} rows total, {} columns",
+            batches.len(),
+            total_rows,
+            num_cols,
+        );
+        for (label, strategy) in &round_robin_strategies {
+            group.bench_function(
+                format!("shuffle_writer: RoundRobin clickstream parquet (strategy={label})"),
+                |b| {
+                    let ctx = SessionContext::new();
+                    let rt = Runtime::new().unwrap();
+                    let exec = create_shuffle_writer_exec(
+                        CompressionCodec::Lz4Frame,
+                        CometPartitioning::RoundRobin(clickstream_num_partitions, strategy.clone()),
+                        batches.clone(),
+                    );
+                    b.iter(|| {
+                        let task_ctx = ctx.task_ctx();
+                        let stream = exec.execute(0, task_ctx).unwrap();
+                        rt.block_on(collect(stream)).unwrap();
+                    });
+                },
+            );
+        }
+    } else {
+        eprintln!("clickstream bench: /tmp/clickstream_data_smoke not found; skipping");
+    }
     group.finish();
 
     // High partition counts stress the per-partition write path (one short-lived
@@ -189,8 +320,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 let exec = create_shuffle_writer_exec(
                     CompressionCodec::None,
                     CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], num_partitions),
-                    8192,
-                    10,
+                    create_batches(8192, 10),
                 );
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
@@ -204,13 +334,66 @@ fn criterion_benchmark(c: &mut Criterion) {
     high_partition_group.finish();
 }
 
+/// Loads every `*.parquet` file under `dir` (recursively, one level deep for
+/// partitioned datasets like `date=…/part-*.parquet`) and returns their record
+/// batches at 8192 rows per batch. Returns `None` if the directory is missing
+/// or contains no parquet files.
+fn load_parquet_dir_batches(dir: &str) -> Option<Vec<RecordBatch>> {
+    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+
+    let root = std::path::Path::new(dir);
+    if !root.exists() {
+        return None;
+    }
+
+    let mut files: Vec<std::path::PathBuf> = Vec::new();
+    let mut stack: Vec<std::path::PathBuf> = vec![root.to_path_buf()];
+    while let Some(p) = stack.pop() {
+        let entries = match std::fs::read_dir(&p) {
+            Ok(e) => e,
+            Err(_) => continue,
+        };
+        for entry in entries.flatten() {
+            let path = entry.path();
+            if path.is_dir() {
+                stack.push(path);
+            } else if path
+                .extension()
+                .and_then(|s| s.to_str())
+                .map(|s| s.eq_ignore_ascii_case("parquet"))
+                .unwrap_or(false)
+            {
+                files.push(path);
+            }
+        }
+    }
+    files.sort();
+    if files.is_empty() {
+        return None;
+    }
+
+    let mut all = Vec::new();
+    for path in files {
+        let file = std::fs::File::open(&path).unwrap_or_else(|e| {
+            panic!("failed to open {}: {e}", path.display());
+        });
+        let reader = ParquetRecordBatchReaderBuilder::try_new(file)
+            .unwrap_or_else(|e| panic!("parquet builder for {}: {e}", path.display()))
+            .with_batch_size(8192)
+            .build()
+            .unwrap_or_else(|e| panic!("parquet build for {}: {e}", path.display()));
+        for batch in reader {
+            all.push(batch.unwrap());
+        }
+    }
+    Some(all)
+}
+
 fn create_shuffle_writer_exec(
     compression_codec: CompressionCodec,
     partitioning: CometPartitioning,
-    rows_per_batch: usize,
-    num_batches: usize,
+    batches: Vec<RecordBatch>,
 ) -> ShuffleWriterExec {
-    let batches = create_batches(rows_per_batch, num_batches);
     let schema = batches[0].schema();
     let partitions = &[batches];
     ShuffleWriterExec::try_new(
@@ -278,7 +461,7 @@ fn schema_encoding_benchmark(c: &mut Criterion) {
 
     for (name, batch) in [
         ("flat", flat_schema_batch(8192)),
-        ("nested", nested_schema_batch(8192)),
+        ("nested", nested_schema_batch(8192, 4, 6)),
     ] {
         let writer =
             ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap();
@@ -307,7 +490,7 @@ fn ipc_context_reuse_benchmark(c: &mut Criterion) {
         for (name, batch) in [
             ("mixed", create_batch(rows, true)),
             ("flat", flat_schema_batch(rows)),
-            ("nested", nested_schema_batch(rows)),
+            ("nested", nested_schema_batch(rows, 4, 6)),
         ] {
             for codec in [
                 CompressionCodec::None,
@@ -370,11 +553,8 @@ fn flat_schema_batch(num_rows: usize) -> RecordBatch {
     RecordBatch::try_new(schema, columns).unwrap()
 }
 
-/// A schema of several deeply nested struct columns.
-fn nested_schema_batch(num_rows: usize) -> RecordBatch {
-    let num_cols = 4;
-    let depth = 6;
-
+/// A schema of `num_cols` struct columns, each nested `depth` levels deep.
+fn nested_schema_batch(num_rows: usize, num_cols: usize, depth: usize) -> RecordBatch {
     let mut fields: Vec<Field> = Vec::with_capacity(num_cols);
     let mut columns: Vec<Arc<dyn Array>> = Vec::with_capacity(num_cols);
     for col in 0..num_cols {
diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs
index 051e18fbdd8..c14f15aebd5 100644
--- a/native/shuffle/src/bin/shuffle_bench.rs
+++ b/native/shuffle/src/bin/shuffle_bench.rs
@@ -45,7 +45,9 @@ use datafusion::physical_plan::common::collect;
 use datafusion::physical_plan::metrics::{MetricValue, MetricsSet};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_comet_shuffle::{CometPartitioning, CompressionCodec, ShuffleWriterExec};
+use datafusion_comet_shuffle::{
+    CometPartitioning, CompressionCodec, RoundRobinStrategy, ShuffleWriterExec,
+};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use std::fs;
 use std::path::{Path, PathBuf};
@@ -583,7 +585,12 @@ fn build_partitioning(
 ) -> CometPartitioning {
     match scheme {
         "single" => CometPartitioning::SinglePartition,
-        "round-robin" => CometPartitioning::RoundRobin(num_partitions, 0),
+        "round-robin" => CometPartitioning::RoundRobin(
+            num_partitions,
+            RoundRobinStrategy::HashAll {
+                max_hash_columns: 0,
+            },
+        ),
         "hash" => {
             let exprs: Vec<Arc<dyn datafusion::physical_expr::PhysicalExpr>> = hash_col_indices
                 .iter()
diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs
index 15912e6481d..7b331ed3eb9 100644
--- a/native/shuffle/src/comet_partitioning.rs
+++ b/native/shuffle/src/comet_partitioning.rs
@@ -19,6 +19,20 @@ use arrow::row::{OwnedRow, RowConverter};
 use datafusion::physical_expr::{LexOrdering, PhysicalExpr};
 use std::sync::Arc;
 
+/// How Comet's RoundRobin partitioning assigns rows to output partitions.
+#[derive(Debug, Clone)]
+pub enum RoundRobinStrategy {
+    /// Hash every row across up to `max_hash_columns` columns (0 means no limit)
+    /// and assign rows individually. Deterministic under retry regardless of upstream
+    /// order preservation, but pays the per-row hash cost.
+    HashAll { max_hash_columns: usize },
+    /// Assign each incoming RecordBatch as a whole to one output partition, based on
+    /// the input partition id and a per-task batch sequence number. Skips per-row
+    /// hashing entirely. Retry-safe only when the upstream operator emits the same
+    /// batches in the same order under retry.
+    WholeBatch,
+}
+
 /// Partitioning scheme for distributing rows across shuffle output partitions.
 #[derive(Debug, Clone)]
 pub enum CometPartitioning {
@@ -32,10 +46,9 @@ pub enum CometPartitioning {
     /// Rows for comparing to 4) OwnedRows that represent the boundaries of each partition, used with
     /// LexOrdering to bin each value in the RecordBatch to a partition.
     RangePartitioning(LexOrdering, usize, Arc<RowConverter>, Vec<OwnedRow>),
-    /// Round robin partitioning. Distributes rows across partitions by sorting them by hash
-    /// (computed from columns) and then assigning partitions sequentially. Args are:
-    /// 1) number of partitions, 2) max columns to hash (0 means no limit).
-    RoundRobin(usize, usize),
+    /// Round robin partitioning. Args are the number of partitions and the strategy that
+    /// decides how rows are assigned. See [`RoundRobinStrategy`] for the trade-offs.
+    RoundRobin(usize, RoundRobinStrategy),
 }
 
 impl CometPartitioning {
diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs
index 1158a2b1e2e..1844ce2da8c 100644
--- a/native/shuffle/src/lib.rs
+++ b/native/shuffle/src/lib.rs
@@ -32,9 +32,18 @@ pub mod spark_unsafe;
 pub(crate) mod writers;
 
 pub use codec_context::ShuffleCodecContext;
-pub use comet_partitioning::CometPartitioning;
+pub use comet_partitioning::{CometPartitioning, RoundRobinStrategy};
 pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated, reset_schema_cache};
 pub use remote_schema::{decode_remote_shuffle_batch, validate_remote_schema};
 pub use schema_align::SchemaAlignExec;
 pub use shuffle_writer::{PartitionOffsets, ShuffleWriterDestination, ShuffleWriterExec};
 pub use writers::{CompressionCodec, ShuffleBlockWriter};
+
+// Bench-only re-exports. `#[doc(hidden)]` marks these as not part of the stable API contract.
+// Consumed by the criterion benches under `native/shuffle/benches/`.
+#[doc(hidden)]
+pub use metrics::ShufflePartitionerMetrics;
+#[doc(hidden)]
+pub use partitioners::{MultiPartitionShuffleRepartitioner, ShufflePartitioner};
+#[doc(hidden)]
+pub use writers::LocalPartitionWriter;
diff --git a/native/shuffle/src/metrics.rs b/native/shuffle/src/metrics.rs
index 855bb950111..10d8106329a 100644
--- a/native/shuffle/src/metrics.rs
+++ b/native/shuffle/src/metrics.rs
@@ -20,7 +20,7 @@ use datafusion::physical_plan::metrics::{
 };
 
 /// Execution metrics for a shuffle partition operation.
-pub(crate) struct ShufflePartitionerMetrics {
+pub struct ShufflePartitionerMetrics {
     /// metrics
     pub(crate) baseline: BaselineMetrics,
 
@@ -55,7 +55,7 @@ pub(crate) struct ShufflePartitionerMetrics {
 }
 
 impl ShufflePartitionerMetrics {
-    pub(crate) fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
+    pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
         Self {
             baseline: BaselineMetrics::new(metrics, partition),
             repart_time: MetricBuilder::new(metrics).subset_time("repart_time", partition),
diff --git a/native/shuffle/src/partitioners/mod.rs b/native/shuffle/src/partitioners/mod.rs
index 8f4239a820e..46d51603ad8 100644
--- a/native/shuffle/src/partitioners/mod.rs
+++ b/native/shuffle/src/partitioners/mod.rs
@@ -22,6 +22,6 @@ mod single_partition;
 mod traits;
 
 pub(crate) use empty_schema::EmptySchemaShufflePartitioner;
-pub(crate) use multi_partition::MultiPartitionShuffleRepartitioner;
+pub use multi_partition::MultiPartitionShuffleRepartitioner;
 pub(crate) use single_partition::SinglePartitionShufflePartitioner;
-pub(crate) use traits::ShufflePartitioner;
+pub use traits::ShufflePartitioner;
diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs
index 37f67eacbd6..017dc240f46 100644
--- a/native/shuffle/src/partitioners/multi_partition.rs
+++ b/native/shuffle/src/partitioners/multi_partition.rs
@@ -19,7 +19,7 @@ use crate::metrics::ShufflePartitionerMetrics;
 use crate::partitioners::partitioned_batch_iterator::PartitionedBatchesProducer;
 use crate::partitioners::ShufflePartitioner;
 use crate::writers::PartitionWriter;
-use crate::{comet_partitioning, CometPartitioning};
+use crate::{comet_partitioning, CometPartitioning, RoundRobinStrategy};
 use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch};
 use datafusion::common::utils::proxy::VecAllocExt;
 use datafusion::common::{DataFusionError, HashSet};
@@ -99,7 +99,7 @@ impl ScratchSpace {
 }
 
 /// A partitioner that uses a hash function to partition data into multiple partitions
-pub(crate) struct MultiPartitionShuffleRepartitioner<T: PartitionWriter> {
+pub struct MultiPartitionShuffleRepartitioner<T: PartitionWriter> {
     buffered_batches: Vec<RecordBatch>,
     partition_indices: Vec<Vec<(u32, u32)>>,
     partition_writer: T,
@@ -122,6 +122,11 @@ pub(crate) struct MultiPartitionShuffleRepartitioner<T: PartitionWriter> {
     /// allocation once rather than once per slice that references it. Cleared whenever the
     /// buffered batches drain (spill / shuffle_write). See `count_new_buffers`.
     pinned_buffers: HashSet<usize>,
+    /// Batch counter for the batch-granular RoundRobin path. Seeded with the input partition
+    /// id so mappers with adjacent partition ids do not concentrate their first batches on the
+    /// same output partition, then incremented once per input batch slice that reaches
+    /// `partitioning_batch`. Unused by other strategies.
+    round_robin_batch_seq: usize,
 }
 
 /// Sum of the capacities of the backing buffers reachable from `batch` whose start address is
@@ -172,7 +177,7 @@ fn count_new_buffers(batch: &RecordBatch, seen: &mut HashSet<usize>) -> usize {
 
 impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
     #[allow(clippy::too_many_arguments)]
-    pub(crate) fn try_new(
+    pub fn try_new(
         partition: usize,
         partition_writer: T,
         partitioning: CometPartitioning,
@@ -191,17 +196,36 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
         // Vectors in the scratch space will be filled with valid values before being used, this
         // initialization code is simply initializing the vectors to the desired size.
         // The initial values are not used.
+        //
+        // `partition_ids` and `partition_row_indices` are only touched by the row-level
+        // strategies (`Hash`, `RangePartitioning`, and hash-all-columns RoundRobin). The
+        // whole-batch RoundRobin path never inspects rows, so allocating them there would be
+        // ~64 KB of dead scratch on every task.
+        let needs_row_scratch = !matches!(
+            &partitioning,
+            CometPartitioning::SinglePartition
+                | CometPartitioning::RoundRobin(_, RoundRobinStrategy::WholeBatch),
+        );
         let scratch = ScratchSpace {
-            hashes_buf: match partitioning {
-                // Allocate hashes_buf for hash and round robin partitioning.
-                // Round robin hashes all columns to achieve even, deterministic distribution.
-                CometPartitioning::Hash(_, _) | CometPartitioning::RoundRobin(_, _) => {
+            hashes_buf: match &partitioning {
+                // Allocate hashes_buf for hash and hash-all-columns round robin partitioning.
+                // Whole-batch round robin does no per-row hashing.
+                CometPartitioning::Hash(_, _)
+                | CometPartitioning::RoundRobin(_, RoundRobinStrategy::HashAll { .. }) => {
                     vec![0; batch_size]
                 }
                 _ => vec![],
             },
-            partition_ids: vec![0; batch_size],
-            partition_row_indices: vec![0; batch_size],
+            partition_ids: if needs_row_scratch {
+                vec![0; batch_size]
+            } else {
+                vec![]
+            },
+            partition_row_indices: if needs_row_scratch {
+                vec![0; batch_size]
+            } else {
+                vec![]
+            },
             partition_starts: vec![0; num_output_partitions + 1],
         };
 
@@ -221,6 +245,10 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
             max_buffer_bytes,
             tracing_enabled,
             pinned_buffers: HashSet::new(),
+            // Seed with the input partition id so batches from mapper i land on
+            // partition (i + k) mod N on the k-th batch, spreading concurrent mappers' first
+            // batches across distinct output partitions.
+            round_robin_batch_seq: partition,
         })
     }
 
@@ -290,7 +318,7 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
 
                 self.buffer_partitioned_batch_may_spill(
                     input,
-                    partition_row_indices,
+                    Some(partition_row_indices),
                     partition_starts,
                 )
                 .await?;
@@ -343,68 +371,85 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
 
                 self.buffer_partitioned_batch_may_spill(
                     input,
-                    partition_row_indices,
+                    Some(partition_row_indices),
                     partition_starts,
                 )
                 .await?;
                 self.scratch = scratch;
             }
-            CometPartitioning::RoundRobin(num_output_partitions, max_hash_columns) => {
-                // Comet implements "round robin" as hash partitioning on columns.
-                // This achieves the same goal as Spark's round robin (even distribution
-                // without semantic grouping) while being deterministic for fault tolerance.
-                //
-                // Note: This produces different partition assignments than Spark's round robin,
-                // which sorts by UnsafeRow binary representation before assigning partitions.
-                // However, both approaches provide even distribution and determinism.
+            CometPartitioning::RoundRobin(num_output_partitions, strategy) => {
+                // Two strategies share the scratch-take / timer / spill / scratch-restore
+                // scaffold. Only the middle "fill partition_row_indices + partition_starts"
+                // step differs. WholeBatch: assign the whole input batch to one partition and
+                // pay no per-row cost. HashAll: hash every row (over up to max_hash_columns
+                // columns) and route rows individually. WholeBatch is retry-safe only when the
+                // upstream operator emits the same batches in the same order under retry
+                // (Comet's `CometNativeScan` and other order-preserving operators do so).
                 let mut scratch = std::mem::take(&mut self.scratch);
-                let (partition_starts, partition_row_indices): (&Vec<u32>, &Vec<u32>) = {
+                let num_rows = input.num_rows();
+                let partition_row_indices: Option<&[u32]> = {
                     let mut timer = self.metrics.repart_time.timer();
 
-                    let num_rows = input.num_rows();
+                    let indices: Option<&[u32]> = match strategy {
+                        RoundRobinStrategy::WholeBatch => {
+                            let target_idx = self.round_robin_batch_seq % *num_output_partitions;
+                            self.round_robin_batch_seq = self.round_robin_batch_seq.wrapping_add(1);
+
+                            // `partition_starts[k]..partition_starts[k+1]` is partition k's
+                            // slice. Slots 0..=target_idx are 0; slots after `target_idx` are
+                            // `num_rows`, so `target_idx` owns the whole [0..num_rows) range.
+                            // `partition_row_indices` is left unmaterialized — the spill path
+                            // treats the target partition's slice as an identity mapping.
+                            let partition_starts = &mut scratch.partition_starts;
+                            partition_starts.clear();
+                            partition_starts.resize(target_idx + 1, 0);
+                            partition_starts.resize(*num_output_partitions + 1, num_rows as u32);
+                            None
+                        }
+                        RoundRobinStrategy::HashAll { max_hash_columns } => {
+                            // Hash-partition rows into pmod(hash, N). This produces different
+                            // partition assignments than Spark's round robin (which sorts by
+                            // UnsafeRow binary representation before assigning partitions), but
+                            // both approaches provide even distribution and determinism.
+                            //
+                            // max_hash_columns of 0 means no limit (hash all columns). Negative
+                            // values are normalized to 0 in the planner.
+                            let num_columns_to_hash = if *max_hash_columns == 0 {
+                                input.num_columns()
+                            } else {
+                                (*max_hash_columns).min(input.num_columns())
+                            };
+                            let columns_to_hash: Vec<ArrayRef> = (0..num_columns_to_hash)
+                                .map(|i| Arc::clone(input.column(i)))
+                                .collect();
+
+                            // Use identical seed as Spark hash partitioning.
+                            let hashes_buf = &mut scratch.hashes_buf[..num_rows];
+                            hashes_buf.fill(42_u32);
+                            create_murmur3_hashes(&columns_to_hash, hashes_buf)?;
+
+                            let partition_ids = &mut scratch.partition_ids[..num_rows];
+                            hashes_buf.iter().enumerate().for_each(|(idx, hash)| {
+                                partition_ids[idx] =
+                                    comet_partitioning::pmod(*hash, *num_output_partitions) as u32;
+                            });
 
-                    // Collect columns for hashing, respecting max_hash_columns limit
-                    // max_hash_columns of 0 means no limit (hash all columns)
-                    // Negative values are normalized to 0 in the planner
-                    let num_columns_to_hash = if *max_hash_columns == 0 {
-                        input.num_columns()
-                    } else {
-                        (*max_hash_columns).min(input.num_columns())
+                            scratch.map_partition_ids_to_starts_and_indices(
+                                *num_output_partitions,
+                                num_rows,
+                            );
+                            Some(scratch.partition_row_indices.as_slice())
+                        }
                     };
-                    let columns_to_hash: Vec<ArrayRef> = (0..num_columns_to_hash)
-                        .map(|i| Arc::clone(input.column(i)))
-                        .collect();
-
-                    // Use identical seed as Spark hash partitioning.
-                    let hashes_buf = &mut scratch.hashes_buf[..num_rows];
-                    hashes_buf.fill(42_u32);
-
-                    // Compute hash for selected columns
-                    create_murmur3_hashes(&columns_to_hash, hashes_buf)?;
-
-                    // Assign partition IDs based on hash (same as hash partitioning)
-                    let partition_ids = &mut scratch.partition_ids[..num_rows];
-                    hashes_buf.iter().enumerate().for_each(|(idx, hash)| {
-                        partition_ids[idx] =
-                            comet_partitioning::pmod(*hash, *num_output_partitions) as u32;
-                    });
-
-                    // We now have partition ids for every input row, map that to partition starts
-                    // and partition indices to eventually write these rows to partition buffers.
-                    scratch
-                        .map_partition_ids_to_starts_and_indices(*num_output_partitions, num_rows);
 
                     timer.stop();
-                    Ok::<(&Vec<u32>, &Vec<u32>), DataFusionError>((
-                        &scratch.partition_starts,
-                        &scratch.partition_row_indices,
-                    ))
-                }?;
+                    indices
+                };
 
                 self.buffer_partitioned_batch_may_spill(
                     input,
                     partition_row_indices,
-                    partition_starts,
+                    &scratch.partition_starts,
                 )
                 .await?;
                 self.scratch = scratch;
@@ -423,7 +468,7 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
     async fn buffer_partitioned_batch_may_spill(
         &mut self,
         input: RecordBatch,
-        partition_row_indices: &[u32],
+        partition_row_indices: Option<&[u32]>,
         partition_starts: &[u32],
     ) -> datafusion::common::Result<()> {
         // Charge both the reservation and the data_size metric for the buffers this batch newly
@@ -434,26 +479,29 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
         let buffered_partition_idx = self.buffered_batches.len() as u32;
         self.buffered_batches.push(input);
 
-        // partition_starts conceptually slices partition_row_indices into smaller slices,
-        // each slice contains the indices of rows in input that will go into the corresponding
-        // partition. The following loop iterates over the slices and put the row indices into
-        // the indices array of the corresponding partition.
+        // `partition_starts` slices the input's rows into per-partition ranges: partition K
+        // owns rows `partition_starts[K]..partition_starts[K + 1]`. When `partition_row_indices`
+        // is `Some(indices)`, `indices[start..end]` are the source-row ids in the input for
+        // partition K (see the hash arms). When it is `None`, the source rows are the identity
+        // range `start..end` itself (whole-batch RoundRobin — one target partition owns all
+        // input rows, and the mapping is trivial so we do not materialize it).
         for (partition_id, (&start, &end)) in partition_starts
             .iter()
             .tuple_windows()
             .enumerate()
             .filter(|(_, (start, end))| start < end)
         {
-            let row_indices = &partition_row_indices[start as usize..end as usize];
-
-            // Put row indices for the current partition into the indices array of that partition.
-            // This indices array will be used for calling interleave_record_batch to produce
-            // shuffled batches.
             let indices = &mut self.partition_indices[partition_id];
             let before_size = indices.allocated_size();
-            indices.reserve(row_indices.len());
-            for row_idx in row_indices {
-                indices.push((buffered_partition_idx, *row_idx));
+            match partition_row_indices {
+                Some(row_indices) => indices.extend(
+                    row_indices[start as usize..end as usize]
+                        .iter()
+                        .map(|&row_idx| (buffered_partition_idx, row_idx)),
+                ),
+                None => {
+                    indices.extend((start..end).map(|row_idx| (buffered_partition_idx, row_idx)))
+                }
             }
             let after_size = indices.allocated_size();
             mem_growth += after_size.saturating_sub(before_size);
@@ -683,7 +731,12 @@ mod tests {
         let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
             0,
             FailingPartitionWriter::default(),
-            CometPartitioning::RoundRobin(2, 1),
+            CometPartitioning::RoundRobin(
+                2,
+                RoundRobinStrategy::HashAll {
+                    max_hash_columns: 1,
+                },
+            ),
             ShufflePartitionerMetrics::new(&metrics_set, 0),
             Arc::clone(&runtime),
             64,
@@ -788,7 +841,12 @@ mod tests {
         let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
             0,
             FailingPartitionWriter::default(),
-            CometPartitioning::RoundRobin(2, 0),
+            CometPartitioning::RoundRobin(
+                2,
+                RoundRobinStrategy::HashAll {
+                    max_hash_columns: 0,
+                },
+            ),
             ShufflePartitionerMetrics::new(&metrics_set, 0),
             Arc::clone(&runtime),
             2,
@@ -798,7 +856,7 @@ mod tests {
         .unwrap();
 
         repartitioner
-            .buffer_partitioned_batch_may_spill(batch.slice(0, 2), &[0, 1], &[0, 1, 2])
+            .buffer_partitioned_batch_may_spill(batch.slice(0, 2), Some(&[0, 1]), &[0, 1, 2])
             .await
             .unwrap();
         assert_eq!(repartitioner.reservation.size(), 0);
@@ -816,7 +874,7 @@ mod tests {
 
         repartitioner.max_buffer_bytes = None;
         repartitioner
-            .buffer_partitioned_batch_may_spill(batch.slice(2, 2), &[0, 1], &[0, 1, 2])
+            .buffer_partitioned_batch_may_spill(batch.slice(2, 2), Some(&[0, 1]), &[0, 1, 2])
             .await
             .unwrap();
         let reservation_before_failure = repartitioner.reservation.size();
@@ -912,7 +970,12 @@ mod tests {
                 let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
                     0,
                     FailingPartitionWriter::default(),
-                    CometPartitioning::RoundRobin(2, 0),
+                    CometPartitioning::RoundRobin(
+                        2,
+                        RoundRobinStrategy::HashAll {
+                            max_hash_columns: 0,
+                        },
+                    ),
                     ShufflePartitionerMetrics::new(&metrics_set, 0),
                     Arc::clone(&runtime),
                     64,
@@ -922,7 +985,7 @@ mod tests {
                 .unwrap();
                 let row_indices = (0..64).collect::<Vec<_>>();
                 repartitioner
-                    .buffer_partitioned_batch_may_spill(batch, &row_indices, &[0, 32, 64])
+                    .buffer_partitioned_batch_may_spill(batch, Some(&row_indices), &[0, 32, 64])
                     .await
                     .unwrap();
                 // The first reservation includes 512 bytes of input and 512 bytes of indices.
@@ -932,7 +995,11 @@ mod tests {
                 repartitioner.partition_writer.fail = fail;
                 repartitioner.partition_writer.consume_before_failure = consume_before_failure;
                 let result = repartitioner
-                    .buffer_partitioned_batch_may_spill(next_batch, &row_indices, &[0, 32, 64])
+                    .buffer_partitioned_batch_may_spill(
+                        next_batch,
+                        Some(&row_indices),
+                        &[0, 32, 64],
+                    )
                     .await;
                 assert_eq!(result.is_err(), fail);
                 // The indices grow by 512 bytes. Only independent input adds another 512.
diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs
index 424fa827d52..adc580b80ed 100644
--- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs
+++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs
@@ -111,6 +111,31 @@ impl<'a> PartitionedBatchIterator<'a> {
             interleave_time,
         }
     }
+
+    /// Returns the source batch when `indices[self.pos..indices_end]` selects every row of a
+    /// single source batch in natural order, so a clone is equivalent to interleaving them.
+    fn whole_source_batch(&self, indices_end: usize) -> Option<RecordBatch> {
+        let indices = &self.indices[self.pos..indices_end];
+        let &(first_batch, first_row) = indices.first()?;
+        let &(last_batch, last_row) = indices.last()?;
+        let source_rows = self.record_batches[first_batch as usize].num_rows();
+        if first_batch != last_batch
+            || first_row != 0
+            || last_row as usize + 1 != source_rows
+            || indices.len() != source_rows
+        {
+            return None;
+        }
+        // Verify the full contiguous invariant. Invariant holds by construction in
+        // `buffer_partitioned_batch_may_spill`, but the check is O(indices.len())
+        // with two integer compares per step and vectorizes easily. Two orders of
+        // magnitude cheaper than a nested-schema interleave.
+        indices
+            .iter()
+            .enumerate()
+            .all(|(i, &(b, r))| b == first_batch && r as usize == i)
+            .then(|| self.record_batches[first_batch as usize].clone())
+    }
 }
 
 impl Iterator for PartitionedBatchIterator<'_> {
@@ -122,6 +147,17 @@ impl Iterator for PartitionedBatchIterator<'_> {
         }
 
         let indices_end = std::cmp::min(self.pos + self.batch_size, self.indices.len());
+
+        // Whole-source-batch fast path: when this chunk covers every row of one source
+        // batch in natural order, cloning the source batch is equivalent to `interleave`
+        // and skips walking every column (and every nested child on wide/nested schemas).
+        // This is the vectorized path used by RoundRobinStrategy::WholeBatch, but the check
+        // is generic and fires whenever indices happen to line up this way.
+        if let Some(batch) = self.whole_source_batch(indices_end) {
+            self.pos = indices_end;
+            return Some(Ok(batch));
+        }
+
         self.chunk_scratch.clear();
         self.chunk_scratch.extend(
             self.indices[self.pos..indices_end]
diff --git a/native/shuffle/src/partitioners/traits.rs b/native/shuffle/src/partitioners/traits.rs
index 904c7d00883..e89fc3dbf62 100644
--- a/native/shuffle/src/partitioners/traits.rs
+++ b/native/shuffle/src/partitioners/traits.rs
@@ -19,7 +19,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion::common::Result;
 
 #[async_trait::async_trait]
-pub(crate) trait ShufflePartitioner: Send {
+pub trait ShufflePartitioner: Send {
     /// Insert a batch into the partitioner
     async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()>;
     /// Write the buffered shuffle data to the partition writer
diff --git a/native/shuffle/src/rss_execution_tests.rs b/native/shuffle/src/rss_execution_tests.rs
index e9269edc9e5..a3d92f79f7f 100644
--- a/native/shuffle/src/rss_execution_tests.rs
+++ b/native/shuffle/src/rss_execution_tests.rs
@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::{
-    read_ipc_compressed, CometPartitioning, CompressionCodec, PartitionOffsets,
+    read_ipc_compressed, CometPartitioning, CompressionCodec, PartitionOffsets, RoundRobinStrategy,
     ShuffleWriterDestination, ShuffleWriterExec,
 };
 use arrow::array::{Array, Int32Array, RecordBatch, RecordBatchOptions};
@@ -231,7 +231,12 @@ fn rss_multi_partition_supports_hash_range_and_round_robin() {
     for partitioning in [
         CometPartitioning::Hash(vec![expression], 4),
         CometPartitioning::RangePartitioning(ordering, 4, Arc::new(converter), boundaries),
-        CometPartitioning::RoundRobin(4, 0),
+        CometPartitioning::RoundRobin(
+            4,
+            RoundRobinStrategy::HashAll {
+                max_hash_columns: 0,
+            },
+        ),
     ] {
         let pusher = Arc::new(RecordingPusher::default());
         let execution = rss_execution(
@@ -271,7 +276,12 @@ fn rss_empty_schema_preserves_row_counts_in_partition_zero() {
     let execution = rss_execution(
         vec![batch.clone(), batch],
         schema,
-        CometPartitioning::RoundRobin(4, 0),
+        CometPartitioning::RoundRobin(
+            4,
+            RoundRobinStrategy::HashAll {
+                max_hash_columns: 0,
+            },
+        ),
         pusher.clone(),
         CompressionCodec::None,
         1024 * 1024,
@@ -299,7 +309,12 @@ fn rss_empty_schema_without_rows_does_not_push_frames() {
     let execution = rss_execution(
         vec![batch],
         schema,
-        CometPartitioning::RoundRobin(4, 0),
+        CometPartitioning::RoundRobin(
+            4,
+            RoundRobinStrategy::HashAll {
+                max_hash_columns: 0,
+            },
+        ),
         pusher.clone(),
         CompressionCodec::None,
         1024 * 1024,
diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs
index ca88509a816..91a9ca3ad6d 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -440,7 +440,7 @@ fn contextualize_shuffle_error(error: DataFusionError, phase: &str) -> DataFusio
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::{read_ipc_compressed, ShuffleBlockWriter, ShuffleCodecContext};
+    use crate::{read_ipc_compressed, RoundRobinStrategy, ShuffleBlockWriter, ShuffleCodecContext};
     use arrow::array::{Array, Int64Array, StringArray, StringBuilder};
     use arrow::datatypes::{DataType, Field, Schema};
     use arrow::record_batch::RecordBatch;
@@ -1143,7 +1143,12 @@ mod test {
                 Arc::new(row_converter),
                 owned_rows,
             ),
-            CometPartitioning::RoundRobin(num_partitions, 0),
+            CometPartitioning::RoundRobin(
+                num_partitions,
+                RoundRobinStrategy::HashAll {
+                    max_hash_columns: 0,
+                },
+            ),
         ] {
             let batches = (0..num_batches).map(|_| batch.clone()).collect::<Vec<_>>();
 
@@ -1211,7 +1216,12 @@ mod test {
                 Arc::new(DataSourceExec::new(Arc::new(
                     MemorySourceConfig::try_new(partitions, batch.schema(), None).unwrap(),
                 ))),
-                CometPartitioning::RoundRobin(num_partitions, 0),
+                CometPartitioning::RoundRobin(
+                    num_partitions,
+                    RoundRobinStrategy::HashAll {
+                        max_hash_columns: 0,
+                    },
+                ),
                 CompressionCodec::Zstd(1),
                 data_file.clone(),
                 false,
@@ -1276,8 +1286,206 @@ mod test {
         let _ = fs::remove_file("/tmp/rr_index_1.out");
     }
 
-    /// Test that batch coalescing in BufBatchWriter reduces output size by
-    /// writing fewer, larger IPC blocks instead of many small ones.
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_round_robin_batch_granular_retry_deterministic() {
+        // A retry that sees the same batches in the same order must produce byte-identical
+        // shuffle output. This is the contract the batch-granular path relies on in place
+        // of per-row content hashing.
+        use arrow::array::{Float64Array, Int64Array, StringArray, StructArray};
+        use std::fs;
+        use std::io::Read;
+
+        let num_rows = 1024;
+        let num_batches = 6;
+        let num_partitions = 8;
+
+        // A schema with a nested struct to exercise the code path this optimization targets.
+        let leaf_dt = DataType::Struct(
+            vec![
+                Field::new("a", DataType::Int64, false),
+                Field::new("b", DataType::Utf8, false),
+                Field::new("c", DataType::Float64, false),
+            ]
+            .into(),
+        );
+        let wrapped_dt = DataType::Struct(vec![Field::new("leaf", leaf_dt.clone(), false)].into());
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("scalar", DataType::Int64, false),
+            Field::new("nested", wrapped_dt.clone(), false),
+        ]));
+
+        let make_batch = |seed: i64| {
+            let scalar = Int64Array::from_iter_values((0..num_rows as i64).map(|i| i + seed));
+            let leaf = StructArray::from(vec![
+                (
+                    Arc::new(Field::new("a", DataType::Int64, false)),
+                    Arc::new(Int64Array::from_iter_values(
+                        (0..num_rows as i64).map(|i| i * 3 + seed),
+                    )) as Arc<dyn Array>,
+                ),
+                (
+                    Arc::new(Field::new("b", DataType::Utf8, false)),
+                    Arc::new(StringArray::from_iter_values(
+                        (0..num_rows).map(|i| format!("v{i}")),
+                    )) as Arc<dyn Array>,
+                ),
+                (
+                    Arc::new(Field::new("c", DataType::Float64, false)),
+                    Arc::new(Float64Array::from_iter_values(
+                        (0..num_rows).map(|i| (i as f64) * 0.5),
+                    )) as Arc<dyn Array>,
+                ),
+            ]);
+            let nested = StructArray::from(vec![(
+                Arc::new(Field::new("leaf", leaf_dt.clone(), false)),
+                Arc::new(leaf) as Arc<dyn Array>,
+            )]);
+            RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![Arc::new(scalar), Arc::new(nested)],
+            )
+            .unwrap()
+        };
+        let batches: Vec<RecordBatch> = (0..num_batches as i64).map(make_batch).collect();
+
+        let dir = tempfile::tempdir().unwrap();
+        let run = |tag: &str| -> (Vec<u8>, Vec<i64>) {
+            let data_file = dir.path().join(format!("{tag}_data.out"));
+            let partitions = std::slice::from_ref(&batches);
+            let exec = ShuffleWriterExec::try_new(
+                Arc::new(DataSourceExec::new(Arc::new(
+                    MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(),
+                ))),
+                CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::WholeBatch),
+                CompressionCodec::Zstd(1),
+                data_file.to_str().unwrap().to_string(),
+                false,
+                1024 * 1024,
+                None,
+            )
+            .unwrap();
+
+            let ctx = SessionContext::new_with_config_rt(
+                SessionConfig::new(),
+                Arc::new(
+                    RuntimeEnvBuilder::new()
+                        .with_memory_limit(10 * 1024 * 1024, 1.0)
+                        .build()
+                        .unwrap(),
+                ),
+            );
+            let stream = exec.execute(0, ctx.task_ctx()).unwrap();
+            Runtime::new().unwrap().block_on(collect(stream)).unwrap();
+
+            let mut data = Vec::new();
+            fs::File::open(&data_file)
+                .unwrap()
+                .read_to_end(&mut data)
+                .unwrap();
+            let offsets = exec
+                .partition_offsets()
+                .expect("local destination publishes offsets")
+                .get()
+                .expect("writer published its partition offsets")
+                .to_vec();
+            (data, offsets)
+        };
+
+        let (data_a, offsets_a) = run("attempt0");
+        let (data_b, offsets_b) = run("attempt1");
+        assert_eq!(
+            data_a, data_b,
+            "batch-granular round robin must produce byte-identical shuffle data on retry"
+        );
+        assert_eq!(
+            offsets_a, offsets_b,
+            "batch-granular round robin must produce identical partition offsets on retry"
+        );
+
+        // Sanity: the assignment must have distributed the 6 batches across the partitions
+        // (specifically, the offsets must show non-trivial partition sizes, not all in one
+        // bucket).
+        assert_eq!(offsets_a.len(), num_partitions + 1);
+        let sizes: Vec<i64> = offsets_a.windows(2).map(|w| w[1] - w[0]).collect();
+        let nonempty = sizes.iter().filter(|&&s| s > 0).count();
+        assert!(
+            nonempty >= 2,
+            "batch-granular round robin should populate at least 2 partitions across 6 batches; got sizes {sizes:?}"
+        );
+
+        // And every batch's rows must show up somewhere in the output.
+        let total_rows: usize = read_all_ipc_batches(&data_a)
+            .iter()
+            .map(|b| b.num_rows())
+            .sum();
+        assert_eq!(total_rows, num_rows * num_batches, "no rows may be dropped");
+    }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_round_robin_batch_granular_whole_batch_per_partition() {
+        // Each input batch must land entirely on a single output partition. Reading back
+        // the IPC blocks from any one partition should give whole batches (num_rows each),
+        // not row-level slices.
+        use std::fs;
+
+        let num_rows = 500;
+        let num_batches = 4;
+        let num_partitions = 8;
+
+        let batch = create_batch(num_rows);
+        let batches: Vec<RecordBatch> = (0..num_batches).map(|_| batch.clone()).collect();
+
+        let dir = tempfile::tempdir().unwrap();
+        let data_file = dir.path().join("data.out");
+
+        let partitions = std::slice::from_ref(&batches);
+        let exec = ShuffleWriterExec::try_new(
+            Arc::new(DataSourceExec::new(Arc::new(
+                MemorySourceConfig::try_new(partitions, batch.schema(), None).unwrap(),
+            ))),
+            CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::WholeBatch),
+            CompressionCodec::Zstd(1),
+            data_file.to_str().unwrap().to_string(),
+            false,
+            1024 * 1024,
+            None,
+        )
+        .unwrap();
+
+        let ctx = SessionContext::new_with_config_rt(
+            SessionConfig::new(),
+            Arc::new(RuntimeEnvBuilder::new().build().unwrap()),
+        );
+        let stream = exec.execute(0, ctx.task_ctx()).unwrap();
+        Runtime::new().unwrap().block_on(collect(stream)).unwrap();
+
+        let data = fs::read(&data_file).unwrap();
+        let offsets = exec
+            .partition_offsets()
+            .expect("local destination publishes offsets")
+            .get()
+            .expect("writer published its partition offsets");
+        assert_eq!(offsets.len(), num_partitions + 1);
+
+        for p in 0..num_partitions {
+            let start = offsets[p] as usize;
+            let end = offsets[p + 1] as usize;
+            if start == end {
+                continue;
+            }
+            let batches = read_all_ipc_batches(&data[start..end]);
+            for b in &batches {
+                assert_eq!(
+                    b.num_rows(),
+                    num_rows,
+                    "each shuffle block in partition {p} must contain a whole input batch"
+                );
+            }
+        }
+    }
+
     #[test]
     #[cfg_attr(miri, ignore)]
     fn test_batch_coalescing_reduces_size() {
@@ -1615,7 +1823,12 @@ mod test {
             Arc::new(DataSourceExec::new(Arc::new(
                 MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(),
             ))),
-            CometPartitioning::RoundRobin(num_partitions, 0),
+            CometPartitioning::RoundRobin(
+                num_partitions,
+                RoundRobinStrategy::HashAll {
+                    max_hash_columns: 0,
+                },
+            ),
             CompressionCodec::Zstd(1),
             data_file.to_str().unwrap().to_string(),
             false,
@@ -1699,7 +1912,12 @@ mod test {
             Arc::new(DataSourceExec::new(Arc::new(
                 MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(),
             ))),
-            CometPartitioning::RoundRobin(num_partitions, 0),
+            CometPartitioning::RoundRobin(
+                num_partitions,
+                RoundRobinStrategy::HashAll {
+                    max_hash_columns: 0,
+                },
+            ),
             CompressionCodec::Zstd(1),
             data_file.to_str().unwrap().to_string(),
             false,
diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs
index d76ff23eff0..48775f00302 100644
--- a/native/shuffle/src/writers/local/local_partition_writer.rs
+++ b/native/shuffle/src/writers/local/local_partition_writer.rs
@@ -77,7 +77,7 @@ enum DataOutput {
 /// Writes shuffle output to a single data file and publishes the byte offset where
 /// each partition begins through [`PartitionOffsets`]. See [`DataOutput`] for how the
 /// single- and multi-partition modes differ.
-pub(crate) struct LocalPartitionWriter {
+pub struct LocalPartitionWriter {
     partition_offsets: Arc<PartitionOffsets>,
     data_output: DataOutput,
     /// Compression state shared by every block this task writes; the per-partition
@@ -97,7 +97,7 @@ pub(crate) struct LocalPartitionWriter {
 }
 
 impl LocalPartitionWriter {
-    pub(crate) fn try_new(
+    pub fn try_new(
         output_data_file: String,
         partition_offsets: Arc<PartitionOffsets>,
         shuffle_block_writer: ShuffleBlockWriter,
diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs
index 4586e46c25b..350a1fd3932 100644
--- a/native/shuffle/src/writers/mod.rs
+++ b/native/shuffle/src/writers/mod.rs
@@ -24,7 +24,7 @@ mod shuffle_block_writer;
 
 pub(crate) use buf_batch_writer::BufBatchWriter;
 pub(crate) use checksum::Checksum;
-pub(crate) use local::local_partition_writer::LocalPartitionWriter;
+pub use local::local_partition_writer::LocalPartitionWriter;
 pub(crate) use partition_writer::PartitionWriter;
 pub(crate) use rss::rss_partition_writer::RssPartitionWriter;
 pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter};
diff --git a/native/shuffle/src/writers/partition_writer.rs b/native/shuffle/src/writers/partition_writer.rs
index 0fcec17c00d..fdba82cc27e 100644
--- a/native/shuffle/src/writers/partition_writer.rs
+++ b/native/shuffle/src/writers/partition_writer.rs
@@ -32,7 +32,7 @@ use arrow::record_batch::RecordBatch;
 /// ascending id order, then a single [`finish_all`](PartitionWriter::finish_all).
 ///
 /// [`LocalPartitionWriter`]: crate::writers::local::local_partition_writer::LocalPartitionWriter
-pub(crate) trait PartitionWriter: Send {
+pub trait PartitionWriter: Send {
     /// Stages the batches from `iter` for partition `pid` without finalizing it.
     ///
     /// Used to stream single-partition output and to stage multi-partition
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index d371f1ba44c..b541787d62d 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -475,6 +475,18 @@ object CometConf extends ShimCometConf {
         "The maximum number of columns to hash for round robin partitioning must be non-negative.")
       .createWithDefault(0)
 
+  val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR: ConfigEntry[Boolean] =
+    conf("spark.comet.shuffle.native.partitioning.roundrobin.batchGranular")
+      .category(CATEGORY_SHUFFLE)
+      .doc(
+        "When true, Comet's native round-robin shuffle assigns each RecordBatch as a whole " +
+          "to one output partition instead of hashing every column of every row. This is much " +
+          "cheaper on wide/nested schemas. Distribution is quasi-even at batch granularity. " +
+          "Retry-safe only when the upstream operator emits the same batches in the same order " +
+          "under retry (Comet's Parquet scan and other order-preserving operators satisfy this).")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.convertFromSparkPlan.enabled")
       .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.convertFromSparkPlan.enabled")
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index fce4291deb6..947701e4a1d 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -421,6 +421,8 @@ class CometNativeShuffleWriter[K, V](
         partitioning.setNumPartitions(effectivePartitionCount)
         partitioning.setMaxHashColumns(
           CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get())
+        partitioning.setBatchGranular(
+          CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.get())
 
         val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder()
         shuffleWriterBuilder.setPartitioning(

From 3aba03dc03f3f0656243c0adfe4a34d8c8d904e1 Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Mon, 21 Sep 2026 09:05:21 -0700
Subject: [PATCH 2/4] feat: fail the job when a positional round-robin shuffle
 task is retried

`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>
---
 .../contributor-guide/native_shuffle.md       |  54 +++-
 .../scala/org/apache/comet/CometConf.scala    |  19 ++
 .../shuffle/CometNativeShuffleInputRDD.scala  |  45 +++-
 .../shuffle/CometNativeShuffleWriter.scala    |  59 +++-
 .../shuffle/CometShuffleExchangeExec.scala    |  26 +-
 .../CometNativeRoundRobinRetrySuite.scala     | 252 ++++++++++++++++++
 6 files changed, 448 insertions(+), 7 deletions(-)
 create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala

diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md
index 8d84aa2f6b8..bddce77791d 100644
--- a/docs/source/contributor-guide/native_shuffle.md
+++ b/docs/source/contributor-guide/native_shuffle.md
@@ -305,7 +305,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 +325,54 @@ 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`, where the
+counter is seeded with the input partition id so concurrent mappers do not all target partition 0
+first. `partition_row_indices` is left unmaterialized and the flush clones the source batch instead
+of interleaving it, so the per-batch cost is one modulo, two `resize` calls, and some `Arc` bumps.
+
+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:
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index b541787d62d..d972f2fe3e0 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -487,6 +487,25 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_FAIL_ON_RETRY: ConfigEntry[Boolean] =
+    conf("spark.comet.shuffle.native.partitioning.roundrobin.batchGranular.failOnRetry")
+      .category(CATEGORY_SHUFFLE)
+      .doc(
+        "Only applies when " +
+          "spark.comet.shuffle.native.partitioning.roundrobin.batchGranular is true. " +
+          "Batch-granular round robin places rows by position rather than by content, so a " +
+          "re-executed map task can send rows to different output partitions than the attempt " +
+          "it replaces, silently dropping and duplicating rows once any of that output has " +
+          "been fetched. When this is true, a native round-robin map task that runs as a retry " +
+          "(a task attempt after the first, or any attempt of a re-submitted stage, which is " +
+          "what an executor loss produces) fails immediately instead of writing output, which " +
+          "fails the stage and therefore the job. Note that speculative attempts also count as " +
+          "retries. Set to false to keep Spark's normal fault tolerance, in which case safety " +
+          "rests on the whole stage being rolled back, which Comet requests by declaring the " +
+          "shuffle input RDD INDETERMINATE whenever its own input is not DETERMINATE.")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.convertFromSparkPlan.enabled")
       .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.convertFromSparkPlan.enabled")
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
index bced55b34df..e7e2abbf1a0 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
@@ -20,7 +20,7 @@
 package org.apache.spark.sql.comet.execution.shuffle
 
 import org.apache.spark._
-import org.apache.spark.rdd.RDD
+import org.apache.spark.rdd.{DeterministicLevel, RDD}
 import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
@@ -33,6 +33,10 @@ import org.apache.comet.CometShuffleBlockIterator
  * [[CometNativeShuffleInputIterator]]. The iterator reports `hasNext = false`;
  * [[CometNativeShuffleWriter]] downcasts it and reads those slots directly to drive the unified
  * `ShuffleWriter(child = childNativeOp)` plan.
+ *
+ * @param positionalRoundRobin
+ *   whether the writer fed by this RDD will place rows by position instead of by content; see
+ *   [[CometShuffleExchangeExec.usesPositionalRoundRobin]] and `getOutputDeterministicLevel`.
  */
 private[shuffle] class CometNativeShuffleInputRDD(
     sc: SparkContext,
@@ -40,7 +44,8 @@ private[shuffle] class CometNativeShuffleInputRDD(
     numPartitionsParam: Int,
     shuffleScanIndices: Set[Int],
     spillMetricNode: CometMetricNode,
-    @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty)
+    @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty,
+    positionalRoundRobin: Boolean = false)
     extends RDD[Product2[Int, ColumnarBatch]](
       sc,
       inputRDDs.map(rdd => new OneToOneDependency(rdd))) {
@@ -57,7 +62,41 @@ private[shuffle] class CometNativeShuffleInputRDD(
       numPartitionsParam,
       shuffleScanIndices,
       spillMetricNode,
-      perPartitionByKey)
+      perPartitionByKey,
+      positionalRoundRobin)
+
+  /**
+   * Batch-granular round robin assigns each Arrow batch to an output partition by position, so
+   * the placement of a row depends on the order and the framing of the batches the upstream
+   * operator produced, not on the row itself. Re-running one map task against a differently
+   * ordered or differently framed input therefore sends rows somewhere else, which duplicates and
+   * drops rows once any of the replaced output has already been fetched.
+   *
+   * Spark states that risk declaratively rather than defending against it per-operator: a
+   * round-robin repartition is wrapped in a `MapPartitionsRDD` with `isOrderSensitive = true`
+   * (Comet's own JVM path does this in `prepareJVMShuffleDependency`), and that RDD reports
+   * `INDETERMINATE` whenever its parent is `UNORDERED`. The DAGScheduler then rolls the whole
+   * stage back, or aborts the job when a result stage has already consumed output, instead of
+   * re-running a single task. The native path has no `MapPartitionsRDD` to carry the flag, so
+   * apply the same rule here.
+   *
+   * The parent level does the discriminating, exactly as it does for Spark. A plain scan is
+   * `DETERMINATE`: it replays identically, so positional assignment is reproducible and ordinary
+   * per-task retry stays cheap. Anything below another exchange is `UNORDERED`, because reduce
+   * tasks see shuffle blocks in arrival order, and that is when positional assignment stops being
+   * reproducible.
+   *
+   * Content-hash round robin (the default) needs none of this: it places rows by hash, so its
+   * output is a pure function of the rows regardless of how they arrive.
+   */
+  override protected def getOutputDeterministicLevel: DeterministicLevel.Value = {
+    val inheritedLevel = super.getOutputDeterministicLevel
+    if (positionalRoundRobin && inheritedLevel != DeterministicLevel.DETERMINATE) {
+      DeterministicLevel.INDETERMINATE
+    } else {
+      inheritedLevel
+    }
+  }
 
   override protected def getPartitions: Array[Partition] =
     (0 until numPartitionsParam).map { i =>
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index 947701e4a1d..7f6bf800645 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -38,7 +38,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.types.StructField
 import org.apache.spark.util.{ThreadUtils, Utils}
 
-import org.apache.comet.{CometConf, CometExecIterator, CometShuffleSizeLimitException}
+import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, CometShuffleSizeLimitException}
 import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde}
 import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator}
 import org.apache.comet.serde.operator.schema2Proto
@@ -108,7 +108,64 @@ class CometNativeShuffleWriter[K, V](
     }
   }
 
+  /**
+   * Refuse to re-execute a map task whose round-robin placement is positional.
+   *
+   * With `RoundRobinStrategy::WholeBatch` an input batch's output partition comes from a per-task
+   * counter, so which partition a row lands in is a function of the order and the framing of the
+   * batches the upstream operator produced rather than of the row itself. An attempt that frames
+   * or orders its input even slightly differently writes a different partitioning of the same
+   * rows. Once any consumer has fetched the output this attempt replaces, the reduce side gets
+   * some rows twice and some not at all, and nothing downstream can detect it.
+   *
+   * [[CometNativeShuffleInputRDD.getOutputDeterministicLevel]] already asks the DAGScheduler to
+   * roll the whole stage back rather than re-run one task, which is Spark's own answer to this.
+   * That answer is only as good as the parent's determinism level being an accurate description
+   * of what the upstream operator replays, so while the strategy is opt-in and unproven the
+   * default is to not rely on it and fail here instead.
+   *
+   * Both counters are needed and neither subsumes the other. `attemptNumber` covers a task re-run
+   * inside the current stage attempt: a task-level failure, speculation, or an executor lost
+   * while the stage was still running. `stageAttemptNumber` covers a re-submitted stage, which is
+   * what a fetch failure against a dead executor's shuffle output produces, and which also
+   * re-runs map tasks that had already succeeded elsewhere.
+   *
+   * There is no Spark API for failing an application from inside a task, and the two exceptions
+   * Spark declines to retry are matched by class name (`NotSerializableException`,
+   * `TaskOutputFileAlreadyExistException`), so this throws an ordinary exception. Because the
+   * condition is sticky - a later attempt only has a higher `attemptNumber` - every remaining
+   * attempt fails here too and the task set aborts after `spark.task.maxFailures`, failing the
+   * stage and the job. Each of those attempts fails before doing any work.
+   */
+  private def failIfRetryingPositionalRoundRobin(): Unit = {
+    if (!CometShuffleExchangeExec.usesPositionalRoundRobin(outputPartitioning)) return
+    if (!CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_FAIL_ON_RETRY.get()) return
+    // `context` is null only in unit tests that drive the writer directly.
+    val taskContext = context
+    if (taskContext == null) return
+    val taskAttempt = taskContext.attemptNumber()
+    val stageAttempt = taskContext.stageAttemptNumber()
+    if (taskAttempt == 0 && stageAttempt == 0) return
+
+    val batchGranularKey =
+      CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.key
+    val failOnRetryKey =
+      CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_FAIL_ON_RETRY.key
+    throw new CometRuntimeException(
+      s"Refusing to re-execute a Comet native round-robin shuffle map task: shuffle $shuffleId, " +
+        s"map partition ${taskContext.partitionId()}, stage ${taskContext.stageId()} " +
+        s"(attempt $stageAttempt), task attempt $taskAttempt. `$batchGranularKey` is enabled, " +
+        "so each Arrow batch is assigned to an output partition by position rather than by " +
+        "hashing its rows, and this attempt may place rows differently than the attempt it " +
+        "replaces, duplicating and dropping rows on the reduce side without any error. Failing " +
+        s"the job instead. Set `$batchGranularKey` to false to use content-hash round robin, " +
+        s"which is safe to re-execute, or set `$failOnRetryKey` to false to allow retries and " +
+        "rely on Spark rolling the stage back.")
+  }
+
   private def writeInternal(inputs: Iterator[Product2[K, V]]): Unit = {
+    failIfRetryingPositionalRoundRobin()
+
     val localOutput = if (remoteDestination.isEmpty) {
       val resolver =
         SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index 2a539643755..b490b7164c4 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -127,7 +127,8 @@ case class CometShuffleExchangeExec(
           ctx.numPartitions,
           ctx.shuffleScanIndices,
           CometMetricNode(metrics, Seq(nativeChildMetricNode)),
-          ctx.perPartitionByKey)
+          ctx.perPartitionByKey,
+          CometShuffleExchangeExec.usesPositionalRoundRobin(outputPartitioning))
       case None =>
         // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets
         // built via the convenience overload below; we just need a real RDD of batches.
@@ -298,6 +299,26 @@ object CometShuffleExchangeExec
     if (shuffleSupported(op).isDefined) Compatible() else Unsupported()
   }
 
+  /**
+   * True when this exchange will run the native round-robin writer in its positional
+   * (batch-granular) mode, where an input `RecordBatch` is assigned whole to one output partition
+   * by a per-task counter rather than by hashing its rows.
+   *
+   * Positional assignment is what makes the strategy cheap and also what makes it unsafe to
+   * re-execute: the placement depends on the order and the framing of the batches the upstream
+   * operator hands over, not on the rows themselves. Two call sites need the same answer, so they
+   * share this predicate rather than each reading the config: [[CometNativeShuffleInputRDD]],
+   * which declares its output indeterminate so the DAGScheduler rolls a stage back rather than
+   * re-running one task into partially consumed output, and [[CometNativeShuffleWriter]], which
+   * refuses to run at all on a retry.
+   *
+   * The equivalent decision on the native side is `PhysicalPlanner::create_partitioning` turning
+   * `batch_granular` into `RoundRobinStrategy::WholeBatch`; keep the two in step.
+   */
+  def usesPositionalRoundRobin(outputPartitioning: Partitioning): Boolean =
+    outputPartitioning.isInstanceOf[RoundRobinPartitioning] &&
+      CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.get()
+
   override def createExec(
       nativeOp: OperatorOuterClass.Operator,
       op: ShuffleExchangeExec): CometNativeExec = {
@@ -777,7 +798,8 @@ object CometShuffleExchangeExec
       Seq(streamRDD),
       rdd.getNumPartitions,
       shuffleScanIndices = Set.empty,
-      spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)))
+      spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)),
+      positionalRoundRobin = usesPositionalRoundRobin(outputPartitioning))
 
     val ctx = NativeExecContext(
       inputs = Seq(streamRDD),
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala
new file mode 100644
index 00000000000..dd7a36045ed
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala
@@ -0,0 +1,252 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.shuffle
+
+import java.util.Properties
+
+import org.apache.spark.{Partition, TaskContext, TaskContextImpl}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.rdd.{DeterministicLevel, RDD}
+import org.apache.spark.sql.{CometTestBase, DataFrame}
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition}
+import org.apache.spark.sql.comet.CometMetricNode
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.{CometConf, CometRuntimeException}
+
+/**
+ * Retry safety for the batch-granular (positional) native round-robin strategy.
+ *
+ * That strategy assigns a whole Arrow batch to an output partition from a per-task counter, so a
+ * row's destination depends on the order and the framing of the batches the upstream operator
+ * produced rather than on the row itself. Re-executing a map task can therefore write a different
+ * partitioning of the same rows, which duplicates and drops rows on the reduce side once any of
+ * the replaced output has been fetched. Two defences are covered here: declaring the shuffle
+ * input RDD indeterminate so the DAGScheduler rolls the stage back instead of re-running one
+ * task, and refusing to run a retried map task at all.
+ *
+ * Lives in the `execution.shuffle` package so it can construct the `private[shuffle]`
+ * [[CometNativeShuffleInputRDD]] and the `private[spark]` [[TaskContextImpl]] directly.
+ */
+class CometNativeRoundRobinRetrySuite extends CometTestBase {
+
+  private val batchGranularKey =
+    CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.key
+  private val failOnRetryKey =
+    CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_FAIL_ON_RETRY.key
+
+  /** An input RDD that reports exactly `level`, standing in for a real upstream subtree. */
+  private def parentWithLevel(level: DeterministicLevel.Value): RDD[AnyRef] =
+    new RDD[AnyRef](spark.sparkContext, Nil) {
+      override protected def getOutputDeterministicLevel: DeterministicLevel.Value = level
+      override protected def getPartitions: Array[Partition] = Array.empty
+      override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] =
+        Iterator.empty
+    }
+
+  private def shuffleInput(
+      parent: RDD[AnyRef],
+      positionalRoundRobin: Boolean): CometNativeShuffleInputRDD =
+    new CometNativeShuffleInputRDD(
+      spark.sparkContext,
+      Seq(parent),
+      0,
+      Set.empty,
+      CometMetricNode(Map.empty),
+      positionalRoundRobin = positionalRoundRobin)
+
+  /** Mirrors `TaskContext.empty()`, which hard-codes both attempt numbers to zero. */
+  private def taskContextFor(stageAttempt: Int, taskAttempt: Int): TaskContext =
+    new TaskContextImpl(
+      0,
+      stageAttempt,
+      0,
+      0L,
+      taskAttempt,
+      1,
+      null,
+      new Properties,
+      null,
+      TaskMetrics.empty,
+      1)
+
+  private def writerFor(
+      outputPartitioning: Partitioning,
+      taskContext: TaskContext): CometNativeShuffleWriter[Int, Any] =
+    new CometNativeShuffleWriter[Int, Any](
+      NativeShuffleSpec(null, CometMetricNode(Map.empty), null),
+      outputPartitioning,
+      Nil,
+      Map.empty,
+      4,
+      0,
+      0L,
+      taskContext,
+      null)
+
+  /**
+   * `write` is driven with a plain iterator rather than a [[CometNativeShuffleInputIterator]], so
+   * anything that gets past the retry guard fails soon after on that. Returning the message lets
+   * a caller assert on which of the two happened without depending on the unrelated failure's
+   * type.
+   */
+  private def writeFailureMessage(writer: CometNativeShuffleWriter[Int, Any]): String = {
+    val failure = intercept[Throwable](writer.write(Iterator.empty))
+    Option(failure.getMessage).getOrElse("")
+  }
+
+  private val refusalPrefix = "Refusing to re-execute a Comet native round-robin"
+
+  test("positional round robin is claimed only for round robin with the config enabled") {
+    val partitionings = Seq(
+      RoundRobinPartitioning(4),
+      SinglePartition,
+      HashPartitioning(Seq(Literal(1)), 4),
+      RangePartitioning(Nil, 4))
+    withSQLConf(batchGranularKey -> "true") {
+      val claimed = partitionings.filter(CometShuffleExchangeExec.usesPositionalRoundRobin)
+      assert(claimed == Seq(RoundRobinPartitioning(4)))
+    }
+    withSQLConf(batchGranularKey -> "false") {
+      assert(!partitionings.exists(CometShuffleExchangeExec.usesPositionalRoundRobin))
+    }
+  }
+
+  test("positional round robin declares indeterminate output unless its parent is determinate") {
+    // A determinate parent (a plain scan) replays identically, so positional assignment is
+    // reproducible and per-task retry stays cheap. Anything below another exchange is unordered,
+    // which is where Spark's own round robin flips to indeterminate too.
+    Seq(
+      DeterministicLevel.DETERMINATE -> DeterministicLevel.DETERMINATE,
+      DeterministicLevel.UNORDERED -> DeterministicLevel.INDETERMINATE,
+      DeterministicLevel.INDETERMINATE -> DeterministicLevel.INDETERMINATE).foreach {
+      case (parentLevel, expected) =>
+        val input = shuffleInput(parentWithLevel(parentLevel), positionalRoundRobin = true)
+        assert(input.outputDeterministicLevel == expected, s"parent was $parentLevel")
+        assert(
+          input.copyForLocalShuffle().outputDeterministicLevel == expected,
+          s"local fallback lost the declaration for parent $parentLevel")
+    }
+  }
+
+  test("content-hash round robin keeps inheriting its parent's determinism") {
+    // Hash placement is a pure function of the rows, so nothing here should be indeterminate on
+    // its own account. This is the default path and must stay as retryable as it is today.
+    Seq(
+      DeterministicLevel.DETERMINATE,
+      DeterministicLevel.UNORDERED,
+      DeterministicLevel.INDETERMINATE).foreach { level =>
+      val input = shuffleInput(parentWithLevel(level), positionalRoundRobin = false)
+      assert(input.outputDeterministicLevel == level)
+    }
+  }
+
+  test("positional round robin refuses to run a retried map task") {
+    withSQLConf(batchGranularKey -> "true") {
+      // A task re-run inside the current stage attempt, and a re-submitted stage, which is what
+      // losing an executor produces. Neither counter subsumes the other.
+      Seq(("task attempt", 0, 1), ("stage attempt", 1, 0), ("both", 2, 3)).foreach {
+        case (label, stageAttempt, taskAttempt) =>
+          val writer =
+            writerFor(RoundRobinPartitioning(4), taskContextFor(stageAttempt, taskAttempt))
+          val failure = intercept[CometRuntimeException](writer.write(Iterator.empty))
+          assert(failure.getMessage.startsWith(refusalPrefix), label)
+          // The message has to say which knob to turn; an operator hitting this mid-job has no
+          // other signal that positional placement is what failed them.
+          assert(failure.getMessage.contains(batchGranularKey), label)
+          assert(failure.getMessage.contains(failOnRetryKey), label)
+      }
+    }
+  }
+
+  test("positional round robin lets a first attempt through") {
+    withSQLConf(batchGranularKey -> "true") {
+      val writer = writerFor(RoundRobinPartitioning(4), taskContextFor(0, 0))
+      assert(!writeFailureMessage(writer).startsWith(refusalPrefix))
+    }
+  }
+
+  test("the retry guard is scoped to positional round robin") {
+    // Hash and range partitioning place rows by content, and content-hash round robin does too,
+    // so a retry of any of those is safe and must not be turned into a job failure.
+    withSQLConf(batchGranularKey -> "true") {
+      Seq(SinglePartition, HashPartitioning(Seq(Literal(1)), 4), RangePartitioning(Nil, 4))
+        .foreach { partitioning =>
+          val writer = writerFor(partitioning, taskContextFor(1, 1))
+          assert(
+            !writeFailureMessage(writer).startsWith(refusalPrefix),
+            s"$partitioning should not be treated as positional round robin")
+        }
+    }
+    withSQLConf(batchGranularKey -> "false") {
+      val writer = writerFor(RoundRobinPartitioning(4), taskContextFor(1, 1))
+      assert(!writeFailureMessage(writer).startsWith(refusalPrefix))
+    }
+  }
+
+  test("failOnRetry=false hands retry handling back to the DAGScheduler") {
+    withSQLConf(batchGranularKey -> "true", failOnRetryKey -> "false") {
+      val writer = writerFor(RoundRobinPartitioning(4), taskContextFor(1, 1))
+      assert(!writeFailureMessage(writer).startsWith(refusalPrefix))
+    }
+  }
+
+  /**
+   * The determinism level the DAGScheduler would read off a planned round-robin exchange.
+   * `ShuffleMapStage.rdd` is the dependency's RDD, so this is exactly what
+   * `Stage.isIndeterminate` sees.
+   */
+  private def roundRobinStageLevel(df: DataFrame): DeterministicLevel.Value = {
+    val plan = df.queryExecution.executedPlan
+    val exchanges = plan.collect {
+      case exchange: CometShuffleExchangeExec
+          if exchange.outputPartitioning.isInstanceOf[RoundRobinPartitioning] =>
+        exchange
+    }
+    assert(exchanges.size == 1, s"expected exactly one round robin exchange in\n$plan")
+    assert(exchanges.head.shuffleType == CometNativeShuffle, s"not native shuffle in\n$plan")
+    exchanges.head.shuffleDependency.rdd.outputDeterministicLevel
+  }
+
+  test("a planned positional round robin exchange carries the declaration to its stage") {
+    withSQLConf(
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+      CometConf.COMET_SHUFFLE_MODE.key -> "native",
+      CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true",
+      batchGranularKey -> "true",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+      withParquetTable((0 until 100).map(i => (i % 10, i)), "tbl") {
+        // Straight over a scan. The scan replays identically, so positional assignment is
+        // reproducible and per-task retry stays as cheap as it is for hash placement.
+        assert(
+          roundRobinStageLevel(sql("SELECT * FROM tbl").repartition(4)) ==
+            DeterministicLevel.DETERMINATE)
+
+        // Below an aggregate's exchange. The reduce side sees shuffle blocks in arrival order, so
+        // a replay can frame the batches differently and the stage has to be rolled back whole.
+        assert(
+          roundRobinStageLevel(sql("SELECT _1, sum(_2) FROM tbl GROUP BY _1").repartition(4)) ==
+            DeterministicLevel.INDETERMINATE)
+      }
+    }
+  }
+}

From be42b31612fdcc4a42f38e0055f37ea25a79494a Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Mon, 21 Sep 2026 10:45:23 -0700
Subject: [PATCH 3/4] vectorized RR partitioning

---
 .github/workflows/pr_build_linux.yml          |   1 +
 .github/workflows/pr_build_macos.yml          |   1 +
 .../contributor-guide/native_shuffle.md       |  62 +++---
 native/core/src/execution/planner.rs          |  12 +-
 native/shuffle/Cargo.toml                     |   3 -
 native/shuffle/README.md                      |  30 +--
 native/shuffle/benches/shuffle_writer.rs      | 186 +++++-------------
 native/shuffle/src/bin/shuffle_bench.rs       |  13 +-
 native/shuffle/src/comet_partitioning.rs      |  28 ++-
 .../src/partitioners/multi_partition.rs       |  83 ++++----
 .../partitioned_batch_iterator.rs             |  29 ++-
 native/shuffle/src/rss_execution_tests.rs     |  21 +-
 native/shuffle/src/shuffle_writer.rs          | 103 +++++++---
 .../scala/org/apache/comet/CometConf.scala    |  29 +--
 .../shuffle/CometNativeShuffleInputRDD.scala  |  28 +--
 .../shuffle/CometNativeShuffleWriter.scala    |  23 +--
 .../shuffle/CometShuffleExchangeExec.scala    |  19 +-
 .../CometNativeRoundRobinRetrySuite.scala     | 108 +++-------
 .../CometNativeShuffleInputRDDSuite.scala     |  52 +++--
 19 files changed, 375 insertions(+), 456 deletions(-)

diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index a9c184a20b0..9c6c4881c39 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.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
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index b47ed5a46f6..b5d7f8816e3 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.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
diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md
index bddce77791d..e040c87a77f 100644
--- a/docs/source/contributor-guide/native_shuffle.md
+++ b/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,7 +163,6 @@ 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)
 
@@ -173,7 +170,6 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet-
    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
@@ -333,10 +328,27 @@ are hashed. `0`, the default, hashes all of them.
 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`, where the
-counter is seeded with the input partition id so concurrent mappers do not all target partition 0
-first. `partition_row_indices` is left unmaterialized and the flush clones the source batch instead
-of interleaving it, so the per-batch cost is one modulo, two `resize` calls, and some `Arc` bumps.
+`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.
@@ -422,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
 
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 22ca8c02d6d..e569bcf0615 100644
--- a/native/core/src/execution/planner.rs
+++ b/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;
@@ -3653,11 +3653,17 @@ impl PhysicalPlanner {
             PartitioningStruct::SinglePartition(_) => Ok(CometPartitioning::SinglePartition),
             PartitioningStruct::RoundRobinPartition(rr_partition) => {
                 let strategy = if rr_partition.batch_granular {
-                    crate::execution::shuffle::RoundRobinStrategy::WholeBatch
+                    // 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 {
                     // Treat negative max_hash_columns as 0 (no limit).
                     let max_hash_columns = rr_partition.max_hash_columns.max(0) as usize;
-                    crate::execution::shuffle::RoundRobinStrategy::HashAll { max_hash_columns }
+                    RoundRobinStrategy::HashAll { max_hash_columns }
                 };
                 Ok(CometPartitioning::RoundRobin(
                     rr_partition.num_partitions as usize,
diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml
index 6d6d85bd2c9..f0ed22ad730 100644
--- a/native/shuffle/Cargo.toml
+++ b/native/shuffle/Cargo.toml
@@ -57,9 +57,6 @@ zstd = { version = "0.13.3", features = ["experimental"] }
 criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] }
 datafusion = { workspace = true, features = ["parquet_encryption", "sql"] }
 itertools = "0.15.0"
-# The shuffle_writer bench reads a real on-disk parquet dataset; the optional dependency above
-# is gated behind shuffle-bench, which benches do not activate.
-parquet = { workspace = true }
 tempfile = "3.26.0"
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 
diff --git a/native/shuffle/README.md b/native/shuffle/README.md
index 0f53604fa37..2a47f643ff2 100644
--- a/native/shuffle/README.md
+++ b/native/shuffle/README.md
@@ -41,21 +41,21 @@ cargo run --release --features shuffle-bench --bin shuffle_bench -- \
 
 ### Options
 
-| Option                | Default                    | Description                                            |
-| --------------------- | -------------------------- | ------------------------------------------------------ |
-| `--input`             | _(required)_               | Path to a Parquet file or directory of Parquet files   |
-| `--partitions`        | `200`                      | Number of output shuffle partitions                    |
-| `--partitioning`      | `hash`                     | Partitioning scheme: `hash`, `single`, `round-robin`   |
-| `--hash-columns`      | `0`                        | Comma-separated column indices to hash on (e.g. `0,3`) |
-| `--codec`             | `lz4`                      | Compression codec: `none`, `lz4`, `zstd`, `snappy`     |
-| `--zstd-level`        | `1`                        | Zstd compression level (1–22)                          |
-| `--batch-size`        | `8192`                     | Batch size for reading Parquet data                    |
-| `--memory-limit`      | _(none)_                   | Memory limit in bytes; triggers spilling when exceeded |
-| `--write-buffer-size` | `1048576`                  | Write buffer size in bytes                             |
-| `--limit`             | `0`                        | Limit rows processed per iteration (0 = no limit)      |
-| `--iterations`        | `1`                        | Number of timed iterations                             |
-| `--warmup`            | `0`                        | Number of warmup iterations before timing              |
-| `--output-dir`        | `/tmp/comet_shuffle_bench` | Directory for temporary shuffle output files           |
+| Option                | Default                    | Description                                                                     |
+| --------------------- | -------------------------- | ------------------------------------------------------------------------------- |
+| `--input`             | _(required)_               | Path to a Parquet file or directory of Parquet files                            |
+| `--partitions`        | `200`                      | Number of output shuffle partitions                                             |
+| `--partitioning`      | `hash`                     | Partitioning scheme: `hash`, `single`, `round-robin`, `round-robin-whole-batch` |
+| `--hash-columns`      | `0`                        | Comma-separated column indices to hash on (e.g. `0,3`)                          |
+| `--codec`             | `lz4`                      | Compression codec: `none`, `lz4`, `zstd`, `snappy`                              |
+| `--zstd-level`        | `1`                        | Zstd compression level (1–22)                                                   |
+| `--batch-size`        | `8192`                     | Batch size for reading Parquet data                                             |
+| `--memory-limit`      | _(none)_                   | Memory limit in bytes; triggers spilling when exceeded                          |
+| `--write-buffer-size` | `1048576`                  | Write buffer size in bytes                                                      |
+| `--limit`             | `0`                        | Limit rows processed per iteration (0 = no limit)                               |
+| `--iterations`        | `1`                        | Number of timed iterations                                                      |
+| `--warmup`            | `0`                        | Number of warmup iterations before timing                                       |
+| `--output-dir`        | `/tmp/comet_shuffle_bench` | Directory for temporary shuffle output files                                    |
 
 ### Profiling with flamegraph
 
diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs
index 0ac4bccb421..5fd0ebb9452 100644
--- a/native/shuffle/benches/shuffle_writer.rs
+++ b/native/shuffle/benches/shuffle_writer.rs
@@ -19,7 +19,7 @@ use arrow::array::builder::{Date32Builder, Decimal128Builder, Int32Builder};
 use arrow::array::{builder::StringBuilder, Array, Int32Array, RecordBatch};
 use arrow::datatypes::{DataType, Field, Schema};
 use arrow::row::{RowConverter, SortField};
-use criterion::{criterion_group, criterion_main, Criterion};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::datasource::source::DataSourceExec;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
@@ -82,10 +82,10 @@ fn criterion_benchmark(c: &mut Criterion) {
                     CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16),
                     create_batches(8192, 10),
                 );
+                let rt = Runtime::new().unwrap();
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
                     let stream = exec.execute(0, task_ctx).unwrap();
-                    let rt = Runtime::new().unwrap();
                     rt.block_on(collect(stream)).unwrap();
                 });
             },
@@ -134,10 +134,10 @@ fn criterion_benchmark(c: &mut Criterion) {
                     partitioning.clone(),
                     create_batches(8192, 10),
                 );
+                let rt = Runtime::new().unwrap();
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
                     let stream = exec.execute(0, task_ctx).unwrap();
-                    let rt = Runtime::new().unwrap();
                     rt.block_on(collect(stream)).unwrap();
                 });
             },
@@ -164,10 +164,10 @@ fn criterion_benchmark(c: &mut Criterion) {
                     CometPartitioning::SinglePartition,
                     create_batches(rows_per_batch, num_batches),
                 );
+                let rt = Runtime::new().unwrap();
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
                     let stream = exec.execute(0, task_ctx).unwrap();
-                    let rt = Runtime::new().unwrap();
                     rt.block_on(collect(stream)).unwrap();
                 });
             },
@@ -184,13 +184,11 @@ fn criterion_benchmark(c: &mut Criterion) {
     let wide_batches: Vec<RecordBatch> = (0..8).map(|_| nested_schema_batch(8192, 40, 2)).collect();
     let wide_schema = wide_batches[0].schema();
     let round_robin_strategies = [
+        ("hash_all_columns", RoundRobinStrategy::default()),
         (
-            "hash_all_columns",
-            RoundRobinStrategy::HashAll {
-                max_hash_columns: 0,
-            },
+            "whole_batch",
+            RoundRobinStrategy::WholeBatch { start_partition: 0 },
         ),
-        ("whole_batch", RoundRobinStrategy::WholeBatch),
     ];
     for (label, strategy) in &round_robin_strategies {
         group.bench_function(
@@ -215,6 +213,10 @@ fn criterion_benchmark(c: &mut Criterion) {
     // Partitioning-only microbench. `insert_batch` runs partition assignment + index
     // buffering but never flushes to disk, so the two RoundRobin strategies can be compared
     // without the IPC encode + disk write cost that dominates the end-to-end bench above.
+    //
+    // The repartitioner is built in the setup closure, not the timed one: `try_new` sizes the
+    // row scratch from the strategy, so timing it would charge `hash_all_columns` ~64 KB of
+    // per-iteration allocation that `whole_batch` does not pay and bias the comparison.
     for (label, strategy) in &round_robin_strategies {
         group.bench_function(
             format!("partitioning_only: RoundRobin nested schema (strategy={label})"),
@@ -229,81 +231,48 @@ fn criterion_benchmark(c: &mut Criterion) {
                 );
                 let dir = tempfile::tempdir().unwrap();
                 let data_path = dir.path().join("data.out").to_str().unwrap().to_string();
-                b.iter(|| {
-                    let block_writer =
-                        ShuffleBlockWriter::try_new(wide_schema.as_ref(), CompressionCodec::None)
-                            .unwrap();
-                    let writer = LocalPartitionWriter::try_new(
-                        data_path.clone(),
-                        Arc::new(PartitionOffsets::default()),
-                        block_writer,
-                        num_partitions,
-                        8192,
-                        1024 * 1024,
-                        Arc::clone(&runtime_env),
-                    )
-                    .unwrap();
-                    let metrics =
-                        ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
-                    let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
-                        0,
-                        writer,
-                        CometPartitioning::RoundRobin(num_partitions, strategy.clone()),
-                        metrics,
-                        Arc::clone(&runtime_env),
-                        8192,
-                        false,
-                        None,
-                    )
-                    .unwrap();
-                    rt.block_on(async {
-                        for batch in &batches {
-                            repartitioner.insert_batch(batch.clone()).await.unwrap();
-                        }
-                    });
-                });
+                b.iter_batched(
+                    || {
+                        let block_writer = ShuffleBlockWriter::try_new(
+                            wide_schema.as_ref(),
+                            CompressionCodec::None,
+                        )
+                        .unwrap();
+                        let writer = LocalPartitionWriter::try_new(
+                            data_path.clone(),
+                            Arc::new(PartitionOffsets::default()),
+                            block_writer,
+                            num_partitions,
+                            8192,
+                            1024 * 1024,
+                            Arc::clone(&runtime_env),
+                        )
+                        .unwrap();
+                        MultiPartitionShuffleRepartitioner::try_new(
+                            0,
+                            writer,
+                            CometPartitioning::RoundRobin(num_partitions, strategy.clone()),
+                            ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+                            Arc::clone(&runtime_env),
+                            8192,
+                            false,
+                            None,
+                        )
+                        .unwrap()
+                    },
+                    |mut repartitioner| {
+                        rt.block_on(async {
+                            for batch in &batches {
+                                repartitioner.insert_batch(batch.clone()).await.unwrap();
+                            }
+                        });
+                    },
+                    BatchSize::LargeInput,
+                );
             },
         );
     }
 
-    // End-to-end bench on a real on-disk parquet dataset at
-    // /tmp/clickstream_data_smoke (Spark-written, snappy-compressed, wide nested
-    // clickstream schema — the shape this optimization targets). Loaded once per
-    // process; each iteration streams the batches through `ShuffleWriterExec`
-    // for both RoundRobin strategies. Skipped if the dataset is absent.
-    let clickstream_batches = load_parquet_dir_batches("/tmp/clickstream_data_smoke");
-    if let Some(batches) = clickstream_batches {
-        let clickstream_num_partitions = 50usize;
-        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
-        let num_cols = batches[0].num_columns();
-        eprintln!(
-            "clickstream bench: {} batches, {} rows total, {} columns",
-            batches.len(),
-            total_rows,
-            num_cols,
-        );
-        for (label, strategy) in &round_robin_strategies {
-            group.bench_function(
-                format!("shuffle_writer: RoundRobin clickstream parquet (strategy={label})"),
-                |b| {
-                    let ctx = SessionContext::new();
-                    let rt = Runtime::new().unwrap();
-                    let exec = create_shuffle_writer_exec(
-                        CompressionCodec::Lz4Frame,
-                        CometPartitioning::RoundRobin(clickstream_num_partitions, strategy.clone()),
-                        batches.clone(),
-                    );
-                    b.iter(|| {
-                        let task_ctx = ctx.task_ctx();
-                        let stream = exec.execute(0, task_ctx).unwrap();
-                        rt.block_on(collect(stream)).unwrap();
-                    });
-                },
-            );
-        }
-    } else {
-        eprintln!("clickstream bench: /tmp/clickstream_data_smoke not found; skipping");
-    }
     group.finish();
 
     // High partition counts stress the per-partition write path (one short-lived
@@ -322,10 +291,10 @@ fn criterion_benchmark(c: &mut Criterion) {
                     CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], num_partitions),
                     create_batches(8192, 10),
                 );
+                let rt = Runtime::new().unwrap();
                 b.iter(|| {
                     let task_ctx = ctx.task_ctx();
                     let stream = exec.execute(0, task_ctx).unwrap();
-                    let rt = Runtime::new().unwrap();
                     rt.block_on(collect(stream)).unwrap();
                 });
             },
@@ -334,61 +303,6 @@ fn criterion_benchmark(c: &mut Criterion) {
     high_partition_group.finish();
 }
 
-/// Loads every `*.parquet` file under `dir` (recursively, one level deep for
-/// partitioned datasets like `date=…/part-*.parquet`) and returns their record
-/// batches at 8192 rows per batch. Returns `None` if the directory is missing
-/// or contains no parquet files.
-fn load_parquet_dir_batches(dir: &str) -> Option<Vec<RecordBatch>> {
-    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
-
-    let root = std::path::Path::new(dir);
-    if !root.exists() {
-        return None;
-    }
-
-    let mut files: Vec<std::path::PathBuf> = Vec::new();
-    let mut stack: Vec<std::path::PathBuf> = vec![root.to_path_buf()];
-    while let Some(p) = stack.pop() {
-        let entries = match std::fs::read_dir(&p) {
-            Ok(e) => e,
-            Err(_) => continue,
-        };
-        for entry in entries.flatten() {
-            let path = entry.path();
-            if path.is_dir() {
-                stack.push(path);
-            } else if path
-                .extension()
-                .and_then(|s| s.to_str())
-                .map(|s| s.eq_ignore_ascii_case("parquet"))
-                .unwrap_or(false)
-            {
-                files.push(path);
-            }
-        }
-    }
-    files.sort();
-    if files.is_empty() {
-        return None;
-    }
-
-    let mut all = Vec::new();
-    for path in files {
-        let file = std::fs::File::open(&path).unwrap_or_else(|e| {
-            panic!("failed to open {}: {e}", path.display());
-        });
-        let reader = ParquetRecordBatchReaderBuilder::try_new(file)
-            .unwrap_or_else(|e| panic!("parquet builder for {}: {e}", path.display()))
-            .with_batch_size(8192)
-            .build()
-            .unwrap_or_else(|e| panic!("parquet build for {}: {e}", path.display()));
-        for batch in reader {
-            all.push(batch.unwrap());
-        }
-    }
-    Some(all)
-}
-
 fn create_shuffle_writer_exec(
     compression_codec: CompressionCodec,
     partitioning: CometPartitioning,
diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs
index c14f15aebd5..6c7feaf8c72 100644
--- a/native/shuffle/src/bin/shuffle_bench.rs
+++ b/native/shuffle/src/bin/shuffle_bench.rs
@@ -72,7 +72,7 @@ struct Args {
     #[arg(long, default_value_t = 200)]
     partitions: usize,
 
-    /// Partitioning scheme: hash, single, round-robin
+    /// Partitioning scheme: hash, single, round-robin, round-robin-whole-batch
     #[arg(long, default_value = "hash")]
     partitioning: String,
 
@@ -585,11 +585,14 @@ fn build_partitioning(
 ) -> CometPartitioning {
     match scheme {
         "single" => CometPartitioning::SinglePartition,
-        "round-robin" => CometPartitioning::RoundRobin(
+        "round-robin" => {
+            CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default())
+        }
+        "round-robin-whole-batch" => CometPartitioning::RoundRobin(
             num_partitions,
-            RoundRobinStrategy::HashAll {
-                max_hash_columns: 0,
-            },
+            // Standalone bench, so there is no Spark map partition id to carry. One mapper
+            // means the start offset makes no difference to the work measured.
+            RoundRobinStrategy::WholeBatch { start_partition: 0 },
         ),
         "hash" => {
             let exprs: Vec<Arc<dyn datafusion::physical_expr::PhysicalExpr>> = hash_col_indices
diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs
index 7b331ed3eb9..16d65654959 100644
--- a/native/shuffle/src/comet_partitioning.rs
+++ b/native/shuffle/src/comet_partitioning.rs
@@ -26,11 +26,29 @@ pub enum RoundRobinStrategy {
     /// and assign rows individually. Deterministic under retry regardless of upstream
     /// order preservation, but pays the per-row hash cost.
     HashAll { max_hash_columns: usize },
-    /// Assign each incoming RecordBatch as a whole to one output partition, based on
-    /// the input partition id and a per-task batch sequence number. Skips per-row
-    /// hashing entirely. Retry-safe only when the upstream operator emits the same
-    /// batches in the same order under retry.
-    WholeBatch,
+    /// Assign each incoming RecordBatch as a whole to one output partition, chosen by a
+    /// per-task batch counter. Skips per-row hashing entirely. Retry-safe only when the
+    /// upstream operator emits the same batches in the same order under retry.
+    ///
+    /// `start_partition` seeds that counter, so mapper i sends its k-th batch to
+    /// `(i + k) % num_partitions`. It must be the Spark map partition id, for two reasons
+    /// that pull in different directions and are both required:
+    ///
+    ///   * 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;
+    ///   * a pure function of the map partition, so a re-executed task reproduces its
+    ///     placement. Spark's own round robin seeds `XORShiftRandom(partitionId)` for the
+    ///     same pair of reasons.
+    WholeBatch { start_partition: usize },
+}
+
+impl Default for RoundRobinStrategy {
+    /// Hash every column, which is what Comet's round robin did before `WholeBatch` existed.
+    fn default() -> Self {
+        Self::HashAll {
+            max_hash_columns: 0,
+        }
+    }
 }
 
 /// Partitioning scheme for distributing rows across shuffle output partitions.
diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs
index 017dc240f46..0de84db6079 100644
--- a/native/shuffle/src/partitioners/multi_partition.rs
+++ b/native/shuffle/src/partitioners/multi_partition.rs
@@ -122,9 +122,8 @@ pub struct MultiPartitionShuffleRepartitioner<T: PartitionWriter> {
     /// allocation once rather than once per slice that references it. Cleared whenever the
     /// buffered batches drain (spill / shuffle_write). See `count_new_buffers`.
     pinned_buffers: HashSet<usize>,
-    /// Batch counter for the batch-granular RoundRobin path. Seeded with the input partition
-    /// id so mappers with adjacent partition ids do not concentrate their first batches on the
-    /// same output partition, then incremented once per input batch slice that reaches
+    /// Batch counter for [`RoundRobinStrategy::WholeBatch`], seeded from that variant's
+    /// `start_partition` and incremented once per input batch slice that reaches
     /// `partitioning_batch`. Unused by other strategies.
     round_robin_batch_seq: usize,
 }
@@ -197,15 +196,18 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
         // initialization code is simply initializing the vectors to the desired size.
         // The initial values are not used.
         //
-        // `partition_ids` and `partition_row_indices` are only touched by the row-level
-        // strategies (`Hash`, `RangePartitioning`, and hash-all-columns RoundRobin). The
-        // whole-batch RoundRobin path never inspects rows, so allocating them there would be
-        // ~64 KB of dead scratch on every task.
-        let needs_row_scratch = !matches!(
-            &partitioning,
-            CometPartitioning::SinglePartition
-                | CometPartitioning::RoundRobin(_, RoundRobinStrategy::WholeBatch),
-        );
+        // Whole-batch RoundRobin is the one strategy that never inspects rows. It needs neither
+        // `partition_ids` nor `partition_row_indices` (~64 KB of dead scratch per task), and it
+        // is the only one that needs a seeded batch counter. Read both off the partitioning here,
+        // before it is moved into the struct below.
+        let whole_batch_start = match &partitioning {
+            CometPartitioning::RoundRobin(
+                _,
+                RoundRobinStrategy::WholeBatch { start_partition },
+            ) => Some(*start_partition),
+            _ => None,
+        };
+        let needs_row_scratch = whole_batch_start.is_none();
         let scratch = ScratchSpace {
             hashes_buf: match &partitioning {
                 // Allocate hashes_buf for hash and hash-all-columns round robin partitioning.
@@ -245,10 +247,9 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
             max_buffer_bytes,
             tracing_enabled,
             pinned_buffers: HashSet::new(),
-            // Seed with the input partition id so batches from mapper i land on
-            // partition (i + k) mod N on the k-th batch, spreading concurrent mappers' first
-            // batches across distinct output partitions.
-            round_robin_batch_seq: partition,
+            // `partition` is the DataFusion partition and is always 0 here, so the seed comes
+            // from the strategy instead. See `RoundRobinStrategy::WholeBatch`.
+            round_robin_batch_seq: whole_batch_start.unwrap_or(0),
         })
     }
 
@@ -378,32 +379,28 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
                 self.scratch = scratch;
             }
             CometPartitioning::RoundRobin(num_output_partitions, strategy) => {
-                // Two strategies share the scratch-take / timer / spill / scratch-restore
-                // scaffold. Only the middle "fill partition_row_indices + partition_starts"
-                // step differs. WholeBatch: assign the whole input batch to one partition and
-                // pay no per-row cost. HashAll: hash every row (over up to max_hash_columns
-                // columns) and route rows individually. WholeBatch is retry-safe only when the
-                // upstream operator emits the same batches in the same order under retry
-                // (Comet's `CometNativeScan` and other order-preserving operators do so).
                 let mut scratch = std::mem::take(&mut self.scratch);
                 let num_rows = input.num_rows();
                 let partition_row_indices: Option<&[u32]> = {
                     let mut timer = self.metrics.repart_time.timer();
 
-                    let indices: Option<&[u32]> = match strategy {
-                        RoundRobinStrategy::WholeBatch => {
+                    let indices = match strategy {
+                        RoundRobinStrategy::WholeBatch { .. } => {
                             let target_idx = self.round_robin_batch_seq % *num_output_partitions;
                             self.round_robin_batch_seq = self.round_robin_batch_seq.wrapping_add(1);
 
-                            // `partition_starts[k]..partition_starts[k+1]` is partition k's
-                            // slice. Slots 0..=target_idx are 0; slots after `target_idx` are
-                            // `num_rows`, so `target_idx` owns the whole [0..num_rows) range.
-                            // `partition_row_indices` is left unmaterialized — the spill path
-                            // treats the target partition's slice as an identity mapping.
+                            // `target_idx` owns the whole `[0..num_rows)` range and every other
+                            // partition an empty one. `partition_row_indices` stays
+                            // unmaterialized, so `None` below tells the spill path to read that
+                            // range as an identity mapping.
+                            //
+                            // Overwrite in place rather than `clear` + `resize`: the vector is
+                            // always `num_output_partitions + 1` long (set in `try_new`, kept by
+                            // `map_partition_ids_to_starts_and_indices`), and `fill` lowers to a
+                            // memset where `resize` goes through `extend_with`.
                             let partition_starts = &mut scratch.partition_starts;
-                            partition_starts.clear();
-                            partition_starts.resize(target_idx + 1, 0);
-                            partition_starts.resize(*num_output_partitions + 1, num_rows as u32);
+                            partition_starts[..=target_idx].fill(0);
+                            partition_starts[target_idx + 1..].fill(num_rows as u32);
                             None
                         }
                         RoundRobinStrategy::HashAll { max_hash_columns } => {
@@ -419,14 +416,12 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
                             } else {
                                 (*max_hash_columns).min(input.num_columns())
                             };
-                            let columns_to_hash: Vec<ArrayRef> = (0..num_columns_to_hash)
-                                .map(|i| Arc::clone(input.column(i)))
-                                .collect();
+                            let columns_to_hash = &input.columns()[..num_columns_to_hash];
 
                             // Use identical seed as Spark hash partitioning.
                             let hashes_buf = &mut scratch.hashes_buf[..num_rows];
                             hashes_buf.fill(42_u32);
-                            create_murmur3_hashes(&columns_to_hash, hashes_buf)?;
+                            create_murmur3_hashes(columns_to_hash, hashes_buf)?;
 
                             let partition_ids = &mut scratch.partition_ids[..num_rows];
                             hashes_buf.iter().enumerate().for_each(|(idx, hash)| {
@@ -841,12 +836,7 @@ mod tests {
         let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
             0,
             FailingPartitionWriter::default(),
-            CometPartitioning::RoundRobin(
-                2,
-                RoundRobinStrategy::HashAll {
-                    max_hash_columns: 0,
-                },
-            ),
+            CometPartitioning::RoundRobin(2, RoundRobinStrategy::default()),
             ShufflePartitionerMetrics::new(&metrics_set, 0),
             Arc::clone(&runtime),
             2,
@@ -970,12 +960,7 @@ mod tests {
                 let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new(
                     0,
                     FailingPartitionWriter::default(),
-                    CometPartitioning::RoundRobin(
-                        2,
-                        RoundRobinStrategy::HashAll {
-                            max_hash_columns: 0,
-                        },
-                    ),
+                    CometPartitioning::RoundRobin(2, RoundRobinStrategy::default()),
                     ShufflePartitionerMetrics::new(&metrics_set, 0),
                     Arc::clone(&runtime),
                     64,
diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs
index adc580b80ed..8385563e745 100644
--- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs
+++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs
@@ -116,25 +116,21 @@ impl<'a> PartitionedBatchIterator<'a> {
     /// single source batch in natural order, so a clone is equivalent to interleaving them.
     fn whole_source_batch(&self, indices_end: usize) -> Option<RecordBatch> {
         let indices = &self.indices[self.pos..indices_end];
-        let &(first_batch, first_row) = indices.first()?;
-        let &(last_batch, last_row) = indices.last()?;
-        let source_rows = self.record_batches[first_batch as usize].num_rows();
-        if first_batch != last_batch
-            || first_row != 0
-            || last_row as usize + 1 != source_rows
-            || indices.len() != source_rows
-        {
+        let &(first_batch, _) = indices.first()?;
+        let source = self.record_batches[first_batch as usize];
+        // O(1) reject, which is all a chunk that cannot take this path pays.
+        if indices.len() != source.num_rows() {
             return None;
         }
-        // Verify the full contiguous invariant. Invariant holds by construction in
-        // `buffer_partitioned_batch_may_spill`, but the check is O(indices.len())
-        // with two integer compares per step and vectorizes easily. Two orders of
-        // magnitude cheaper than a nested-schema interleave.
+        // Given the length check, `b == first_batch && r == i` for every element implies the
+        // rest: one source batch, starting at row 0, ending at its last row. `WholeBatch`
+        // guarantees it by construction, but `HashAll` can land a chunk here by chance, so the
+        // scan has to run in release builds too. It short-circuits on the first scattered pair.
         indices
             .iter()
             .enumerate()
             .all(|(i, &(b, r))| b == first_batch && r as usize == i)
-            .then(|| self.record_batches[first_batch as usize].clone())
+            .then(|| source.clone())
     }
 }
 
@@ -148,11 +144,8 @@ impl Iterator for PartitionedBatchIterator<'_> {
 
         let indices_end = std::cmp::min(self.pos + self.batch_size, self.indices.len());
 
-        // Whole-source-batch fast path: when this chunk covers every row of one source
-        // batch in natural order, cloning the source batch is equivalent to `interleave`
-        // and skips walking every column (and every nested child on wide/nested schemas).
-        // This is the vectorized path used by RoundRobinStrategy::WholeBatch, but the check
-        // is generic and fires whenever indices happen to line up this way.
+        // Generic, so `HashAll` benefits opportunistically, but this is the path that makes
+        // `RoundRobinStrategy::WholeBatch` cheap.
         if let Some(batch) = self.whole_source_batch(indices_end) {
             self.pos = indices_end;
             return Some(Ok(batch));
diff --git a/native/shuffle/src/rss_execution_tests.rs b/native/shuffle/src/rss_execution_tests.rs
index a3d92f79f7f..f18409d0baf 100644
--- a/native/shuffle/src/rss_execution_tests.rs
+++ b/native/shuffle/src/rss_execution_tests.rs
@@ -231,12 +231,7 @@ fn rss_multi_partition_supports_hash_range_and_round_robin() {
     for partitioning in [
         CometPartitioning::Hash(vec![expression], 4),
         CometPartitioning::RangePartitioning(ordering, 4, Arc::new(converter), boundaries),
-        CometPartitioning::RoundRobin(
-            4,
-            RoundRobinStrategy::HashAll {
-                max_hash_columns: 0,
-            },
-        ),
+        CometPartitioning::RoundRobin(4, RoundRobinStrategy::default()),
     ] {
         let pusher = Arc::new(RecordingPusher::default());
         let execution = rss_execution(
@@ -276,12 +271,7 @@ fn rss_empty_schema_preserves_row_counts_in_partition_zero() {
     let execution = rss_execution(
         vec![batch.clone(), batch],
         schema,
-        CometPartitioning::RoundRobin(
-            4,
-            RoundRobinStrategy::HashAll {
-                max_hash_columns: 0,
-            },
-        ),
+        CometPartitioning::RoundRobin(4, RoundRobinStrategy::default()),
         pusher.clone(),
         CompressionCodec::None,
         1024 * 1024,
@@ -309,12 +299,7 @@ fn rss_empty_schema_without_rows_does_not_push_frames() {
     let execution = rss_execution(
         vec![batch],
         schema,
-        CometPartitioning::RoundRobin(
-            4,
-            RoundRobinStrategy::HashAll {
-                max_hash_columns: 0,
-            },
-        ),
+        CometPartitioning::RoundRobin(4, RoundRobinStrategy::default()),
         pusher.clone(),
         CompressionCodec::None,
         1024 * 1024,
diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs
index 91a9ca3ad6d..f0f6f05164f 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -1143,12 +1143,7 @@ mod test {
                 Arc::new(row_converter),
                 owned_rows,
             ),
-            CometPartitioning::RoundRobin(
-                num_partitions,
-                RoundRobinStrategy::HashAll {
-                    max_hash_columns: 0,
-                },
-            ),
+            CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default()),
         ] {
             let batches = (0..num_batches).map(|_| batch.clone()).collect::<Vec<_>>();
 
@@ -1216,12 +1211,7 @@ mod test {
                 Arc::new(DataSourceExec::new(Arc::new(
                     MemorySourceConfig::try_new(partitions, batch.schema(), None).unwrap(),
                 ))),
-                CometPartitioning::RoundRobin(
-                    num_partitions,
-                    RoundRobinStrategy::HashAll {
-                        max_hash_columns: 0,
-                    },
-                ),
+                CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default()),
                 CompressionCodec::Zstd(1),
                 data_file.clone(),
                 false,
@@ -1357,7 +1347,10 @@ mod test {
                 Arc::new(DataSourceExec::new(Arc::new(
                     MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(),
                 ))),
-                CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::WholeBatch),
+                CometPartitioning::RoundRobin(
+                    num_partitions,
+                    RoundRobinStrategy::WholeBatch { start_partition: 0 },
+                ),
                 CompressionCodec::Zstd(1),
                 data_file.to_str().unwrap().to_string(),
                 false,
@@ -1422,6 +1415,69 @@ mod test {
         assert_eq!(total_rows, num_rows * num_batches, "no rows may be dropped");
     }
 
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_round_robin_batch_granular_starts_at_map_partition() {
+        // Every mapper seeding 0 would send its first batch to partition 0, so a stage whose
+        // tasks each emit fewer batches than there are output partitions would leave the tail
+        // partitions empty everywhere. `start_partition` carries the Spark map partition id, so
+        // mapper i starts at partition i and the stage covers the full range.
+        let num_partitions = 8;
+        let batch = create_batch(100);
+        let batches: Vec<RecordBatch> = (0..2).map(|_| batch.clone()).collect();
+        let dir = tempfile::tempdir().unwrap();
+
+        let nonempty_partitions = |start_partition: usize| -> Vec<usize> {
+            let data_file = dir.path().join(format!("start{start_partition}.out"));
+            let exec = ShuffleWriterExec::try_new(
+                Arc::new(DataSourceExec::new(Arc::new(
+                    MemorySourceConfig::try_new(
+                        std::slice::from_ref(&batches),
+                        batch.schema(),
+                        None,
+                    )
+                    .unwrap(),
+                ))),
+                CometPartitioning::RoundRobin(
+                    num_partitions,
+                    RoundRobinStrategy::WholeBatch { start_partition },
+                ),
+                CompressionCodec::Zstd(1),
+                data_file.to_str().unwrap().to_string(),
+                false,
+                1024 * 1024,
+                None,
+            )
+            .unwrap();
+
+            let ctx = SessionContext::new_with_config_rt(
+                SessionConfig::new(),
+                Arc::new(RuntimeEnvBuilder::new().build().unwrap()),
+            );
+            let stream = exec.execute(0, ctx.task_ctx()).unwrap();
+            Runtime::new().unwrap().block_on(collect(stream)).unwrap();
+
+            let offsets = exec
+                .partition_offsets()
+                .expect("local destination publishes offsets")
+                .get()
+                .expect("writer published its partition offsets")
+                .to_vec();
+            (0..num_partitions)
+                .filter(|&p| offsets[p + 1] > offsets[p])
+                .collect()
+        };
+
+        assert_eq!(nonempty_partitions(0), vec![0, 1]);
+        assert_eq!(
+            nonempty_partitions(3),
+            vec![3, 4],
+            "mapper 3 must not start where mapper 0 does"
+        );
+        // A map partition id above the output partition count wraps rather than panicking.
+        assert_eq!(nonempty_partitions(num_partitions + 1), vec![1, 2]);
+    }
+
     #[test]
     #[cfg_attr(miri, ignore)]
     fn test_round_robin_batch_granular_whole_batch_per_partition() {
@@ -1445,7 +1501,10 @@ mod test {
             Arc::new(DataSourceExec::new(Arc::new(
                 MemorySourceConfig::try_new(partitions, batch.schema(), None).unwrap(),
             ))),
-            CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::WholeBatch),
+            CometPartitioning::RoundRobin(
+                num_partitions,
+                RoundRobinStrategy::WholeBatch { start_partition: 0 },
+            ),
             CompressionCodec::Zstd(1),
             data_file.to_str().unwrap().to_string(),
             false,
@@ -1486,6 +1545,8 @@ mod test {
         }
     }
 
+    /// Test that batch coalescing in BufBatchWriter reduces output size by
+    /// writing fewer, larger IPC blocks instead of many small ones.
     #[test]
     #[cfg_attr(miri, ignore)]
     fn test_batch_coalescing_reduces_size() {
@@ -1823,12 +1884,7 @@ mod test {
             Arc::new(DataSourceExec::new(Arc::new(
                 MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(),
             ))),
-            CometPartitioning::RoundRobin(
-                num_partitions,
-                RoundRobinStrategy::HashAll {
-                    max_hash_columns: 0,
-                },
-            ),
+            CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default()),
             CompressionCodec::Zstd(1),
             data_file.to_str().unwrap().to_string(),
             false,
@@ -1912,12 +1968,7 @@ mod test {
             Arc::new(DataSourceExec::new(Arc::new(
                 MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(),
             ))),
-            CometPartitioning::RoundRobin(
-                num_partitions,
-                RoundRobinStrategy::HashAll {
-                    max_hash_columns: 0,
-                },
-            ),
+            CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default()),
             CompressionCodec::Zstd(1),
             data_file.to_str().unwrap().to_string(),
             false,
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index d972f2fe3e0..ca945bdd61d 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -478,12 +478,13 @@ object CometConf extends ShimCometConf {
   val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.native.partitioning.roundrobin.batchGranular")
       .category(CATEGORY_SHUFFLE)
-      .doc(
-        "When true, Comet's native round-robin shuffle assigns each RecordBatch as a whole " +
-          "to one output partition instead of hashing every column of every row. This is much " +
-          "cheaper on wide/nested schemas. Distribution is quasi-even at batch granularity. " +
-          "Retry-safe only when the upstream operator emits the same batches in the same order " +
-          "under retry (Comet's Parquet scan and other order-preserving operators satisfy this).")
+      .doc("When true, Comet's native round-robin shuffle assigns each RecordBatch as a whole " +
+        "to one output partition instead of hashing every column of every row. This is much " +
+        "cheaper on wide/nested schemas. Distribution is quasi-even at batch granularity. " +
+        "Retry-safe only when the upstream operator emits the same batches in the same order " +
+        "under retry (Comet's Parquet scan and other order-preserving operators satisfy this). " +
+        s"Has no effect unless ${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key} " +
+        "is also true.")
       .booleanConf
       .createWithDefault(false)
 
@@ -492,17 +493,17 @@ object CometConf extends ShimCometConf {
       .category(CATEGORY_SHUFFLE)
       .doc(
         "Only applies when " +
-          "spark.comet.shuffle.native.partitioning.roundrobin.batchGranular is true. " +
+          s"${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.key} is true. " +
           "Batch-granular round robin places rows by position rather than by content, so a " +
           "re-executed map task can send rows to different output partitions than the attempt " +
           "it replaces, silently dropping and duplicating rows once any of that output has " +
-          "been fetched. When this is true, a native round-robin map task that runs as a retry " +
-          "(a task attempt after the first, or any attempt of a re-submitted stage, which is " +
-          "what an executor loss produces) fails immediately instead of writing output, which " +
-          "fails the stage and therefore the job. Note that speculative attempts also count as " +
-          "retries. Set to false to keep Spark's normal fault tolerance, in which case safety " +
-          "rests on the whole stage being rolled back, which Comet requests by declaring the " +
-          "shuffle input RDD INDETERMINATE whenever its own input is not DETERMINATE.")
+          "been fetched. When true, such a map task fails immediately instead of writing " +
+          "output, which fails the stage and therefore the job. A retry here means any task " +
+          "attempt after the first (speculative attempts included) or any attempt of a " +
+          "re-submitted stage, which is what an executor loss produces. Set to false to keep " +
+          "Spark's normal fault tolerance, leaving safety to whole-stage rollback, which Comet " +
+          "requests by declaring the shuffle input RDD INDETERMINATE whenever its own input " +
+          "is not DETERMINATE.")
       .booleanConf
       .createWithDefault(true)
 
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
index e7e2abbf1a0..16827727022 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala
@@ -66,28 +66,16 @@ private[shuffle] class CometNativeShuffleInputRDD(
       positionalRoundRobin)
 
   /**
-   * Batch-granular round robin assigns each Arrow batch to an output partition by position, so
-   * the placement of a row depends on the order and the framing of the batches the upstream
-   * operator produced, not on the row itself. Re-running one map task against a differently
-   * ordered or differently framed input therefore sends rows somewhere else, which duplicates and
-   * drops rows once any of the replaced output has already been fetched.
-   *
-   * Spark states that risk declaratively rather than defending against it per-operator: a
-   * round-robin repartition is wrapped in a `MapPartitionsRDD` with `isOrderSensitive = true`
+   * Spark handles the retry hazard of positional round robin declaratively rather than
+   * per-operator: it wraps the repartition in a `MapPartitionsRDD` with `isOrderSensitive = true`
    * (Comet's own JVM path does this in `prepareJVMShuffleDependency`), and that RDD reports
-   * `INDETERMINATE` whenever its parent is `UNORDERED`. The DAGScheduler then rolls the whole
-   * stage back, or aborts the job when a result stage has already consumed output, instead of
-   * re-running a single task. The native path has no `MapPartitionsRDD` to carry the flag, so
-   * apply the same rule here.
-   *
-   * The parent level does the discriminating, exactly as it does for Spark. A plain scan is
-   * `DETERMINATE`: it replays identically, so positional assignment is reproducible and ordinary
-   * per-task retry stays cheap. Anything below another exchange is `UNORDERED`, because reduce
-   * tasks see shuffle blocks in arrival order, and that is when positional assignment stops being
-   * reproducible.
+   * `INDETERMINATE` whenever its parent is `UNORDERED`, which makes the DAGScheduler roll the
+   * whole stage back instead of re-running one task. The native path has no `MapPartitionsRDD` to
+   * carry the flag, so apply the same rule here.
    *
-   * Content-hash round robin (the default) needs none of this: it places rows by hash, so its
-   * output is a pure function of the rows regardless of how they arrive.
+   * Letting the parent level discriminate is what keeps a plain scan on the cheap per-task retry
+   * path. See the `WholeBatch` retry-safety section of
+   * `docs/source/contributor-guide/native_shuffle.md`.
    */
   override protected def getOutputDeterministicLevel: DeterministicLevel.Value = {
     val inheritedLevel = super.getOutputDeterministicLevel
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index 7f6bf800645..7c0278b57ae 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -109,20 +109,13 @@ class CometNativeShuffleWriter[K, V](
   }
 
   /**
-   * Refuse to re-execute a map task whose round-robin placement is positional.
+   * Refuse to re-execute a map task whose round-robin placement is positional. See the
+   * `WholeBatch` retry-safety section of `docs/source/contributor-guide/native_shuffle.md`.
    *
-   * With `RoundRobinStrategy::WholeBatch` an input batch's output partition comes from a per-task
-   * counter, so which partition a row lands in is a function of the order and the framing of the
-   * batches the upstream operator produced rather than of the row itself. An attempt that frames
-   * or orders its input even slightly differently writes a different partitioning of the same
-   * rows. Once any consumer has fetched the output this attempt replaces, the reduce side gets
-   * some rows twice and some not at all, and nothing downstream can detect it.
-   *
-   * [[CometNativeShuffleInputRDD.getOutputDeterministicLevel]] already asks the DAGScheduler to
-   * roll the whole stage back rather than re-run one task, which is Spark's own answer to this.
-   * That answer is only as good as the parent's determinism level being an accurate description
-   * of what the upstream operator replays, so while the strategy is opt-in and unproven the
-   * default is to not rely on it and fail here instead.
+   * [[CometNativeShuffleInputRDD.getOutputDeterministicLevel]] is the other half of the defence.
+   * This half does not trust the parent's declared determinism level to describe what the
+   * upstream operator actually replays, so while the strategy is opt-in and unproven the default
+   * is to fail rather than rely on rollback.
    *
    * Both counters are needed and neither subsumes the other. `attemptNumber` covers a task re-run
    * inside the current stage attempt: a task-level failure, speculation, or an executor lost
@@ -133,14 +126,14 @@ class CometNativeShuffleWriter[K, V](
    * There is no Spark API for failing an application from inside a task, and the two exceptions
    * Spark declines to retry are matched by class name (`NotSerializableException`,
    * `TaskOutputFileAlreadyExistException`), so this throws an ordinary exception. Because the
-   * condition is sticky - a later attempt only has a higher `attemptNumber` - every remaining
+   * condition is sticky (a later attempt only has a higher `attemptNumber`) every remaining
    * attempt fails here too and the task set aborts after `spark.task.maxFailures`, failing the
    * stage and the job. Each of those attempts fails before doing any work.
    */
   private def failIfRetryingPositionalRoundRobin(): Unit = {
     if (!CometShuffleExchangeExec.usesPositionalRoundRobin(outputPartitioning)) return
     if (!CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_FAIL_ON_RETRY.get()) return
-    // `context` is null only in unit tests that drive the writer directly.
+    // Nullable for the same reason as the `Option(context)` guard on `cancellationWatch` above.
     val taskContext = context
     if (taskContext == null) return
     val taskAttempt = taskContext.attemptNumber()
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index b490b7164c4..e5629d032c7 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -302,21 +302,20 @@ object CometShuffleExchangeExec
   /**
    * True when this exchange will run the native round-robin writer in its positional
    * (batch-granular) mode, where an input `RecordBatch` is assigned whole to one output partition
-   * by a per-task counter rather than by hashing its rows.
+   * by a per-task counter rather than by hashing its rows. See the `WholeBatch` retry-safety
+   * section of `docs/source/contributor-guide/native_shuffle.md` for why that is unsafe to
+   * re-execute.
    *
-   * Positional assignment is what makes the strategy cheap and also what makes it unsafe to
-   * re-execute: the placement depends on the order and the framing of the batches the upstream
-   * operator hands over, not on the rows themselves. Two call sites need the same answer, so they
-   * share this predicate rather than each reading the config: [[CometNativeShuffleInputRDD]],
-   * which declares its output indeterminate so the DAGScheduler rolls a stage back rather than
-   * re-running one task into partially consumed output, and [[CometNativeShuffleWriter]], which
-   * refuses to run at all on a retry.
+   * Must stay in step with `PhysicalPlanner::create_partitioning`, which turns the
+   * `batch_granular` proto field into `RoundRobinStrategy::WholeBatch`. Nothing enforces that.
    *
-   * The equivalent decision on the native side is `PhysicalPlanner::create_partitioning` turning
-   * `batch_granular` into `RoundRobinStrategy::WholeBatch`; keep the two in step.
+   * The `numPartitions > 1` guard mirrors `isRoundRobin` in `prepareJVMShuffleDependency`. With a
+   * single output partition every row lands in the same place, so there is no placement to get
+   * wrong, and native routes that case to `SinglePartitionShufflePartitioner` regardless.
    */
   def usesPositionalRoundRobin(outputPartitioning: Partitioning): Boolean =
     outputPartitioning.isInstanceOf[RoundRobinPartitioning] &&
+      outputPartitioning.numPartitions > 1 &&
       CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.get()
 
   override def createExec(
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala
index dd7a36045ed..e70fe7b121e 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeRoundRobinRetrySuite.scala
@@ -21,9 +21,9 @@ package org.apache.spark.sql.comet.execution.shuffle
 
 import java.util.Properties
 
-import org.apache.spark.{Partition, TaskContext, TaskContextImpl}
+import org.apache.spark.{TaskContext, TaskContextImpl}
 import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.rdd.{DeterministicLevel, RDD}
+import org.apache.spark.rdd.DeterministicLevel
 import org.apache.spark.sql.{CometTestBase, DataFrame}
 import org.apache.spark.sql.catalyst.expressions.Literal
 import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition}
@@ -39,9 +39,9 @@ import org.apache.comet.{CometConf, CometRuntimeException}
  * row's destination depends on the order and the framing of the batches the upstream operator
  * produced rather than on the row itself. Re-executing a map task can therefore write a different
  * partitioning of the same rows, which duplicates and drops rows on the reduce side once any of
- * the replaced output has been fetched. Two defences are covered here: declaring the shuffle
- * input RDD indeterminate so the DAGScheduler rolls the stage back instead of re-running one
- * task, and refusing to run a retried map task at all.
+ * the replaced output has been fetched. This suite covers the writer refusing to run a retried
+ * map task. The other defence, declaring the shuffle input RDD indeterminate so the DAGScheduler
+ * rolls the stage back, is covered in [[CometNativeShuffleInputRDDSuite]].
  *
  * Lives in the `execution.shuffle` package so it can construct the `private[shuffle]`
  * [[CometNativeShuffleInputRDD]] and the `private[spark]` [[TaskContextImpl]] directly.
@@ -53,26 +53,6 @@ class CometNativeRoundRobinRetrySuite extends CometTestBase {
   private val failOnRetryKey =
     CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_FAIL_ON_RETRY.key
 
-  /** An input RDD that reports exactly `level`, standing in for a real upstream subtree. */
-  private def parentWithLevel(level: DeterministicLevel.Value): RDD[AnyRef] =
-    new RDD[AnyRef](spark.sparkContext, Nil) {
-      override protected def getOutputDeterministicLevel: DeterministicLevel.Value = level
-      override protected def getPartitions: Array[Partition] = Array.empty
-      override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] =
-        Iterator.empty
-    }
-
-  private def shuffleInput(
-      parent: RDD[AnyRef],
-      positionalRoundRobin: Boolean): CometNativeShuffleInputRDD =
-    new CometNativeShuffleInputRDD(
-      spark.sparkContext,
-      Seq(parent),
-      0,
-      Set.empty,
-      CometMetricNode(Map.empty),
-      positionalRoundRobin = positionalRoundRobin)
-
   /** Mirrors `TaskContext.empty()`, which hard-codes both attempt numbers to zero. */
   private def taskContextFor(stageAttempt: Int, taskAttempt: Int): TaskContext =
     new TaskContextImpl(
@@ -115,9 +95,12 @@ class CometNativeRoundRobinRetrySuite extends CometTestBase {
 
   private val refusalPrefix = "Refusing to re-execute a Comet native round-robin"
 
-  test("positional round robin is claimed only for round robin with the config enabled") {
+  test("usesPositionalRoundRobin claims only multi-partition round robin with the config on") {
     val partitionings = Seq(
       RoundRobinPartitioning(4),
+      // One output partition puts every row in the same place, so there is no placement to get
+      // wrong. `isRoundRobin` in `prepareJVMShuffleDependency` excludes it for the same reason.
+      RoundRobinPartitioning(1),
       SinglePartition,
       HashPartitioning(Seq(Literal(1)), 4),
       RangePartitioning(Nil, 4))
@@ -130,35 +113,6 @@ class CometNativeRoundRobinRetrySuite extends CometTestBase {
     }
   }
 
-  test("positional round robin declares indeterminate output unless its parent is determinate") {
-    // A determinate parent (a plain scan) replays identically, so positional assignment is
-    // reproducible and per-task retry stays cheap. Anything below another exchange is unordered,
-    // which is where Spark's own round robin flips to indeterminate too.
-    Seq(
-      DeterministicLevel.DETERMINATE -> DeterministicLevel.DETERMINATE,
-      DeterministicLevel.UNORDERED -> DeterministicLevel.INDETERMINATE,
-      DeterministicLevel.INDETERMINATE -> DeterministicLevel.INDETERMINATE).foreach {
-      case (parentLevel, expected) =>
-        val input = shuffleInput(parentWithLevel(parentLevel), positionalRoundRobin = true)
-        assert(input.outputDeterministicLevel == expected, s"parent was $parentLevel")
-        assert(
-          input.copyForLocalShuffle().outputDeterministicLevel == expected,
-          s"local fallback lost the declaration for parent $parentLevel")
-    }
-  }
-
-  test("content-hash round robin keeps inheriting its parent's determinism") {
-    // Hash placement is a pure function of the rows, so nothing here should be indeterminate on
-    // its own account. This is the default path and must stay as retryable as it is today.
-    Seq(
-      DeterministicLevel.DETERMINATE,
-      DeterministicLevel.UNORDERED,
-      DeterministicLevel.INDETERMINATE).foreach { level =>
-      val input = shuffleInput(parentWithLevel(level), positionalRoundRobin = false)
-      assert(input.outputDeterministicLevel == level)
-    }
-  }
-
   test("positional round robin refuses to run a retried map task") {
     withSQLConf(batchGranularKey -> "true") {
       // A task re-run inside the current stage attempt, and a re-submitted stage, which is what
@@ -177,36 +131,24 @@ class CometNativeRoundRobinRetrySuite extends CometTestBase {
     }
   }
 
-  test("positional round robin lets a first attempt through") {
-    withSQLConf(batchGranularKey -> "true") {
-      val writer = writerFor(RoundRobinPartitioning(4), taskContextFor(0, 0))
-      assert(!writeFailureMessage(writer).startsWith(refusalPrefix))
-    }
-  }
-
-  test("the retry guard is scoped to positional round robin") {
-    // Hash and range partitioning place rows by content, and content-hash round robin does too,
-    // so a retry of any of those is safe and must not be turned into a job failure.
-    withSQLConf(batchGranularKey -> "true") {
-      Seq(SinglePartition, HashPartitioning(Seq(Literal(1)), 4), RangePartitioning(Nil, 4))
-        .foreach { partitioning =>
-          val writer = writerFor(partitioning, taskContextFor(1, 1))
-          assert(
-            !writeFailureMessage(writer).startsWith(refusalPrefix),
-            s"$partitioning should not be treated as positional round robin")
+  test("the retry guard fires only for positional round robin on a retried attempt") {
+    // Everything here must reach the writer. Hash and range partitioning place rows by content,
+    // content-hash round robin does too, a first attempt is not a retry, and failOnRetry=false
+    // hands retry handling back to the DAGScheduler.
+    Seq(
+      ("first attempt", "true", "true", RoundRobinPartitioning(4): Partitioning, 0, 0),
+      ("batchGranular off", "false", "true", RoundRobinPartitioning(4), 1, 1),
+      ("failOnRetry off", "true", "false", RoundRobinPartitioning(4), 1, 1),
+      ("single output partition", "true", "true", RoundRobinPartitioning(1), 1, 1),
+      ("hash partitioning", "true", "true", HashPartitioning(Seq(Literal(1)), 4), 1, 1),
+      ("range partitioning", "true", "true", RangePartitioning(Nil, 4), 1, 1),
+      ("single partition", "true", "true", SinglePartition, 1, 1)).foreach {
+      case (label, granular, failOnRetry, partitioning, stageAttempt, taskAttempt) =>
+        withSQLConf(batchGranularKey -> granular, failOnRetryKey -> failOnRetry) {
+          val writer = writerFor(partitioning, taskContextFor(stageAttempt, taskAttempt))
+          assert(!writeFailureMessage(writer).startsWith(refusalPrefix), label)
         }
     }
-    withSQLConf(batchGranularKey -> "false") {
-      val writer = writerFor(RoundRobinPartitioning(4), taskContextFor(1, 1))
-      assert(!writeFailureMessage(writer).startsWith(refusalPrefix))
-    }
-  }
-
-  test("failOnRetry=false hands retry handling back to the DAGScheduler") {
-    withSQLConf(batchGranularKey -> "true", failOnRetryKey -> "false") {
-      val writer = writerFor(RoundRobinPartitioning(4), taskContextFor(1, 1))
-      assert(!writeFailureMessage(writer).startsWith(refusalPrefix))
-    }
   }
 
   /**
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
index 99d9fe93e0c..92a06486e82 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala
@@ -45,28 +45,56 @@ import org.apache.comet.serde.OperatorOuterClass.Operator
  */
 class CometNativeShuffleInputRDDSuite extends CometTestBase {
 
+  /** An input RDD that reports exactly `level`, standing in for a real upstream subtree. */
+  private def parentWithLevel(level: DeterministicLevel.Value): RDD[AnyRef] =
+    new RDD[AnyRef](spark.sparkContext, Nil) {
+      override protected def getOutputDeterministicLevel: DeterministicLevel.Value = level
+      override protected def getPartitions: Array[Partition] = Array.empty
+      override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] =
+        Iterator.empty
+    }
+
+  private def inputWithParent(
+      level: DeterministicLevel.Value,
+      positionalRoundRobin: Boolean = false): CometNativeShuffleInputRDD =
+    new CometNativeShuffleInputRDD(
+      spark.sparkContext,
+      Seq(parentWithLevel(level)),
+      0,
+      Set.empty,
+      CometMetricNode(Map.empty),
+      positionalRoundRobin = positionalRoundRobin)
+
   test("native shuffle input preserves its parents' determinism") {
+    // Content-hash round robin, like every other partitioning, places rows by hash, so its
+    // output is a pure function of the rows and nothing here is indeterminate on its own account.
     Seq(
       DeterministicLevel.DETERMINATE,
       DeterministicLevel.UNORDERED,
       DeterministicLevel.INDETERMINATE).foreach { level =>
-      val parent = new RDD[AnyRef](spark.sparkContext, Nil) {
-        override protected def getOutputDeterministicLevel: DeterministicLevel.Value = level
-        override protected def getPartitions: Array[Partition] = Array.empty
-        override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] =
-          Iterator.empty
-      }
-      val input = new CometNativeShuffleInputRDD(
-        spark.sparkContext,
-        Seq(parent),
-        0,
-        Set.empty,
-        CometMetricNode(Map.empty))
+      val input = inputWithParent(level)
       assert(input.outputDeterministicLevel == level)
       assert(input.copyForLocalShuffle().outputDeterministicLevel == level)
     }
   }
 
+  test("positional round robin declares indeterminate output unless its parent is determinate") {
+    // A determinate parent (a plain scan) replays identically, so positional assignment is
+    // reproducible and per-task retry stays cheap. Anything below another exchange is unordered,
+    // which is where Spark's own round robin flips to indeterminate too.
+    Seq(
+      DeterministicLevel.DETERMINATE -> DeterministicLevel.DETERMINATE,
+      DeterministicLevel.UNORDERED -> DeterministicLevel.INDETERMINATE,
+      DeterministicLevel.INDETERMINATE -> DeterministicLevel.INDETERMINATE).foreach {
+      case (parentLevel, expected) =>
+        val input = inputWithParent(parentLevel, positionalRoundRobin = true)
+        assert(input.outputDeterministicLevel == expected, s"parent was $parentLevel")
+        assert(
+          input.copyForLocalShuffle().outputDeterministicLevel == expected,
+          s"local fallback lost the declaration for parent $parentLevel")
+    }
+  }
+
   test("local shuffle input is an independent sibling with the same partition inputs") {
     val upstream = new RDD[AnyRef](spark.sparkContext, Nil) {
       override protected def getPartitions: Array[Partition] = Array.tabulate(2) { i =>

From 6f595108c7db6b702d9208537ac7906aa739efd1 Mon Sep 17 00:00:00 2001
From: comphead <comphead@ukr.net>
Date: Mon, 21 Sep 2026 12:42:58 -0700
Subject: [PATCH 4/4] vectorized RR partitioning

---
 native/shuffle/src/partitioners/multi_partition.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs
index 0de84db6079..7028bff9f9e 100644
--- a/native/shuffle/src/partitioners/multi_partition.rs
+++ b/native/shuffle/src/partitioners/multi_partition.rs
@@ -20,7 +20,7 @@ use crate::partitioners::partitioned_batch_iterator::PartitionedBatchesProducer;
 use crate::partitioners::ShufflePartitioner;
 use crate::writers::PartitionWriter;
 use crate::{comet_partitioning, CometPartitioning, RoundRobinStrategy};
-use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch};
+use arrow::array::{Array, ArrayData, RecordBatch};
 use datafusion::common::utils::proxy::VecAllocExt;
 use datafusion::common::{DataFusionError, HashSet};
 use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
@@ -668,7 +668,7 @@ impl<T: PartitionWriter> Debug for MultiPartitionShuffleRepartitioner<T> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::Int64Array;
+    use arrow::array::{ArrayRef, Int64Array};
     use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
 
     #[derive(Default)]