stream ROW_NUMBER output by input batch

docs/source/contributor-guide/debugging.md

@@ -260,21 +260,28 @@ so `spark.comet.debug.enabled` and `COMET_LOG_LEVEL` do not affect it.
 #### Example: dumping window operator output
 
 In `WindowBuilder::build` in `native/core/src/execution/planner/window.rs`,
-the window operator builds `window_agg`, either a `BoundedWindowAggExec` or a
-`WindowAggExec` depending on whether every window expression runs with bounded
-memory. It then optionally wraps `window_agg` in a casting `ProjectionExec` and
-builds the `SparkPlan` from either the projection or `window_agg`. To see the raw
+the window operator builds `window_agg`. A window containing only `ROW_NUMBER`
+uses `RowNumberExec`, which generates one rank array per input batch and carries
+its final partition key and counter across batches. Other functions use
+`BoundedWindowAggExec` or `WindowAggExec`, depending on whether every expression
+runs with bounded memory. The builder then optionally wraps `window_agg` in a casting
+`ProjectionExec` and builds the `SparkPlan` from either the projection or `window_agg`. To see the raw
 window output, wrap `window_agg` right after it is built, before the cast
 projection:
 
 ```rust
 let window_agg: Arc<dyn ExecutionPlan> = if all_bounded {
-    Arc::new(BoundedWindowAggExec::try_new(
+    let window = BoundedWindowAggExec::try_new(
         window_expr,
         Arc::clone(&child.native_plan),
         InputOrderMode::Sorted,
         !partition_exprs.is_empty(),
-    )?)
+    )?;
+    if RowNumberExec::supports(&window) {
+        Arc::new(RowNumberExec::try_new(window)?)
+    } else {
+        Arc::new(window)
+    }
 } else {
     Arc::new(WindowAggExec::try_new(
         window_expr,

docs/source/contributor-guide/memory_management.md

@@ -489,6 +489,25 @@ Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservat
 
 An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses.
 
+### ROW_NUMBER windows
+
+A window containing only `ROW_NUMBER` uses `RowNumberExec` in Rust. It generates one rank
+array per input batch and retains only the final partition key and counter between batches.
+The key's Arrow buffers are charged to the `ROW_NUMBER partition key` consumer. Top-level
+views and dictionaries are rebuilt from the selected value so unrelated source bytes are released.
+Nested values may retain shared buffers, which are included in the reservation. A continuing
+partition reuses its key. Completion, cancellation, and errors release the reservation.
+
+Replacement keys reuse the reservation, so only a change in retained bytes adjusts the pool.
+The key has already been allocated and cannot spill, so its reservation uses `resize`.
+Comet's unified pools carry any bytes Spark does not grant as overcommit. This avoids
+failing the task solely because the pool cannot admit the key, but retained keys can
+exceed the pool budget.
+
+The reservation covers retained key state. Per-batch comparison, key-copy workspace, and emitted
+rank arrays follow the allocation model described in [The accounting gap](#the-accounting-gap).
+Other window functions keep their existing DataFusion executors and memory behaviour.
+
 ### Native writers
 
 Both native writers reserve what they hold between batches through a single consumer per task,

native/core/src/execution/planner.rs

@@ -4460,6 +4460,9 @@ mod tests {
         };
         let (mut scans, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap();
         assert_eq!(planned.additional_native_plans.len(), 1);
+        assert!(planned.additional_native_plans[0]
+            .downcast_ref::<datafusion_comet_operators::RowNumberExec>()
+            .is_some());
 
         let rows = 100;
         let mut inputs = [vec![

native/core/src/execution/planner/window.rs

@@ -40,7 +40,7 @@ use datafusion::{
         ExecutionPlan, InputOrderMode,
     },
 };
-use datafusion_comet_operators::{PartitionedRankLimitExec, WindowFnKind};
+use datafusion_comet_operators::{PartitionedRankLimitExec, RowNumberExec, WindowFnKind};
 use datafusion_comet_proto::{
     spark_expression::{
         self, agg_expr::ExprStruct as AggExprStruct, expr::ExprStruct, literal::Value, AggExpr,
@@ -132,12 +132,17 @@ impl OperatorBuilder for WindowBuilder {
         let window_expr = window_expr?;
         let all_bounded = window_expr.iter().all(|e| e.uses_bounded_memory());
         let window_agg: Arc<dyn ExecutionPlan> = if all_bounded {
-            Arc::new(BoundedWindowAggExec::try_new(
+            let window = BoundedWindowAggExec::try_new(
                 window_expr,
                 Arc::clone(&child.native_plan),
                 InputOrderMode::Sorted,
                 !partition_exprs.is_empty(),
-            )?)
+            )?;
+            if RowNumberExec::supports(&window) {
+                Arc::new(RowNumberExec::try_new(window)?)
+            } else {
+                Arc::new(window)
+            }
         } else {
             Arc::new(WindowAggExec::try_new(
                 window_expr,

native/operators/benches/window.rs

@@ -52,6 +52,7 @@ use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
 use datafusion::physical_plan::windows::{create_window_expr, BoundedWindowAggExec};
 use datafusion::physical_plan::{ExecutionPlan, InputOrderMode};
 use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_comet_operators::RowNumberExec;
 use futures::StreamExt;
 use tokio::runtime::Runtime;
 
@@ -299,7 +300,9 @@ fn run(runtime: &Runtime, plan: &Arc<dyn ExecutionPlan>, ctx: &Arc<TaskContext>)
         let mut stream = plan.execute(0, Arc::clone(ctx)).unwrap();
         let mut rows = 0;
         while let Some(batch) = stream.next().await {
-            rows += batch.unwrap().num_rows();
+            let batch = batch.unwrap();
+            rows += batch.num_rows();
+            std::hint::black_box(batch);
         }
         rows
     });
@@ -380,6 +383,71 @@ fn criterion_benchmark(c: &mut Criterion) {
     }
     group.finish();
 
+    // Compare the same ROW_NUMBER expressions and sorted batches through the generic
+    // window and the batch-at-a-time executor. Validate every rank before measuring.
+    let mut group = c.benchmark_group("row_number_executor");
+    group
+        .sample_size(10)
+        .measurement_time(Duration::from_secs(3));
+    group.throughput(Throughput::Elements(ROWS as u64));
+    for rows_per_partition in [1usize, 4, 64, 16384, ROWS] {
+        let s = shape(Func::RowNumber, Arg::Int64, rows_per_partition, 0);
+        let schema = s.schema();
+        let batches = (0..BATCHES)
+            .map(|i| s.batch(&schema, i))
+            .collect::<Vec<_>>();
+        let build_plan = |streaming| -> Arc<dyn ExecutionPlan> {
+            let generic = s.plan(&session, &batches);
+            if streaming {
+                Arc::new(
+                    RowNumberExec::try_new(
+                        generic
+                            .downcast_ref::<BoundedWindowAggExec>()
+                            .unwrap()
+                            .clone(),
+                    )
+                    .unwrap(),
+                )
+            } else {
+                generic
+            }
+        };
+        for (name, streaming) in [("bounded", false), ("streaming", true)] {
+            let plan = build_plan(streaming);
+            runtime.block_on(async {
+                let mut stream = plan.execute(0, Arc::clone(&task_ctx)).unwrap();
+                let mut row = 0;
+                while let Some(batch) = stream.next().await {
+                    let batch = batch.unwrap();
+                    let ranks = batch
+                        .columns()
+                        .last()
+                        .unwrap()
+                        .as_any()
+                        .downcast_ref::<arrow::array::UInt64Array>()
+                        .unwrap();
+                    for rank in ranks.values() {
+                        assert_eq!(*rank, (row % rows_per_partition + 1) as u64);
+                        row += 1;
+                    }
+                }
+                assert_eq!(row, ROWS);
+            });
+            group.bench_function(BenchmarkId::new(name, rows_per_partition), |b| {
+                // Each execution registers metrics. Keep plan creation and disposal outside timing.
+                b.iter_batched(
+                    || build_plan(streaming),
+                    |plan| {
+                        run(&runtime, &plan, &task_ctx);
+                        plan
+                    },
+                    criterion::BatchSize::PerIteration,
+                );
+            });
+        }
+    }
+    group.finish();
+
     // The argument type alone, at a group size where per-row work dominates.
     let mut group = c.benchmark_group("window_argument_type");
     group

native/operators/src/lib.rs

@@ -33,6 +33,7 @@ mod metric_owning;
 mod projection;
 mod range;
 mod rank_limit;
+mod row_number;
 mod sample;
 
 pub use expand::ExpandExec;
@@ -42,4 +43,5 @@ pub use membership_key_reduction::MembershipKeyReductionExec;
 pub use projection::CometProjectionExec;
 pub use range::{range_exec, range_sort_expr};
 pub use rank_limit::{PartitionedRankLimitExec, WindowFnKind};
+pub use row_number::RowNumberExec;
 pub use sample::SampleExec;