From 4a19543cea4bc3396e2218b890f85a6840435bc9 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 6 Oct 2026 01:36:14 +0000 Subject: [PATCH 1/2] 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 b5083a3b77463..f4e8be3efeb32 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 current_row_values = get_row_at_idx(range_columns, idx)?; let search_start = if SIDE { last_range.start @@ -733,6 +808,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 c09552be0a58b..f6170a30f1cee 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1178,11 +1178,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, @@ -1196,15 +1193,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(_, _) => { @@ -1217,8 +1205,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, @@ -1260,9 +1248,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 a6b64b98de0e3..1e3c7d396a4c1 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 f0d7123ed4da0819d75e09dc4ae2c18df9a9249d Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 7 Oct 2026 03:57:57 +0000 Subject: [PATCH 2/2] 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 f4e8be3efeb32..9dd1913ee503c 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") @@ -813,6 +818,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,