From a02ceeb3c6ff20f08009d7ab65ba77c7f60bbf44 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 6 Oct 2026 01:36:14 +0000 Subject: [PATCH 1/4] fix: handle nested RANGE CURRENT ROW peers --- datafusion/core/tests/window_nested_range.rs | 251 ++++++++++++++++++ datafusion/expr/src/window_state.rs | 143 +++++++++- .../optimizer/src/analyzer/type_coercion.rs | 25 +- datafusion/sqllogictest/test_files/window.slt | 27 +- 4 files changed, 410 insertions(+), 36 deletions(-) create mode 100644 datafusion/core/tests/window_nested_range.rs diff --git a/datafusion/core/tests/window_nested_range.rs b/datafusion/core/tests/window_nested_range.rs new file mode 100644 index 0000000000000..b3a5624311370 --- /dev/null +++ b/datafusion/core/tests/window_nested_range.rs @@ -0,0 +1,251 @@ +// 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. + +//! Nested RANGE peers through public physical execution APIs. +//! Construct plans directly to test peer boundaries independently of SQL type admission. + +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, AsArray, Int32Array, Int64Array, ListArray, StructArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Field, Int64Type, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec}; +use datafusion::physical_plan::windows::{ + BoundedWindowAggExec, WindowAggExec, create_window_expr, +}; +use datafusion::physical_plan::{ + ExecutionPlan, InputOrderMode, SendableRecordBatchStream, collect, +}; +use datafusion::prelude::SessionContext; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, +}; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::{FutureExt, StreamExt}; + +fn nested_batches() -> Result> { + // Three nested keys: a NULL container, a nested NULL, and nested 1. + // Hidden values under the two NULL containers deliberately differ. + let values = [Some(10), Some(99), None, None, None, Some(1)]; + let values: ArrayRef = Arc::new(Int32Array::from(values.to_vec())); + let lists: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + OffsetBuffer::from_lengths([1; 6]), + Arc::clone(&values), + None, + )); + let structs: ArrayRef = Arc::new(StructArray::new( + vec![Arc::new(Field::new("item", DataType::Int32, true))].into(), + vec![values], + None, + )); + let nulls = NullBuffer::from(vec![false, false, true, true, true, true]); + let list_field = Arc::new(Field::new("item", lists.data_type().clone(), true)); + let keys: [ArrayRef; 3] = [ + Arc::new(ListArray::new( + Arc::clone(&list_field), + OffsetBuffer::from_lengths([1; 6]), + Arc::clone(&lists), + Some(nulls.clone()), + )), + Arc::new(ListArray::new( + Arc::new(Field::new("item", structs.data_type().clone(), true)), + OffsetBuffer::from_lengths([1; 6]), + structs, + Some(nulls.clone()), + )), + Arc::new(StructArray::new( + vec![list_field].into(), + vec![lists], + Some(nulls), + )), + ]; + keys.into_iter() + .map(|keys| { + // The tie separates the third nested NULL from the preceding peers. + let tie: ArrayRef = Arc::new(Int32Array::from(vec![0, 0, 0, 0, 1, 0])); + Ok(RecordBatch::try_from_iter(vec![ + ("key", keys), + ("tie", tie), + ( + "value", + Arc::new(Int64Array::from(vec![1, 2, 4, 8, 16, 32])) as ArrayRef, + ), + ])?) + }) + .collect() +} + +#[tokio::test] +async fn nested_range_current_row_physical_operators() -> Result<()> { + use WindowFrameBound::{CurrentRow, Following, Preceding}; + + for batch in nested_batches()? { + let schema = batch.schema(); + let order_by = vec![ + PhysicalSortExpr::new_default(col("key", &schema)?), + PhysicalSortExpr::new_default(col("tie", &schema)?), + ]; + // Repeated peers cross input batch boundaries. + let batches = (0..6).map(|row| batch.slice(row, 1)).collect(); + let source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ + LexOrdering::new(order_by.clone()).unwrap(), + ])?; + let input: Arc = DataSourceExec::from_data_source(source); + for (start, end, expected) in [ + ( + Preceding(ScalarValue::UInt64(None)), + CurrentRow, + [3, 3, 15, 15, 31, 63], + ), + (CurrentRow, CurrentRow, [3, 3, 12, 12, 16, 32]), + ( + CurrentRow, + Following(ScalarValue::UInt64(None)), + [63, 63, 60, 60, 48, 32], + ), + ] { + for bounded in [false, true] { + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("value", &schema)?], + &[], + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Range, + start.clone(), + end.clone(), + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let plan: Arc = if bounded { + Arc::new(BoundedWindowAggExec::try_new( + vec![expr], + Arc::clone(&input), + InputOrderMode::Sorted, + false, + )?) + } else { + Arc::new(WindowAggExec::try_new( + vec![expr], + Arc::clone(&input), + false, + )?) + }; + let output_schema = plan.schema(); + let output = collect(plan, SessionContext::new().task_ctx()).await?; + let output = concat_batches(&output_schema, &output)?; + let actual = output.column(3).as_primitive::(); + assert_eq!( + actual.iter().collect::>(), + expected.map(Some), + "{}, {start:?} to {end:?}, bounded={bounded}", + batch.column(0).data_type(), + ); + } + } + } + Ok(()) +} + +#[derive(Debug)] +struct OpenEndedPartition { + batch: RecordBatch, +} + +impl PartitionStream for OpenEndedPartition { + fn schema(&self) -> &SchemaRef { + self.batch.schema_ref() + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + let batches = (0..self.batch.num_rows()) + .map(|row| Ok(self.batch.slice(row, 1))) + .collect::>(); + Box::pin(RecordBatchStreamAdapter::new( + self.batch.schema(), + futures::stream::iter(batches).chain(futures::stream::pending()), + )) + } +} + +#[tokio::test] +async fn sorted_nested_range_emits_completed_peers_before_eof() -> Result<()> { + let batch = nested_batches()?.remove(0); + let schema = batch.schema(); + let order_by = vec![ + PhysicalSortExpr::new_default(col("key", &schema)?), + PhysicalSortExpr::new_default(col("tie", &schema)?), + ]; + let source = Arc::new(StreamingTableExec::try_new( + Arc::clone(&schema), + vec![Arc::new(OpenEndedPartition { batch })], + None, + vec![LexOrdering::new(order_by.clone()).unwrap()], + true, + None, + )?); + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("value", &schema)?], + &[], + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + schema, + false, + false, + None, + )?; + let plan = + BoundedWindowAggExec::try_new(vec![expr], source, InputOrderMode::Sorted, false)?; + let mut stream = plan.execute(0, SessionContext::new().task_ctx())?; + let actual = tokio::time::timeout(Duration::from_secs(5), async { + let mut sums = Vec::new(); + while sums.len() < 5 { + let batch = stream.next().await.unwrap()?; + sums.extend(batch.column(3).as_primitive::().iter()); + } + Ok::<_, datafusion_common::DataFusionError>(sums) + }) + .await + .expect("completed peer groups should emit before EOF")?; + assert_eq!(actual, [Some(3), Some(3), Some(15), Some(15), Some(31)]); + // The final peer group cannot finish until another key or EOF arrives. + assert!(stream.next().now_or_never().is_none()); + Ok(()) +} diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index acba75268b50d..a55aa3d4de7c9 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -21,6 +21,7 @@ use std::{collections::VecDeque, ops::Range, sync::Arc}; use crate::{WindowFrame, WindowFrameBound, WindowFrameUnits}; +use arrow::compute::kernels::sort::{LexicographicalComparator, SortColumn}; use arrow::{ array::ArrayRef, compute::{SortOptions, concat, concat_batches}, @@ -70,7 +71,10 @@ impl WindowAggState { match self.window_frame_ctx.as_mut() { // Rows have no state do nothing Some(WindowFrameContext::Rows(_)) => {} - Some(WindowFrameContext::Range { .. }) => {} + Some(WindowFrameContext::Range { state, .. }) => { + // The comparator refers to the old, unpruned ORDER BY arrays. + state.peer_comparator = None; + } Some(WindowFrameContext::Groups { state, .. }) => { let mut n_group_to_del = 0; for (_, end_idx) in &state.group_end_indices { @@ -343,12 +347,61 @@ impl PartitionBatchState { #[derive(Debug, Default, Clone)] pub struct WindowFrameStateRange { sort_options: Vec, + peer_comparator: Option, +} + +/// Cache comparators because dictionary logical-null masks can scan entire arrays. +#[derive(Clone)] +struct RangePeerComparator { + columns: Vec, + comparator: Arc, +} + +impl std::fmt::Debug for RangePeerComparator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RangePeerComparator") + .field("num_columns", &self.columns.len()) + .finish_non_exhaustive() + } } impl WindowFrameStateRange { /// Create a new object to store the search state. fn new(sort_options: Vec) -> Self { - Self { sort_options } + Self { + sort_options, + peer_comparator: None, + } + } + + fn peer_comparator( + &mut self, + range_columns: &[ArrayRef], + ) -> Result<&LexicographicalComparator> { + let needs_rebuild = self.peer_comparator.as_ref().is_none_or(|cached| { + cached.columns.len() != range_columns.len() + || cached + .columns + .iter() + .zip(range_columns) + .any(|(old, new)| !Arc::ptr_eq(old, new)) + }); + if needs_rebuild { + // Sliced arrays have different row indices even when they share buffers. + let sort_columns = range_columns + .iter() + .zip(&self.sort_options) + .map(|(values, options)| SortColumn { + values: Arc::clone(values), + options: Some(*options), + }) + .collect::>(); + self.peer_comparator = Some(RangePeerComparator { + columns: range_columns.to_vec(), + comparator: Arc::new(LexicographicalComparator::try_new(&sort_columns)?), + }); + } + Ok(&self.peer_comparator.as_ref().unwrap().comparator) } /// This function calculates beginning/ending indices for the frame of the current row. @@ -440,6 +493,28 @@ impl WindowFrameStateRange { delta: Option<&ScalarValue>, length: usize, ) -> Result { + if delta.is_none() + && range_columns + .iter() + .any(|column| column.data_type().is_nested()) + { + // CURRENT ROW needs peer equality. Scalar ordering does not support + // all nested keys; Arrow's sort comparator also handles dictionary nulls. + let mut boundary = if SIDE { + last_range.start + } else { + last_range.end.max(idx) + }; + let comparator = self.peer_comparator(range_columns)?; + while boundary < length { + let is_peer = comparator.compare(boundary, idx).is_eq(); + if if SIDE { is_peer } else { !is_peer } { + break; + } + boundary += 1; + } + return Ok(boundary); + } let mut end_range = get_row_at_idx(range_columns, idx)?; let search_start = if SIDE { last_range.start @@ -729,6 +804,70 @@ mod tests { use arrow::array::{Float64Array, UInt64Array}; use arrow::datatypes::{Field, Schema}; + #[test] + fn range_peer_comparator_reuses_arrays_and_invalidates_after_pruning() -> Result<()> { + use arrow::array::ListArray; + use arrow::datatypes::Int64Type; + + let frame = Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::CurrentRow, + WindowFrameBound::CurrentRow, + )); + let columns: [ArrayRef; 1] = + [Arc::new(ListArray::from_iter_primitive::( + [1, 1, 2, 2].into_iter().map(|v| Some(vec![Some(v)])), + ))]; + let mut range = WindowFrameStateRange::new(vec![SortOptions::default()]); + assert_eq!( + range.calculate_range(&frame, &(0..0), &columns, 4, 0)?, + 0..2 + ); + let first = Arc::downgrade(&range.peer_comparator.as_ref().unwrap().comparator); + assert_eq!( + range.calculate_range(&frame, &(0..2), &columns, 4, 1)?, + 0..2 + ); + assert!(Arc::ptr_eq( + &first.upgrade().unwrap(), + &range.peer_comparator.as_ref().unwrap().comparator, + )); + + let columns = [columns[0].slice(1, 3)]; + assert_eq!( + range.calculate_range(&frame, &(0..0), &columns, 3, 0)?, + 0..1 + ); + assert!(first.upgrade().is_none()); + let retained = Arc::downgrade(&columns[0]); + let cached = Arc::downgrade(&range.peer_comparator.as_ref().unwrap().comparator); + let mut state = WindowAggState::new(&DataType::Int64)?; + state.window_frame_range = + range.calculate_range(&frame, &(0..1), &columns, 3, 1)?; + assert_eq!(state.window_frame_range, 1..3); + state.last_calculated_index = 1; + state.window_frame_ctx = Some(WindowFrameContext::Range { + window_frame: frame, + state: range, + }); + let pruned = [columns[0].slice(1, 2)]; + drop(columns); + assert!(retained.upgrade().is_some()); + state.prune_state(1); + assert!(retained.upgrade().is_none()); + assert!(cached.upgrade().is_none()); + assert_eq!( + state.window_frame_ctx.as_mut().unwrap().calculate_range( + &pruned, + &state.window_frame_range, + 2, + 0, + )?, + 0..2 + ); + Ok(()) + } + fn get_test_data() -> (Vec, Vec) { let range_columns: Vec = vec![Arc::new(Float64Array::from(vec![ 5.0, 7.0, 8.0, 8.0, 9., 10., 10., 10., 11., diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index e31dc65173474..f50fdd63598da 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1325,11 +1325,8 @@ fn extract_window_frame_target_type(col_type: &DataType) -> Option { /// run over an ORDER BY column of `col_type` even though the type has no /// arithmetic for finite offsets. /// -/// Such a frame only compares rows to find peers, so the type must compare -/// the same way in the RANGE peer check (`ScalarValue::partial_cmp`) as in -/// the sort that produced the input order. That holds for durations and -/// intervals; it does not for structs and maps, whose `ScalarValue` -/// comparison differs from the sorter's, so they stay unsupported. +/// This fallback admits durations and intervals, including encoded values. +/// Structs and maps remain outside its supported types. fn supports_free_range_frame(col_type: &DataType) -> bool { match col_type { DataType::Duration(_) | DataType::Interval(_) => true, @@ -1343,15 +1340,6 @@ fn supports_free_range_frame(col_type: &DataType) -> bool { /// Whether `col_type` is a list, possibly behind dictionary or run-end /// encoding. -/// -/// Lists are kept out of the free-range fallback because their peer -/// comparison and the sort do not agree on element NULLs: `compare_rows` -/// applies the NULLS FIRST / NULLS LAST option to the top-level value only and -/// then calls `ScalarValue::partial_cmp`, whose `partial_cmp_list` always -/// orders a NULL element after a non-NULL one (Postgres semantics), while the -/// sorter's `make_comparator` applies the option to the elements as well. Under -/// `ORDER BY d, l NULLS FIRST` with tied `d`, the sort puts `[NULL]` before -/// `[1]` and the peer check orders them the other way round. fn is_list_type(col_type: &DataType) -> bool { match col_type { DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) => { @@ -1364,8 +1352,8 @@ fn is_list_type(col_type: &DataType) -> bool { } /// Errors if any ORDER BY expression has a type not supported in a free RANGE -/// frame: a type with neither an offset target nor a sound peer comparison, or -/// a list (see `is_list_type`). +/// frame: a type with neither an offset target nor support in the free-range +/// fallback, or a list (see `is_list_type`). fn check_free_range_order_by_types( expressions: &[Sort], schema: &DFSchema, @@ -1407,9 +1395,8 @@ fn coerce_window_frame( } target_type } - // A free range frame has no offsets to coerce, so ORDER BY - // types without arithmetic are fine as long as their peer - // comparison is sound (see `supports_free_range_frame`). + // A free range frame has no offsets to coerce. Check whether + // its ORDER BY types are supported by the free-range fallback. None if window_frame.free_range() => { check_free_range_order_by_types(expressions, schema)?; return Ok(window_frame); diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 7c92514e723bd..328db4f2c17fe 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7187,8 +7187,7 @@ statement error Error during planning: RANGE window frames are not supported for select count(*) over (order by d, m) from (values (arrow_cast(1, 'Duration(Second)'), map(['a'], [1])), (arrow_cast(1, 'Duration(Second)'), map(['a'], [2]))) t(d, m) # the trailing key check also runs when the first ORDER BY type supports offset -# arithmetic; otherwise these maps are treated as peers because their keys are -# equal even though their values differ +# arithmetic, so a map after an integer is rejected as well statement error Error during planning: RANGE window frames are not supported for ORDER BY type Map select count(*) over (order by i, m) from (values (1, map(['a'], [1])), (1, map(['a'], [2]))) t(i, m) @@ -7207,23 +7206,21 @@ select count(*) over (order by l) from (values ([1]), ([2])) t(l) 1 2 -# but a list whose elements are nested (struct, list) plans and then fails at -# execution: the peer comparison in the RANGE frame goes through -# `ScalarValue::partial_cmp`, which cannot compare nested list elements, while -# the sort that produced the input uses `make_comparator` and can. Documented -# here so the fix to the scalar comparison can flip these expectations. -query error DataFusion error: Internal error: Uncomparable values: List\(\[\{c0: 1\}\]\), List\(\[\{c0: 1\}\]\) +# Lists containing structs or lists also form RANGE peer groups. +query I select count(*) over (order by l) from (values ([struct(1)]), ([struct(2)])) t(l) +---- +1 +2 -query error DataFusion error: Internal error: Uncomparable values: List\(\[\[1\]\]\), List\(\[\[1\]\]\) +query I select count(*) over (order by l) from (values ([[1]]), ([[2]])) t(l) +---- +1 +2 -# after a duration key the free-range check looks at every ORDER BY type, and -# a list is rejected at planning even though it sorts fine: the RANGE peer -# check (`compare_rows`, then `ScalarValue::partial_cmp`) always orders a NULL -# list element after a non-NULL one, while the sort applies NULLS FIRST to the -# elements as well, so with tied durations `[NULL]` and `[1]` would be sorted -# one way and compared the other +# after a duration key the free-range check looks at every ORDER BY type. +# Trailing list keys remain rejected by this check. query error Error during planning: RANGE window frames are not supported for ORDER BY type List select count(*) over (order by d, l nulls first) from (values (arrow_cast(1, 'Duration(Second)'), [NULL]), (arrow_cast(1, 'Duration(Second)'), [1])) t(d, l) From 3ee0478a4e8de73d663bc6951d0d09ac13b08cce Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 7 Oct 2026 03:57:57 +0000 Subject: [PATCH 2/4] fix: preserve range window unwind safety --- datafusion/expr/src/window_state.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index a55aa3d4de7c9..81922bc6e3c19 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -357,6 +357,11 @@ struct RangePeerComparator { comparator: Arc, } +// The cached Arrow arrays and comparison closures are immutable. A comparison +// that unwinds cannot leave this cache partially updated. +impl std::panic::UnwindSafe for RangePeerComparator {} +impl std::panic::RefUnwindSafe for RangePeerComparator {} + impl std::fmt::Debug for RangePeerComparator { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("RangePeerComparator") @@ -809,6 +814,9 @@ mod tests { use arrow::array::ListArray; use arrow::datatypes::Int64Type; + fn assert_unwind_safe() {} + assert_unwind_safe::(); + let frame = Arc::new(WindowFrame::new_bounds( WindowFrameUnits::Range, WindowFrameBound::CurrentRow, From 99100274dd139574854e40393474caac5c97e16b Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 22 Oct 2026 14:58:38 +0000 Subject: [PATCH 3/4] fix: address nested range window review --- datafusion/core/tests/execution/mod.rs | 1 + .../{ => execution}/window_nested_range.rs | 176 ++++++++++++------ datafusion/expr/src/window_state.rs | 56 +++++- .../optimizer/src/analyzer/type_coercion.rs | 9 +- .../physical-expr/src/window/window_expr.rs | 44 +++-- datafusion/sqllogictest/test_files/window.slt | 8 + 6 files changed, 214 insertions(+), 80 deletions(-) rename datafusion/core/tests/{ => execution}/window_nested_range.rs (59%) diff --git a/datafusion/core/tests/execution/mod.rs b/datafusion/core/tests/execution/mod.rs index f33ef87aa3023..24a3ff6c7f587 100644 --- a/datafusion/core/tests/execution/mod.rs +++ b/datafusion/core/tests/execution/mod.rs @@ -19,3 +19,4 @@ mod coop; mod datasource_split; mod logical_plan; mod register_arrow; +mod window_nested_range; diff --git a/datafusion/core/tests/window_nested_range.rs b/datafusion/core/tests/execution/window_nested_range.rs similarity index 59% rename from datafusion/core/tests/window_nested_range.rs rename to datafusion/core/tests/execution/window_nested_range.rs index b3a5624311370..a7fd780d281ef 100644 --- a/datafusion/core/tests/window_nested_range.rs +++ b/datafusion/core/tests/execution/window_nested_range.rs @@ -35,7 +35,8 @@ use datafusion::physical_plan::windows::{ BoundedWindowAggExec, WindowAggExec, create_window_expr, }; use datafusion::physical_plan::{ - ExecutionPlan, InputOrderMode, SendableRecordBatchStream, collect, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, SendableRecordBatchStream, + collect, }; use datafusion::prelude::SessionContext; use datafusion_common::{Result, ScalarValue}; @@ -92,6 +93,11 @@ fn nested_batches() -> Result> { Ok(RecordBatch::try_from_iter(vec![ ("key", keys), ("tie", tie), + // In Linear mode, the final row advances another partition's peers. + ( + "p", + Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0, 1])) as ArrayRef, + ), ( "value", Arc::new(Int64Array::from(vec![1, 2, 4, 8, 16, 32])) as ArrayRef, @@ -131,12 +137,25 @@ async fn nested_range_current_row_physical_operators() -> Result<()> { [63, 63, 60, 60, 48, 32], ), ] { - for bounded in [false, true] { + for mode in [ + None, + Some(InputOrderMode::Sorted), + Some(InputOrderMode::Linear), + ] { + // The planner uses WindowAggExec for UNBOUNDED FOLLOWING. + if mode.is_some() && end.is_unbounded() { + continue; + } + let partition_by = if mode == Some(InputOrderMode::Linear) { + vec![col("p", &schema)?] + } else { + vec![] + }; let expr = create_window_expr( &WindowFunctionDefinition::AggregateUDF(sum_udaf()), "sum".to_string(), &[col("value", &schema)?], - &[], + &partition_by, &order_by, Arc::new(WindowFrame::new_bounds( WindowFrameUnits::Range, @@ -148,11 +167,11 @@ async fn nested_range_current_row_physical_operators() -> Result<()> { false, None, )?; - let plan: Arc = if bounded { + let plan: Arc = if let Some(mode) = &mode { Arc::new(BoundedWindowAggExec::try_new( vec![expr], Arc::clone(&input), - InputOrderMode::Sorted, + mode.clone(), false, )?) } else { @@ -162,14 +181,31 @@ async fn nested_range_current_row_physical_operators() -> Result<()> { false, )?) }; + assert!( + input + .equivalence_properties() + .ordering_satisfy_requirement( + plan.required_input_ordering()[0] + .clone() + .unwrap() + .into_single(), + )? + ); let output_schema = plan.schema(); let output = collect(plan, SessionContext::new().task_ctx()).await?; let output = concat_batches(&output_schema, &output)?; - let actual = output.column(3).as_primitive::(); + let actual = output + .column(batch.num_columns()) + .as_primitive::(); + let mut expected = expected; + if mode == Some(InputOrderMode::Linear) { + // The final row is in its own PARTITION BY group. + expected[5] = 32; + } assert_eq!( actual.iter().collect::>(), expected.map(Some), - "{}, {start:?} to {end:?}, bounded={bounded}", + "{}, {start:?} to {end:?}, mode={mode:?}", batch.column(0).data_type(), ); } @@ -200,52 +236,86 @@ impl PartitionStream for OpenEndedPartition { } #[tokio::test] -async fn sorted_nested_range_emits_completed_peers_before_eof() -> Result<()> { - let batch = nested_batches()?.remove(0); - let schema = batch.schema(); - let order_by = vec![ - PhysicalSortExpr::new_default(col("key", &schema)?), - PhysicalSortExpr::new_default(col("tie", &schema)?), - ]; - let source = Arc::new(StreamingTableExec::try_new( - Arc::clone(&schema), - vec![Arc::new(OpenEndedPartition { batch })], - None, - vec![LexOrdering::new(order_by.clone()).unwrap()], - true, - None, - )?); - let expr = create_window_expr( - &WindowFunctionDefinition::AggregateUDF(sum_udaf()), - "sum".to_string(), - &[col("value", &schema)?], - &[], - &order_by, - Arc::new(WindowFrame::new_bounds( - WindowFrameUnits::Range, - WindowFrameBound::Preceding(ScalarValue::UInt64(None)), - WindowFrameBound::CurrentRow, - )), - schema, - false, - false, - None, - )?; - let plan = - BoundedWindowAggExec::try_new(vec![expr], source, InputOrderMode::Sorted, false)?; - let mut stream = plan.execute(0, SessionContext::new().task_ctx())?; - let actual = tokio::time::timeout(Duration::from_secs(5), async { - let mut sums = Vec::new(); - while sums.len() < 5 { - let batch = stream.next().await.unwrap()?; - sums.extend(batch.column(3).as_primitive::().iter()); +async fn nested_range_emits_completed_peers_before_eof() -> Result<()> { + for batch in nested_batches()? { + for mode in [InputOrderMode::Sorted, InputOrderMode::Linear] { + let schema = batch.schema(); + let order_by = vec![ + PhysicalSortExpr::new_default(col("key", &schema)?), + PhysicalSortExpr::new_default(col("tie", &schema)?), + ]; + let source: Arc = Arc::new(StreamingTableExec::try_new( + Arc::clone(&schema), + vec![Arc::new(OpenEndedPartition { + batch: batch.clone(), + })], + None, + vec![LexOrdering::new(order_by.clone()).unwrap()], + true, + None, + )?); + let partition_by = if mode == InputOrderMode::Linear { + vec![col("p", &schema)?] + } else { + vec![] + }; + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("value", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + schema, + false, + false, + None, + )?; + let plan = BoundedWindowAggExec::try_new( + vec![expr], + Arc::clone(&source), + mode.clone(), + false, + )?; + assert!( + source + .equivalence_properties() + .ordering_satisfy_requirement( + plan.required_input_ordering()[0] + .clone() + .unwrap() + .into_single(), + )? + ); + let mut stream = plan.execute(0, SessionContext::new().task_ctx())?; + let actual = tokio::time::timeout(Duration::from_secs(5), async { + let mut sums = Vec::new(); + while sums.len() < 5 { + let output = stream.next().await.unwrap()?; + sums.extend( + output + .column(batch.num_columns()) + .as_primitive::() + .iter(), + ); + } + Ok::<_, datafusion_common::DataFusionError>(sums) + }) + .await + .expect("completed peer groups should emit before EOF")?; + assert_eq!( + actual, + [Some(3), Some(3), Some(15), Some(15), Some(31)], + "{}, mode={mode:?}", + batch.column(0).data_type() + ); + // The final peer group cannot finish until another key or EOF arrives. + assert!(stream.next().now_or_never().is_none()); } - Ok::<_, datafusion_common::DataFusionError>(sums) - }) - .await - .expect("completed peer groups should emit before EOF")?; - assert_eq!(actual, [Some(3), Some(3), Some(15), Some(15), Some(31)]); - // The final peer group cannot finish until another key or EOF arrives. - assert!(stream.next().now_or_never().is_none()); + } Ok(()) } diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index 81922bc6e3c19..53d09a1bfdbe5 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -72,7 +72,8 @@ impl WindowAggState { // Rows have no state do nothing Some(WindowFrameContext::Rows(_)) => {} Some(WindowFrameContext::Range { state, .. }) => { - // The comparator refers to the old, unpruned ORDER BY arrays. + // Release the old ORDER BY arrays early. Replacement arrays would + // also invalidate the cache by identity when it is next used. state.peer_comparator = None; } Some(WindowFrameContext::Groups { state, .. }) => { @@ -350,14 +351,16 @@ pub struct WindowFrameStateRange { peer_comparator: Option, } -/// Cache comparators because dictionary logical-null masks can scan entire arrays. +/// Reuse comparators to avoid per-row setup work, including logical-null masks +/// that can take O(n) to build for dictionary arrays. #[derive(Clone)] struct RangePeerComparator { columns: Vec, comparator: Arc, } -// The cached Arrow arrays and comparison closures are immutable. A comparison +// Preserve the public WindowFrameStateRange type's existing unwind-safety traits. +// The cached Arrow arrays and comparison closures are immutable, so a comparison // that unwinds cannot leave this cache partially updated. impl std::panic::UnwindSafe for RangePeerComparator {} impl std::panic::RefUnwindSafe for RangePeerComparator {} @@ -503,17 +506,21 @@ impl WindowFrameStateRange { .iter() .any(|column| column.data_type().is_nested()) { - // CURRENT ROW needs peer equality. Scalar ordering does not support - // all nested keys; Arrow's sort comparator also handles dictionary nulls. + // CURRENT ROW needs peer equality. ScalarValue equality can distinguish + // encodings of the same nested value: a list with a null dictionary + // key and one with a key referencing null. Use Arrow's logical + // comparison, which also matches the ordering used to sort these keys. let mut boundary = if SIDE { last_range.start } else { + // Evaluation may start partway through a partition. An equality + // scan cannot skip preceding non-peers to reach the current row. last_range.end.max(idx) }; let comparator = self.peer_comparator(range_columns)?; while boundary < length { let is_peer = comparator.compare(boundary, idx).is_eq(); - if if SIDE { is_peer } else { !is_peer } { + if is_peer == SIDE { break; } boundary += 1; @@ -810,7 +817,7 @@ mod tests { use arrow::datatypes::{Field, Schema}; #[test] - fn range_peer_comparator_reuses_arrays_and_invalidates_after_pruning() -> Result<()> { + fn range_peer_comparator_reuses_arrays_and_releases_after_pruning() -> Result<()> { use arrow::array::ListArray; use arrow::datatypes::Int64Type; @@ -846,9 +853,7 @@ mod tests { range.calculate_range(&frame, &(0..0), &columns, 3, 0)?, 0..1 ); - assert!(first.upgrade().is_none()); let retained = Arc::downgrade(&columns[0]); - let cached = Arc::downgrade(&range.peer_comparator.as_ref().unwrap().comparator); let mut state = WindowAggState::new(&DataType::Int64)?; state.window_frame_range = range.calculate_range(&frame, &(0..1), &columns, 3, 1)?; @@ -863,7 +868,6 @@ mod tests { assert!(retained.upgrade().is_some()); state.prune_state(1); assert!(retained.upgrade().is_none()); - assert!(cached.upgrade().is_none()); assert_eq!( state.window_frame_ctx.as_mut().unwrap().calculate_range( &pruned, @@ -876,6 +880,38 @@ mod tests { Ok(()) } + #[test] + fn range_peers_compare_nested_dictionary_nulls_logically() -> Result<()> { + use arrow::array::{DictionaryArray, Int8Array, Int32Array, ListArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Int8Type; + + // Both lists contain NULL: one dictionary key is null, the other + // references a null value. ScalarValue equality distinguishes them. + let values: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![None, Some(0)]), + Arc::new(Int32Array::from(vec![None])), + )?); + let columns: [ArrayRef; 1] = [Arc::new(ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::from_lengths([1, 1]), + values, + None, + ))]; + let frame = Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::CurrentRow, + WindowFrameBound::CurrentRow, + )); + let mut state = WindowFrameContext::new(frame, vec![SortOptions::default()]); + let mut range = 0..0; + for idx in 0..2 { + range = state.calculate_range(&columns, &range, 2, idx)?; + assert_eq!(range, 0..2); + } + Ok(()) + } + fn get_test_data() -> (Vec, Vec) { let range_columns: Vec = vec![Arc::new(Float64Array::from(vec![ 5.0, 7.0, 8.0, 8.0, 9., 10., 10., 10., 11., diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index f50fdd63598da..bd3aea11ae8db 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1326,7 +1326,8 @@ fn extract_window_frame_target_type(col_type: &DataType) -> Option { /// arithmetic for finite offsets. /// /// This fallback admits durations and intervals, including encoded values. -/// Structs and maps remain outside its supported types. +/// Admission remains conservative: Arrow peer comparison also supports nested +/// types, but admitting structs and maps requires separate validation. fn supports_free_range_frame(col_type: &DataType) -> bool { match col_type { DataType::Duration(_) | DataType::Interval(_) => true, @@ -1339,7 +1340,8 @@ fn supports_free_range_frame(col_type: &DataType) -> bool { } /// Whether `col_type` is a list, possibly behind dictionary or run-end -/// encoding. +/// encoding. Free-range fallback still excludes lists conservatively, although +/// Arrow can now find peers for supported list keys. fn is_list_type(col_type: &DataType) -> bool { match col_type { DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) => { @@ -1353,7 +1355,8 @@ fn is_list_type(col_type: &DataType) -> bool { /// Errors if any ORDER BY expression has a type not supported in a free RANGE /// frame: a type with neither an offset target nor support in the free-range -/// fallback, or a list (see `is_list_type`). +/// fallback, or a list (see `is_list_type`). These conservative admission rules +/// can be relaxed separately from the nested peer-comparison fix. fn check_free_range_order_by_types( expressions: &[Sort], schema: &DFSchema, diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 5f6bbc46c27ea..4b2e8258fd626 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -23,7 +23,7 @@ use std::sync::Arc; use crate::PhysicalExpr; use arrow::array::BooleanArray; -use arrow::array::{Array, ArrayRef, new_empty_array}; +use arrow::array::{Array, ArrayRef, make_comparator, new_empty_array}; use arrow::compute::SortOptions; use arrow::compute::filter as arrow_filter; use arrow::compute::kernels::sort::SortColumn; @@ -31,7 +31,6 @@ use arrow::datatypes::FieldRef; use arrow::record_batch::RecordBatch; use datafusion_common::cast::as_boolean_array; use datafusion_common::hash_utils::RandomState; -use datafusion_common::utils::compare_rows; use datafusion_common::{ Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err, internal_err, @@ -353,6 +352,10 @@ pub trait AggregateWindowExpr: WindowExpr { let mut row_wise_results: Vec = Vec::with_capacity(length.saturating_sub(idx)); let is_causal = self.get_window_frame().is_causal(); + // These ordering arrays do not change during this evaluation. Compare + // their last/most-recent rows lazily once, since comparator setup can + // scan entire dictionary arrays to compute logical nulls. + let mut row_ahead = None; while idx < length { // Start search from the last_range. This squeezes searched range. let cur_range = @@ -367,6 +370,7 @@ pub trait AggregateWindowExpr: WindowExpr { most_recent_row_order_bys, self.order_by(), idx, + &mut row_ahead, )? { break; @@ -422,6 +426,7 @@ pub(crate) fn filter_arrays( /// columns. /// * `sort_exprs`: Defines the lexicographical ordering in question. /// * `idx`: The current index in the window frame. +/// * `row_ahead`: The last/most-recent row comparison, reused within one evaluation. /// /// # Returns /// @@ -432,6 +437,7 @@ pub(crate) fn is_end_bound_safe( most_recent_order_bys: Option<&[ArrayRef]>, sort_exprs: &[PhysicalSortExpr], idx: usize, + row_ahead: &mut Option, ) -> Result { if sort_exprs.is_empty() { // Early return if no sort expressions are present: @@ -448,6 +454,7 @@ pub(crate) fn is_end_bound_safe( most_recent_order_bys.map(|items| &items[0]), &sort_exprs[0].options, idx, + row_ahead, ), WindowFrameContext::Groups { window_frame, @@ -458,6 +465,7 @@ pub(crate) fn is_end_bound_safe( &order_bys[0], most_recent_order_bys.map(|items| &items[0]), &sort_exprs[0].options, + row_ahead, ), } } @@ -505,13 +513,14 @@ fn is_end_bound_safe_for_range( most_recent_ob_col: Option<&ArrayRef>, sort_options: &SortOptions, idx: usize, + row_ahead: &mut Option, ) -> Result { match end_bound { // A frame ending in `N PRECEDING` can only reach the end of the buffer // when the current row's frame ends at its last peer, as for a NULL or // NaN ORDER BY value, so it is safe once a row past its peers arrives. WindowFrameBound::Preceding(_) | WindowFrameBound::CurrentRow => { - is_row_ahead(orderby_col, most_recent_ob_col, sort_options) + is_row_ahead(orderby_col, most_recent_ob_col, sort_options, row_ahead) } WindowFrameBound::Following(delta) => { let Some(most_recent_ob_col) = most_recent_ob_col else { @@ -556,23 +565,24 @@ fn is_end_bound_safe_for_groups( orderby_col: &ArrayRef, most_recent_ob_col: Option<&ArrayRef>, sort_options: &SortOptions, + row_ahead: &mut Option, ) -> Result { match end_bound { WindowFrameBound::Preceding(value) => { let zero = ScalarValue::new_zero(&value.data_type())?; if value.eq(&zero) { - is_row_ahead(orderby_col, most_recent_ob_col, sort_options) + is_row_ahead(orderby_col, most_recent_ob_col, sort_options, row_ahead) } else { Ok(true) } } WindowFrameBound::CurrentRow => { - is_row_ahead(orderby_col, most_recent_ob_col, sort_options) + is_row_ahead(orderby_col, most_recent_ob_col, sort_options, row_ahead) } WindowFrameBound::Following(ScalarValue::UInt64(Some(offset))) => { let delta = state.group_end_indices.len() - state.current_group_idx; if delta == (*offset as usize) + 1 { - is_row_ahead(orderby_col, most_recent_ob_col, sort_options) + is_row_ahead(orderby_col, most_recent_ob_col, sort_options, row_ahead) } else { Ok(false) } @@ -581,23 +591,27 @@ fn is_end_bound_safe_for_groups( } } -/// This utility function checks whether `current_cols` is ahead of the `old_cols` -/// in terms of `sort_options`. +/// Checks whether the most recent input row is ahead of the partition's last row. +/// Both arrays stay unchanged during one evaluation, so reuse the comparison. fn is_row_ahead( old_col: &ArrayRef, current_col: Option<&ArrayRef>, sort_options: &SortOptions, + row_ahead: &mut Option, ) -> Result { + if let Some(ahead) = *row_ahead { + return Ok(ahead); + } let Some(current_col) = current_col else { return Ok(false); }; if old_col.is_empty() || current_col.is_empty() { return Ok(false); } - let last_value = ScalarValue::try_from_array(old_col, old_col.len() - 1)?; - let current_value = ScalarValue::try_from_array(current_col, 0)?; - let cmp = compare_rows(&[current_value], &[last_value], &[*sort_options])?; - Ok(cmp.is_gt()) + let cmp = make_comparator(old_col.as_ref(), current_col.as_ref(), *sort_options)?; + let ahead = cmp(old_col.len() - 1, 0).is_lt(); + *row_ahead = Some(ahead); + Ok(ahead) } /// Get order by expression results inside `order_by_columns`. @@ -783,7 +797,8 @@ mod tests { &SortOptions { descending: false, nulls_first: false - } + }, + &mut None, )?); assert!(!is_row_ahead( &old_values, @@ -791,7 +806,8 @@ mod tests { &SortOptions { descending: false, nulls_first: false - } + }, + &mut None, )?); Ok(()) diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 328db4f2c17fe..607d5f75368c6 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7206,6 +7206,14 @@ select count(*) over (order by l) from (values ([1]), ([2])) t(l) 1 2 +# NULL elements must form the same peers as Arrow's list sort order. +query I +select count(*) over (order by l nulls first) from (values ([NULL::int]), ([1]), ([NULL::int])) t(l) +---- +2 +2 +3 + # Lists containing structs or lists also form RANGE peer groups. query I select count(*) over (order by l) from (values ([struct(1)]), ([struct(2)])) t(l) From 39358f523fbc7f48a02c55606c3c66fc08a23990 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 23 Oct 2026 23:49:43 +0000 Subject: [PATCH 4/4] fix: simplify nested RANGE peer caching and test streaming order Keep the comparator in the frame state and classify nested keys once per range calculation. Exercise tied keys across batches in SQL and check descending order and null placement before EOF in Linear mode. --- .../tests/execution/window_nested_range.rs | 296 ++++++++++-------- datafusion/expr/src/window_state.rs | 134 ++++---- datafusion/sqllogictest/test_files/window.slt | 55 +++- 3 files changed, 279 insertions(+), 206 deletions(-) diff --git a/datafusion/core/tests/execution/window_nested_range.rs b/datafusion/core/tests/execution/window_nested_range.rs index a7fd780d281ef..f2343167eb10e 100644 --- a/datafusion/core/tests/execution/window_nested_range.rs +++ b/datafusion/core/tests/execution/window_nested_range.rs @@ -24,7 +24,7 @@ use std::time::Duration; use arrow::array::{ArrayRef, AsArray, Int32Array, Int64Array, ListArray, StructArray}; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::compute::concat_batches; -use arrow::datatypes::{DataType, Field, Int64Type, SchemaRef}; +use arrow::datatypes::{DataType, Field, Int32Type, Int64Type, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -107,23 +107,86 @@ fn nested_batches() -> Result> { .collect() } +fn nested_ordering(schema: &SchemaRef) -> Result { + Ok(LexOrdering::from([ + PhysicalSortExpr::new_default(col("key", schema)?), + PhysicalSortExpr::new_default(col("tie", schema)?), + ])) +} + +fn window_plan( + input: &Arc, + mode: Option<&InputOrderMode>, + start: WindowFrameBound, + end: WindowFrameBound, +) -> Result> { + let schema = input.schema(); + let partition_by = if mode == Some(&InputOrderMode::Linear) { + vec![col("p", &schema)?] + } else { + vec![] + }; + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("value", &schema)?], + &partition_by, + input.output_ordering().unwrap(), + Arc::new(WindowFrame::new_bounds(WindowFrameUnits::Range, start, end)), + schema, + false, + false, + None, + )?; + let plan: Arc = match mode { + Some(mode) => Arc::new(BoundedWindowAggExec::try_new( + vec![expr], + Arc::clone(input), + mode.clone(), + false, + )?), + None => Arc::new(WindowAggExec::try_new( + vec![expr], + Arc::clone(input), + false, + )?), + }; + assert!( + input + .equivalence_properties() + .ordering_satisfy_requirement( + plan.required_input_ordering()[0] + .clone() + .unwrap() + .into_single(), + )? + ); + Ok(plan) +} + #[tokio::test] async fn nested_range_current_row_physical_operators() -> Result<()> { use WindowFrameBound::{CurrentRow, Following, Preceding}; for batch in nested_batches()? { let schema = batch.schema(); - let order_by = vec![ - PhysicalSortExpr::new_default(col("key", &schema)?), - PhysicalSortExpr::new_default(col("tie", &schema)?), - ]; // Repeated peers cross input batch boundaries. - let batches = (0..6).map(|row| batch.slice(row, 1)).collect(); + let batches = (0..batch.num_rows()) + .map(|row| batch.slice(row, 1)) + .collect(); let source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)? - .try_with_sort_information(vec![ - LexOrdering::new(order_by.clone()).unwrap(), - ])?; + .try_with_sort_information(vec![nested_ordering(&schema)?])?; let input: Arc = DataSourceExec::from_data_source(source); + // window.slt covers list keys in SQL-planned sorted windows. + let modes: &[Option] = if batch.column(0).data_type().is_list() { + &[Some(InputOrderMode::Linear)] + } else { + &[ + None, + Some(InputOrderMode::Sorted), + Some(InputOrderMode::Linear), + ] + }; for (start, end, expected) in [ ( Preceding(ScalarValue::UInt64(None)), @@ -137,60 +200,12 @@ async fn nested_range_current_row_physical_operators() -> Result<()> { [63, 63, 60, 60, 48, 32], ), ] { - for mode in [ - None, - Some(InputOrderMode::Sorted), - Some(InputOrderMode::Linear), - ] { + for mode in modes.iter().map(Option::as_ref) { // The planner uses WindowAggExec for UNBOUNDED FOLLOWING. if mode.is_some() && end.is_unbounded() { continue; } - let partition_by = if mode == Some(InputOrderMode::Linear) { - vec![col("p", &schema)?] - } else { - vec![] - }; - let expr = create_window_expr( - &WindowFunctionDefinition::AggregateUDF(sum_udaf()), - "sum".to_string(), - &[col("value", &schema)?], - &partition_by, - &order_by, - Arc::new(WindowFrame::new_bounds( - WindowFrameUnits::Range, - start.clone(), - end.clone(), - )), - Arc::clone(&schema), - false, - false, - None, - )?; - let plan: Arc = if let Some(mode) = &mode { - Arc::new(BoundedWindowAggExec::try_new( - vec![expr], - Arc::clone(&input), - mode.clone(), - false, - )?) - } else { - Arc::new(WindowAggExec::try_new( - vec![expr], - Arc::clone(&input), - false, - )?) - }; - assert!( - input - .equivalence_properties() - .ordering_satisfy_requirement( - plan.required_input_ordering()[0] - .clone() - .unwrap() - .into_single(), - )? - ); + let plan = window_plan(&input, mode, start.clone(), end.clone())?; let output_schema = plan.schema(); let output = collect(plan, SessionContext::new().task_ctx()).await?; let output = concat_batches(&output_schema, &output)?; @@ -198,7 +213,7 @@ async fn nested_range_current_row_physical_operators() -> Result<()> { .column(batch.num_columns()) .as_primitive::(); let mut expected = expected; - if mode == Some(InputOrderMode::Linear) { + if mode == Some(&InputOrderMode::Linear) { // The final row is in its own PARTITION BY group. expected[5] = 32; } @@ -235,87 +250,96 @@ impl PartitionStream for OpenEndedPartition { } } +async fn assert_emits_before_eof( + batch: &RecordBatch, + order_by: LexOrdering, + mode: InputOrderMode, + expected: &[Option], +) -> Result<()> { + let source: Arc = Arc::new(StreamingTableExec::try_new( + batch.schema(), + vec![Arc::new(OpenEndedPartition { + batch: batch.clone(), + })], + None, + vec![order_by], + true, + None, + )?); + let plan = window_plan( + &source, + Some(&mode), + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )?; + let mut stream = plan.execute(0, SessionContext::new().task_ctx())?; + let actual = tokio::time::timeout(Duration::from_secs(5), async { + let mut sums = Vec::new(); + while sums.len() < expected.len() { + let output = stream.next().await.unwrap()?; + sums.extend( + output + .column(batch.num_columns()) + .as_primitive::() + .iter(), + ); + } + Ok::<_, datafusion_common::DataFusionError>(sums) + }) + .await + .expect("completed peer groups should emit before EOF")?; + assert_eq!( + actual, + expected, + "{}, mode={mode:?}", + batch.column(0).data_type() + ); + // The final peer group cannot finish until another key or EOF arrives. + assert!(stream.next().now_or_never().is_none()); + Ok(()) +} + #[tokio::test] async fn nested_range_emits_completed_peers_before_eof() -> Result<()> { for batch in nested_batches()? { for mode in [InputOrderMode::Sorted, InputOrderMode::Linear] { - let schema = batch.schema(); - let order_by = vec![ - PhysicalSortExpr::new_default(col("key", &schema)?), - PhysicalSortExpr::new_default(col("tie", &schema)?), - ]; - let source: Arc = Arc::new(StreamingTableExec::try_new( - Arc::clone(&schema), - vec![Arc::new(OpenEndedPartition { - batch: batch.clone(), - })], - None, - vec![LexOrdering::new(order_by.clone()).unwrap()], - true, - None, - )?); - let partition_by = if mode == InputOrderMode::Linear { - vec![col("p", &schema)?] - } else { - vec![] - }; - let expr = create_window_expr( - &WindowFunctionDefinition::AggregateUDF(sum_udaf()), - "sum".to_string(), - &[col("value", &schema)?], - &partition_by, - &order_by, - Arc::new(WindowFrame::new_bounds( - WindowFrameUnits::Range, - WindowFrameBound::Preceding(ScalarValue::UInt64(None)), - WindowFrameBound::CurrentRow, - )), - schema, - false, - false, - None, - )?; - let plan = BoundedWindowAggExec::try_new( - vec![expr], - Arc::clone(&source), - mode.clone(), - false, - )?; - assert!( - source - .equivalence_properties() - .ordering_satisfy_requirement( - plan.required_input_ordering()[0] - .clone() - .unwrap() - .into_single(), - )? - ); - let mut stream = plan.execute(0, SessionContext::new().task_ctx())?; - let actual = tokio::time::timeout(Duration::from_secs(5), async { - let mut sums = Vec::new(); - while sums.len() < 5 { - let output = stream.next().await.unwrap()?; - sums.extend( - output - .column(batch.num_columns()) - .as_primitive::() - .iter(), - ); - } - Ok::<_, datafusion_common::DataFusionError>(sums) - }) - .await - .expect("completed peer groups should emit before EOF")?; - assert_eq!( - actual, - [Some(3), Some(3), Some(15), Some(15), Some(31)], - "{}, mode={mode:?}", - batch.column(0).data_type() - ); - // The final peer group cannot finish until another key or EOF arrives. - assert!(stream.next().now_or_never().is_none()); + assert_emits_before_eof( + &batch, + nested_ordering(batch.schema_ref())?, + mode, + &[Some(3), Some(3), Some(15), Some(15), Some(31)], + ) + .await?; } } Ok(()) } + +#[tokio::test] +async fn nested_range_desc_nulls_last_emits_before_eof() -> Result<()> { + // Both [1] and a top-level NULL sort after [2] under DESC NULLS LAST. + // Each stream stops after that row arrives in a different partition. + for last_key in [Some([Some(1)]), None] { + let keys = ListArray::from_iter_primitive::([ + Some([Some(2)]), + Some([Some(2)]), + last_key, + ]); + let batch = RecordBatch::try_from_iter([ + ("key", Arc::new(keys) as ArrayRef), + ("p", Arc::new(Int32Array::from(vec![0, 0, 1]))), + ("value", Arc::new(Int64Array::from(vec![1, 2, 4]))), + ])?; + let order = PhysicalSortExpr::new_default(col("key", batch.schema_ref())?) + .desc() + .nulls_last(); + assert_emits_before_eof( + &batch, + [order].into(), + InputOrderMode::Linear, + &[Some(3), Some(3)], + ) + .await?; + } + Ok(()) +} diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index 53d09a1bfdbe5..514fdc839c577 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -345,7 +345,7 @@ impl PartitionBatchState { /// ranges of data while processing RANGE frames. /// Attribute `sort_options` stores the column ordering specified by the ORDER /// BY clause. This information is used to calculate the range. -#[derive(Debug, Default, Clone)] +#[derive(Debug, Default)] pub struct WindowFrameStateRange { sort_options: Vec, peer_comparator: Option, @@ -353,10 +353,28 @@ pub struct WindowFrameStateRange { /// Reuse comparators to avoid per-row setup work, including logical-null masks /// that can take O(n) to build for dictionary arrays. -#[derive(Clone)] struct RangePeerComparator { columns: Vec, - comparator: Arc, + comparator: LexicographicalComparator, +} + +impl RangePeerComparator { + fn matches_columns(&self, columns: &[ArrayRef]) -> bool { + // Sliced arrays have different row indices even when they share buffers. + self.columns.len() == columns.len() + && self + .columns + .iter() + .zip(columns) + .all(|(old, new)| Arc::ptr_eq(old, new)) + } +} + +impl Clone for WindowFrameStateRange { + fn clone(&self) -> Self { + // A clone builds its own comparator when it first sees nested keys. + Self::new(self.sort_options.clone()) + } } // Preserve the public WindowFrameStateRange type's existing unwind-safety traits. @@ -382,34 +400,40 @@ impl WindowFrameStateRange { } } - fn peer_comparator( + fn peer_boundary( &mut self, range_columns: &[ArrayRef], - ) -> Result<&LexicographicalComparator> { - let needs_rebuild = self.peer_comparator.as_ref().is_none_or(|cached| { - cached.columns.len() != range_columns.len() - || cached - .columns + idx: usize, + mut boundary: usize, + length: usize, + ) -> Result { + let cached = match &self.peer_comparator { + Some(cached) if cached.matches_columns(range_columns) => cached, + _ => { + let sort_columns = range_columns .iter() - .zip(range_columns) - .any(|(old, new)| !Arc::ptr_eq(old, new)) - }); - if needs_rebuild { - // Sliced arrays have different row indices even when they share buffers. - let sort_columns = range_columns - .iter() - .zip(&self.sort_options) - .map(|(values, options)| SortColumn { - values: Arc::clone(values), - options: Some(*options), + .zip(&self.sort_options) + .map(|(values, options)| SortColumn { + values: Arc::clone(values), + options: Some(*options), + }) + .collect::>(); + self.peer_comparator.insert(RangePeerComparator { + columns: range_columns.to_vec(), + comparator: LexicographicalComparator::try_new(&sort_columns)?, }) - .collect::>(); - self.peer_comparator = Some(RangePeerComparator { - columns: range_columns.to_vec(), - comparator: Arc::new(LexicographicalComparator::try_new(&sort_columns)?), - }); + } + }; + // ScalarValue equality can distinguish encodings of the same nested value: + // a list with a null dictionary key and one with a key referencing null. + while boundary < length { + let is_peer = cached.comparator.compare(boundary, idx).is_eq(); + if is_peer == START { + break; + } + boundary += 1; } - Ok(&self.peer_comparator.as_ref().unwrap().comparator) + Ok(boundary) } /// This function calculates beginning/ending indices for the frame of the current row. @@ -425,6 +449,12 @@ impl WindowFrameStateRange { length: usize, idx: usize, ) -> Result> { + let has_current_row = window_frame.start_bound == WindowFrameBound::CurrentRow + || window_frame.end_bound == WindowFrameBound::CurrentRow; + let nested_current_row = has_current_row + && range_columns + .iter() + .any(|column| column.data_type().is_nested()); let start = match window_frame.start_bound { WindowFrameBound::Preceding(ref n) => { if n.is_null() { @@ -440,6 +470,9 @@ impl WindowFrameStateRange { )? } } + WindowFrameBound::CurrentRow if nested_current_row => { + self.peer_boundary::(range_columns, idx, last_range.start, length)? + } WindowFrameBound::CurrentRow => self.calculate_index_of_row::( range_columns, last_range, @@ -465,6 +498,16 @@ impl WindowFrameStateRange { Some(n), length, )?, + WindowFrameBound::CurrentRow if nested_current_row => { + // An equality scan cannot skip preceding non-peers if evaluation + // starts partway through a partition. + self.peer_boundary::( + range_columns, + idx, + last_range.end.max(idx), + length, + )? + } WindowFrameBound::CurrentRow => self.calculate_index_of_row::( range_columns, last_range, @@ -494,39 +537,13 @@ impl WindowFrameStateRange { /// called twice, in succession, to get window frame start and end indices (with `SIDE` /// supplied as true and false, respectively). fn calculate_index_of_row( - &mut self, + &self, range_columns: &[ArrayRef], last_range: &Range, idx: usize, delta: Option<&ScalarValue>, length: usize, ) -> Result { - if delta.is_none() - && range_columns - .iter() - .any(|column| column.data_type().is_nested()) - { - // CURRENT ROW needs peer equality. ScalarValue equality can distinguish - // encodings of the same nested value: a list with a null dictionary - // key and one with a key referencing null. Use Arrow's logical - // comparison, which also matches the ordering used to sort these keys. - let mut boundary = if SIDE { - last_range.start - } else { - // Evaluation may start partway through a partition. An equality - // scan cannot skip preceding non-peers to reach the current row. - last_range.end.max(idx) - }; - let comparator = self.peer_comparator(range_columns)?; - while boundary < length { - let is_peer = comparator.compare(boundary, idx).is_eq(); - if is_peer == SIDE { - break; - } - boundary += 1; - } - return Ok(boundary); - } let mut end_range = get_row_at_idx(range_columns, idx)?; let search_start = if SIDE { last_range.start @@ -817,7 +834,7 @@ mod tests { use arrow::datatypes::{Field, Schema}; #[test] - fn range_peer_comparator_reuses_arrays_and_releases_after_pruning() -> Result<()> { + fn nested_range_peers_across_cloning_and_pruning() -> Result<()> { use arrow::array::ListArray; use arrow::datatypes::Int64Type; @@ -838,15 +855,12 @@ mod tests { range.calculate_range(&frame, &(0..0), &columns, 4, 0)?, 0..2 ); - let first = Arc::downgrade(&range.peer_comparator.as_ref().unwrap().comparator); assert_eq!( - range.calculate_range(&frame, &(0..2), &columns, 4, 1)?, + range + .clone() + .calculate_range(&frame, &(0..2), &columns, 4, 1)?, 0..2 ); - assert!(Arc::ptr_eq( - &first.upgrade().unwrap(), - &range.peer_comparator.as_ref().unwrap().comparator, - )); let columns = [columns[0].slice(1, 3)]; assert_eq!( diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 607d5f75368c6..42e189eebefca 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7214,18 +7214,53 @@ select count(*) over (order by l nulls first) from (values ([NULL::int]), ([1]), 2 3 -# Lists containing structs or lists also form RANGE peer groups. -query I -select count(*) over (order by l) from (values ([struct(1)]), ([struct(2)])) t(l) ----- -1 -2 +# With one row per batch, nested peers cross batch boundaries. The tie column +# separates rows with the same nested key. +statement ok +set datafusion.execution.batch_size = 1; -query I -select count(*) over (order by l) from (values ([[1]]), ([[2]])) t(l) +statement ok +create table nested_range_peers as +select * from (values + (NULL, NULL, 0, 1), + (NULL, NULL, 0, 2), + ([[NULL::int]], [struct(NULL::int)], 0, 4), + ([[NULL::int]], [struct(NULL::int)], 0, 8), + ([[NULL::int]], [struct(NULL::int)], 1, 16), + ([[1]], [struct(1::int)], 0, 32) +) t(lists, structs, tie, v); + +query IIII +select count(*) over (order by lists nulls first, tie), + count(*) over (order by structs nulls first, tie), + sum(v) over (order by lists nulls first, tie range between current row and current row), + sum(v) over (order by structs nulls first, tie range between current row and current row) +from nested_range_peers order by v +---- +2 2 3 3 +2 2 3 3 +4 4 12 12 +4 4 12 12 +5 5 16 16 +6 6 32 32 + +query II +select sum(v) over (order by lists nulls first, tie range between current row and unbounded following), + sum(v) over (order by structs nulls first, tie range between current row and unbounded following) +from nested_range_peers order by v ---- -1 -2 +63 63 +63 63 +60 60 +60 60 +48 48 +32 32 + +statement ok +drop table nested_range_peers; + +statement ok +reset datafusion.execution.batch_size; # after a duration key the free-range check looks at every ORDER BY type. # Trailing list keys remain rejected by this check.