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..c19ebba6e26 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/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, 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..5eb91d91eaf 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,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::>(); + 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 + } + }; + 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::() + .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 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..008e30eaba2 --- /dev/null +++ b/native/operators/src/row_number.rs @@ -0,0 +1,338 @@ +// 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::{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; +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, + 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. +/// 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. +/// +/// 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, + metrics: ExecutionPlanMetricsSet, +} + +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; + }; + 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 InputOrderMode::Sorted and 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 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 = 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, partition_by), + move |(mut input, mut state, metrics, partition_by)| { + 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, partition_by)))) + } + }, + ); + 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) { + // 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![]; + 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::(); + // The key already exists and cannot spill, so account for it unconditionally. + self.reservation.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..2c44a42c97e --- /dev/null +++ b/native/operators/src/row_number/tests.rs @@ -0,0 +1,868 @@ +// 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 std::sync::Mutex; + +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, Int8Type, Schema}; +use datafusion::common::ScalarValue; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +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, +}; +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 { + 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(offset..offset + rows as i64)), + ], + ) + .unwrap() +} + +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 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(); + 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() + .map(|&name| col(name, &schema).unwrap()) + .collect::>(); + 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, + input_order_mode, + !partition_columns.is_empty(), + ) + .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 +} + +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![ + 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() { + 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_with_offset(second, 4)], + 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; +} + +#[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]))); + 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); +} + +#[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::Grow(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::Grow(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::Grow(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 replacement_growth_can_exceed_the_pool_limit() { + 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 small = pool.reserved(); + assert!(small > 0 && small < 1024); + let mut calls = pool.take_calls(); + + 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_eq!(pool.reserved(), 0); + assert_eq!(pool.take_calls(), vec![PoolCall::Shrink(large)]); +} + +#[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] + ); + 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] +async fn completion_and_cancellation_release_the_boundary_key() { + 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() { + 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/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..fd379ee71f9 --- /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 CURRENT ROW AND 1 FOLLOWING) 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 642a73e337a..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