diff --git a/datafusion/core/tests/execution/mod.rs b/datafusion/core/tests/execution/mod.rs index 24a3ff6c7f587..ac595de2ab4c2 100644 --- a/datafusion/core/tests/execution/mod.rs +++ b/datafusion/core/tests/execution/mod.rs @@ -19,4 +19,5 @@ mod coop; mod datasource_split; mod logical_plan; mod register_arrow; +mod window_nested_groups_rank; mod window_nested_range; diff --git a/datafusion/core/tests/execution/window_nested_groups_rank.rs b/datafusion/core/tests/execution/window_nested_groups_rank.rs new file mode 100644 index 0000000000000..61fb096a5138e --- /dev/null +++ b/datafusion/core/tests/execution/window_nested_groups_rank.rs @@ -0,0 +1,92 @@ +// 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. + +//! GROUPS frames and bounded rank peers for nested dictionary nulls (#26689). + +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, DictionaryArray, Int8Array, Int32Array, Int64Array, ListArray, UInt64Array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{Field, Int8Type}; +use arrow::record_batch::RecordBatch; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::Result; + +fn nested_dict_null_batch() -> Result { + let children: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![None, Some(0)]), + Arc::new(Int32Array::from(vec![None])), + )?); + let key: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", children.data_type().clone(), true)), + OffsetBuffer::from_lengths([1, 1]), + children, + None, + )); + Ok(RecordBatch::try_from_iter([("k", key)])?) +} + +async fn ctx_with_batch() -> Result { + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + ctx.register_batch("t", nested_dict_null_batch()?)?; + Ok(ctx) +} + +#[tokio::test] +async fn groups_current_row_counts_nested_dictionary_null_peers() -> Result<()> { + let ctx = ctx_with_batch().await?; + let batches = ctx + .sql("SELECT count(*) OVER (ORDER BY k GROUPS BETWEEN CURRENT ROW AND CURRENT ROW) AS c FROM t") + .await? + .collect() + .await?; + assert_eq!(batches.len(), 1); + let col = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("count OVER returns Int64"); + assert_eq!(col.values(), &[2, 2]); + Ok(()) +} + +#[tokio::test] +async fn bounded_rank_keeps_nested_dictionary_null_peers() -> Result<()> { + let ctx = ctx_with_batch().await?; + let batches = ctx + .sql("SELECT rank() OVER w AS r, dense_rank() OVER w AS d FROM t WINDOW w AS (ORDER BY k)") + .await? + .collect() + .await?; + assert_eq!(batches.len(), 1); + let rank = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let dense = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(rank.values(), &[1, 1]); + assert_eq!(dense.values(), &[1, 1]); + Ok(()) +} diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index f301c91253df3..48635707cf5a4 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -90,6 +90,8 @@ impl WindowAggState { .iter_mut() .for_each(|(_, start_idx)| *start_idx -= n_prune); state.current_group_idx -= n_group_to_del; + // Replacement ORDER BY arrays invalidate the cache by identity. + state.peer_comparator = None; } None => {} } @@ -204,7 +206,7 @@ impl WindowFrameContext { }, WindowFrameUnits::Groups => WindowFrameContext::Groups { window_frame, - state: WindowFrameStateGroups::default(), + state: WindowFrameStateGroups::new(sort_options), }, } } @@ -234,9 +236,8 @@ impl WindowFrameContext { length, idx, ), - // Sort options is not used in GROUPS mode calculations as the - // inequality of two rows indicates a group change, and ordering - // or position of NULLs do not impact inequality. + // Sort options feed the nested peer comparator. Non-nested keys + // still use ScalarValue equality; null ordering does not affect it. WindowFrameContext::Groups { window_frame, state, @@ -624,7 +625,7 @@ impl WindowFrameStateRange { /// This structure encapsulates all the state information we require as we /// scan groups of data while processing window frames. -#[derive(Debug, Default, Clone)] +#[derive(Debug, Default)] pub struct WindowFrameStateGroups { /// A tuple containing group values and the row index where the group ends. /// Example: [[1, 1], [1, 1], [2, 1], [2, 1], ...] would correspond to @@ -632,9 +633,98 @@ pub struct WindowFrameStateGroups { pub group_end_indices: VecDeque<(Vec, usize)>, /// The group index to which the row index belongs. pub current_group_idx: usize, + /// ORDER BY sort options used to build the nested peer comparator. + sort_options: Vec, + /// Cached Arrow comparator for nested ORDER BY keys. + peer_comparator: Option, +} + +impl Clone for WindowFrameStateGroups { + fn clone(&self) -> Self { + // A clone builds its own comparator when it first sees nested keys. + Self::new(self.sort_options.clone()) + } } impl WindowFrameStateGroups { + /// Create a new object to store the search state. + fn new(sort_options: Vec) -> Self { + Self { + group_end_indices: VecDeque::new(), + current_group_idx: 0, + sort_options, + peer_comparator: None, + } + } + + fn ensure_peer_comparator( + &mut self, + range_columns: &[ArrayRef], + ) -> Result<&LexicographicalComparator> { + let needs_rebuild = match &self.peer_comparator { + Some(cached) => !cached.matches_columns(range_columns), + None => true, + }; + if needs_rebuild { + let sort_columns = range_columns + .iter() + .enumerate() + .map(|(i, values)| SortColumn { + values: Arc::clone(values), + options: Some(self.sort_options.get(i).copied().unwrap_or_default()), + }) + .collect::>(); + self.peer_comparator = Some(RangePeerComparator { + columns: range_columns.to_vec(), + comparator: LexicographicalComparator::try_new(&sort_columns)?, + }); + } + Ok(&self.peer_comparator.as_ref().unwrap().comparator) + } + + /// Exclusive end of the group containing `peer_idx`, searching from `start`. + fn find_group_end( + &mut self, + range_columns: &[ArrayRef], + peer_idx: usize, + start: usize, + length: usize, + ) -> Result { + let nested = range_columns + .iter() + .any(|column| column.data_type().is_nested()); + if nested { + let comparator = self.ensure_peer_comparator(range_columns)?; + let mut end = start; + while end < length && comparator.compare(peer_idx, end).is_eq() { + end += 1; + } + Ok(end) + } else { + let group_row = get_row_at_idx(range_columns, peer_idx)?; + search_in_slice(range_columns, &group_row, check_equality, start, length) + } + } + + /// Whether rows `peer_idx` and `idx` belong to the same GROUPS peer set. + fn is_same_group( + &mut self, + range_columns: &[ArrayRef], + peer_idx: usize, + idx: usize, + group_row: &[ScalarValue], + ) -> Result { + let nested = range_columns + .iter() + .any(|column| column.data_type().is_nested()); + if nested { + let comparator = self.ensure_peer_comparator(range_columns)?; + Ok(comparator.compare(peer_idx, idx).is_eq()) + } else { + Ok(get_row_at_idx(range_columns, idx)? == group_row) + } + } + fn calculate_range( &mut self, window_frame: &Arc, @@ -723,38 +813,40 @@ impl WindowFrameStateGroups { } else { 0 }; - let mut group_start = 0; - let last_group = self.group_end_indices.back_mut(); - if let Some((group_row, group_end)) = last_group { - if *group_end < length { - let new_group_row = get_row_at_idx(range_columns, *group_end)?; - // If last/current group keys are the same, we extend the last group: - if new_group_row.eq(group_row) { - // Update the end boundary of the group (search right boundary): - *group_end = search_in_slice( - range_columns, - group_row, - check_equality, - *group_end, - length, - )?; + + // Optionally extend the trailing group when new rows still belong to it. + // Nested ORDER BY keys need Arrow logical comparison: ScalarValue + // equality can distinguish a null dictionary key from a key that + // references a null value. + let extend_info = self + .group_end_indices + .back() + .map(|(row, end)| (row.clone(), *end)); + if let Some((group_row, group_end)) = extend_info { + if group_end < length && group_end > 0 { + let peer_idx = group_end - 1; + if self.is_same_group(range_columns, peer_idx, group_end, &group_row)? { + let new_end = + self.find_group_end(range_columns, peer_idx, group_end, length)?; + if let Some((_, end)) = self.group_end_indices.back_mut() { + *end = new_end; + } } } - // Start searching from the last group boundary: - group_start = *group_end; } + let mut group_start = self + .group_end_indices + .back() + .map(|(_, end)| *end) + .unwrap_or(0); + // Advance groups until `idx` is inside a group: while idx >= group_start { - let group_row = get_row_at_idx(range_columns, group_start)?; - // Find end boundary of the group (search right boundary): - let group_end = search_in_slice( - range_columns, - &group_row, - check_equality, - group_start, - length, - )?; + let peer_idx = group_start; + let group_row = get_row_at_idx(range_columns, peer_idx)?; + let group_end = + self.find_group_end(range_columns, peer_idx, group_start, length)?; self.group_end_indices.push_back((group_row, group_end)); group_start = group_end; } @@ -773,17 +865,12 @@ impl WindowFrameStateGroups { self.current_group_idx + delta }; - // Extend `group_start_indices` until it includes at least `group_idx`: + // Extend groups until the frame boundary group is materialized: while self.group_end_indices.len() <= group_idx && group_start < length { - let group_row = get_row_at_idx(range_columns, group_start)?; - // Find end boundary of the group (search right boundary): - let group_end = search_in_slice( - range_columns, - &group_row, - check_equality, - group_start, - length, - )?; + let peer_idx = group_start; + let group_row = get_row_at_idx(range_columns, peer_idx)?; + let group_end = + self.find_group_end(range_columns, peer_idx, group_start, length)?; self.group_end_indices.push_back((group_row, group_end)); group_start = group_end; } @@ -927,6 +1014,38 @@ mod tests { Ok(()) } + #[test] + fn groups_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::Groups, + 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, "row {idx} should see both peers"); + } + 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/functions-window/src/rank.rs b/datafusion/functions-window/src/rank.rs index d18bd0748917e..963ce6dad7ef2 100644 --- a/datafusion/functions-window/src/rank.rs +++ b/datafusion/functions-window/src/rank.rs @@ -22,6 +22,9 @@ use arrow::datatypes::FieldRef; use datafusion_common::arrow::array::ArrayRef; use datafusion_common::arrow::array::{Float64Array, UInt64Array}; use datafusion_common::arrow::compute::SortOptions; +use datafusion_common::arrow::compute::kernels::sort::{ + LexicographicalComparator, SortColumn, +}; use datafusion_common::arrow::datatypes::DataType; use datafusion_common::arrow::datatypes::Field; use datafusion_common::utils::get_row_at_idx; @@ -284,8 +287,9 @@ impl PartitionEvaluator for RankEvaluator { let last_rank_data = get_row_at_idx(range_columns, row_idx)?; let new_rank_encountered = if let Some(state_last_rank_data) = &self.state.last_rank_data { - // if rank data changes, new rank is encountered - state_last_rank_data != &last_rank_data + // Nested ORDER BY keys can encode the same value with different + // physical layouts; ScalarValue inequality would split peers. + !rows_logically_equal(state_last_rank_data, &last_rank_data)? } else { // First rank seen true @@ -369,6 +373,39 @@ impl PartitionEvaluator for RankEvaluator { } } +/// Compare two ORDER BY rows with Arrow logical equality when any value is +/// nested. Physical [`PartialEq`] can distinguish a null dictionary key from a +/// key that references a null value. +fn rows_logically_equal(left: &[ScalarValue], right: &[ScalarValue]) -> Result { + if left.len() != right.len() { + return Ok(false); + } + if left == right { + return Ok(true); + } + if !left + .iter() + .chain(right.iter()) + .any(|value| value.data_type().is_nested()) + { + return Ok(false); + } + let columns = left + .iter() + .zip(right.iter()) + .map(|(a, b)| ScalarValue::iter_to_array([a.clone(), b.clone()])) + .collect::>>()?; + let sort_columns = columns + .iter() + .map(|values| SortColumn { + values: Arc::clone(values), + options: Some(SortOptions::default()), + }) + .collect::>(); + let comparator = LexicographicalComparator::try_new(&sort_columns)?; + Ok(comparator.compare(0, 1).is_eq()) +} + #[cfg(test)] mod tests { use super::*; @@ -453,4 +490,43 @@ mod tests { Ok(()) } + + #[test] + fn nested_dictionary_null_peers_share_rank() -> Result<()> { + use arrow::array::{DictionaryArray, Int8Array, Int32Array, ListArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Int8Type}; + use std::sync::Arc; + + // Both lists contain NULL via different dictionary encodings. + let children: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![None, Some(0)]), + Arc::new(Int32Array::from(vec![None])), + )?); + let key: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", children.data_type().clone(), true)), + OffsetBuffer::from_lengths([1, 1]), + children, + None, + )); + + for (expr, expected) in [ + (Rank::basic(), vec![1u64, 1u64]), + (Rank::dense_rank(), vec![1u64, 1u64]), + ] { + let mut evaluator = + expr.partition_evaluator(PartitionEvaluatorArgs::default())?; + let mut out = Vec::new(); + for idx in 0..2 { + let value = + evaluator.evaluate(std::slice::from_ref(&key), &(idx..idx + 1))?; + out.push(match value { + ScalarValue::UInt64(Some(v)) => v, + other => panic!("unexpected {other:?}"), + }); + } + assert_eq!(out, expected, "{}", expr.name()); + } + Ok(()) + } }