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/execution/window_nested_range.rs b/datafusion/core/tests/execution/window_nested_range.rs new file mode 100644 index 0000000000000..f2343167eb10e --- /dev/null +++ b/datafusion/core/tests/execution/window_nested_range.rs @@ -0,0 +1,345 @@ +// 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, Int32Type, 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, ExecutionPlanProperties, 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), + // 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, + ), + ])?) + }) + .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(); + // Repeated peers cross input batch boundaries. + 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![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)), + 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 mode in modes.iter().map(Option::as_ref) { + // The planner uses WindowAggExec for UNBOUNDED FOLLOWING. + if mode.is_some() && end.is_unbounded() { + continue; + } + 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)?; + 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:?}, mode={mode:?}", + 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()), + )) + } +} + +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] { + 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 acba75268b50d..514fdc839c577 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,11 @@ 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, .. }) => { + // 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, .. }) => { let mut n_group_to_del = 0; for (_, end_idx) in &state.group_end_indices { @@ -340,15 +345,95 @@ 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, +} + +/// Reuse comparators to avoid per-row setup work, including logical-null masks +/// that can take O(n) to build for dictionary arrays. +struct RangePeerComparator { + columns: Vec, + 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. +// 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 {} + +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_boundary( + &mut self, + range_columns: &[ArrayRef], + 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(&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)?, + }) + } + }; + // 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(boundary) } /// This function calculates beginning/ending indices for the frame of the current row. @@ -364,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() { @@ -379,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, @@ -404,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, @@ -433,7 +537,7 @@ 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, @@ -729,6 +833,99 @@ mod tests { use arrow::array::{Float64Array, UInt64Array}; use arrow::datatypes::{Field, Schema}; + #[test] + fn nested_range_peers_across_cloning_and_pruning() -> Result<()> { + 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, + 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 + ); + assert_eq!( + range + .clone() + .calculate_range(&frame, &(0..2), &columns, 4, 1)?, + 0..2 + ); + + let columns = [columns[0].slice(1, 3)]; + assert_eq!( + range.calculate_range(&frame, &(0..0), &columns, 3, 0)?, + 0..1 + ); + let retained = Arc::downgrade(&columns[0]); + 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_eq!( + state.window_frame_ctx.as_mut().unwrap().calculate_range( + &pruned, + &state.window_frame_range, + 2, + 0, + )?, + 0..2 + ); + 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 e31dc65173474..bd3aea11ae8db 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1325,11 +1325,9 @@ 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. +/// 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, @@ -1342,16 +1340,8 @@ 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. +/// 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(_, _) => { @@ -1364,8 +1354,9 @@ 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`). 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, @@ -1407,9 +1398,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/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 7c92514e723bd..42e189eebefca 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,64 @@ 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\}\]\) -select count(*) over (order by l) from (values ([struct(1)]), ([struct(2)])) t(l) - -query error DataFusion error: Internal error: Uncomparable values: List\(\[\[1\]\]\), List\(\[\[1\]\]\) -select count(*) over (order by l) from (values ([[1]]), ([[2]])) t(l) - -# 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 +# 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 + +# 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; + +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 +---- +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. 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)