From a50622e823441b1dc173d2ff78d297b2f9248d1e Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Tue, 3 Nov 2026 09:02:16 -0700 Subject: [PATCH 1/9] stream ROW_NUMBER output by input batch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generate ROW_NUMBER results per input batch, carrying the final partition key and counter across batches. This avoids per-group slicing and window state allocations, reducing overhead for many small partitions. Preserve existing plan contracts and result casts while accounting for retained key buffers. Add regression tests for boundaries and cleanup, plus benchmarks comparing the streaming and generic executors. Benchmark timing: | Rows per partition | Generic | Streaming | Speedup | |---:|---:|---:|---:| | 1 | 243.08 ms | 3.38 ms | 71.8× | | 4 | 66.53 ms | 2.75 ms | 24.2× | | 64 | 9.82 ms | 2.34 ms | 4.2× | | 16,384 | 6.39 ms | 2.32 ms | 2.8× | | 262,144 | 7.22 ms | 2.36 ms | 3.1× | --- docs/source/contributor-guide/debugging.md | 19 +- .../contributor-guide/memory_management.md | 13 + native/core/src/execution/planner.rs | 3 + native/core/src/execution/planner/window.rs | 11 +- native/operators/benches/window.rs | 55 ++- native/operators/src/lib.rs | 2 + native/operators/src/row_number.rs | 326 +++++++++++++ native/operators/src/row_number/tests.rs | 435 ++++++++++++++++++ .../comet/exec/CometWindowExecSuite.scala | 35 ++ 9 files changed, 889 insertions(+), 10 deletions(-) create mode 100644 native/operators/src/row_number.rs create mode 100644 native/operators/src/row_number/tests.rs diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md index 8795854193d..92550933157 100644 --- a/docs/source/contributor-guide/debugging.md +++ b/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 = 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, diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 4e7248f2ba3..70275d766be 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -489,6 +489,19 @@ 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. + +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, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 603bba872bc..2a6a32f4a33 100644 --- a/native/core/src/execution/planner.rs +++ b/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::() + .is_some()); let rows = 100; let mut inputs = [vec![ diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 2d15fdc86e1..674e6f13377 100644 --- a/native/core/src/execution/planner/window.rs +++ b/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 = 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, diff --git a/native/operators/benches/window.rs b/native/operators/benches/window.rs index 131bbd0067b..df1e6e61680 100644 --- a/native/operators/benches/window.rs +++ b/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, ctx: &Arc) 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,56 @@ 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::>(); + let generic = s.plan(&session, &batches); + let streaming: Arc = Arc::new( + RowNumberExec::try_new( + generic + .downcast_ref::() + .unwrap() + .clone(), + ) + .unwrap(), + ); + for (name, plan) in [("bounded", generic), ("streaming", 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::() + .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| { + b.iter(|| run(&runtime, &plan, &task_ctx)); + }); + } + } + 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 diff --git a/native/operators/src/lib.rs b/native/operators/src/lib.rs index 05a0324567c..7a58749fb7f 100644 --- a/native/operators/src/lib.rs +++ b/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; diff --git a/native/operators/src/row_number.rs b/native/operators/src/row_number.rs new file mode 100644 index 00000000000..623919837a3 --- /dev/null +++ b/native/operators/src/row_number.rs @@ -0,0 +1,326 @@ +// 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. + +//! Batch-at-a-time ROW_NUMBER for input sorted by its normalised partition keys. + +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{ArrayRef, RecordBatch, UInt64Array}; +use arrow::compute::{partition, take, SortOptions}; +use arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, plan_err, Result, ScalarValue, Statistics}; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion::execution::TaskContext; +use datafusion::functions_window::row_number::RowNumber; +use datafusion::physical_expr::window::StandardWindowExpr; +use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{ + CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions, +}; +use datafusion::physical_plan::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, +}; +use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + PlanProperties, SendableRecordBatchStream, +}; +use futures::StreamExt; + +/// Generates every ROW_NUMBER result in one pass over each input batch. +/// +/// The unexecuted window supplies the schema, ordering, distribution and statistics contracts. +/// Its expressions must all be the built-in ROW_NUMBER over the same partition and sort keys. +/// The caller must normalise floating-point partition keys before building the window, just as +/// it does for ordinary native windows. Other window functions keep their existing executor. +#[derive(Debug)] +pub struct RowNumberExec { + window: BoundedWindowAggExec, + metrics: ExecutionPlanMetricsSet, +} + +impl RowNumberExec { + pub fn supports(window: &BoundedWindowAggExec) -> bool { + let Some(first) = window.window_expr().first() else { + return false; + }; + window.window_expr().iter().all(|expr| { + let Some(standard) = expr.as_any().downcast_ref::() else { + return false; + }; + let Some(udf) = standard + .get_standard_func_expr() + .as_any() + .downcast_ref::() + else { + return false; + }; + udf.fun().inner().downcast_ref::().is_some() + && expr.expressions().is_empty() + && expr.partition_by() == first.partition_by() + && expr.order_by() == first.order_by() + }) + } + + pub fn try_new(window: BoundedWindowAggExec) -> Result { + if !Self::supports(&window) { + return plan_err!( + "RowNumberExec requires only built-in ROW_NUMBER expressions with matching keys" + ); + } + Ok(Self { + window, + metrics: ExecutionPlanMetricsSet::new(), + }) + } +} + +impl DisplayAs for RowNumberExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometRowNumberExec: ")?; + self.window.fmt_as(t, f) + } +} + +impl ExecutionPlan for RowNumberExec { + fn name(&self) -> &str { + "CometRowNumberExec" + } + + fn properties(&self) -> &Arc { + self.window.properties() + } + + fn children(&self) -> Vec<&Arc> { + self.window.children() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.window.apply_expressions(f) + } + + fn required_input_ordering(&self) -> Vec> { + self.window.required_input_ordering() + } + + fn required_input_distribution(&self) -> Vec { + self.window + .input_distribution_requirements() + .into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + self.window.input_distribution_requirements() + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + let replaced = Arc::new(self.window.clone()).replace_children(children, options)?; + let Some(window) = replaced.downcast_ref::() else { + return internal_err!("ROW_NUMBER child replacement changed its window type"); + }; + Ok(Arc::new(Self::try_new(window.clone())?)) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self + .window + .input() + .execute(partition, Arc::clone(&context))?; + let partition_by = self.window.window_expr()[0].partition_by().to_vec(); + let output_count = self.window.window_expr().len(); + let schema = self.schema(); + let state = RowNumberState { + key: vec![], + rank: 0, + reservation: MemoryConsumer::new("ROW_NUMBER partition key") + .register(context.memory_pool()), + }; + let metrics = BaselineMetrics::new(&self.metrics, partition); + let stream = futures::stream::try_unfold( + (input, state, metrics), + move |(mut input, mut state, metrics)| { + let partition_by = partition_by.clone(); + let schema = Arc::clone(&schema); + async move { + let Some(batch) = input.next().await else { + return Ok(None); + }; + let batch = batch?; + let output = { + let _timer = metrics.elapsed_compute().timer(); + state.evaluate(batch, &partition_by, schema, output_count)? + }; + (&output).record_output(&metrics); + Ok(Some((output, (input, state, metrics)))) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + self.window.child_stats_requests(partition) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + self.window.statistics_from_inputs(input_stats, args) + } +} + +struct RowNumberState { + key: Vec, + rank: u64, + reservation: MemoryReservation, +} + +impl RowNumberState { + fn evaluate( + &mut self, + batch: RecordBatch, + partition_by: &[Arc], + schema: SchemaRef, + output_count: usize, + ) -> Result { + let rows = batch.num_rows(); + let mut ranks = Vec::with_capacity(rows); + if rows != 0 { + let keys = partition_by + .iter() + .map(|expr| expr.evaluate(&batch)?.into_array(rows)) + .collect::>>()?; + let mut continues = keys.len() == self.key.len(); + for (key, previous) in keys.iter().zip(&self.key) { + let compare = arrow::array::make_comparator( + previous.as_ref(), + key.as_ref(), + SortOptions::default(), + )?; + continues &= compare(0, 0).is_eq(); + } + let ranges = if keys.is_empty() { + std::iter::once(0..rows).collect::>() + } else { + partition(&keys)?.ranges() + }; + for range in &ranges { + let start = if range.start == 0 && continues { + self.rank + } else { + 0 + }; + ranks.extend((1..=range.len() as u64).map(|offset| start.wrapping_add(offset))); + } + self.rank = *ranks.last().unwrap(); + if !keys.is_empty() && (!continues || ranges.len() != 1) { + // The old boundary key is no longer needed. Do not keep both copies alive + // when a wide key changes. A continued partition needs no new key copy. + self.key = vec![]; + self.reservation.free(); + let index = UInt64Array::from(vec![(rows - 1) as u64]); + let key = keys + .iter() + .map(|array| { + if matches!( + array.data_type(), + DataType::Dictionary(_, _) | DataType::Utf8View | DataType::BinaryView + ) { + // A view or dictionary take retains its source buffers. Rebuild + // the logical value so unrelated bytes do not stay alive. + ScalarValue::try_from_array(array.as_ref(), rows - 1)?.to_array() + } else { + take(array.as_ref(), &index, None).map_err(Into::into) + } + }) + .collect::>>()?; + // Nested views and dictionaries can still share backing buffers after take. + // Charge the buffers retained, rather than assuming a one-row key is small. + let size = key + .iter() + .map(|array| array.get_array_memory_size()) + .sum::() + + key.capacity() * std::mem::size_of::(); + self.reservation.try_resize(size)?; + self.key = key; + } + } + let ranks: ArrayRef = Arc::new(UInt64Array::from(ranks)); + let columns = batch + .columns() + .iter() + .cloned() + .chain(std::iter::repeat_n(ranks, output_count)) + .collect(); + Ok(RecordBatch::try_new(schema, columns)?) + } +} + +#[cfg(test)] +mod tests; diff --git a/native/operators/src/row_number/tests.rs b/native/operators/src/row_number/tests.rs new file mode 100644 index 00000000000..b7260049073 --- /dev/null +++ b/native/operators/src/row_number/tests.rs @@ -0,0 +1,435 @@ +// 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. + +use super::*; + +use arrow::array::{ + Array, DictionaryArray, Int64Array, Int8Array, ListArray, StringArray, StringViewArray, + StructArray, +}; +use arrow::buffer::NullBuffer; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Field, Int64Type, Schema}; +use datafusion::common::ScalarValue; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion::execution::FunctionRegistry; +use datafusion::logical_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, +}; +use datafusion::physical_expr::expressions::col; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_plan::windows::create_window_expr; +use datafusion::physical_plan::{collect, ExecutionPlanProperties, InputOrderMode}; +use datafusion::prelude::SessionContext; + +fn batch(keys: ArrayRef) -> RecordBatch { + let rows = keys.len(); + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("k", keys.data_type().clone(), true), + Field::new("i", DataType::Int64, false), + ])), + vec![keys, Arc::new(Int64Array::from_iter_values(0..rows as i64))], + ) + .unwrap() +} + +fn window(batches: &[RecordBatch], partitioned: bool, functions: &[&str]) -> BoundedWindowAggExec { + let schema = batches[0].schema(); + let sort = |name| PhysicalSortExpr::new(col(name, &schema).unwrap(), SortOptions::default()); + let ordering = if partitioned { + vec![sort("k"), sort("i")] + } else { + vec![sort("i")] + }; + let source = MemorySourceConfig::try_new(&[batches.to_vec()], Arc::clone(&schema), None) + .unwrap() + .try_with_sort_information(vec![LexOrdering::new(ordering).unwrap()]) + .unwrap(); + let input = DataSourceExec::from_data_source(source); + let ctx = SessionContext::new(); + let keys = if partitioned { + vec![col("k", &schema).unwrap()] + } else { + vec![] + }; + let frame = Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )); + let expressions = functions + .iter() + .enumerate() + .map(|(index, function)| { + create_window_expr( + &WindowFunctionDefinition::WindowUDF(ctx.udwf(function).unwrap()), + format!("result_{index}"), + &[], + &keys, + &[sort("i")], + Arc::clone(&frame), + Arc::clone(&schema), + false, + false, + None, + ) + .unwrap() + }) + .collect(); + BoundedWindowAggExec::try_new(expressions, input, InputOrderMode::Sorted, partitioned).unwrap() +} + +async fn compare( + batches: &[RecordBatch], + partitioned: bool, + functions: &[&str], +) -> Vec { + let generic = window(batches, partitioned, functions); + let streaming = Arc::new(RowNumberExec::try_new(generic.clone()).unwrap()); + let ctx = SessionContext::new(); + let expected = collect(Arc::new(generic), ctx.task_ctx()).await.unwrap(); + let actual = collect(Arc::clone(&streaming) as _, ctx.task_ctx()) + .await + .unwrap(); + assert_eq!( + concat_batches(&streaming.schema(), &actual).unwrap(), + concat_batches(&streaming.schema(), &expected).unwrap() + ); + assert_eq!( + actual.iter().map(RecordBatch::num_rows).collect::>(), + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>() + ); + actual +} + +#[tokio::test] +async fn partitions_continue_across_slices_and_empty_batches() { + let input = batch(Arc::new(Int64Array::from(vec![ + None, + None, + Some(1), + Some(1), + Some(1), + Some(2), + Some(2), + Some(3), + ]))); + let batches = vec![ + input.slice(0, 0), + input.slice(0, 3), + input.slice(3, 0), + input.slice(3, 1), + input.slice(4, 4), + input.slice(8, 0), + ]; + let output = compare(&batches, true, &["row_number"]).await; + let ranks = output + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!(ranks, vec![1, 2, 1, 2, 3, 1, 2, 1]); +} + +#[tokio::test] +async fn global_windows_share_alias_output_and_input_columns() { + let input = batch(Arc::new(Int64Array::from_iter_values(0..23))); + let batches = vec![input.slice(0, 5), input.slice(5, 7), input.slice(12, 11)]; + let actual = compare(&batches, false, &["row_number", "row_number"]).await; + for (output, input) in actual.iter().zip(&batches) { + assert!(Arc::ptr_eq(output.column(0), input.column(0))); + assert!(Arc::ptr_eq(output.column(2), output.column(3))); + assert_eq!(output.column(2).null_count(), 0); + } + let last = actual[2] + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(last.values().as_ref(), &(13..=23).collect::>()); +} + +#[tokio::test] +async fn empty_input_preserves_schema() { + compare( + &[batch(Arc::new(Int64Array::from(Vec::::new())))], + true, + &["row_number"], + ) + .await; +} + +#[tokio::test] +async fn dictionary_keys_compare_values_across_different_codebooks() { + use arrow::datatypes::Int8Type; + let first = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![None, Some(0), Some(1), Some(1)]), + Arc::new(StringArray::from(vec![None, Some("a")])), + ) + .unwrap(), + ); + let second = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(1), Some(0), Some(0)]), + Arc::new(StringArray::from(vec![Some("b"), Some("a")])), + ) + .unwrap(), + ); + let output = compare(&[batch(first), batch(second)], true, &["row_number"]).await; + let ranks = output[1] + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ranks.values().as_ref(), &[3, 1, 2]); +} + +#[tokio::test] +async fn nested_keys_and_nullable_children_match_generic_windows() { + let keys = Arc::new(ListArray::from_iter_primitive::(vec![ + None, + Some(vec![]), + Some(vec![None]), + Some(vec![None]), + Some(vec![Some(2)]), + Some(vec![Some(2)]), + ])); + let input = batch(keys); + compare( + &[input.slice(0, 3), input.slice(3, 3)], + true, + &["row_number"], + ) + .await; + let keys = Arc::new(StructArray::new( + vec![Field::new("value", DataType::Int64, true)].into(), + vec![Arc::new(Int64Array::from(vec![ + None, + None, + Some(1), + Some(1), + ]))], + Some(NullBuffer::from(vec![false, false, true, true])), + )); + let input = batch(keys); + compare( + &[input.slice(0, 1), input.slice(1, 2), input.slice(3, 1)], + true, + &["row_number"], + ) + .await; +} + +#[test] +fn other_functions_cannot_enter_the_specialised_executor() { + let input = batch(Arc::new(Int64Array::from(vec![1, 1, 2]))); + let mixed = window(std::slice::from_ref(&input), true, &["row_number", "rank"]); + assert!(!RowNumberExec::supports(&mixed)); + assert!(RowNumberExec::try_new(mixed).is_err()); + assert!(!RowNumberExec::supports(&window(&[input], true, &["rank"]))); +} + +#[test] +fn plan_rewrites_keep_the_window_contracts() { + let input = batch(Arc::new(Int64Array::from(vec![1, 1, 2]))); + let generic = window(&[input], true, &["row_number"]); + let expected_order = generic.required_input_ordering(); + let expected_partitions = generic.properties().output_partitioning().partition_count(); + let plan = Arc::new(RowNumberExec::try_new(generic.clone()).unwrap()); + let child = Arc::clone(generic.input()); + for replacement in [ + Arc::clone(&plan) + .replace_children( + vec![Arc::clone(&child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(), + Arc::clone(&plan) + .replace_children( + vec![Arc::clone(&child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + .unwrap(), + ] { + assert_eq!(replacement.schema(), generic.schema()); + assert_eq!(replacement.required_input_ordering(), expected_order); + assert!(matches!( + replacement.input_distribution_requirements().into_per_child().as_slice(), + [Distribution::KeyPartitioned(keys)] if keys == &generic.partition_keys() + )); + assert_eq!( + replacement.output_partitioning().partition_count(), + expected_partitions + ); + assert!(matches!( + replacement.cardinality_effect(), + CardinalityEffect::Equal + )); + assert_eq!(replacement.maintains_input_order(), vec![true]); + } +} + +#[test] +fn retained_view_buffers_are_charged_and_a_continuing_key_is_reused() { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let reservation = MemoryConsumer::new("test boundary").register(&pool); + let mut state = RowNumberState { + key: vec![], + rank: 0, + reservation, + }; + let keys: ArrayRef = Arc::new(StringViewArray::from(vec!["x".repeat(16 * 1024); 3])); + let input = batch(keys); + let plan = window(std::slice::from_ref(&input), true, &["row_number"]); + let expressions = plan.window_expr()[0].partition_by(); + state + .evaluate(input.clone(), expressions, plan.schema(), 1) + .unwrap(); + let retained = pool.reserved(); + assert!(retained >= 16 * 1024); + let key = Arc::clone(&state.key[0]); + state + .evaluate(input, expressions, plan.schema(), 1) + .unwrap(); + assert!(Arc::ptr_eq(&key, &state.key[0])); + assert_eq!(pool.reserved(), retained); + drop(key); + drop(state); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn refused_boundary_reservations_are_released() { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1)); + let reservation = MemoryConsumer::new("test refusal").register(&pool); + let mut state = RowNumberState { + key: vec![], + rank: 0, + reservation, + }; + let input = batch(Arc::new(StringArray::from(vec!["some partition key"]))); + let plan = window(std::slice::from_ref(&input), true, &["row_number"]); + let result = state.evaluate( + input, + plan.window_expr()[0].partition_by(), + plan.schema(), + 1, + ); + assert!(result.is_err()); + assert!(state.key.is_empty()); + assert_eq!(pool.reserved(), 0); +} + +#[tokio::test] +async fn completion_and_cancellation_release_the_boundary_key() { + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + let input = batch(Arc::new(Int64Array::from(vec![1, 1, 1, 2]))); + let plan = RowNumberExec::try_new(window( + &[input.slice(0, 2), input.slice(2, 2)], + true, + &["row_number"], + )) + .unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc() + .unwrap(); + let ctx = SessionContext::new_with_config_rt(Default::default(), runtime); + let mut stream = plan.execute(0, ctx.task_ctx()).unwrap(); + assert_eq!(stream.next().await.unwrap().unwrap().num_rows(), 2); + assert!(pool.reserved() > 0); + drop(stream); + assert_eq!(pool.reserved(), 0); + let mut stream = plan.execute(0, ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + batch.unwrap(); + } + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn a_dictionary_boundary_does_not_retain_unused_values() { + use arrow::datatypes::Int8Type; + let values: ArrayRef = Arc::new(StringArray::from(vec![ + "small", + &"unused".repeat(64 * 1024), + ])); + let keys: ArrayRef = Arc::new( + DictionaryArray::::try_new(Int8Array::from(vec![0, 0]), values).unwrap(), + ); + let input = batch(keys); + let plan = window(std::slice::from_ref(&input), true, &["row_number"]); + let pool: Arc = Arc::new(GreedyMemoryPool::new(4096)); + let mut state = RowNumberState { + key: vec![], + rank: 0, + reservation: MemoryConsumer::new("dictionary boundary").register(&pool), + }; + state + .evaluate( + input, + plan.window_expr()[0].partition_by(), + plan.schema(), + 1, + ) + .unwrap(); + assert!(pool.reserved() < 4096); + drop(state); + assert_eq!(pool.reserved(), 0); +} + +#[test] +fn a_view_boundary_does_not_retain_an_unrelated_large_buffer() { + let keys: ArrayRef = Arc::new(StringViewArray::from(vec![ + "a".repeat(1024 * 1024), + "z".repeat(13), + ])); + let input = batch(keys); + let plan = window(std::slice::from_ref(&input), true, &["row_number"]); + let pool: Arc = Arc::new(GreedyMemoryPool::new(64 * 1024)); + let mut state = RowNumberState { + key: vec![], + rank: 0, + reservation: MemoryConsumer::new("view boundary").register(&pool), + }; + state + .evaluate( + input, + plan.window_expr()[0].partition_by(), + plan.schema(), + 1, + ) + .unwrap(); + assert!(pool.reserved() < 64 * 1024); + drop(state); + assert_eq!(pool.reserved(), 0); +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 642a73e337a..c54fca71215 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -1287,6 +1287,41 @@ class CometWindowExecSuite extends CometTestBase { } } + test("window: streaming ROW_NUMBER across input batches") { + withTempDir { dir => + (0 until 97) + .map(i => (if (i < 9) null else s"key-${i / 13}", i, "payload-" * 32)) + .toDF("k", "i", "payload") + .coalesce(1) + .write + .parquet(dir.toString) + spark.read.parquet(dir.toString).createOrReplaceTempView("streaming_row_number") + for (batchSize <- Seq(3, 16, 128); ansi <- Seq("false", "true")) { + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> batchSize.toString, + SQLConf.ANSI_ENABLED.key -> ansi) { + for (partition <- Seq("", "PARTITION BY k")) { + val df = sql(s""" + SELECT k, i, payload, + ROW_NUMBER() OVER ($partition ORDER BY i) AS first_rank, + ROW_NUMBER() OVER ($partition ORDER BY i) AS second_rank + FROM streaming_row_number + """) + checkSparkAnswerAndOperator(df) + assertCometWindowExecExists(df.queryExecution.executedPlan) + } + // A mixed window must retain the generic executor and its frame behaviour. + checkSparkAnswerAndOperator(sql(""" + SELECT k, i, + ROW_NUMBER() OVER (PARTITION BY k ORDER BY i) AS rn, + COUNT(*) OVER (PARTITION BY k ORDER BY i) AS n + FROM streaming_row_number + """)) + } + } + } + } + test("window: RANK with PARTITION BY and ORDER BY") { withTempDir { dir => (0 until 30) From 2ab0d3e99f746c5adda629228af0bdadff4f8a11 Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Tue, 3 Nov 2026 09:58:13 -0700 Subject: [PATCH 2/9] fix: require sorted input for streaming ROW_NUMBER Reject Linear and PartiallySorted windows because a partition can recur after another key, which the streaming evaluator cannot track. Cover both modes within and across batches and verify that the generic executor still produces the expected row numbers. --- native/operators/src/row_number.rs | 11 ++- native/operators/src/row_number/tests.rs | 90 +++++++++++++++++++++--- 2 files changed, 88 insertions(+), 13 deletions(-) diff --git a/native/operators/src/row_number.rs b/native/operators/src/row_number.rs index 623919837a3..09cf43eab9e 100644 --- a/native/operators/src/row_number.rs +++ b/native/operators/src/row_number.rs @@ -41,14 +41,15 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, - PlanProperties, SendableRecordBatchStream, + InputOrderMode, PlanProperties, SendableRecordBatchStream, }; use futures::StreamExt; /// Generates every ROW_NUMBER result in one pass over each input batch. /// /// The unexecuted window supplies the schema, ordering, distribution and statistics contracts. -/// Its expressions must all be the built-in ROW_NUMBER over the same partition and sort keys. +/// Input must use `InputOrderMode::Sorted` so each partition is contiguous. Its expressions must +/// all be the built-in ROW_NUMBER over the same partition and sort keys. /// The caller must normalise floating-point partition keys before building the window, just as /// it does for ordinary native windows. Other window functions keep their existing executor. #[derive(Debug)] @@ -59,6 +60,9 @@ pub struct RowNumberExec { impl RowNumberExec { pub fn supports(window: &BoundedWindowAggExec) -> bool { + if window.input_order_mode != InputOrderMode::Sorted { + return false; + } let Some(first) = window.window_expr().first() else { return false; }; @@ -83,7 +87,8 @@ impl RowNumberExec { pub fn try_new(window: BoundedWindowAggExec) -> Result { if !Self::supports(&window) { return plan_err!( - "RowNumberExec requires only built-in ROW_NUMBER expressions with matching keys" + "RowNumberExec requires InputOrderMode::Sorted and only built-in ROW_NUMBER \ + expressions with matching keys" ); } Ok(Self { diff --git a/native/operators/src/row_number/tests.rs b/native/operators/src/row_number/tests.rs index b7260049073..1c0221b547f 100644 --- a/native/operators/src/row_number/tests.rs +++ b/native/operators/src/row_number/tests.rs @@ -51,24 +51,41 @@ fn batch(keys: ArrayRef) -> RecordBatch { } fn window(batches: &[RecordBatch], partitioned: bool, functions: &[&str]) -> BoundedWindowAggExec { + window_with_order_mode( + batches, + if partitioned { &["k"] } else { &[] }, + functions, + InputOrderMode::Sorted, + ) +} + +fn window_with_order_mode( + batches: &[RecordBatch], + partition_columns: &[&str], + functions: &[&str], + input_order_mode: InputOrderMode, +) -> BoundedWindowAggExec { let schema = batches[0].schema(); let sort = |name| PhysicalSortExpr::new(col(name, &schema).unwrap(), SortOptions::default()); - let ordering = if partitioned { - vec![sort("k"), sort("i")] - } else { - vec![sort("i")] + let mut ordering: Vec<_> = match &input_order_mode { + InputOrderMode::Sorted => partition_columns.iter().map(|&name| sort(name)).collect(), + InputOrderMode::PartiallySorted(indices) => indices + .iter() + .map(|&index| sort(partition_columns[index])) + .collect(), + InputOrderMode::Linear => vec![], }; + ordering.push(sort("i")); let source = MemorySourceConfig::try_new(&[batches.to_vec()], Arc::clone(&schema), None) .unwrap() .try_with_sort_information(vec![LexOrdering::new(ordering).unwrap()]) .unwrap(); let input = DataSourceExec::from_data_source(source); let ctx = SessionContext::new(); - let keys = if partitioned { - vec![col("k", &schema).unwrap()] - } else { - vec![] - }; + let keys = partition_columns + .iter() + .map(|&name| col(name, &schema).unwrap()) + .collect::>(); let frame = Arc::new(WindowFrame::new_bounds( WindowFrameUnits::Rows, WindowFrameBound::Preceding(ScalarValue::UInt64(None)), @@ -93,7 +110,13 @@ fn window(batches: &[RecordBatch], partitioned: bool, functions: &[&str]) -> Bou .unwrap() }) .collect(); - BoundedWindowAggExec::try_new(expressions, input, InputOrderMode::Sorted, partitioned).unwrap() + BoundedWindowAggExec::try_new( + expressions, + input, + input_order_mode, + !partition_columns.is_empty(), + ) + .unwrap() } async fn compare( @@ -248,6 +271,53 @@ async fn nested_keys_and_nullable_children_match_generic_windows() { .await; } +#[tokio::test] +async fn interleaved_partitions_require_the_generic_executor() { + let input = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("prefix", DataType::Int64, false), + Field::new("k", DataType::Int64, false), + Field::new("i", DataType::Int64, false), + ])), + vec![ + Arc::new(Int64Array::from(vec![0, 0, 0, 1, 1, 1])) as ArrayRef, + Arc::new(Int64Array::from(vec![1, 2, 1, 2, 1, 2])), + Arc::new(Int64Array::from_iter_values(0..6)), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + for mode in [ + InputOrderMode::Linear, + InputOrderMode::PartiallySorted(vec![0]), + ] { + for batch_size in [1, 2, input.num_rows()] { + let batches = (0..input.num_rows()) + .step_by(batch_size) + .map(|start| input.slice(start, batch_size.min(input.num_rows() - start))) + .collect::>(); + // Only the prefix is grouped. The complete partition keys recur within each prefix. + let generic = + window_with_order_mode(&batches, &["prefix", "k"], &["row_number"], mode.clone()); + let mut columns = input.columns().to_vec(); + columns.push(Arc::new(UInt64Array::from(vec![1, 1, 2, 1, 1, 2]))); + let expected = RecordBatch::try_new(generic.schema(), columns).unwrap(); + let actual = collect(Arc::new(generic.clone()), ctx.task_ctx()) + .await + .unwrap(); + assert_eq!( + concat_batches(&expected.schema(), &actual).unwrap(), + expected + ); + assert!( + !RowNumberExec::supports(&generic), + "{mode:?} with batch size {batch_size} must retain the generic executor" + ); + assert!(RowNumberExec::try_new(generic).is_err()); + } + } +} + #[test] fn other_functions_cannot_enter_the_specialised_executor() { let input = batch(Arc::new(Int64Array::from(vec![1, 1, 2]))); From 95da6179b64d1064c29a07c8ee56d5230c855127 Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Tue, 3 Nov 2026 14:29:49 -0700 Subject: [PATCH 3/9] fix: restore ROW_NUMBER ordering checks and boundary coverage Validate the child's declared partition ordering before opening its stream. Cover exact batch-boundary resets and final-key updates for primitive and dictionary keys, and allow the Spark fixture to overwrite its temporary path. --- native/operators/src/row_number.rs | 8 +- native/operators/src/row_number/tests.rs | 178 +++++++++++++++++- .../comet/exec/CometWindowExecSuite.scala | 1 + 3 files changed, 179 insertions(+), 8 deletions(-) diff --git a/native/operators/src/row_number.rs b/native/operators/src/row_number.rs index 09cf43eab9e..cb0ecbf2629 100644 --- a/native/operators/src/row_number.rs +++ b/native/operators/src/row_number.rs @@ -24,7 +24,7 @@ use arrow::array::{ArrayRef, RecordBatch, UInt64Array}; use arrow::compute::{partition, take, SortOptions}; use arrow::datatypes::{DataType, SchemaRef}; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{internal_err, plan_err, Result, ScalarValue, Statistics}; +use datafusion::common::{exec_err, internal_err, plan_err, Result, ScalarValue, Statistics}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::TaskContext; use datafusion::functions_window::row_number::RowNumber; @@ -184,11 +184,15 @@ impl ExecutionPlan for RowNumberExec { partition: usize, context: Arc, ) -> Result { + let partition_by = self.window.window_expr()[0].partition_by(); + if self.window.partition_by_sort_keys()?.len() != partition_by.len() { + return exec_err!("All partition by columns should have an ordering in Sorted mode."); + } let input = self .window .input() .execute(partition, Arc::clone(&context))?; - let partition_by = self.window.window_expr()[0].partition_by().to_vec(); + let partition_by = partition_by.to_vec(); let output_count = self.window.window_expr().len(); let schema = self.schema(); let state = RowNumberState { diff --git a/native/operators/src/row_number/tests.rs b/native/operators/src/row_number/tests.rs index 1c0221b547f..4f96e3e6b0f 100644 --- a/native/operators/src/row_number/tests.rs +++ b/native/operators/src/row_number/tests.rs @@ -23,7 +23,7 @@ use arrow::array::{ }; use arrow::buffer::NullBuffer; use arrow::compute::concat_batches; -use arrow::datatypes::{DataType, Field, Int64Type, Schema}; +use arrow::datatypes::{DataType, Field, Int64Type, Int8Type, Schema}; use datafusion::common::ScalarValue; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -39,13 +39,20 @@ use datafusion::physical_plan::{collect, ExecutionPlanProperties, InputOrderMode use datafusion::prelude::SessionContext; fn batch(keys: ArrayRef) -> RecordBatch { + batch_with_offset(keys, 0) +} + +fn batch_with_offset(keys: ArrayRef, offset: i64) -> RecordBatch { let rows = keys.len(); RecordBatch::try_new( Arc::new(Schema::new(vec![ Field::new("k", keys.data_type().clone(), true), Field::new("i", DataType::Int64, false), ])), - vec![keys, Arc::new(Int64Array::from_iter_values(0..rows as i64))], + vec![ + keys, + Arc::new(Int64Array::from_iter_values(offset..offset + rows as i64)), + ], ) .unwrap() } @@ -80,7 +87,22 @@ fn window_with_order_mode( .unwrap() .try_with_sort_information(vec![LexOrdering::new(ordering).unwrap()]) .unwrap(); - let input = DataSourceExec::from_data_source(source); + window_on_input( + DataSourceExec::from_data_source(source), + partition_columns, + functions, + input_order_mode, + ) +} + +fn window_on_input( + input: Arc, + partition_columns: &[&str], + functions: &[&str], + input_order_mode: InputOrderMode, +) -> BoundedWindowAggExec { + let schema = input.schema(); + let sort = |name| PhysicalSortExpr::new(col(name, &schema).unwrap(), SortOptions::default()); let ctx = SessionContext::new(); let keys = partition_columns .iter() @@ -145,6 +167,147 @@ async fn compare( actual } +fn row_numbers(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect() +} + +fn assert_ordering_rejected(generic: BoundedWindowAggExec, message: &str) { + let streaming = RowNumberExec::try_new(generic.clone()).unwrap(); + let context = Arc::new(TaskContext::default()); + for plan in [ + &generic as &dyn ExecutionPlan, + &streaming as &dyn ExecutionPlan, + ] { + let error = plan + .execute(0, Arc::clone(&context)) + .err() + .expect("window must reject missing partition ordering"); + assert!(matches!( + &error, + datafusion::common::DataFusionError::Execution(_) + )); + assert!(error.to_string().contains(message), "{error}"); + } +} + +#[test] +fn missing_partition_order_is_rejected() { + for (keys, partition_columns, declared_columns) in [ + (vec![1, 2, 1, 2], vec!["k"], vec![]), + (vec![1, 1, 2, 2], vec!["k"], vec![]), + (vec![1, 1, 2, 2], vec!["k", "i"], vec!["k"]), + ] { + let input = batch(Arc::new(Int64Array::from(keys))); + let schema = input.schema(); + let ordering = declared_columns + .iter() + .map(|name| PhysicalSortExpr::new_default(col(name, &schema).unwrap())) + .collect::>(); + let source = MemorySourceConfig::try_new(&[vec![input]], schema, None) + .unwrap() + .try_with_sort_information(LexOrdering::new(ordering).into_iter().collect()) + .unwrap(); + let generic = window_on_input( + DataSourceExec::from_data_source(source), + &partition_columns, + &["row_number"], + InputOrderMode::Sorted, + ); + assert_ordering_rejected(generic, "Expects PARTITION BY expression to be ordered"); + } +} + +#[test] +fn sorted_mode_requires_all_cached_partition_keys() { + let input = batch(Arc::new(Int64Array::from(vec![1, 1, 2, 2]))); + let mut generic = + window_with_order_mode(&[input], &["k"], &["row_number"], InputOrderMode::Linear); + // The mode is public, but changing it does not rebuild the cached partition-key indices. + generic.input_order_mode = InputOrderMode::Sorted; + assert_ordering_rejected( + generic, + "All partition by columns should have an ordering in Sorted mode", + ); +} + +#[tokio::test] +async fn global_window_needs_no_declared_partition_order() { + let input = batch(Arc::new(Int64Array::from(vec![1, 2, 1, 2]))); + let source = + MemorySourceConfig::try_new_exec(&[vec![input.clone()]], input.schema(), None).unwrap(); + let generic = window_on_input(source, &[], &["row_number"], InputOrderMode::Sorted); + let streaming = Arc::new(RowNumberExec::try_new(generic.clone()).unwrap()); + let ctx = SessionContext::new(); + let expected = collect(Arc::new(generic), ctx.task_ctx()).await.unwrap(); + let actual = collect(Arc::clone(&streaming) as _, ctx.task_ctx()) + .await + .unwrap(); + assert_eq!( + concat_batches(&streaming.schema(), &actual).unwrap(), + concat_batches(&streaming.schema(), &expected).unwrap() + ); + assert_eq!(row_numbers(&actual), vec![1, 2, 3, 4]); +} + +#[tokio::test] +async fn partitions_reset_at_batch_boundaries() { + let input = batch(Arc::new(Int64Array::from(vec![1, 1, 2, 2]))); + let actual = compare( + &[input.slice(0, 2), input.slice(2, 2)], + true, + &["row_number"], + ) + .await; + assert_eq!(row_numbers(&actual), vec![1, 2, 1, 2]); +} + +#[tokio::test] +async fn dictionary_partitions_reset_at_batch_boundaries() { + let dictionary = |value: &str| -> ArrayRef { + Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![0, 0]), + Arc::new(StringArray::from(vec![value])), + ) + .unwrap(), + ) + }; + let actual = compare( + &[ + batch_with_offset(dictionary("a"), 0), + batch_with_offset(dictionary("b"), 2), + ], + true, + &["row_number"], + ) + .await; + assert_eq!(row_numbers(&actual), vec![1, 2, 1, 2]); +} + +#[tokio::test] +async fn mixed_batch_retains_its_final_partition_key() { + let input = batch(Arc::new(Int64Array::from(vec![1, 1, 2, 2, 3, 3]))); + let actual = compare( + &[input.slice(0, 1), input.slice(1, 4), input.slice(5, 1)], + true, + &["row_number"], + ) + .await; + assert_eq!(row_numbers(&actual), vec![1, 2, 1, 2, 1, 2]); +} + #[tokio::test] async fn partitions_continue_across_slices_and_empty_batches() { let input = batch(Arc::new(Int64Array::from(vec![ @@ -211,7 +374,6 @@ async fn empty_input_preserves_schema() { #[tokio::test] async fn dictionary_keys_compare_values_across_different_codebooks() { - use arrow::datatypes::Int8Type; let first = Arc::new( DictionaryArray::::try_new( Int8Array::from(vec![None, Some(0), Some(1), Some(1)]), @@ -226,7 +388,12 @@ async fn dictionary_keys_compare_values_across_different_codebooks() { ) .unwrap(), ); - let output = compare(&[batch(first), batch(second)], true, &["row_number"]).await; + let output = compare( + &[batch(first), batch_with_offset(second, 4)], + true, + &["row_number"], + ) + .await; let ranks = output[1] .column(2) .as_any() @@ -448,7 +615,6 @@ async fn completion_and_cancellation_release_the_boundary_key() { #[test] fn a_dictionary_boundary_does_not_retain_unused_values() { - use arrow::datatypes::Int8Type; let values: ArrayRef = Arc::new(StringArray::from(vec![ "small", &"unused".repeat(64 * 1024), diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index c54fca71215..fb052f164aa 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -1294,6 +1294,7 @@ class CometWindowExecSuite extends CometTestBase { .toDF("k", "i", "payload") .coalesce(1) .write + .mode("overwrite") .parquet(dir.toString) spark.read.parquet(dir.toString).createOrReplaceTempView("streaming_row_number") for (batchSize <- Seq(3, 16, 128); ansi <- Seq("false", "true")) { From 3a8eae322aa59ff7f8487360404b3b57f24d80e6 Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Tue, 3 Nov 2026 14:38:28 -0700 Subject: [PATCH 4/9] test: broaden streaming ROW_NUMBER key coverage Move batch-boundary checks into a SQL fixture covering seven nullable partition-key types and three batch sizes. Exercise floating keys with ROW_NUMBER and retain global, repeated-result and mixed-frame controls. --- .../windows/streaming_row_number.sql | 117 ++++++++++++++++++ .../comet/exec/CometWindowExecSuite.scala | 47 ++----- 2 files changed, 124 insertions(+), 40 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/windows/streaming_row_number.sql diff --git a/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql b/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql new file mode 100644 index 00000000000..be66b0dd661 --- /dev/null +++ b/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql @@ -0,0 +1,117 @@ +-- 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. + +-- ConfigMatrix: spark.comet.batchSize=3,16,128 +-- Config: spark.sql.shuffle.partitions=1 + +-- The seven nullable key types describe the same groups: 33 NULL rows followed +-- by groups of 17 rows. Groups cross every configured input batch size. +statement +CREATE TABLE streaming_row_number_keys( + id BIGINT, + k_int INT, + k_bigint BIGINT, + k_decimal DECIMAL(12, 2), + k_date DATE, + k_timestamp TIMESTAMP, + k_string STRING, + k_binary BINARY, + payload STRING +) USING parquet + +statement +INSERT INTO streaming_row_number_keys +SELECT id, + g, + CAST(g AS BIGINT), + CAST(g AS DECIMAL(12, 2)), + DATE_ADD(DATE '2000-01-01', g), + CAST(DATE_ADD(DATE '2000-01-01', g) AS TIMESTAMP), + CONCAT('key-', CAST(g AS STRING)), + CAST(CONCAT('key-', CAST(g AS STRING)) AS BINARY), + CONCAT('payload-', CAST(id AS STRING)) +FROM ( + SELECT id, + CASE WHEN id < 33 THEN CAST(NULL AS INT) + ELSE CAST((id - 33) DIV 17 AS INT) END AS g + FROM range(385) +) generated_keys + +query +SELECT id, k_int, + ROW_NUMBER() OVER (PARTITION BY k_int ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +query +SELECT id, k_bigint, + ROW_NUMBER() OVER (PARTITION BY k_bigint ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +query +SELECT id, k_decimal, + ROW_NUMBER() OVER (PARTITION BY k_decimal ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +query +SELECT id, k_date, + ROW_NUMBER() OVER (PARTITION BY k_date ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +query +SELECT id, k_timestamp, + ROW_NUMBER() OVER (PARTITION BY k_timestamp ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +query +SELECT id, k_string, + ROW_NUMBER() OVER (PARTITION BY k_string ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +query +SELECT id, k_binary, + ROW_NUMBER() OVER (PARTITION BY k_binary ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +-- The global counter crosses batches without a retained partition key. +query +SELECT id, ROW_NUMBER() OVER (ORDER BY id) AS rn +FROM streaming_row_number_keys +ORDER BY id + +-- Repeated results and carried columns survive batch-at-a-time execution. +query +SELECT id, k_string, payload, + ROW_NUMBER() OVER (PARTITION BY k_string ORDER BY id) AS first_rank, + ROW_NUMBER() OVER (PARTITION BY k_string ORDER BY id) AS second_rank +FROM streaming_row_number_keys +ORDER BY id + +-- A mixed window retains the generic executor and its frame behaviour. +query +SELECT id, k_string, + ROW_NUMBER() OVER (PARTITION BY k_string ORDER BY id) AS rn, + COUNT(*) OVER (PARTITION BY k_string ORDER BY id + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS n +FROM streaming_row_number_keys +ORDER BY id diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index fb052f164aa..1241e121e58 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -336,7 +336,9 @@ class CometWindowExecSuite extends CometTestBase { "RANK()" -> false, "PERCENT_RANK()" -> false, "NTILE(2)" -> false, - "RANK()" -> true) + "RANK()" -> true, + "ROW_NUMBER()" -> false, + "ROW_NUMBER()" -> true) } { test( s"window: floating partition keys ($partitionColumn, $function, group limit=$groupLimit)") { @@ -360,6 +362,7 @@ class CometWindowExecSuite extends CometTestBase { SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "1", CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "3", CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.key -> groupLimit.toString) { withTempView("floating_window_partitions") { data @@ -367,9 +370,9 @@ class CometWindowExecSuite extends CometTestBase { .createOrReplaceTempView("floating_window_partitions") for (partitionBy <- Seq(partitionColumn, s"secondary, $partitionColumn")) { - // Spark normalizes floating partition keys before constructing the child sort. - // Native sort and window expressions must still agree on these keys, including - // for bounded RANK and non-bounded PERCENT_RANK/NTILE without WindowGroupLimit. + // Spark normalises floating partition keys before constructing the child sort. + // Native sort and window expressions must still agree across batches, including + // streaming ROW_NUMBER, bounded RANK, and non-bounded PERCENT_RANK/NTILE. val windowQuery = s""" |SELECT id, $function OVER (PARTITION BY $partitionBy ORDER BY id) AS rnk |FROM floating_window_partitions @@ -1287,42 +1290,6 @@ class CometWindowExecSuite extends CometTestBase { } } - test("window: streaming ROW_NUMBER across input batches") { - withTempDir { dir => - (0 until 97) - .map(i => (if (i < 9) null else s"key-${i / 13}", i, "payload-" * 32)) - .toDF("k", "i", "payload") - .coalesce(1) - .write - .mode("overwrite") - .parquet(dir.toString) - spark.read.parquet(dir.toString).createOrReplaceTempView("streaming_row_number") - for (batchSize <- Seq(3, 16, 128); ansi <- Seq("false", "true")) { - withSQLConf( - CometConf.COMET_BATCH_SIZE.key -> batchSize.toString, - SQLConf.ANSI_ENABLED.key -> ansi) { - for (partition <- Seq("", "PARTITION BY k")) { - val df = sql(s""" - SELECT k, i, payload, - ROW_NUMBER() OVER ($partition ORDER BY i) AS first_rank, - ROW_NUMBER() OVER ($partition ORDER BY i) AS second_rank - FROM streaming_row_number - """) - checkSparkAnswerAndOperator(df) - assertCometWindowExecExists(df.queryExecution.executedPlan) - } - // A mixed window must retain the generic executor and its frame behaviour. - checkSparkAnswerAndOperator(sql(""" - SELECT k, i, - ROW_NUMBER() OVER (PARTITION BY k ORDER BY i) AS rn, - COUNT(*) OVER (PARTITION BY k ORDER BY i) AS n - FROM streaming_row_number - """)) - } - } - } - } - test("window: RANK with PARTITION BY and ORDER BY") { withTempDir { dir => (0 until 30) From 60062fcc2c97f96281b4fec45997216a1ffb6e80 Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Tue, 3 Nov 2026 14:41:03 -0700 Subject: [PATCH 5/9] perf: retain ROW_NUMBER expressions in stream state Move the partition expressions through the unfolding stream state so every batch reuses the same vector. Avoid allocating a replacement vector and cloning its expression references on each iteration, including completion. --- native/operators/src/row_number.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/native/operators/src/row_number.rs b/native/operators/src/row_number.rs index cb0ecbf2629..a3e6130c208 100644 --- a/native/operators/src/row_number.rs +++ b/native/operators/src/row_number.rs @@ -203,9 +203,8 @@ impl ExecutionPlan for RowNumberExec { }; let metrics = BaselineMetrics::new(&self.metrics, partition); let stream = futures::stream::try_unfold( - (input, state, metrics), - move |(mut input, mut state, metrics)| { - let partition_by = partition_by.clone(); + (input, state, metrics, partition_by), + move |(mut input, mut state, metrics, partition_by)| { let schema = Arc::clone(&schema); async move { let Some(batch) = input.next().await else { @@ -217,7 +216,7 @@ impl ExecutionPlan for RowNumberExec { state.evaluate(batch, &partition_by, schema, output_count)? }; (&output).record_output(&metrics); - Ok(Some((output, (input, state, metrics)))) + Ok(Some((output, (input, state, metrics, partition_by)))) } }, ); From 101e4c1698a35835a34da027824e89f2f93ddf2a Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Tue, 3 Nov 2026 15:58:25 -0700 Subject: [PATCH 6/9] PRFB: reservation churn Removed unnecesary free/try_resize churn. I've stuck to try_resize for now while I think about it more. --- .../contributor-guide/memory_management.md | 5 + native/operators/benches/window.rs | 37 +++- native/operators/src/row_number.rs | 9 +- native/operators/src/row_number/tests.rs | 197 +++++++++++++++++- 4 files changed, 232 insertions(+), 16 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 70275d766be..cfd182a81bb 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -498,6 +498,11 @@ views and dictionaries are rebuilt from the selected value so unrelated source b 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 operator cannot spill its key. If the pool refuses the initial reservation or later growth, +the task fails and releases the state. This can fail a task that the generic bounded window +executor would run with unreserved state. + 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. diff --git a/native/operators/benches/window.rs b/native/operators/benches/window.rs index df1e6e61680..5eb91d91eaf 100644 --- a/native/operators/benches/window.rs +++ b/native/operators/benches/window.rs @@ -396,17 +396,24 @@ fn criterion_benchmark(c: &mut Criterion) { let batches = (0..BATCHES) .map(|i| s.batch(&schema, i)) .collect::>(); - let generic = s.plan(&session, &batches); - let streaming: Arc = Arc::new( - RowNumberExec::try_new( + let build_plan = |streaming| -> Arc { + let generic = s.plan(&session, &batches); + if streaming { + Arc::new( + RowNumberExec::try_new( + generic + .downcast_ref::() + .unwrap() + .clone(), + ) + .unwrap(), + ) + } else { generic - .downcast_ref::() - .unwrap() - .clone(), - ) - .unwrap(), - ); - for (name, plan) in [("bounded", generic), ("streaming", streaming)] { + } + }; + 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; @@ -427,7 +434,15 @@ fn criterion_benchmark(c: &mut Criterion) { assert_eq!(row, ROWS); }); group.bench_function(BenchmarkId::new(name, rows_per_partition), |b| { - b.iter(|| run(&runtime, &plan, &task_ctx)); + // 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, + ); }); } } diff --git a/native/operators/src/row_number.rs b/native/operators/src/row_number.rs index a3e6130c208..92c21be0548 100644 --- a/native/operators/src/row_number.rs +++ b/native/operators/src/row_number.rs @@ -52,6 +52,10 @@ use futures::StreamExt; /// all be the built-in ROW_NUMBER over the same partition and sort keys. /// The caller must normalise floating-point partition keys before building the window, just as /// it does for ordinary native windows. Other window functions keep their existing executor. +/// +/// This avoids the per-partition overhead tracked in +/// . Reassess this executor when upstream +/// changes provide equivalent results, memory accounting and performance. #[derive(Debug)] pub struct RowNumberExec { window: BoundedWindowAggExec, @@ -288,10 +292,9 @@ impl RowNumberState { } self.rank = *ranks.last().unwrap(); if !keys.is_empty() && (!continues || ranges.len() != 1) { - // The old boundary key is no longer needed. Do not keep both copies alive - // when a wide key changes. A continued partition needs no new key copy. + // Drop the old key before copying its replacement, but retain the reservation + // so only the change in retained bytes adjusts the pool. self.key = vec![]; - self.reservation.free(); let index = UInt64Array::from(vec![(rows - 1) as u64]); let key = keys .iter() diff --git a/native/operators/src/row_number/tests.rs b/native/operators/src/row_number/tests.rs index 4f96e3e6b0f..f68ba9967e7 100644 --- a/native/operators/src/row_number/tests.rs +++ b/native/operators/src/row_number/tests.rs @@ -17,6 +17,8 @@ use super::*; +use std::sync::Mutex; + use arrow::array::{ Array, DictionaryArray, Int64Array, Int8Array, ListArray, StringArray, StringViewArray, StructArray, @@ -27,7 +29,8 @@ use arrow::datatypes::{DataType, Field, Int64Type, Int8Type, Schema}; use datafusion::common::ScalarValue; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; -use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryLimit, MemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::{ WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, @@ -563,6 +566,197 @@ fn retained_view_buffers_are_charged_and_a_continuing_key_is_reused() { assert_eq!(pool.reserved(), 0); } +#[derive(Debug, PartialEq, Eq)] +enum PoolCall { + Grow(usize), + TryGrow(usize), + Shrink(usize), +} + +#[derive(Debug)] +struct RecordingMemoryPool { + inner: GreedyMemoryPool, + calls: Mutex>, +} + +impl RecordingMemoryPool { + fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + calls: Mutex::new(vec![]), + } + } + + fn take_calls(&self) -> Vec { + std::mem::take(&mut *self.calls.lock().unwrap()) + } +} + +impl std::fmt::Display for RecordingMemoryPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "RecordingMemoryPool({})", self.inner) + } +} + +impl MemoryPool for RecordingMemoryPool { + fn name(&self) -> &str { + "RecordingMemoryPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.calls.lock().unwrap().push(PoolCall::Grow(additional)); + self.inner.grow(reservation, additional); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.calls + .lock() + .unwrap() + .push(PoolCall::TryGrow(additional)); + self.inner.try_grow(reservation, additional) + } + + fn shrink(&self, reservation: &MemoryReservation, size: usize) { + self.calls.lock().unwrap().push(PoolCall::Shrink(size)); + self.inner.shrink(reservation, size); + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +fn task_context_with_pool(pool: Arc) -> Arc { + Arc::new( + TaskContext::default().with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc() + .unwrap(), + ), + ) +} + +#[tokio::test] +async fn changing_keys_reuse_the_boundary_reservation() { + let batches = [ + batch_with_offset(Arc::new(Int64Array::from(vec![1, 1])), 0), + batch_with_offset(Arc::new(Int64Array::from(vec![2, 2])), 2), + ]; + let plan = RowNumberExec::try_new(window(&batches, true, &["row_number"])).unwrap(); + let pool = Arc::new(RecordingMemoryPool::new(1024)); + let mut stream = plan + .execute(0, task_context_with_pool(Arc::clone(&pool) as _)) + .unwrap(); + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + let retained = pool.reserved(); + assert!(retained > 0); + assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(retained)]); + + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + assert_eq!(pool.reserved(), retained); + assert!(pool.take_calls().is_empty()); + assert!(stream.next().await.is_none()); + assert_eq!(pool.reserved(), 0); + assert_eq!(pool.take_calls(), vec![PoolCall::Shrink(retained)]); +} + +#[tokio::test] +async fn boundary_reservations_follow_key_size_changes() { + let batches = ["a".repeat(32), "b".repeat(8192), "c".repeat(32)] + .into_iter() + .enumerate() + .map(|(index, key)| { + batch_with_offset(Arc::new(StringArray::from(vec![key; 2])), index as i64 * 2) + }) + .collect::>(); + let plan = RowNumberExec::try_new(window(&batches, true, &["row_number"])).unwrap(); + let pool = Arc::new(RecordingMemoryPool::new(64 * 1024)); + let mut stream = plan + .execute(0, task_context_with_pool(Arc::clone(&pool) as _)) + .unwrap(); + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + let small = pool.reserved(); + assert!(small > 0); + assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(small)]); + + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + let large = pool.reserved(); + assert!(large > small); + assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(large - small)]); + + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + assert_eq!(pool.reserved(), small); + assert_eq!(pool.take_calls(), vec![PoolCall::Shrink(large - small)]); + drop(stream); + assert_eq!(pool.reserved(), 0); + assert_eq!(pool.take_calls(), vec![PoolCall::Shrink(small)]); +} + +#[tokio::test] +async fn refused_replacement_releases_the_boundary_reservation() { + let batches = ["a".repeat(32), "b".repeat(8192)] + .into_iter() + .enumerate() + .map(|(index, key)| { + batch_with_offset(Arc::new(StringArray::from(vec![key; 2])), index as i64 * 2) + }) + .collect::>(); + let plan = RowNumberExec::try_new(window(&batches, true, &["row_number"])).unwrap(); + let pool = Arc::new(RecordingMemoryPool::new(1024)); + let mut stream = plan + .execute(0, task_context_with_pool(Arc::clone(&pool) as _)) + .unwrap(); + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + let retained = pool.reserved(); + assert!(retained > 0); + assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(retained)]); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!(matches!( + error, + datafusion::common::DataFusionError::ResourcesExhausted(_) + )); + assert_eq!(pool.reserved(), 0); + assert!(matches!( + pool.take_calls().as_slice(), + [PoolCall::TryGrow(additional), PoolCall::Shrink(released)] + if *additional > 1024 && *released == retained + )); + drop(stream); + assert!(pool.take_calls().is_empty()); +} + #[test] fn refused_boundary_reservations_are_released() { let pool: Arc = Arc::new(GreedyMemoryPool::new(1)); @@ -587,7 +781,6 @@ fn refused_boundary_reservations_are_released() { #[tokio::test] async fn completion_and_cancellation_release_the_boundary_key() { - use datafusion::execution::runtime_env::RuntimeEnvBuilder; let input = batch(Arc::new(Int64Array::from(vec![1, 1, 1, 2]))); let plan = RowNumberExec::try_new(window( &[input.slice(0, 2), input.slice(2, 2)], From 63761ef5f9ace7fdb9864f669e8ebe191f71124b Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Wed, 4 Nov 2026 08:33:44 -0700 Subject: [PATCH 7/9] use resize rather than try_resize --- .../contributor-guide/memory_management.md | 9 ++- native/operators/src/row_number.rs | 3 +- native/operators/src/row_number/tests.rs | 78 ++++++++++--------- 3 files changed, 48 insertions(+), 42 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index cfd182a81bb..b59ebf9b1d1 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -487,7 +487,7 @@ Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservat less, as overcommit (see [The unified pools](#the-unified-pools)). - `shrink(n)` returns bytes to the pool. -An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses. +An operator that never reserves memory is invisible to the pool no matter how much memory it uses. ### ROW_NUMBER windows @@ -499,9 +499,10 @@ Nested values may retain shared buffers, which are included in the reservation. 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 operator cannot spill its key. If the pool refuses the initial reservation or later growth, -the task fails and releases the state. This can fail a task that the generic bounded window -executor would run with unreserved state. +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). diff --git a/native/operators/src/row_number.rs b/native/operators/src/row_number.rs index 92c21be0548..008e30eaba2 100644 --- a/native/operators/src/row_number.rs +++ b/native/operators/src/row_number.rs @@ -318,7 +318,8 @@ impl RowNumberState { .map(|array| array.get_array_memory_size()) .sum::() + key.capacity() * std::mem::size_of::(); - self.reservation.try_resize(size)?; + // The key already exists and cannot spill, so account for it unconditionally. + self.reservation.resize(size); self.key = key; } } diff --git a/native/operators/src/row_number/tests.rs b/native/operators/src/row_number/tests.rs index f68ba9967e7..2c44a42c97e 100644 --- a/native/operators/src/row_number/tests.rs +++ b/native/operators/src/row_number/tests.rs @@ -666,7 +666,7 @@ async fn changing_keys_reuse_the_boundary_reservation() { ); let retained = pool.reserved(); assert!(retained > 0); - assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(retained)]); + assert_eq!(pool.take_calls(), vec![PoolCall::Grow(retained)]); assert_eq!( row_numbers(&[stream.next().await.unwrap().unwrap()]), @@ -699,7 +699,7 @@ async fn boundary_reservations_follow_key_size_changes() { ); let small = pool.reserved(); assert!(small > 0); - assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(small)]); + assert_eq!(pool.take_calls(), vec![PoolCall::Grow(small)]); assert_eq!( row_numbers(&[stream.next().await.unwrap().unwrap()]), @@ -707,7 +707,7 @@ async fn boundary_reservations_follow_key_size_changes() { ); let large = pool.reserved(); assert!(large > small); - assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(large - small)]); + assert_eq!(pool.take_calls(), vec![PoolCall::Grow(large - small)]); assert_eq!( row_numbers(&[stream.next().await.unwrap().unwrap()]), @@ -721,7 +721,7 @@ async fn boundary_reservations_follow_key_size_changes() { } #[tokio::test] -async fn refused_replacement_releases_the_boundary_reservation() { +async fn replacement_growth_can_exceed_the_pool_limit() { let batches = ["a".repeat(32), "b".repeat(8192)] .into_iter() .enumerate() @@ -738,45 +738,49 @@ async fn refused_replacement_releases_the_boundary_reservation() { row_numbers(&[stream.next().await.unwrap().unwrap()]), vec![1, 2] ); - let retained = pool.reserved(); - assert!(retained > 0); - assert_eq!(pool.take_calls(), vec![PoolCall::TryGrow(retained)]); + let small = pool.reserved(); + assert!(small > 0 && small < 1024); + let mut calls = pool.take_calls(); - let error = stream.next().await.unwrap().unwrap_err(); - assert!(matches!( - error, - datafusion::common::DataFusionError::ResourcesExhausted(_) - )); - assert_eq!(pool.reserved(), 0); - assert!(matches!( - pool.take_calls().as_slice(), - [PoolCall::TryGrow(additional), PoolCall::Shrink(released)] - if *additional > 1024 && *released == retained - )); + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] + ); + let large = pool.reserved(); + assert!(large >= 8192); + calls.extend(pool.take_calls()); + assert_eq!( + calls, + vec![PoolCall::Grow(small), PoolCall::Grow(large - small)] + ); drop(stream); - assert!(pool.take_calls().is_empty()); + assert_eq!(pool.reserved(), 0); + assert_eq!(pool.take_calls(), vec![PoolCall::Shrink(large)]); } -#[test] -fn refused_boundary_reservations_are_released() { - let pool: Arc = Arc::new(GreedyMemoryPool::new(1)); - let reservation = MemoryConsumer::new("test refusal").register(&pool); - let mut state = RowNumberState { - key: vec![], - rank: 0, - reservation, - }; - let input = batch(Arc::new(StringArray::from(vec!["some partition key"]))); - let plan = window(std::slice::from_ref(&input), true, &["row_number"]); - let result = state.evaluate( - input, - plan.window_expr()[0].partition_by(), - plan.schema(), - 1, +#[tokio::test] +async fn initial_nested_key_can_exceed_the_pool_limit() { + let keys: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("value", DataType::Utf8View, true)].into(), + vec![Arc::new(StringViewArray::from(vec!["x".repeat(8192); 2]))], + None, + )); + let input = batch(keys); + let plan = RowNumberExec::try_new(window(&[input], true, &["row_number"])).unwrap(); + let pool = Arc::new(RecordingMemoryPool::new(1)); + let mut stream = plan + .execute(0, task_context_with_pool(Arc::clone(&pool) as _)) + .unwrap(); + assert_eq!( + row_numbers(&[stream.next().await.unwrap().unwrap()]), + vec![1, 2] ); - assert!(result.is_err()); - assert!(state.key.is_empty()); + let retained = pool.reserved(); + assert!(retained >= 8192); + assert_eq!(pool.take_calls(), vec![PoolCall::Grow(retained)]); + assert!(stream.next().await.is_none()); assert_eq!(pool.reserved(), 0); + assert_eq!(pool.take_calls(), vec![PoolCall::Shrink(retained)]); } #[tokio::test] From 91e3357e318c6fa5b66a2313fccd57852f66b3f0 Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Wed, 4 Nov 2026 08:35:06 -0700 Subject: [PATCH 8/9] remove extra diff --- docs/source/contributor-guide/memory_management.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index b59ebf9b1d1..c19ebba6e26 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -487,7 +487,7 @@ Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservat less, as overcommit (see [The unified pools](#the-unified-pools)). - `shrink(n)` returns bytes to the pool. -An operator that never reserves memory is invisible to the pool no matter how much memory it uses. +An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses. ### ROW_NUMBER windows From 2b9254fd8ca42de67fd867aa86ec2b26e1ac021d Mon Sep 17 00:00:00 2001 From: Nat Wilson Date: Wed, 4 Nov 2026 10:39:21 -0700 Subject: [PATCH 9/9] fix SQL test --- .../test/resources/sql-tests/windows/streaming_row_number.sql | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql b/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql index be66b0dd661..fd379ee71f9 100644 --- a/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql +++ b/spark/src/test/resources/sql-tests/windows/streaming_row_number.sql @@ -112,6 +112,6 @@ query SELECT id, k_string, ROW_NUMBER() OVER (PARTITION BY k_string ORDER BY id) AS rn, COUNT(*) OVER (PARTITION BY k_string ORDER BY id - ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS n + ROWS BETWEEN CURRENT ROW AND 1 FOLLOWING) AS n FROM streaming_row_number_keys ORDER BY id