From add9da343061372d3348ff27462e2a5282cd9578 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:17:23 -0700 Subject: [PATCH 1/9] Add group selection for preserving reads --- .../expr-common/src/groups_accumulator.rs | 77 ++++++++++++++++++- datafusion/expr/src/lib.rs | 4 +- 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index 5c01418e04ce7..5e14f83f97af3 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -18,7 +18,7 @@ //! Vectorized [`GroupsAccumulator`] use arrow::array::{ArrayRef, BooleanArray}; -use datafusion_common::{Result, utils::split_vec_min_alloc}; +use datafusion_common::{Result, exec_err, utils::split_vec_min_alloc}; /// Describes how many rows should be emitted during grouping. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -50,6 +50,57 @@ impl EmitTo { } } +/// Selects groups for a non-destructive grouped aggregation read. +/// +/// Unlike [`EmitTo`], this selection does not remove groups or change their +/// indices. [`Self::Indices`] preserves the requested order and supports +/// duplicate indices. +/// +/// Indices are trusted to be valid by preserving read APIs. Call +/// [`Self::validate`] first when they do not come from a source that guarantees +/// they are in bounds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GroupSelection<'a> { + /// Select all groups in group-index order. + All, + /// Select groups in the order specified by their group indices. + Indices(&'a [usize]), +} + +impl<'a> GroupSelection<'a> { + /// Validates that every selected index is less than `total_num_groups`. + /// + /// [`Self::len`] and [`Self::iter`] do not call this method implicitly. + pub fn validate(&self, total_num_groups: usize) -> Result<()> { + if let Self::Indices(indices) = self + && let Some(index) = indices.iter().find(|&&index| index >= total_num_groups) + { + return exec_err!( + "Group index {index} is out of bounds for {total_num_groups} groups" + ); + } + Ok(()) + } + + /// Returns the number of selected groups without validating the selection. + pub fn len(&self, total_num_groups: usize) -> usize { + match self { + Self::All => total_num_groups, + Self::Indices(indices) => indices.len(), + } + } + + /// Returns the selected group indices in output order without validating + /// the selection. + pub fn iter(self, total_num_groups: usize) -> impl Iterator + 'a { + let (all, indices): (_, &'a [usize]) = match self { + Self::All => (0..total_num_groups, &[]), + Self::Indices(indices) => (0..0, indices), + }; + all.chain(indices.iter().copied()) + } +} + /// `GroupsAccumulator` implements a single aggregate (e.g. AVG) and /// stores the state for *all* groups internally. /// @@ -247,7 +298,7 @@ pub trait GroupsAccumulator: Send + std::any::Any { #[cfg(test)] mod tests { - use super::EmitTo; + use super::{EmitTo, GroupSelection}; /// When `n` is small relative to `len`, the old `split_off(n) + swap` pattern had /// two allocation problems: @@ -293,4 +344,26 @@ mod tests { original_capacity, ); } + + #[test] + fn group_selection_order_duplicates_and_explicit_validation() { + let values = [10, 20, 30, 40]; + let selected = GroupSelection::Indices(&[3, 1, 3]) + .iter(values.len()) + .map(|index| values[index]) + .collect::>(); + assert_eq!(selected, vec![40, 20, 40]); + + let selected = GroupSelection::All + .iter(values.len()) + .map(|index| values[index]) + .collect::>(); + assert_eq!(selected, values); + + let invalid = GroupSelection::Indices(&[4]); + assert_eq!(invalid.len(values.len()), 1); + assert_eq!(invalid.iter(values.len()).collect::>(), vec![4]); + let error = invalid.validate(values.len()).unwrap_err(); + assert!(error.to_string().contains("out of bounds")); + } } diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 75041c701454a..5b0526efbf3f6 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -103,7 +103,9 @@ pub use datafusion_doc::{ }; pub use datafusion_expr_common::accumulator::Accumulator; pub use datafusion_expr_common::columnar_value::ColumnarValue; -pub use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; +pub use datafusion_expr_common::groups_accumulator::{ + EmitTo, GroupSelection, GroupsAccumulator, +}; pub use datafusion_expr_common::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; pub use datafusion_expr_common::signature::{ From fa459250f20490a6ec99f7dfe84dc9c6ea0f8e51 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:17:49 -0700 Subject: [PATCH 2/9] Add preserving grouped read APIs --- .../expr-common/src/groups_accumulator.rs | 45 ++++++++++++++++++- .../src/aggregates/group_values/mod.rs | 24 +++++++++- 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index 5e14f83f97af3..92957d186fab1 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -18,7 +18,7 @@ //! Vectorized [`GroupsAccumulator`] use arrow::array::{ArrayRef, BooleanArray}; -use datafusion_common::{Result, exec_err, utils::split_vec_min_alloc}; +use datafusion_common::{Result, exec_err, not_impl_err, utils::split_vec_min_alloc}; /// Describes how many rows should be emitted during grouping. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -205,6 +205,28 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// `n`. See [`EmitTo::First`] for more details. fn evaluate(&mut self, emit_to: EmitTo) -> Result; + /// Returns final aggregate values without changing the logical state or + /// group indices. + /// + /// Rows are returned in the order specified by `selection`. Implementations + /// may mutate internal caches or builders, but repeated calls and later + /// updates must observe the same logical accumulator state. + /// + /// Every index in [`GroupSelection::Indices`] must refer to an existing + /// group. Call [`GroupSelection::validate`] first if this is not guaranteed + /// by the source of the indices. Invalid indices may cause a panic. + fn evaluate_preserving( + &mut self, + _selection: GroupSelection<'_>, + ) -> Result { + not_impl_err!("Preserving grouped evaluation is not implemented") + } + + /// Returns `true` if [`Self::evaluate_preserving`] is implemented. + fn supports_evaluate_preserving(&self) -> bool { + false + } + /// Returns the intermediate aggregate state for this accumulator, /// used for multi-phase grouping, resetting its internal state. /// @@ -223,6 +245,27 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// [`Accumulator::state`]: crate::accumulator::Accumulator::state fn state(&mut self, emit_to: EmitTo) -> Result>; + /// Returns intermediate aggregate state without changing the logical state + /// or group indices. + /// + /// Each returned array has one row per selected group, in the order + /// specified by `selection`. + /// + /// Every index in [`GroupSelection::Indices`] must refer to an existing + /// group. Call [`GroupSelection::validate`] first if this is not guaranteed + /// by the source of the indices. Invalid indices may cause a panic. + fn state_preserving( + &mut self, + _selection: GroupSelection<'_>, + ) -> Result> { + not_impl_err!("Preserving grouped state is not implemented") + } + + /// Returns `true` if [`Self::state_preserving`] is implemented. + fn supports_state_preserving(&self) -> bool { + false + } + /// Merges intermediate state (the output from [`Self::state`]) /// into this accumulator's current state. /// diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 1101d535311e4..b3e97d8270792 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -24,9 +24,9 @@ use arrow::array::types::{ }; use arrow::array::{ArrayRef, downcast_primitive}; use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; -use datafusion_common::Result; +use datafusion_common::{Result, not_impl_err}; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; pub mod multi_group_by; @@ -113,6 +113,26 @@ pub trait GroupValues: Send { /// Emits the group values fn emit(&mut self, emit_to: EmitTo) -> Result>; + /// Materializes selected group values without changing the stored values or + /// their group indices. + /// + /// Rows are returned in the order specified by `selection`. + /// + /// Every index in [`GroupSelection::Indices`] must refer to an existing + /// group. Call [`GroupSelection::validate`] first if this is not guaranteed + /// by the source of the indices. Invalid indices may cause a panic. + fn values_preserving( + &mut self, + _selection: GroupSelection<'_>, + ) -> Result> { + not_impl_err!("Preserving group values are not implemented") + } + + /// Returns `true` if [`Self::values_preserving`] is implemented. + fn supports_values_preserving(&self) -> bool { + false + } + /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage) fn clear_shrink(&mut self, num_rows: usize); } From 2e4901ce6356033c2788c1149e394d5c15319685 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:18:07 -0700 Subject: [PATCH 3/9] Implement preserving reads for group values --- .../physical-expr-common/src/binary_map.rs | 52 +++++++- .../src/binary_view_map.rs | 66 +++++++++- .../src/aggregates/group_values/mod.rs | 120 ++++++++++++++++++ .../group_values/multi_group_by/boolean.rs | 10 ++ .../group_values/multi_group_by/bytes.rs | 32 +++++ .../group_values/multi_group_by/bytes_view.rs | 50 ++++++++ .../multi_group_by/fixed_size_binary.rs | 44 +++++++ .../group_values/multi_group_by/mod.rs | 79 +++++++++++- .../group_values/multi_group_by/primitive.rs | 15 +++ .../group_values/multi_group_by/row_backed.rs | 36 ++++++ .../aggregates/group_values/null_builder.rs | 21 +++ .../src/aggregates/group_values/row.rs | 77 ++++++++++- .../group_values/single_group_by/boolean.rs | 32 ++++- .../group_values/single_group_by/bytes.rs | 14 +- .../single_group_by/bytes_view.rs | 14 +- .../group_values/single_group_by/primitive.rs | 33 ++++- 16 files changed, 682 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 44ca35c7f8708..0bd7f97fb76b3 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -19,8 +19,8 @@ //! StringArray / LargeStringArray / BinaryArray / LargeBinaryArray. use arrow::array::{ - Array, ArrayRef, GenericBinaryArray, GenericStringArray, NullBufferBuilder, - OffsetSizeTrait, + Array, ArrayRef, BufferBuilder, GenericBinaryArray, GenericStringArray, + NullBufferBuilder, OffsetSizeTrait, cast::AsArray, types::{ByteArrayType, GenericBinaryType, GenericStringType}, }; @@ -29,6 +29,7 @@ use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::{Result, exec_err}; use std::any::type_name; use std::fmt::Debug; use std::mem::{size_of, swap}; @@ -526,6 +527,53 @@ where } } + /// Copies keys at `indices` into an array without changing this map. + /// + /// Keys are returned in index order, and duplicate indices are supported. + pub fn keys(&self, indices: I) -> Result + where + I: IntoIterator, + { + let indices = indices.into_iter(); + let mut output_offsets = Vec::with_capacity(indices.size_hint().0 + 1); + let mut output_values = BufferBuilder::::new(0); + let mut output_nulls = NullBufferBuilder::new(indices.size_hint().0); + output_offsets.push(O::default()); + let null_index = self.null.map(|(_, index)| index); + let num_keys = self.offsets.len() - 1; + + for index in indices { + if index >= num_keys { + return exec_err!( + "Key index {index} is out of bounds for {num_keys} keys" + ); + } + let start = self.offsets[index].as_usize(); + let end = self.offsets[index + 1].as_usize(); + output_values.append_slice(&self.buffer.as_slice()[start..end]); + if O::from_usize(output_values.len()).is_none() { + return exec_err!("Offset overflow while copying byte map keys"); + } + output_offsets.push(O::usize_as(output_values.len())); + output_nulls.append(index != null_index.unwrap_or(usize::MAX)); + } + + // SAFETY: offsets are constructed from the length of `output_values`. + let offsets = + unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(output_offsets)) }; + let values = output_values.finish(); + let nulls = output_nulls.finish(); + Ok(match self.output_type { + OutputType::Binary => Arc::new(unsafe { + GenericBinaryArray::new_unchecked(offsets, values, nulls) + }), + OutputType::Utf8 => Arc::new(unsafe { + GenericStringArray::new_unchecked(offsets, values, nulls) + }), + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + }) + } + /// Total number of entries (including null, if present) pub fn len(&self) -> usize { self.non_null_len() + self.null.map(|_| 1).unwrap_or(0) diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9d4b556393a24..2467b6e7bda09 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -20,12 +20,16 @@ use crate::binary_map::OutputType; use arrow::array::NullBufferBuilder; use arrow::array::cast::AsArray; -use arrow::array::{Array, ArrayRef, BinaryViewArray, ByteView, make_view}; +use arrow::array::{ + Array, ArrayRef, BinaryViewArray, BinaryViewBuilder, ByteView, StringViewBuilder, + make_view, +}; use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType}; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::{Result, exec_err}; use std::fmt::Debug; use std::mem::size_of; use std::sync::Arc; @@ -375,6 +379,25 @@ where } } + /// Returns the bytes for the key at `index`, irrespective of nullness. + fn key(&self, index: usize) -> &[u8] { + let view = &self.views[index]; + let byte_view = ByteView::from(*view); + let length = byte_view.length as usize; + if length <= 12 { + // SAFETY: `view` is a valid inline view with `length` bytes. + unsafe { BinaryViewArray::inline_value(view, length) } + } else { + let buffer_index = byte_view.buffer_index as usize; + let offset = byte_view.offset as usize; + if buffer_index < self.completed.len() { + &self.completed[buffer_index][offset..offset + length] + } else { + &self.in_progress[offset..offset + length] + } + } + } + /// Converts this set into a `StringViewArray`, or `BinaryViewArray`, /// containing each distinct value /// that was inserted. This is done without copying the values. @@ -406,6 +429,47 @@ where } } + /// Copies keys at `indices` into an array without changing this map. + /// + /// Keys are returned in index order, and duplicate indices are supported. + pub fn keys(&self, indices: I) -> Result + where + I: IntoIterator, + { + let indices = indices.into_iter(); + let capacity = indices.size_hint().0; + let num_keys = self.views.len(); + let null_index = self.null.map(|(_, index)| index); + + macro_rules! build_keys { + ($builder:ty, $value:expr) => {{ + let mut builder = <$builder>::with_capacity(capacity); + for index in indices { + if index >= num_keys { + return exec_err!( + "Key index {index} is out of bounds for {num_keys} keys" + ); + } + if null_index == Some(index) { + builder.append_null(); + } else { + builder.append_value($value(self.key(index))); + } + } + Arc::new(builder.finish()) as ArrayRef + }}; + } + + Ok(match self.output_type { + OutputType::BinaryView => build_keys!(BinaryViewBuilder, |value| value), + OutputType::Utf8View => build_keys!(StringViewBuilder, |value| unsafe { + // Inputs to an Utf8View map are validated UTF-8. + std::str::from_utf8_unchecked(value) + }), + _ => unreachable!("Utf8/Binary should use `ArrowBytesMap`"), + }) + } + /// Append an already-computed inline view (len <= 12) directly, bypassing /// buffer allocation. /// diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index b3e97d8270792..162f6c1de617b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -232,3 +232,123 @@ pub fn new_group_values( Ok(Box::new(GroupValuesRows::try_new(schema)?)) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{ArrayRef, AsArray, Int32Array, StringArray, StringViewArray}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use datafusion_expr::{EmitTo, GroupSelection}; + + use super::new_group_values; + use crate::aggregates::order::GroupOrdering; + + #[test] + fn preserving_values_keep_group_indices_valid() { + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + DataType::Int32, + true, + )])); + let mut group_values = new_group_values(schema, &GroupOrdering::None).unwrap(); + assert!(group_values.supports_values_preserving()); + + let input = Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + Some(10), + None, + Some(30), + ])) as ArrayRef; + let mut groups = vec![]; + group_values.intern(&[input], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 0, 2, 3]); + + let selection = GroupSelection::Indices(&[3, 0, 2, 0]); + let expected = Int32Array::from(vec![Some(30), Some(10), None, Some(10)]); + for _ in 0..2 { + let actual = group_values.values_preserving(selection).unwrap(); + assert_eq!(actual[0].as_primitive::(), &expected); + } + + let input = + Arc::new(Int32Array::from(vec![Some(20), Some(40), None])) as ArrayRef; + group_values.intern(&[input], &mut groups).unwrap(); + assert_eq!(groups, vec![1, 4, 2]); + + let expected = + Int32Array::from(vec![Some(10), Some(20), None, Some(30), Some(40)]); + let actual = group_values.values_preserving(GroupSelection::All).unwrap(); + assert_eq!(actual[0].as_primitive::(), &expected); + + let error = GroupSelection::Indices(&[5]) + .validate(group_values.len()) + .unwrap_err(); + assert!(error.to_string().contains("out of bounds")); + + let actual = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(actual[0].as_primitive::(), &expected); + } + + #[test] + fn preserving_variable_width_values() { + for data_type in [DataType::Utf8, DataType::Utf8View] { + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + data_type.clone(), + true, + )])); + let mut group_values = + new_group_values(schema, &GroupOrdering::None).unwrap(); + let input: ArrayRef = match data_type { + DataType::Utf8 => Arc::new(StringArray::from(vec![ + Some("a"), + None, + Some("a long value that is not inline"), + Some("a"), + ])), + DataType::Utf8View => Arc::new(StringViewArray::from(vec![ + Some("a"), + None, + Some("a long value that is not inline"), + Some("a"), + ])), + _ => unreachable!(), + }; + let mut groups = vec![]; + group_values.intern(&[input], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0]); + + let selected = group_values + .values_preserving(GroupSelection::Indices(&[2, 1, 0, 2])) + .unwrap(); + let expected = vec![ + Some("a long value that is not inline"), + None, + Some("a"), + Some("a long value that is not inline"), + ]; + match data_type { + DataType::Utf8 => assert_eq!( + selected[0].as_string::(), + &StringArray::from(expected) + ), + DataType::Utf8View => assert_eq!( + selected[0].as_string_view(), + &StringViewArray::from(expected) + ), + _ => unreachable!(), + } + + // Reading did not remove values or alter interned indices. + let input: ArrayRef = match data_type { + DataType::Utf8 => Arc::new(StringArray::from(vec!["new"])), + DataType::Utf8View => Arc::new(StringViewArray::from(vec!["new"])), + _ => unreachable!(), + }; + group_values.intern(&[input], &mut groups).unwrap(); + assert_eq!(groups, vec![3]); + } + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs index 5fdbe434f9f30..b307fdb1b6c2c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs @@ -22,6 +22,7 @@ use crate::aggregates::group_values::multi_group_by::{GroupColumn, nulls_equal_t use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::{Array as _, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder}; use datafusion_common::Result; +use datafusion_expr::GroupSelection; /// An implementation of [`GroupColumn`] for booleans /// @@ -177,6 +178,15 @@ impl GroupColumn for BooleanGroupValueBuilder { Arc::new(arr) } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + let mut values = BooleanBufferBuilder::new(selection.len(self.buffer.len())); + for index in selection.iter(self.buffer.len()) { + values.append(self.buffer.get_bit(index)); + } + let nulls = self.nulls.build_preserving(selection, self.buffer.len())?; + Ok(Arc::new(BooleanArray::new(values.finish(), nulls))) + } + fn take_n(&mut self, n: usize) -> ArrayRef { let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index c83b1da4049bc..9c666ce019c2b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -28,6 +28,7 @@ use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, exec_datafusion_err}; +use datafusion_expr::GroupSelection; use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType}; use std::mem::size_of; use std::sync::Arc; @@ -373,6 +374,37 @@ where } } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + let selected_len = selection.len(self.len()); + let mut buffer = BufferBuilder::::new(0); + let mut offsets = Vec::with_capacity(selected_len + 1); + let mut nulls = MaybeNullBufferBuilder::new(); + offsets.push(O::default()); + + for index in selection.iter(self.len()) { + let is_null = self.nulls.is_null(index); + nulls.append(is_null); + if !is_null { + buffer.append_slice(self.value(index)); + } + offsets.push(O::usize_as(buffer.len())); + } + + // SAFETY: offsets are constructed from the length of `buffer`. + let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; + let values = buffer.finish(); + let nulls = nulls.build(); + Ok(match self.output_type { + OutputType::Binary => Arc::new(unsafe { + GenericBinaryArray::new_unchecked(offsets, values, nulls) + }), + OutputType::Utf8 => Arc::new(unsafe { + GenericStringArray::new_unchecked(offsets, values, nulls) + }), + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + }) + } + fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(self.len() >= n); let null_buffer = self.nulls.take_n(n); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs index 8625772e2c995..ecbddb7f431c4 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs @@ -21,11 +21,13 @@ use crate::aggregates::group_values::multi_group_by::{ use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::{ Array, ArrayRef, AsArray, BooleanBufferBuilder, ByteView, GenericByteViewArray, + make_view, }; use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::ByteViewType; use datafusion_common::Result; use datafusion_common::utils::split_vec_min_alloc; +use datafusion_expr::GroupSelection; use std::marker::PhantomData; use std::mem::{replace, size_of}; use std::sync::Arc; @@ -327,6 +329,50 @@ impl ByteViewGroupValueBuilder { } } + /// Returns the bytes stored at `index`, irrespective of nullness. + fn value(&self, index: usize) -> &[u8] { + let view = &self.views[index]; + let byte_view = ByteView::from(*view); + let length = byte_view.length as usize; + if length <= 12 { + // SAFETY: `view` is a valid inline view with `length` bytes. + unsafe { GenericByteViewArray::::inline_value(view, length) } + } else { + let buffer_index = byte_view.buffer_index as usize; + let offset = byte_view.offset as usize; + if buffer_index < self.completed.len() { + &self.completed[buffer_index][offset..offset + length] + } else { + &self.in_progress[offset..offset + length] + } + } + } + + fn values_preserving_inner(&self, selection: GroupSelection<'_>) -> Result { + let mut selected = Self::new().with_max_block_size(self.max_block_size); + for index in selection.iter(self.len()) { + let is_null = self.nulls.is_null(index); + selected.nulls.append(is_null); + if is_null { + selected.views.push(0); + continue; + } + + let value = self.value(index); + let view = if value.len() <= 12 { + make_view(value, 0, 0) + } else { + selected.ensure_in_progress_big_enough(value.len()); + let buffer_index = selected.completed.len() as u32; + let offset = selected.in_progress.len() as u32; + selected.in_progress.extend_from_slice(value); + make_view(value, buffer_index, offset) + }; + selected.views.push(view); + } + Ok(selected.build_inner()) + } + fn build_inner(self) -> ArrayRef { let Self { views, @@ -601,6 +647,10 @@ impl GroupColumn for ByteViewGroupValueBuilder { Self::build_inner(*self) } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + self.values_preserving_inner(selection) + } + fn take_n(&mut self, n: usize) -> ArrayRef { self.take_n_inner(n) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs index 589083c8f7ce2..a8ecf47c9baa6 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs @@ -26,6 +26,7 @@ use arrow::buffer::{Buffer, NullBuffer}; use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, exec_datafusion_err}; +use datafusion_expr::GroupSelection; use std::sync::Arc; /// An implementation of [`GroupColumn`] for `FixedSizeBinary` values @@ -224,6 +225,23 @@ impl GroupColumn for FixedSizeBinaryGroupValueBuilder { Self::build_array(byte_width, buffer, nulls.build(), len) } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.len).is_ok()); + let len = selection.len(self.len); + let mut values = Vec::with_capacity(len * self.byte_width); + let mut nulls = MaybeNullBufferBuilder::new(); + for index in selection.iter(self.len) { + nulls.append(self.nulls.is_null(index)); + values.extend_from_slice(self.value(index)); + } + Ok(Self::build_array( + self.byte_width, + values, + nulls.build(), + len, + )) + } + fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(self.len >= n); @@ -241,6 +259,7 @@ mod tests { use crate::aggregates::group_values::multi_group_by::fixed_size_binary::FixedSizeBinaryGroupValueBuilder; use arrow::array::{ArrayRef, BooleanBufferBuilder, FixedSizeBinaryArray}; + use datafusion_expr::GroupSelection; use super::GroupColumn; @@ -474,6 +493,31 @@ mod tests { assert_eq!(builder.len(), 0); } + #[test] + fn test_fixed_size_binary_values_preserving() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let input = make_array( + vec![Some(b"aa".as_slice()), None, Some(b"bb".as_slice())], + 2, + ); + builder.vectorized_append(&input, &[0, 1, 2]).unwrap(); + + let output = builder + .values_preserving(GroupSelection::Indices(&[2, 0, 1, 2])) + .unwrap(); + let expected = make_array( + vec![ + Some(b"bb".as_slice()), + Some(b"aa".as_slice()), + None, + Some(b"bb".as_slice()), + ], + 2, + ); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 3); + } + #[test] fn test_fixed_size_binary_build() { let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 5b474f3bae075..418d0780730b7 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -33,7 +33,7 @@ use crate::aggregates::group_values::multi_group_by::{ fixed_size_binary::FixedSizeBinaryGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, row_backed::RowsGroupColumn, }; -use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder, new_empty_array}; use arrow::compute::cast; use arrow::datatypes::{ BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, @@ -50,7 +50,7 @@ use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; use hashbrown::hash_table::HashTable; @@ -109,6 +109,13 @@ pub trait GroupColumn: Send + Sync { /// Builds a new array from all of the stored rows fn build(self: Box) -> ArrayRef; + /// Builds a new array from selected stored rows without changing this + /// column. Rows are returned in selection order. The caller must ensure all + /// selected indices are in bounds. + fn values_preserving(&self, _selection: GroupSelection<'_>) -> Result { + not_impl_err!("Preserving group column values are not implemented") + } + /// Builds a new array from the first `n` stored rows, shifting the /// remaining rows to the start of the builder fn take_n(&mut self, n: usize) -> ArrayRef; @@ -1285,6 +1292,30 @@ impl GroupValues for GroupValuesColumn { Ok(output) } + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + debug_assert!(selection.validate(self.len()).is_ok()); + if self.group_values.is_empty() { + return Ok(self + .schema + .fields() + .iter() + .map(|field| new_empty_array(field.data_type())) + .collect()); + } + + self.group_values + .iter() + .map(|column| column.values_preserving(selection)) + .collect() + } + + fn supports_values_preserving(&self) -> bool { + true + } + fn clear_shrink(&mut self, num_rows: usize) { // Reset to a fresh column-builder vector. The schema was validated // in `try_new`, so rebuilding cannot fail unless something else @@ -1333,12 +1364,15 @@ mod tests { use arrow::array::{ Array, ArrayRef, DurationMicrosecondArray, FixedSizeBinaryArray, Float16Array, Int32Array, Int64Array, PrimitiveArray, RecordBatch, StringArray, - StringViewArray, + StringViewArray, UInt32Array, }; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; + use arrow::{ + compute::{concat_batches, take}, + util::pretty::pretty_format_batches, + }; use datafusion_common::utils::proxy::HashTableAllocExt; - use datafusion_expr::EmitTo; + use datafusion_expr::{EmitTo, GroupSelection}; use crate::aggregates::group_values::{ GroupValues, multi_group_by::GroupValuesColumn, @@ -1990,6 +2024,41 @@ mod tests { check_result(&actual_batch, &expected_batch); } + #[test] + fn test_preserving_selected_vectorized_group_values() { + let data_set = VectorizedTestDataSet::new(); + let mut group_values = + GroupValuesColumn::::try_new(data_set.schema()).unwrap(); + data_set.load_to_group_values(&mut group_values); + + let selection = [16, 0, 4, 0]; + let actual = group_values + .values_preserving(GroupSelection::Indices(&selection)) + .unwrap(); + let indices = UInt32Array::from_iter_values(selection.map(|index| index as u32)); + let mut destructive_group_values = + GroupValuesColumn::::try_new(data_set.schema()).unwrap(); + data_set.load_to_group_values(&mut destructive_group_values); + let all = destructive_group_values.emit(EmitTo::All).unwrap(); + let expected = all + .iter() + .map(|column| take(column.as_ref(), &indices, None).unwrap()) + .collect::>(); + let expected = RecordBatch::try_new(data_set.schema(), expected).unwrap(); + let actual = RecordBatch::try_new(data_set.schema(), actual).unwrap(); + assert_eq!(actual, expected); + + // A repeated preserving read returns the same rows and leaves all groups. + let repeated = group_values + .values_preserving(GroupSelection::Indices(&selection)) + .unwrap(); + assert_eq!( + RecordBatch::try_new(data_set.schema(), repeated).unwrap(), + expected + ); + assert_eq!(group_values.len(), data_set.expected_batch.num_rows()); + } + #[test] fn test_emit_first_n_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index 148c5697dea3b..d3a036ab66a42 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -31,6 +31,7 @@ use arrow::util::bit_util::apply_bitwise_binary_op; use datafusion_common::Result; use datafusion_common::utils::split_vec_min_alloc; use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::GroupSelection; use std::iter; use std::sync::Arc; @@ -278,6 +279,20 @@ where Arc::new(arr.with_data_type(data_type)) } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + let values: Vec = selection + .iter(self.group_values.len()) + .map(|index| self.group_values[index]) + .collect(); + let nulls = self + .nulls + .build_preserving(selection, self.group_values.len())?; + Ok(Arc::new( + PrimitiveArray::::new(ScalarBuffer::from(values), nulls) + .with_data_type(self.data_type.clone()), + )) + } + fn take_n(&mut self, n: usize) -> ArrayRef { let first_n = split_vec_min_alloc(&mut self.group_values, n); let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index 1445a81f2189b..6067a18ac5b0f 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -53,6 +53,7 @@ use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::datatypes::DataType; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::{DataFusionError, Result}; +use datafusion_expr::GroupSelection; /// A [`GroupColumn`] that stores group values for a single column in the arrow /// [row format], backed by a single-field [`RowConverter`]. @@ -294,6 +295,14 @@ impl GroupColumn for RowsGroupColumn { self.rows_to_array(&self.group_values) } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.group_values.num_rows()).is_ok()); + let rows = selection + .iter(self.group_values.num_rows()) + .map(|index| self.group_values.row(index)); + Ok(self.rows_to_array(rows)) + } + fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(n <= self.group_values.num_rows()); @@ -553,6 +562,33 @@ mod tests { assert!(!col.equal_to(0, &input, 1)); } + #[test] + fn struct_values_preserving() { + let dt = DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2)])); + let input: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("a", DataType::Int32, true)].into(), + vec![values], + None, + )); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let output = col + .values_preserving(GroupSelection::Indices(&[1, 0, 1])) + .unwrap(); + let output = output.as_any().downcast_ref::().unwrap(); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(), + &Int32Array::from(vec![Some(2), Some(1), Some(2)]) + ); + assert_eq!(col.len(), 2); + } + #[test] fn supports_type_matches_row_converter_impl() { assert!(RowsGroupColumn::supports_type(&DataType::FixedSizeList( diff --git a/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs b/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs index 6a84d685b6c79..602017369532b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs @@ -17,6 +17,8 @@ use arrow::array::NullBufferBuilder; use arrow::buffer::NullBuffer; +use datafusion_common::Result; +use datafusion_expr::GroupSelection; /// Builder for an (optional) null mask /// @@ -72,6 +74,25 @@ impl MaybeNullBufferBuilder { self.nulls.finish() } + /// Returns a null buffer for `selection` without changing this builder. + pub fn build_preserving( + &self, + selection: GroupSelection<'_>, + total_num_values: usize, + ) -> Result> { + let selected_len = selection.len(total_num_values); + if self.nulls.as_slice().is_none() { + return Ok(None); + } + + debug_assert_eq!(self.nulls.len(), total_num_values); + let mut selected = NullBufferBuilder::new(selected_len); + for index in selection.iter(total_num_values) { + selected.append(self.nulls.is_valid(index)); + } + Ok(selected.finish()) + } + /// Returns a NullBuffer representing the first `n` rows accumulated so far /// shifting any remaining down by `n` pub fn take_n(&mut self, n: usize) -> Option { diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index cbd7a609c5caa..970fb4d6a22dc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -29,7 +29,7 @@ use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::normalize_float_zero; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; use hashbrown::hash_table::HashTable; use log::debug; use std::mem::size_of; @@ -255,6 +255,35 @@ impl GroupValues for GroupValuesRows { Ok(output) } + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + let empty_rows; + let group_values = if let Some(group_values) = self.group_values.as_ref() { + group_values + } else { + empty_rows = self.row_converter.empty_rows(0, 0); + &empty_rows + }; + debug_assert!(selection.validate(group_values.num_rows()).is_ok()); + let rows = selection + .iter(group_values.num_rows()) + .map(|index| group_values.row(index)); + let mut output = self.row_converter.convert_rows(rows)?; + + // TODO: Materialize dictionaries in group keys + // https://github.com/apache/datafusion/issues/7647 + for (field, array) in self.schema.fields.iter().zip(&mut output) { + *array = encode_array_if_necessary(array, field.data_type())?; + } + Ok(output) + } + + fn supports_values_preserving(&self) -> bool { + true + } + fn clear_shrink(&mut self, num_rows: usize) { self.group_values = self.group_values.take().map(|mut rows| { rows.clear(); @@ -412,3 +441,49 @@ pub(crate) fn encode_array_if_necessary( (_, _) => Ok(Arc::::clone(array)), } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, ListArray}; + use arrow::datatypes::{Field, Int32Type, Schema}; + + #[test] + fn preserving_nested_row_values() -> Result<()> { + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + DataType::List(field), + true, + )])); + let mut group_values = GroupValuesRows::try_new(schema)?; + let input = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3)]), + Some(vec![Some(1), Some(2)]), + ])) as ArrayRef; + let mut groups = vec![]; + group_values.intern(&[input], &mut groups)?; + assert_eq!(groups, vec![0, 1, 2, 0]); + + let selection = GroupSelection::Indices(&[2, 0, 1, 2]); + let expected = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(3)]), + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3)]), + ]); + for _ in 0..2 { + let actual = group_values.values_preserving(selection)?; + assert_eq!(actual[0].as_list::(), &expected); + } + + let input = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(4)]), + ])) as ArrayRef; + group_values.intern(&[input], &mut groups)?; + assert_eq!(groups, vec![3]); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs index e993c0c53d199..8ddb9abd85444 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs @@ -21,7 +21,7 @@ use arrow::array::{ ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, }; use datafusion_common::Result; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; use std::{mem::size_of, sync::Arc}; #[derive(Debug)] @@ -145,6 +145,36 @@ impl GroupValues for GroupValuesBoolean { Ok(vec![Arc::new(BooleanArray::new(values, nulls)) as _]) } + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + let num_groups = self.len(); + debug_assert!(selection.validate(num_groups).is_ok()); + let mut values = BooleanBufferBuilder::new(selection.len(num_groups)); + let mut nulls = NullBufferBuilder::new(selection.len(num_groups)); + for index in selection.iter(num_groups) { + if self.null_group == Some(index) { + values.append(false); + nulls.append_null(); + } else { + debug_assert!( + self.false_group == Some(index) || self.true_group == Some(index) + ); + values.append(self.true_group == Some(index)); + nulls.append_non_null(); + } + } + Ok(vec![Arc::new(BooleanArray::new( + values.finish(), + nulls.finish(), + ))]) + } + + fn supports_values_preserving(&self) -> bool { + true + } + fn clear_shrink(&mut self, _num_rows: usize) { self.false_group = None; self.true_group = None; diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index b881a51b25474..33faad8dae4d4 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -21,7 +21,7 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use datafusion_common::Result; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; /// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values @@ -120,6 +120,18 @@ impl GroupValues for GroupValuesBytes { Ok(vec![group_values]) } + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + debug_assert!(selection.validate(self.len()).is_ok()); + Ok(vec![self.map.keys(selection.iter(self.len()))?]) + } + + fn supports_values_preserving(&self) -> bool { + true + } + fn clear_shrink(&mut self, _num_rows: usize) { // in theory we could potentially avoid this reallocation and clear the // contents of the maps, but for now we just reset the map from the beginning diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 7a56f7c52c11a..5814e8289d5f2 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -17,7 +17,7 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef}; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; use std::mem::size_of; @@ -122,6 +122,18 @@ impl GroupValues for GroupValuesBytesView { Ok(vec![group_values]) } + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> datafusion_common::Result> { + debug_assert!(selection.validate(self.len()).is_ok()); + Ok(vec![self.map.keys(selection.iter(self.len()))?]) + } + + fn supports_values_preserving(&self) -> bool { + true + } + fn clear_shrink(&mut self, _num_rows: usize) { // in theory we could potentially avoid this reallocation and clear the // contents of the maps, but for now we just reset the map from the beginning diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index e254aebcfd7ce..d5fa0521c0c9f 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -26,7 +26,7 @@ use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::utils::split_vec_min_alloc; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::EmitTo; +use datafusion_expr::{EmitTo, GroupSelection}; use half::f16; use hashbrown::hash_table::HashTable; #[cfg(not(feature = "force_hash_collisions"))] @@ -239,6 +239,37 @@ where Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) } + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + debug_assert!(selection.validate(self.values.len()).is_ok()); + let values: Vec = selection + .iter(self.values.len()) + .map(|index| self.values[index]) + .collect(); + let nulls = if let Some(null_group) = self.null_group { + let mut nulls = NullBufferBuilder::new(values.len()); + for index in selection.iter(self.values.len()) { + if index == null_group { + nulls.append_null(); + } else { + nulls.append_non_null(); + } + } + nulls.finish() + } else { + None + }; + let array = PrimitiveArray::::new(values.into(), nulls) + .with_data_type(self.data_type.clone()); + Ok(vec![Arc::new(array)]) + } + + fn supports_values_preserving(&self) -> bool { + true + } + fn clear_shrink(&mut self, num_rows: usize) { self.values.clear(); self.values.shrink_to(num_rows); From 06d254a2d1cd7145ede3e70f246bf908b7720120 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:18:16 -0700 Subject: [PATCH 4/9] Implement preserving reads for grouped accumulators --- .../src/aggregate/count_distinct/groups.rs | 20 +- .../src/aggregate/groups_accumulator.rs | 32 ++- .../groups_accumulator/accumulate.rs | 29 +- .../aggregate/groups_accumulator/bool_op.rs | 31 +- .../aggregate/groups_accumulator/prim_op.rs | 73 ++++- datafusion/functions-aggregate/src/average.rs | 187 ++++++++++--- .../functions-aggregate/src/correlation.rs | 157 +++++++---- datafusion/functions-aggregate/src/count.rs | 63 ++++- .../src/min_max/min_max_bytes.rs | 264 +++++++++++------- .../src/min_max/min_max_struct.rs | 77 +++-- datafusion/functions-aggregate/src/stddev.rs | 25 +- .../functions-aggregate/src/string_agg.rs | 27 +- .../functions-aggregate/src/variance.rs | 85 +++++- 13 files changed, 830 insertions(+), 240 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 10aa21c3acad2..001487b71e539 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -23,7 +23,9 @@ use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ArrowPrimitiveType, Field}; use datafusion_common::HashSet; use datafusion_common::hash_utils::RandomState; -use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; +use datafusion_expr_common::groups_accumulator::{ + EmitTo, GroupSelection, GroupsAccumulator, +}; use std::hash::Hash; use std::mem::size_of; use std::sync::Arc; @@ -103,6 +105,22 @@ where Ok(Arc::new(Int64Array::from(counts))) } + fn evaluate_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> datafusion_common::Result { + debug_assert!(selection.validate(self.counts.len()).is_ok()); + let counts = selection + .iter(self.counts.len()) + .map(|index| self.counts[index]) + .collect::>(); + Ok(Arc::new(Int64Array::from(counts))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result> { let num_emitted = match emit_to { EmitTo::All => self.counts.len(), diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index b5610419166df..ac148deabdd1e 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -34,7 +34,9 @@ use arrow::{ }; use datafusion_common::{Result, ScalarValue, arrow_datafusion_err}; use datafusion_expr_common::accumulator::Accumulator; -use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; +use datafusion_expr_common::groups_accumulator::{ + EmitTo, GroupSelection, GroupsAccumulator, +}; /// An adapter that implements [`GroupsAccumulator`] for any [`Accumulator`] /// @@ -335,6 +337,34 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { result } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.states.len()).is_ok()); + let selected_len = selection.len(self.states.len()); + if selected_len == 0 { + // ScalarValue::iter_to_array needs at least one value to infer the + // output type, so evaluate a temporary empty accumulator. + let mut accumulator = (self.factory)()?; + return Ok(ScalarValue::iter_to_array([accumulator.evaluate()?])?.slice(0, 0)); + } + + let mut results = Vec::with_capacity(selected_len); + for group_index in selection.iter(self.states.len()) { + let (result, size_pre, size_post) = { + let state = &mut self.states[group_index]; + let size_pre = state.size(); + let result = state.accumulator.evaluate()?; + (result, size_pre, state.size()) + }; + self.adjust_allocation(size_pre, size_post); + results.push(result); + } + ScalarValue::iter_to_array(results) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + // filtered_null_mask(opt_filter, &values); fn state(&mut self, emit_to: EmitTo) -> Result> { let vec_size_pre = self.states.allocated_size(); diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs index 09e1df4eae70c..c86a795f3f37a 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -24,7 +24,8 @@ use arrow::buffer::NullBuffer; use arrow::datatypes::ArrowPrimitiveType; use crate::aggregate::groups_accumulator::nulls::filter_to_validity; -use datafusion_expr_common::groups_accumulator::EmitTo; +use datafusion_common::Result; +use datafusion_expr_common::groups_accumulator::{EmitTo, GroupSelection}; /// If the input has nulls, then the accumulator must potentially /// handle each input null value specially (e.g. for `SUM` to mark the @@ -289,6 +290,32 @@ impl NullState { } } + /// Creates a [`NullBuffer`] for `selection` without changing this state. + /// + /// Indices in `selection` must be less than `total_num_groups`. This method + /// does not validate them. + pub fn build_preserving( + &self, + selection: GroupSelection<'_>, + total_num_groups: usize, + ) -> Result> { + let selected_len = selection.len(total_num_groups); + match &self.seen_values { + SeenValues::All { num_values } => { + debug_assert_eq!(*num_values, total_num_groups); + Ok(None) + } + SeenValues::Some { values } => { + debug_assert_eq!(values.len(), total_num_groups); + let mut selected = BooleanBufferBuilder::new(selected_len); + for index in selection.iter(total_num_groups) { + selected.append(values.get_bit(index)); + } + Ok(Some(NullBuffer::new(selected.finish()))) + } + } + } + /// Creates the a [`NullBuffer`] representing which group_indices /// should have null values (because they never saw any values) /// for the `emit_to` rows. diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs index 77bb7598e2747..665610716207b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs @@ -21,7 +21,9 @@ use crate::aggregate::groups_accumulator::nulls::filtered_null_mask; use arrow::array::{ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder}; use arrow::buffer::BooleanBuffer; use datafusion_common::Result; -use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; +use datafusion_expr_common::groups_accumulator::{ + EmitTo, GroupSelection, GroupsAccumulator, +}; use super::accumulate::NullState; @@ -124,10 +126,37 @@ where Ok(Arc::new(values)) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.values.len()).is_ok()); + let mut values = BooleanBufferBuilder::new(selection.len(self.values.len())); + for index in selection.iter(self.values.len()) { + values.append(self.values.get_bit(index)); + } + let nulls = self + .null_state + .build_preserving(selection, self.values.len())?; + Ok(Arc::new(BooleanArray::new(values.finish(), nulls))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> Result> { self.evaluate(emit_to).map(|arr| vec![arr]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|arr| vec![arr]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index c5d74978664c9..12c679525edd3 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -24,7 +24,9 @@ use arrow::compute; use arrow::datatypes::ArrowPrimitiveType; use arrow::datatypes::DataType; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; +use datafusion_expr_common::groups_accumulator::{ + EmitTo, GroupSelection, GroupsAccumulator, +}; use super::accumulate::NullState; @@ -123,10 +125,39 @@ where Ok(Arc::new(values)) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.values.len()).is_ok()); + let values: Vec = selection + .iter(self.values.len()) + .map(|index| self.values[index]) + .collect(); + let nulls = self + .null_state + .build_preserving(selection, self.values.len())?; + let values = PrimitiveArray::::new(values.into(), nulls) + .with_data_type(self.data_type.clone()); + Ok(Arc::new(values)) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> Result> { self.evaluate(emit_to).map(|arr| vec![arr]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|arr| vec![arr]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], @@ -193,3 +224,43 @@ where self.values.capacity() * size_of::() + self.null_state.size() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::Int64Type; + + #[test] + fn preserving_reads_do_not_change_accumulator_state() -> Result<()> { + let mut accumulator = PrimitiveGroupsAccumulator::::new( + &DataType::Int64, + |current, value| *current += value, + ); + let values = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])); + accumulator.update_batch(&[values], &[0, 1, 2], None, 4)?; + + let selection = GroupSelection::Indices(&[3, 0, 1, 1]); + let expected = Int64Array::from(vec![None, Some(1), None, None]); + for _ in 0..2 { + let actual = accumulator.evaluate_preserving(selection)?; + assert_eq!(actual.as_primitive::(), &expected); + let state = accumulator.state_preserving(selection)?; + assert_eq!(state[0].as_primitive::(), &expected); + } + + // Group indices and unselected state remain valid after repeated reads. + let values = Arc::new(Int64Array::from(vec![5, 7])); + accumulator.update_batch(&[values], &[1, 3], None, 4)?; + let expected = Int64Array::from(vec![Some(1), Some(5), Some(3), Some(7)]); + let actual = accumulator.evaluate_preserving(GroupSelection::All)?; + assert_eq!(actual.as_primitive::(), &expected); + + // A destructive read still sees all state after preserving reads. + let actual = accumulator.evaluate(EmitTo::All)?; + assert_eq!(actual.as_primitive::(), &expected); + assert!(accumulator.supports_evaluate_preserving()); + assert!(accumulator.supports_state_preserving()); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index e5030bf39e409..72efd75b60a6c 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -22,6 +22,7 @@ use arrow::array::{ BooleanArray, PrimitiveArray, PrimitiveBuilder, UInt64Array, }; +use arrow::buffer::NullBuffer; use arrow::compute::{DecimalCast, sum}; use arrow::datatypes::{ ArrowNativeType, DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE, @@ -38,7 +39,7 @@ use datafusion_common::{ use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Coercion, Documentation, EmitTo, Expr, + Accumulator, AggregateUDFImpl, Coercion, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, ReversedUDAF, Signature, TypeSignature, TypeSignatureClass, Volatility, }; @@ -959,6 +960,58 @@ where _phantom: PhantomData, } } + + fn evaluate_values( + &self, + counts: Vec, + sums: Vec, + nulls: Option, + ) -> Result { + if let Some(nulls) = &nulls { + assert_eq!(nulls.len(), sums.len()); + } + assert_eq!(counts.len(), sums.len()); + + // Don't evaluate averages with null inputs to avoid errors on null values. + let array: PrimitiveArray = if let Some(nulls) = &nulls + && nulls.null_count() > 0 + { + let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) + .with_data_type(self.return_data_type.clone()); + let iter = sums.into_iter().zip(counts).zip(nulls.iter()); + + for ((sum, count), is_valid) in iter { + if is_valid { + builder.append_value((self.avg_fn)(sum, count)?) + } else { + builder.append_null(); + } + } + builder.finish() + } else { + let averages: Vec = sums + .into_iter() + .zip(counts) + .map(|(sum, count)| (self.avg_fn)(sum, count)) + .collect::>>()?; + PrimitiveArray::new(averages.into(), nulls) + .with_data_type(self.return_data_type.clone()) + }; + + Ok(Arc::new(array)) + } + + fn state_values( + &self, + counts: Vec, + sums: Vec, + nulls: Option, + ) -> Vec { + let counts = UInt64Array::new(counts.into(), nulls.clone()); + let sums = PrimitiveArray::::new(sums.into(), nulls) + .with_data_type(self.sum_data_type.clone()); + vec![Arc::new(counts), Arc::new(sums)] + } } impl GroupsAccumulator for AvgGroupsAccumulator @@ -1003,57 +1056,58 @@ where let counts = emit_to.take_needed(&mut self.counts); let sums = emit_to.take_needed(&mut self.sums); let nulls = self.null_state.build(emit_to); + self.evaluate_values(counts, sums, nulls) + } - if let Some(nulls) = &nulls { - assert_eq!(nulls.len(), sums.len()); - } - assert_eq!(counts.len(), sums.len()); - - // don't evaluate averages with null inputs to avoid errors on null values - - let array: PrimitiveArray = if let Some(nulls) = &nulls - && nulls.null_count() > 0 - { - let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) - .with_data_type(self.return_data_type.clone()); - let iter = sums.into_iter().zip(counts).zip(nulls.iter()); - - for ((sum, count), is_valid) in iter { - if is_valid { - builder.append_value((self.avg_fn)(sum, count)?) - } else { - builder.append_null(); - } - } - builder.finish() - } else { - let averages: Vec = sums - .into_iter() - .zip(counts) - .map(|(sum, count)| (self.avg_fn)(sum, count)) - .collect::>>()?; - PrimitiveArray::new(averages.into(), nulls) // no copy - .with_data_type(self.return_data_type.clone()) - }; + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.counts.len()).is_ok()); + let counts = selection + .iter(self.counts.len()) + .map(|index| self.counts[index]) + .collect(); + let sums = selection + .iter(self.sums.len()) + .map(|index| self.sums[index]) + .collect(); + let nulls = self + .null_state + .build_preserving(selection, self.sums.len())?; + self.evaluate_values(counts, sums, nulls) + } - Ok(Arc::new(array)) + fn supports_evaluate_preserving(&self) -> bool { + true } // return arrays for sums and counts fn state(&mut self, emit_to: EmitTo) -> Result> { let nulls = self.null_state.build(emit_to); - let counts = emit_to.take_needed(&mut self.counts); - let counts = UInt64Array::new(counts.into(), nulls.clone()); // zero copy - let sums = emit_to.take_needed(&mut self.sums); - let sums = PrimitiveArray::::new(sums.into(), nulls) // zero copy - .with_data_type(self.sum_data_type.clone()); + Ok(self.state_values(counts, sums, nulls)) + } - Ok(vec![ - Arc::new(counts) as ArrayRef, - Arc::new(sums) as ArrayRef, - ]) + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + debug_assert!(selection.validate(self.counts.len()).is_ok()); + let counts = selection + .iter(self.counts.len()) + .map(|index| self.counts[index]) + .collect(); + let sums = selection + .iter(self.sums.len()) + .map(|index| self.sums[index]) + .collect(); + let nulls = self + .null_state + .build_preserving(selection, self.sums.len())?; + Ok(self.state_values(counts, sums, nulls)) + } + + fn supports_state_preserving(&self) -> bool { + true } fn merge_batch( @@ -1394,4 +1448,53 @@ mod tests { Ok(()) } + + #[test] + fn average_groups_preserving_reads() -> Result<()> { + let mut accumulator = AvgGroupsAccumulator::::new( + &DataType::Float64, + &DataType::Float64, + |sum, count| Ok(sum / count as f64), + ); + let values = Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + None, + Some(8.0), + ])); + accumulator.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::Indices(&[2, 0, 3, 2]); + let expected = Float64Array::from(vec![Some(8.0), Some(3.0), None, Some(8.0)]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + + let state = accumulator.state_preserving(selection)?; + assert_eq!( + state[0].as_primitive::(), + &UInt64Array::from(vec![Some(1), Some(2), None, Some(1)]) + ); + assert_eq!( + state[1].as_primitive::(), + &Float64Array::from(vec![Some(8.0), Some(6.0), None, Some(8.0)]) + ); + + let values = Arc::new(Float64Array::from(vec![10.0, 6.0])); + accumulator.update_batch(&[values], &[1, 3], None, 4)?; + let expected = + Float64Array::from(vec![Some(3.0), Some(10.0), Some(8.0), Some(6.0)]); + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::All)? + .as_primitive::(), + &expected + ); + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index b9bc57dfa989c..767c840fc0eef 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -31,7 +31,7 @@ use arrow::{ array::ArrayRef, datatypes::{DataType, Field}, }; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate_multiple; use log::debug; @@ -305,10 +305,67 @@ pub struct CorrelationGroupsAccumulator { sum_yy: Vec, } +fn copy_selected(selection: GroupSelection<'_>, values: &[T]) -> Vec { + selection + .iter(values.len()) + .map(|index| values[index]) + .collect() +} + impl CorrelationGroupsAccumulator { pub fn new() -> Self { Default::default() } + + fn evaluate_values( + counts: &[u64], + sum_xs: &[f64], + sum_ys: &[f64], + sum_xys: &[f64], + sum_xxs: &[f64], + sum_yys: &[f64], + ) -> ArrayRef { + let n = counts.len(); + let mut values = Vec::with_capacity(n); + let mut nulls = NullBufferBuilder::new(n); + + for i in 0..n { + let count = counts[i]; + let sum_x = sum_xs[i]; + let sum_y = sum_ys[i]; + let sum_xy = sum_xys[i]; + let sum_xx = sum_xxs[i]; + let sum_yy = sum_yys[i]; + + // If both inputs are NaN, return NaN. If only one input is NaN, + // or there are too few values, return NULL. + if sum_x.is_nan() && sum_y.is_nan() { + values.push(f64::NAN); + nulls.append_non_null(); + continue; + } else if count < 2 || sum_x.is_nan() || sum_y.is_nan() { + values.push(0.0); + nulls.append_null(); + continue; + } + + let mean_x = sum_x / count as f64; + let mean_y = sum_y / count as f64; + let numerator = sum_xy - sum_x * mean_y; + let denominator = + ((sum_xx - sum_x * mean_x) * (sum_yy - sum_y * mean_y)).sqrt(); + + if denominator == 0.0 { + values.push(0.0); + nulls.append_null(); + } else { + values.push(numerator / denominator); + nulls.append_non_null(); + } + } + + Arc::new(Float64Array::new(values.into(), nulls.finish())) + } } /// Specialized version of `accumulate_multiple` for correlation's merge_batch @@ -409,65 +466,30 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - // Drain the state vectors for the groups being emitted - let counts = emit_to.take_needed(&mut self.count); - let sum_xs = emit_to.take_needed(&mut self.sum_x); - let sum_ys = emit_to.take_needed(&mut self.sum_y); - let sum_xys = emit_to.take_needed(&mut self.sum_xy); - let sum_xxs = emit_to.take_needed(&mut self.sum_xx); - let sum_yys = emit_to.take_needed(&mut self.sum_yy); - - let n = counts.len(); - let mut values = Vec::with_capacity(n); - let mut nulls = NullBufferBuilder::new(n); - - // Notes for `Null` handling: - // - If the `count` state of a group is 0, no valid records are accumulated - // for this group, so the aggregation result is `Null`. - // - Correlation can't be calculated when a group only has 1 record, or when - // the `denominator` state is 0. In these cases, the final aggregation - // result should be `Null` (according to PostgreSQL's behavior). - // - However, if any of the accumulated values contain NaN, the result should - // be NaN regardless of the count (even for single-row groups). - for i in 0..n { - let count = counts[i]; - let sum_x = sum_xs[i]; - let sum_y = sum_ys[i]; - let sum_xy = sum_xys[i]; - let sum_xx = sum_xxs[i]; - let sum_yy = sum_yys[i]; - - // If BOTH sum_x AND sum_y are NaN, then both input values are NaN → return NaN - // If only ONE of them is NaN, then only one input value is NaN → return NULL - if sum_x.is_nan() && sum_y.is_nan() { - // Both inputs are NaN → return NaN - values.push(f64::NAN); - nulls.append_non_null(); - continue; - } else if count < 2 || sum_x.is_nan() || sum_y.is_nan() { - // Only one input is NaN → return NULL - values.push(0.0); - nulls.append_null(); - continue; - } - - let mean_x = sum_x / count as f64; - let mean_y = sum_y / count as f64; - - let numerator = sum_xy - sum_x * mean_y; - let denominator = - ((sum_xx - sum_x * mean_x) * (sum_yy - sum_y * mean_y)).sqrt(); + Ok(Self::evaluate_values( + &emit_to.take_needed(&mut self.count), + &emit_to.take_needed(&mut self.sum_x), + &emit_to.take_needed(&mut self.sum_y), + &emit_to.take_needed(&mut self.sum_xy), + &emit_to.take_needed(&mut self.sum_xx), + &emit_to.take_needed(&mut self.sum_yy), + )) + } - if denominator == 0.0 { - values.push(0.0); - nulls.append_null(); - } else { - values.push(numerator / denominator); - nulls.append_non_null(); - } - } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.count.len()).is_ok()); + Ok(Self::evaluate_values( + ©_selected(selection, &self.count), + ©_selected(selection, &self.sum_x), + ©_selected(selection, &self.sum_y), + ©_selected(selection, &self.sum_xy), + ©_selected(selection, &self.sum_xx), + ©_selected(selection, &self.sum_yy), + )) + } - Ok(Arc::new(Float64Array::new(values.into(), nulls.finish()))) + fn supports_evaluate_preserving(&self) -> bool { + true } fn state(&mut self, emit_to: EmitTo) -> Result> { @@ -539,6 +561,25 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { Arc::new(Float64Array::from(sum_yy)), ]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + debug_assert!(selection.validate(self.count.len()).is_ok()); + Ok(vec![ + Arc::new(UInt64Array::from(copy_selected(selection, &self.count))), + Arc::new(Float64Array::from(copy_selected(selection, &self.sum_x))), + Arc::new(Float64Array::from(copy_selected(selection, &self.sum_y))), + Arc::new(Float64Array::from(copy_selected(selection, &self.sum_xy))), + Arc::new(Float64Array::from(copy_selected(selection, &self.sum_xx))), + Arc::new(Float64Array::from(copy_selected(selection, &self.sum_yy))), + ]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 1e72d8ac3d5b1..19844f28bb3ae 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -35,9 +35,9 @@ use datafusion_common::{ stats::Precision, utils::expr::COUNT_STAR_EXPANSION, }; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupsAccumulator, - ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility, - WindowFunctionDefinition, + Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupSelection, + GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, + TypeSignature, Volatility, WindowFunctionDefinition, expr::WindowFunction, function::{AccumulatorArgs, StateFieldsArgs}, utils::format_state_name, @@ -707,6 +707,19 @@ impl GroupsAccumulator for CountGroupsAccumulator { Ok(Arc::new(array)) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.counts.len()).is_ok()); + let counts = selection + .iter(self.counts.len()) + .map(|index| self.counts[index]) + .collect::>(); + Ok(Arc::new(Int64Array::from(counts))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + // return arrays for counts fn state(&mut self, emit_to: EmitTo) -> Result> { let counts = emit_to.take_needed(&mut self.counts); @@ -714,6 +727,17 @@ impl GroupsAccumulator for CountGroupsAccumulator { Ok(vec![Arc::new(counts) as ArrayRef]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|array| vec![array]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + /// Converts an input batch directly to a state batch /// /// The state of `COUNT` is always a single Int64Array: @@ -913,7 +937,7 @@ mod tests { use super::*; use arrow::{ - array::{DictionaryArray, Int32Array, NullArray, StringArray}, + array::{DictionaryArray, Int32Array, Int64Array, NullArray, StringArray}, datatypes::{DataType, Field, Int32Type, Schema}, }; use datafusion_expr::function::AccumulatorArgs; @@ -958,6 +982,37 @@ mod tests { Ok(()) } + #[test] + fn count_groups_preserving_reads() -> Result<()> { + let mut accumulator = CountGroupsAccumulator::new(); + let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(2), Some(3)])); + accumulator.update_batch(&[values], &[0, 1, 0, 2], None, 4)?; + + let selection = GroupSelection::Indices(&[2, 0, 3, 2]); + let expected = Int64Array::from(vec![1, 2, 0, 1]); + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + assert_eq!( + accumulator.state_preserving(selection)?[0].as_primitive::(), + &expected + ); + + let values = Arc::new(Int32Array::from(vec![4])); + accumulator.update_batch(&[values], &[3], None, 4)?; + let expected = Int64Array::from(vec![2, 0, 1, 1]); + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::Indices(&[0, 1, 2, 3]))? + .as_primitive::(), + &expected + ); + Ok(()) + } + #[test] fn test_nested_dictionary() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index efeaea314c4f5..9a3043e2b08c0 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -22,7 +22,7 @@ use arrow::array::{ use arrow::datatypes::DataType; use datafusion_common::hash_map::Entry; use datafusion_common::{HashMap, Result, internal_err}; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; use std::mem::size_of; use std::sync::Arc; @@ -62,6 +62,100 @@ impl MinMaxBytesAccumulator { is_min: false, } } + + fn build_array<'a>( + &self, + min_maxes: impl Iterator>, + num_values: usize, + data_capacity: usize, + ) -> Result { + let result: ArrayRef = match self.inner.data_type { + DataType::Utf8 => { + let mut builder = StringBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder.append_value(unsafe { + std::str::from_utf8_unchecked(value) + }), + } + } + Arc::new(builder.finish()) + } + DataType::LargeUtf8 => { + let mut builder = + LargeStringBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder.append_value(unsafe { + std::str::from_utf8_unchecked(value) + }), + } + } + Arc::new(builder.finish()) + } + DataType::Utf8View => { + let block_size = capacity_to_view_block_size(data_capacity); + let mut builder = StringViewBuilder::with_capacity(num_values) + .with_fixed_block_size(block_size); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder.append_value(unsafe { + std::str::from_utf8_unchecked(value) + }), + } + } + Arc::new(builder.finish()) + } + DataType::Binary => { + let mut builder = BinaryBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), + } + } + Arc::new(builder.finish()) + } + DataType::LargeBinary => { + let mut builder = + LargeBinaryBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), + } + } + Arc::new(builder.finish()) + } + DataType::BinaryView => { + let block_size = capacity_to_view_block_size(data_capacity); + let mut builder = BinaryViewBuilder::with_capacity(num_values) + .with_fixed_block_size(block_size); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), + } + } + Arc::new(builder.finish()) + } + _ => { + return internal_err!( + "Unexpected data type for MinMaxBytesAccumulator: {:?}", + self.inner.data_type + ); + } + }; + + assert_eq!(&self.inner.data_type, result.data_type()); + Ok(result) + } } impl GroupsAccumulator for MinMaxBytesAccumulator { @@ -203,101 +297,29 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { fn evaluate(&mut self, emit_to: EmitTo) -> Result { let (data_capacity, min_maxes) = self.inner.emit_to(emit_to); + self.build_array( + min_maxes.iter().map(|value| value.as_deref()), + min_maxes.len(), + data_capacity, + ) + } - // Convert the Vec of bytes to a vec of Strings (at no cost) - fn bytes_to_str( - min_maxes: Vec>>, - ) -> impl Iterator> { - min_maxes.into_iter().map(|opt| { - opt.map(|bytes| { - // Safety: only called on data added from update_batch which ensures - // the input type matched the output type - unsafe { String::from_utf8_unchecked(bytes) } - }) - }) - } - - let result: ArrayRef = match self.inner.data_type { - DataType::Utf8 => { - let mut builder = - StringBuilder::with_capacity(min_maxes.len(), data_capacity); - for opt in bytes_to_str(min_maxes) { - match opt { - None => builder.append_null(), - Some(s) => builder.append_value(s.as_str()), - } - } - Arc::new(builder.finish()) - } - DataType::LargeUtf8 => { - let mut builder = - LargeStringBuilder::with_capacity(min_maxes.len(), data_capacity); - for opt in bytes_to_str(min_maxes) { - match opt { - None => builder.append_null(), - Some(s) => builder.append_value(s.as_str()), - } - } - Arc::new(builder.finish()) - } - DataType::Utf8View => { - let block_size = capacity_to_view_block_size(data_capacity); - - let mut builder = StringViewBuilder::with_capacity(min_maxes.len()) - .with_fixed_block_size(block_size); - for opt in bytes_to_str(min_maxes) { - match opt { - None => builder.append_null(), - Some(s) => builder.append_value(s.as_str()), - } - } - Arc::new(builder.finish()) - } - DataType::Binary => { - let mut builder = - BinaryBuilder::with_capacity(min_maxes.len(), data_capacity); - for opt in min_maxes { - match opt { - None => builder.append_null(), - Some(s) => builder.append_value(s.as_ref() as &[u8]), - } - } - Arc::new(builder.finish()) - } - DataType::LargeBinary => { - let mut builder = - LargeBinaryBuilder::with_capacity(min_maxes.len(), data_capacity); - for opt in min_maxes { - match opt { - None => builder.append_null(), - Some(s) => builder.append_value(s.as_ref() as &[u8]), - } - } - Arc::new(builder.finish()) - } - DataType::BinaryView => { - let block_size = capacity_to_view_block_size(data_capacity); - - let mut builder = BinaryViewBuilder::with_capacity(min_maxes.len()) - .with_fixed_block_size(block_size); - for opt in min_maxes { - match opt { - None => builder.append_null(), - Some(s) => builder.append_value(s.as_ref() as &[u8]), - } - } - Arc::new(builder.finish()) - } - _ => { - return internal_err!( - "Unexpected data type for MinMaxBytesAccumulator: {:?}", - self.inner.data_type - ); - } - }; + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + let num_groups = self.inner.min_max.len(); + debug_assert!(selection.validate(num_groups).is_ok()); + let num_values = selection.len(num_groups); + let data_capacity = selection + .iter(num_groups) + .filter_map(|index| self.inner.min_max[index].as_ref().map(Vec::len)) + .sum(); + let min_maxes = selection + .iter(num_groups) + .map(|index| self.inner.min_max[index].as_deref()); + self.build_array(min_maxes, num_values, data_capacity) + } - assert_eq!(&self.inner.data_type, result.data_type()); - Ok(result) + fn supports_evaluate_preserving(&self) -> bool { + true } fn state(&mut self, emit_to: EmitTo) -> Result> { @@ -305,6 +327,17 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { self.evaluate(emit_to).map(|arr| vec![arr]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|array| vec![array]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], @@ -504,3 +537,46 @@ impl MinMaxBytesState { self.total_data_bytes + self.min_max.len() * size_of::>>() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::StringArray; + + #[test] + fn preserving_selected_min_values() -> Result<()> { + let mut accumulator = MinMaxBytesAccumulator::new_min(DataType::Utf8); + let values = Arc::new(StringArray::from(vec![ + Some("z"), + Some("b"), + None, + Some("c"), + Some("a"), + Some("x"), + ])); + accumulator.update_batch(&[values], &[0, 0, 1, 2, 2, 3], None, 4)?; + + let selection = GroupSelection::Indices(&[3, 0, 1, 2, 3]); + let expected = + StringArray::from(vec![Some("x"), Some("b"), None, Some("a"), Some("x")]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_string::(), + &expected + ); + } + + let values = Arc::new(StringArray::from(vec!["aa", "w"])); + accumulator.update_batch(&[values], &[0, 3], None, 4)?; + let expected = StringArray::from(vec![Some("aa"), None, Some("a"), Some("w")]); + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::All)? + .as_string::(), + &expected + ); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs index d1bac4e2f90db..949d156cb60f0 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs @@ -27,7 +27,7 @@ use datafusion_common::{ Result, internal_err, scalar::{copy_array_data, partial_cmp_struct}, }; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; use datafusion_common::utils::split_vec_min_alloc; @@ -59,6 +59,33 @@ impl MinMaxStructAccumulator { is_min: false, } } + + fn build_array<'a>( + &self, + min_maxes: impl Iterator>, + num_values: usize, + ) -> Result { + let fields = match &self.inner.data_type { + DataType::Struct(fields) => fields, + _ => return internal_err!("Data type is not a struct"), + }; + let null_array = StructArray::new_null(fields.clone(), 1); + let min_maxes_data: Vec = min_maxes + .map(|value| match value { + Some(value) => value.to_data(), + None => null_array.to_data(), + }) + .collect(); + let min_maxes_refs: Vec<&ArrayData> = min_maxes_data.iter().collect(); + let mut copy = MutableArrayData::new(min_maxes_refs, true, num_values); + + for (index, item) in min_maxes_data.iter().enumerate() { + copy.try_extend(index, 0, item.len())?; + } + let result = copy.freeze(); + assert_eq!(&self.inner.data_type, result.data_type()); + Ok(Arc::new(StructArray::from(result))) + } } impl GroupsAccumulator for MinMaxStructAccumulator { @@ -102,27 +129,24 @@ impl GroupsAccumulator for MinMaxStructAccumulator { fn evaluate(&mut self, emit_to: EmitTo) -> Result { let (_, min_maxes) = self.inner.emit_to(emit_to); - let fields = match &self.inner.data_type { - DataType::Struct(fields) => fields, - _ => return internal_err!("Data type is not a struct"), - }; - let null_array = StructArray::new_null(fields.clone(), 1); - let min_maxes_data: Vec = min_maxes - .iter() - .map(|v| match v { - Some(v) => v.to_data(), - None => null_array.to_data(), - }) - .collect(); - let min_maxes_refs: Vec<&ArrayData> = min_maxes_data.iter().collect(); - let mut copy = MutableArrayData::new(min_maxes_refs, true, min_maxes_data.len()); + self.build_array( + min_maxes.iter().map(|value| value.as_ref()), + min_maxes.len(), + ) + } - for (i, item) in min_maxes_data.iter().enumerate() { - copy.try_extend(i, 0, item.len())?; - } - let result = copy.freeze(); - assert_eq!(&self.inner.data_type, result.data_type()); - Ok(Arc::new(StructArray::from(result))) + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + let num_groups = self.inner.min_max.len(); + debug_assert!(selection.validate(num_groups).is_ok()); + let num_values = selection.len(num_groups); + let min_maxes = selection + .iter(num_groups) + .map(|index| self.inner.min_max[index].as_ref()); + self.build_array(min_maxes, num_values) + } + + fn supports_evaluate_preserving(&self) -> bool { + true } fn state(&mut self, emit_to: EmitTo) -> Result> { @@ -130,6 +154,17 @@ impl GroupsAccumulator for MinMaxStructAccumulator { self.evaluate(emit_to).map(|arr| vec![arr]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|array| vec![array]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs index 15511bf4a565f..d7b7f45d86d77 100644 --- a/datafusion/functions-aggregate/src/stddev.rs +++ b/datafusion/functions-aggregate/src/stddev.rs @@ -30,8 +30,8 @@ use datafusion_common::{Result, internal_err, not_impl_err}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, GroupsAccumulator, Signature, - Volatility, + Accumulator, AggregateUDFImpl, Documentation, GroupSelection, GroupsAccumulator, + Signature, Volatility, }; use datafusion_functions_aggregate_common::stats::StatsType; use datafusion_macros::user_doc; @@ -341,6 +341,16 @@ impl GroupsAccumulator for StddevGroupsAccumulator { Ok(Arc::new(Float64Array::new(variances.into(), Some(nulls)))) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + let (mut variances, nulls) = self.variance.variance_preserving(selection)?; + variances.iter_mut().for_each(|value| *value = value.sqrt()); + Ok(Arc::new(Float64Array::new(variances.into(), Some(nulls)))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: datafusion_expr::EmitTo) -> Result> { self.variance.state(emit_to) } @@ -352,6 +362,17 @@ impl GroupsAccumulator for StddevGroupsAccumulator { ) -> Result> { self.variance.convert_to_state(values, opt_filter) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.variance.state_preserving(selection) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn size(&self) -> usize { self.variance.size() } diff --git a/datafusion/functions-aggregate/src/string_agg.rs b/datafusion/functions-aggregate/src/string_agg.rs index 3fe2b0a186ae3..d7530f81a78f8 100644 --- a/datafusion/functions-aggregate/src/string_agg.rs +++ b/datafusion/functions-aggregate/src/string_agg.rs @@ -32,8 +32,8 @@ use datafusion_common::{ use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature, - TypeSignature, Volatility, + Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupSelection, + GroupsAccumulator, Signature, TypeSignature, Volatility, }; use datafusion_functions_aggregate_common::accumulator::StateFieldsArgs; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; @@ -405,10 +405,33 @@ impl GroupsAccumulator for StringAggGroupsAccumulator { Ok(result) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + debug_assert!(selection.validate(self.values.len()).is_ok()); + let values = selection + .iter(self.values.len()) + .map(|index| self.values[index].as_deref()); + Ok(Arc::new(LargeStringArray::from_iter(values))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> Result> { self.evaluate(emit_to).map(|arr| vec![arr]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|array| vec![array]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index b8e52f849a7cc..d25ea6596ba5a 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -27,8 +27,8 @@ use arrow::{ use datafusion_common::cast::{as_float64_array, as_uint64_array}; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, GroupsAccumulator, Signature, - Volatility, + Accumulator, AggregateUDFImpl, Documentation, GroupSelection, GroupsAccumulator, + Signature, Volatility, function::{AccumulatorArgs, StateFieldsArgs}, utils::format_state_name, }; @@ -476,16 +476,11 @@ impl VarianceGroupsAccumulator { }); } - pub fn variance( - &mut self, - emit_to: datafusion_expr::EmitTo, + fn variance_values( + &self, + mut counts: Vec, + m2s: Vec, ) -> (Vec, NullBuffer) { - let mut counts = emit_to.take_needed(&mut self.counts); - // means are only needed for updating m2s and are not needed for the final result. - // But we still need to take them to ensure the internal state is consistent. - let _ = emit_to.take_needed(&mut self.means); - let m2s = emit_to.take_needed(&mut self.m2s); - if let StatsType::Sample = self.stats_type { counts.iter_mut().for_each(|count| { *count = count.saturating_sub(1); @@ -493,12 +488,40 @@ impl VarianceGroupsAccumulator { } let nulls = NullBuffer::from_iter(counts.iter().map(|&count| count != 0)); let variance = m2s - .iter() + .into_iter() .zip(counts) .map(|(m2, count)| m2 / count as f64) .collect(); (variance, nulls) } + + pub fn variance( + &mut self, + emit_to: datafusion_expr::EmitTo, + ) -> (Vec, NullBuffer) { + let counts = emit_to.take_needed(&mut self.counts); + // Means are only needed for updating m2s, but still need to be removed + // to keep the internal vectors aligned. + let _ = emit_to.take_needed(&mut self.means); + let m2s = emit_to.take_needed(&mut self.m2s); + self.variance_values(counts, m2s) + } + + pub fn variance_preserving( + &self, + selection: GroupSelection<'_>, + ) -> Result<(Vec, NullBuffer)> { + debug_assert!(selection.validate(self.counts.len()).is_ok()); + let counts = selection + .iter(self.counts.len()) + .map(|index| self.counts[index]) + .collect(); + let m2s = selection + .iter(self.m2s.len()) + .map(|index| self.m2s[index]) + .collect(); + Ok(self.variance_values(counts, m2s)) + } } impl GroupsAccumulator for VarianceGroupsAccumulator { @@ -571,6 +594,15 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { Ok(Arc::new(Float64Array::new(variances.into(), Some(nulls)))) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + let (variances, nulls) = self.variance_preserving(selection)?; + Ok(Arc::new(Float64Array::new(variances.into(), Some(nulls)))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: datafusion_expr::EmitTo) -> Result> { let counts = emit_to.take_needed(&mut self.counts); let means = emit_to.take_needed(&mut self.means); @@ -616,6 +648,35 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { Arc::new(Float64Array::new(m2s.into(), None)), ]) } + + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + debug_assert!(selection.validate(self.counts.len()).is_ok()); + let counts = selection + .iter(self.counts.len()) + .map(|index| self.counts[index]) + .collect::>(); + let means = selection + .iter(self.means.len()) + .map(|index| self.means[index]) + .collect::>(); + let m2s = selection + .iter(self.m2s.len()) + .map(|index| self.m2s[index]) + .collect::>(); + Ok(vec![ + Arc::new(UInt64Array::new(counts.into(), None)), + Arc::new(Float64Array::new(means.into(), None)), + Arc::new(Float64Array::new(m2s.into(), None)), + ]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn size(&self) -> usize { self.m2s.capacity() * size_of::() + self.means.capacity() * size_of::() From f21cf8c9eb9e0cd2206a7a69e08b5223e38e7c96 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:08:33 -0700 Subject: [PATCH 5/9] Validate grouped read selections at construction --- .../expr-common/src/groups_accumulator.rs | 154 +++++++++++------- .../src/aggregate/count_distinct/groups.rs | 4 +- .../src/aggregate/groups_accumulator.rs | 6 +- .../groups_accumulator/accumulate.rs | 12 +- .../aggregate/groups_accumulator/bool_op.rs | 10 +- .../aggregate/groups_accumulator/prim_op.rs | 16 +- datafusion/functions-aggregate/src/average.rs | 38 ++--- .../functions-aggregate/src/correlation.rs | 10 +- datafusion/functions-aggregate/src/count.rs | 10 +- .../src/min_max/min_max_bytes.rs | 12 +- .../src/min_max/min_max_struct.rs | 6 +- .../functions-aggregate/src/string_agg.rs | 6 +- .../functions-aggregate/src/variance.rs | 24 ++- .../src/aggregates/group_values/mod.rs | 33 ++-- .../group_values/multi_group_by/boolean.rs | 7 +- .../group_values/multi_group_by/bytes.rs | 6 +- .../group_values/multi_group_by/bytes_view.rs | 3 +- .../multi_group_by/fixed_size_binary.rs | 10 +- .../group_values/multi_group_by/mod.rs | 15 +- .../group_values/multi_group_by/primitive.rs | 7 +- .../group_values/multi_group_by/row_backed.rs | 8 +- .../aggregates/group_values/null_builder.rs | 7 +- .../src/aggregates/group_values/row.rs | 8 +- .../group_values/single_group_by/boolean.rs | 8 +- .../group_values/single_group_by/bytes.rs | 4 +- .../single_group_by/bytes_view.rs | 4 +- .../group_values/single_group_by/primitive.rs | 10 +- 27 files changed, 230 insertions(+), 208 deletions(-) diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index 92957d186fab1..c682d7caf1b68 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -53,49 +53,83 @@ impl EmitTo { /// Selects groups for a non-destructive grouped aggregation read. /// /// Unlike [`EmitTo`], this selection does not remove groups or change their -/// indices. [`Self::Indices`] preserves the requested order and supports -/// duplicate indices. +/// indices. Selections created by [`Self::try_from_indices`] preserve the +/// requested order and support duplicate indices. /// -/// Indices are trusted to be valid by preserving read APIs. Call -/// [`Self::validate`] first when they do not come from a source that guarantees -/// they are in bounds. +/// A selection is validated once when it is constructed and can then be reused +/// for the group values and accumulators participating in the same snapshot. +/// Construct a new selection if the number or indexing of those groups changes. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum GroupSelection<'a> { - /// Select all groups in group-index order. - All, - /// Select groups in the order specified by their group indices. - Indices(&'a [usize]), +pub struct GroupSelection<'a> { + total_num_groups: usize, + indices: Option<&'a [usize]>, } impl<'a> GroupSelection<'a> { - /// Validates that every selected index is less than `total_num_groups`. + /// Selects all `total_num_groups` groups in group-index order. + pub fn all(total_num_groups: usize) -> Self { + Self { + total_num_groups, + indices: None, + } + } + + /// Selects groups in the order specified by `indices`. /// - /// [`Self::len`] and [`Self::iter`] do not call this method implicitly. - pub fn validate(&self, total_num_groups: usize) -> Result<()> { - if let Self::Indices(indices) = self - && let Some(index) = indices.iter().find(|&&index| index >= total_num_groups) - { + /// Returns an error if an index is not less than `total_num_groups`. Empty + /// selections are valid, and duplicate indices are preserved. + pub fn try_from_indices( + indices: &'a [usize], + total_num_groups: usize, + ) -> Result { + if let Some(index) = indices.iter().find(|&&index| index >= total_num_groups) { return exec_err!( "Group index {index} is out of bounds for {total_num_groups} groups" ); } - Ok(()) + Ok(Self { + total_num_groups, + indices: Some(indices), + }) } - /// Returns the number of selected groups without validating the selection. - pub fn len(&self, total_num_groups: usize) -> usize { - match self { - Self::All => total_num_groups, - Self::Indices(indices) => indices.len(), + /// Returns the group count against which this selection was constructed. + pub fn total_num_groups(self) -> usize { + self.total_num_groups + } + + /// Ensures this selection is being applied to the same number of groups + /// against which it was constructed. + /// + /// Preserving-read implementations should call this method with their + /// stored group count before using [`Self::iter`]. This check is `O(1)`; + /// the selected indices were already checked by [`Self::try_from_indices`]. + pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> { + if actual_num_groups != self.total_num_groups { + return exec_err!( + "Group selection was constructed for {} groups but applied to {actual_num_groups} groups", + self.total_num_groups + ); } + Ok(()) } - /// Returns the selected group indices in output order without validating - /// the selection. - pub fn iter(self, total_num_groups: usize) -> impl Iterator + 'a { - let (all, indices): (_, &'a [usize]) = match self { - Self::All => (0..total_num_groups, &[]), - Self::Indices(indices) => (0..0, indices), + /// Returns the number of selected groups. + pub fn len(self) -> usize { + self.indices + .map_or(self.total_num_groups, |indices| indices.len()) + } + + /// Returns `true` if no groups are selected. + pub fn is_empty(self) -> bool { + self.len() == 0 + } + + /// Returns the selected group indices in output order. + pub fn iter(self) -> impl Iterator + 'a { + let (all, indices): (_, &'a [usize]) = match self.indices { + None => (0..self.total_num_groups, &[]), + Some(indices) => (0..0, indices), }; all.chain(indices.iter().copied()) } @@ -208,13 +242,12 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// Returns final aggregate values without changing the logical state or /// group indices. /// - /// Rows are returned in the order specified by `selection`. Implementations - /// may mutate internal caches or builders, but repeated calls and later - /// updates must observe the same logical accumulator state. + /// Rows are returned in the order specified by `selection`. An empty + /// selection returns a correctly typed array with no rows. /// - /// Every index in [`GroupSelection::Indices`] must refer to an existing - /// group. Call [`GroupSelection::validate`] first if this is not guaranteed - /// by the source of the indices. Invalid indices may cause a panic. + /// This method requires exclusive access because implementations may mutate + /// internal caches or builders. However, repeated calls and later updates + /// must observe the same logical accumulator state. fn evaluate_preserving( &mut self, _selection: GroupSelection<'_>, @@ -249,11 +282,12 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// or group indices. /// /// Each returned array has one row per selected group, in the order - /// specified by `selection`. + /// specified by `selection`. An empty selection returns the normal number + /// of correctly typed state arrays, each with no rows. /// - /// Every index in [`GroupSelection::Indices`] must refer to an existing - /// group. Call [`GroupSelection::validate`] first if this is not guaranteed - /// by the source of the indices. Invalid indices may cause a panic. + /// This method requires exclusive access because implementations may mutate + /// internal caches or builders. However, repeated calls and later updates + /// must observe the same logical accumulator state. fn state_preserving( &mut self, _selection: GroupSelection<'_>, @@ -389,24 +423,34 @@ mod tests { } #[test] - fn group_selection_order_duplicates_and_explicit_validation() { + fn group_selection_is_validated_once_and_reusable() { let values = [10, 20, 30, 40]; - let selected = GroupSelection::Indices(&[3, 1, 3]) - .iter(values.len()) - .map(|index| values[index]) - .collect::>(); - assert_eq!(selected, vec![40, 20, 40]); - - let selected = GroupSelection::All - .iter(values.len()) - .map(|index| values[index]) - .collect::>(); - assert_eq!(selected, values); - - let invalid = GroupSelection::Indices(&[4]); - assert_eq!(invalid.len(values.len()), 1); - assert_eq!(invalid.iter(values.len()).collect::>(), vec![4]); - let error = invalid.validate(values.len()).unwrap_err(); + let selected = + GroupSelection::try_from_indices(&[3, 1, 3], values.len()).unwrap(); + assert_eq!(selected.total_num_groups(), values.len()); + assert_eq!(selected.len(), 3); + assert_eq!( + selected + .iter() + .map(|index| values[index]) + .collect::>(), + vec![40, 20, 40] + ); + selected.validate_num_groups(values.len()).unwrap(); + let error = selected.validate_num_groups(values.len() - 1).unwrap_err(); + assert!(error.to_string().contains("constructed for 4 groups")); + + let all = GroupSelection::all(values.len()); + assert_eq!( + all.iter().map(|index| values[index]).collect::>(), + values + ); + + let empty = GroupSelection::try_from_indices(&[], values.len()).unwrap(); + assert!(empty.is_empty()); + assert!(empty.iter().next().is_none()); + + let error = GroupSelection::try_from_indices(&[4], values.len()).unwrap_err(); assert!(error.to_string().contains("out of bounds")); } } diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 001487b71e539..8acee0554a264 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -109,9 +109,9 @@ where &mut self, selection: GroupSelection<'_>, ) -> datafusion_common::Result { - debug_assert!(selection.validate(self.counts.len()).is_ok()); + selection.validate_num_groups(self.counts.len())?; let counts = selection - .iter(self.counts.len()) + .iter() .map(|index| self.counts[index]) .collect::>(); Ok(Arc::new(Int64Array::from(counts))) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index ac148deabdd1e..afe376f87cfee 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -338,8 +338,8 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.states.len()).is_ok()); - let selected_len = selection.len(self.states.len()); + selection.validate_num_groups(self.states.len())?; + let selected_len = selection.len(); if selected_len == 0 { // ScalarValue::iter_to_array needs at least one value to infer the // output type, so evaluate a temporary empty accumulator. @@ -348,7 +348,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { } let mut results = Vec::with_capacity(selected_len); - for group_index in selection.iter(self.states.len()) { + for group_index in selection.iter() { let (result, size_pre, size_post) = { let state = &mut self.states[group_index]; let size_pre = state.size(); diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs index c86a795f3f37a..5aeaa07aac138 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -291,24 +291,20 @@ impl NullState { } /// Creates a [`NullBuffer`] for `selection` without changing this state. - /// - /// Indices in `selection` must be less than `total_num_groups`. This method - /// does not validate them. pub fn build_preserving( &self, selection: GroupSelection<'_>, - total_num_groups: usize, ) -> Result> { - let selected_len = selection.len(total_num_groups); + let selected_len = selection.len(); match &self.seen_values { SeenValues::All { num_values } => { - debug_assert_eq!(*num_values, total_num_groups); + selection.validate_num_groups(*num_values)?; Ok(None) } SeenValues::Some { values } => { - debug_assert_eq!(values.len(), total_num_groups); + selection.validate_num_groups(values.len())?; let mut selected = BooleanBufferBuilder::new(selected_len); - for index in selection.iter(total_num_groups) { + for index in selection.iter() { selected.append(values.get_bit(index)); } Ok(Some(NullBuffer::new(selected.finish()))) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs index 665610716207b..77b5efc06acb6 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs @@ -127,14 +127,12 @@ where } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.values.len()).is_ok()); - let mut values = BooleanBufferBuilder::new(selection.len(self.values.len())); - for index in selection.iter(self.values.len()) { + selection.validate_num_groups(self.values.len())?; + let mut values = BooleanBufferBuilder::new(selection.len()); + for index in selection.iter() { values.append(self.values.get_bit(index)); } - let nulls = self - .null_state - .build_preserving(selection, self.values.len())?; + let nulls = self.null_state.build_preserving(selection)?; Ok(Arc::new(BooleanArray::new(values.finish(), nulls))) } diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index 12c679525edd3..0453124528a0f 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -126,14 +126,10 @@ where } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.values.len()).is_ok()); - let values: Vec = selection - .iter(self.values.len()) - .map(|index| self.values[index]) - .collect(); - let nulls = self - .null_state - .build_preserving(selection, self.values.len())?; + selection.validate_num_groups(self.values.len())?; + let values: Vec = + selection.iter().map(|index| self.values[index]).collect(); + let nulls = self.null_state.build_preserving(selection)?; let values = PrimitiveArray::::new(values.into(), nulls) .with_data_type(self.data_type.clone()); Ok(Arc::new(values)) @@ -240,7 +236,7 @@ mod tests { let values = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])); accumulator.update_batch(&[values], &[0, 1, 2], None, 4)?; - let selection = GroupSelection::Indices(&[3, 0, 1, 1]); + let selection = GroupSelection::try_from_indices(&[3, 0, 1, 1], 4)?; let expected = Int64Array::from(vec![None, Some(1), None, None]); for _ in 0..2 { let actual = accumulator.evaluate_preserving(selection)?; @@ -253,7 +249,7 @@ mod tests { let values = Arc::new(Int64Array::from(vec![5, 7])); accumulator.update_batch(&[values], &[1, 3], None, 4)?; let expected = Int64Array::from(vec![Some(1), Some(5), Some(3), Some(7)]); - let actual = accumulator.evaluate_preserving(GroupSelection::All)?; + let actual = accumulator.evaluate_preserving(GroupSelection::all(4))?; assert_eq!(actual.as_primitive::(), &expected); // A destructive read still sees all state after preserving reads. diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index 72efd75b60a6c..1292a3c3d60f4 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -1060,18 +1060,11 @@ where } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.counts.len()).is_ok()); - let counts = selection - .iter(self.counts.len()) - .map(|index| self.counts[index]) - .collect(); - let sums = selection - .iter(self.sums.len()) - .map(|index| self.sums[index]) - .collect(); - let nulls = self - .null_state - .build_preserving(selection, self.sums.len())?; + debug_assert_eq!(self.counts.len(), self.sums.len()); + selection.validate_num_groups(self.counts.len())?; + let counts = selection.iter().map(|index| self.counts[index]).collect(); + let sums = selection.iter().map(|index| self.sums[index]).collect(); + let nulls = self.null_state.build_preserving(selection)?; self.evaluate_values(counts, sums, nulls) } @@ -1091,18 +1084,11 @@ where &mut self, selection: GroupSelection<'_>, ) -> Result> { - debug_assert!(selection.validate(self.counts.len()).is_ok()); - let counts = selection - .iter(self.counts.len()) - .map(|index| self.counts[index]) - .collect(); - let sums = selection - .iter(self.sums.len()) - .map(|index| self.sums[index]) - .collect(); - let nulls = self - .null_state - .build_preserving(selection, self.sums.len())?; + debug_assert_eq!(self.counts.len(), self.sums.len()); + selection.validate_num_groups(self.counts.len())?; + let counts = selection.iter().map(|index| self.counts[index]).collect(); + let sums = selection.iter().map(|index| self.sums[index]).collect(); + let nulls = self.null_state.build_preserving(selection)?; Ok(self.state_values(counts, sums, nulls)) } @@ -1464,7 +1450,7 @@ mod tests { ])); accumulator.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; - let selection = GroupSelection::Indices(&[2, 0, 3, 2]); + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; let expected = Float64Array::from(vec![Some(8.0), Some(3.0), None, Some(8.0)]); for _ in 0..2 { assert_eq!( @@ -1491,7 +1477,7 @@ mod tests { Float64Array::from(vec![Some(3.0), Some(10.0), Some(8.0), Some(6.0)]); assert_eq!( accumulator - .evaluate_preserving(GroupSelection::All)? + .evaluate_preserving(GroupSelection::all(4))? .as_primitive::(), &expected ); diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index 767c840fc0eef..e509e0a943267 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -306,10 +306,8 @@ pub struct CorrelationGroupsAccumulator { } fn copy_selected(selection: GroupSelection<'_>, values: &[T]) -> Vec { - selection - .iter(values.len()) - .map(|index| values[index]) - .collect() + debug_assert_eq!(selection.total_num_groups(), values.len()); + selection.iter().map(|index| values[index]).collect() } impl CorrelationGroupsAccumulator { @@ -477,7 +475,7 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.count.len()).is_ok()); + selection.validate_num_groups(self.count.len())?; Ok(Self::evaluate_values( ©_selected(selection, &self.count), ©_selected(selection, &self.sum_x), @@ -565,7 +563,7 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { &mut self, selection: GroupSelection<'_>, ) -> Result> { - debug_assert!(selection.validate(self.count.len()).is_ok()); + selection.validate_num_groups(self.count.len())?; Ok(vec![ Arc::new(UInt64Array::from(copy_selected(selection, &self.count))), Arc::new(Float64Array::from(copy_selected(selection, &self.sum_x))), diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 19844f28bb3ae..f69b8d403655c 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -708,9 +708,9 @@ impl GroupsAccumulator for CountGroupsAccumulator { } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.counts.len()).is_ok()); + selection.validate_num_groups(self.counts.len())?; let counts = selection - .iter(self.counts.len()) + .iter() .map(|index| self.counts[index]) .collect::>(); Ok(Arc::new(Int64Array::from(counts))) @@ -988,7 +988,7 @@ mod tests { let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(2), Some(3)])); accumulator.update_batch(&[values], &[0, 1, 0, 2], None, 4)?; - let selection = GroupSelection::Indices(&[2, 0, 3, 2]); + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; let expected = Int64Array::from(vec![1, 2, 0, 1]); assert_eq!( accumulator @@ -1006,7 +1006,9 @@ mod tests { let expected = Int64Array::from(vec![2, 0, 1, 1]); assert_eq!( accumulator - .evaluate_preserving(GroupSelection::Indices(&[0, 1, 2, 3]))? + .evaluate_preserving( + GroupSelection::try_from_indices(&[0, 1, 2, 3], 4,)? + )? .as_primitive::(), &expected ); diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index 9a3043e2b08c0..31062e7918e41 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -306,14 +306,14 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { let num_groups = self.inner.min_max.len(); - debug_assert!(selection.validate(num_groups).is_ok()); - let num_values = selection.len(num_groups); + selection.validate_num_groups(num_groups)?; + let num_values = selection.len(); let data_capacity = selection - .iter(num_groups) + .iter() .filter_map(|index| self.inner.min_max[index].as_ref().map(Vec::len)) .sum(); let min_maxes = selection - .iter(num_groups) + .iter() .map(|index| self.inner.min_max[index].as_deref()); self.build_array(min_maxes, num_values, data_capacity) } @@ -556,7 +556,7 @@ mod tests { ])); accumulator.update_batch(&[values], &[0, 0, 1, 2, 2, 3], None, 4)?; - let selection = GroupSelection::Indices(&[3, 0, 1, 2, 3]); + let selection = GroupSelection::try_from_indices(&[3, 0, 1, 2, 3], 4)?; let expected = StringArray::from(vec![Some("x"), Some("b"), None, Some("a"), Some("x")]); for _ in 0..2 { @@ -573,7 +573,7 @@ mod tests { let expected = StringArray::from(vec![Some("aa"), None, Some("a"), Some("w")]); assert_eq!( accumulator - .evaluate_preserving(GroupSelection::All)? + .evaluate_preserving(GroupSelection::all(4))? .as_string::(), &expected ); diff --git a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs index 949d156cb60f0..1fe551ef488f4 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs @@ -137,10 +137,10 @@ impl GroupsAccumulator for MinMaxStructAccumulator { fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { let num_groups = self.inner.min_max.len(); - debug_assert!(selection.validate(num_groups).is_ok()); - let num_values = selection.len(num_groups); + selection.validate_num_groups(num_groups)?; + let num_values = selection.len(); let min_maxes = selection - .iter(num_groups) + .iter() .map(|index| self.inner.min_max[index].as_ref()); self.build_array(min_maxes, num_values) } diff --git a/datafusion/functions-aggregate/src/string_agg.rs b/datafusion/functions-aggregate/src/string_agg.rs index d7530f81a78f8..10af87b338e6b 100644 --- a/datafusion/functions-aggregate/src/string_agg.rs +++ b/datafusion/functions-aggregate/src/string_agg.rs @@ -406,10 +406,8 @@ impl GroupsAccumulator for StringAggGroupsAccumulator { } fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.values.len()).is_ok()); - let values = selection - .iter(self.values.len()) - .map(|index| self.values[index].as_deref()); + selection.validate_num_groups(self.values.len())?; + let values = selection.iter().map(|index| self.values[index].as_deref()); Ok(Arc::new(LargeStringArray::from_iter(values))) } diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index d25ea6596ba5a..390bf61e3a9f8 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -511,15 +511,11 @@ impl VarianceGroupsAccumulator { &self, selection: GroupSelection<'_>, ) -> Result<(Vec, NullBuffer)> { - debug_assert!(selection.validate(self.counts.len()).is_ok()); - let counts = selection - .iter(self.counts.len()) - .map(|index| self.counts[index]) - .collect(); - let m2s = selection - .iter(self.m2s.len()) - .map(|index| self.m2s[index]) - .collect(); + debug_assert_eq!(self.counts.len(), self.means.len()); + debug_assert_eq!(self.counts.len(), self.m2s.len()); + selection.validate_num_groups(self.counts.len())?; + let counts = selection.iter().map(|index| self.counts[index]).collect(); + let m2s = selection.iter().map(|index| self.m2s[index]).collect(); Ok(self.variance_values(counts, m2s)) } } @@ -653,17 +649,19 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { &mut self, selection: GroupSelection<'_>, ) -> Result> { - debug_assert!(selection.validate(self.counts.len()).is_ok()); + debug_assert_eq!(self.counts.len(), self.means.len()); + debug_assert_eq!(self.counts.len(), self.m2s.len()); + selection.validate_num_groups(self.counts.len())?; let counts = selection - .iter(self.counts.len()) + .iter() .map(|index| self.counts[index]) .collect::>(); let means = selection - .iter(self.means.len()) + .iter() .map(|index| self.means[index]) .collect::>(); let m2s = selection - .iter(self.m2s.len()) + .iter() .map(|index| self.m2s[index]) .collect::>(); Ok(vec![ diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 162f6c1de617b..016223170bd26 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -116,11 +116,11 @@ pub trait GroupValues: Send { /// Materializes selected group values without changing the stored values or /// their group indices. /// - /// Rows are returned in the order specified by `selection`. + /// Rows are returned in the order specified by `selection`. An empty + /// selection returns one correctly typed empty array per group-value column. /// - /// Every index in [`GroupSelection::Indices`] must refer to an existing - /// group. Call [`GroupSelection::validate`] first if this is not guaranteed - /// by the source of the indices. Invalid indices may cause a panic. + /// This method requires exclusive access because implementations may mutate + /// internal caches or builders, even though stored values are unchanged. fn values_preserving( &mut self, _selection: GroupSelection<'_>, @@ -265,13 +265,23 @@ mod tests { group_values.intern(&[input], &mut groups).unwrap(); assert_eq!(groups, vec![0, 1, 0, 2, 3]); - let selection = GroupSelection::Indices(&[3, 0, 2, 0]); + let selection = + GroupSelection::try_from_indices(&[3, 0, 2, 0], group_values.len()).unwrap(); let expected = Int32Array::from(vec![Some(30), Some(10), None, Some(10)]); for _ in 0..2 { let actual = group_values.values_preserving(selection).unwrap(); assert_eq!(actual[0].as_primitive::(), &expected); } + let empty = group_values + .values_preserving( + GroupSelection::try_from_indices(&[], group_values.len()).unwrap(), + ) + .unwrap(); + assert_eq!(empty.len(), 1); + assert_eq!(empty[0].data_type(), &DataType::Int32); + assert!(empty[0].is_empty()); + let input = Arc::new(Int32Array::from(vec![Some(20), Some(40), None])) as ArrayRef; group_values.intern(&[input], &mut groups).unwrap(); @@ -279,12 +289,13 @@ mod tests { let expected = Int32Array::from(vec![Some(10), Some(20), None, Some(30), Some(40)]); - let actual = group_values.values_preserving(GroupSelection::All).unwrap(); + let actual = group_values + .values_preserving(GroupSelection::all(group_values.len())) + .unwrap(); assert_eq!(actual[0].as_primitive::(), &expected); - let error = GroupSelection::Indices(&[5]) - .validate(group_values.len()) - .unwrap_err(); + let error = + GroupSelection::try_from_indices(&[5], group_values.len()).unwrap_err(); assert!(error.to_string().contains("out of bounds")); let actual = group_values.emit(EmitTo::All).unwrap(); @@ -321,7 +332,9 @@ mod tests { assert_eq!(groups, vec![0, 1, 2, 0]); let selected = group_values - .values_preserving(GroupSelection::Indices(&[2, 1, 0, 2])) + .values_preserving( + GroupSelection::try_from_indices(&[2, 1, 0, 2], 3).unwrap(), + ) .unwrap(); let expected = vec![ Some("a long value that is not inline"), diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs index b307fdb1b6c2c..0311b9ba8f992 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs @@ -179,11 +179,12 @@ impl GroupColumn for BooleanGroupValueBuilder { } fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { - let mut values = BooleanBufferBuilder::new(selection.len(self.buffer.len())); - for index in selection.iter(self.buffer.len()) { + selection.validate_num_groups(self.buffer.len())?; + let mut values = BooleanBufferBuilder::new(selection.len()); + for index in selection.iter() { values.append(self.buffer.get_bit(index)); } - let nulls = self.nulls.build_preserving(selection, self.buffer.len())?; + let nulls = self.nulls.build_preserving(selection)?; Ok(Arc::new(BooleanArray::new(values.finish(), nulls))) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index 9c666ce019c2b..a76a06d2e814d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -375,13 +375,13 @@ where } fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { - let selected_len = selection.len(self.len()); + selection.validate_num_groups(self.len())?; let mut buffer = BufferBuilder::::new(0); - let mut offsets = Vec::with_capacity(selected_len + 1); + let mut offsets = Vec::with_capacity(selection.len() + 1); let mut nulls = MaybeNullBufferBuilder::new(); offsets.push(O::default()); - for index in selection.iter(self.len()) { + for index in selection.iter() { let is_null = self.nulls.is_null(index); nulls.append(is_null); if !is_null { diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs index ecbddb7f431c4..a2d0d26443fcb 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs @@ -349,8 +349,9 @@ impl ByteViewGroupValueBuilder { } fn values_preserving_inner(&self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.len())?; let mut selected = Self::new().with_max_block_size(self.max_block_size); - for index in selection.iter(self.len()) { + for index in selection.iter() { let is_null = self.nulls.is_null(index); selected.nulls.append(is_null); if is_null { diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs index a8ecf47c9baa6..2c420670a1ddc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs @@ -226,11 +226,11 @@ impl GroupColumn for FixedSizeBinaryGroupValueBuilder { } fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.len).is_ok()); - let len = selection.len(self.len); + selection.validate_num_groups(self.len)?; + let len = selection.len(); let mut values = Vec::with_capacity(len * self.byte_width); let mut nulls = MaybeNullBufferBuilder::new(); - for index in selection.iter(self.len) { + for index in selection.iter() { nulls.append(self.nulls.is_null(index)); values.extend_from_slice(self.value(index)); } @@ -503,7 +503,9 @@ mod tests { builder.vectorized_append(&input, &[0, 1, 2]).unwrap(); let output = builder - .values_preserving(GroupSelection::Indices(&[2, 0, 1, 2])) + .values_preserving( + GroupSelection::try_from_indices(&[2, 0, 1, 2], 3).unwrap(), + ) .unwrap(); let expected = make_array( vec![ diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 418d0780730b7..22f28b0543aa1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -110,8 +110,7 @@ pub trait GroupColumn: Send + Sync { fn build(self: Box) -> ArrayRef; /// Builds a new array from selected stored rows without changing this - /// column. Rows are returned in selection order. The caller must ensure all - /// selected indices are in bounds. + /// column. Rows are returned in selection order. fn values_preserving(&self, _selection: GroupSelection<'_>) -> Result { not_impl_err!("Preserving group column values are not implemented") } @@ -1296,7 +1295,7 @@ impl GroupValues for GroupValuesColumn { &mut self, selection: GroupSelection<'_>, ) -> Result> { - debug_assert!(selection.validate(self.len()).is_ok()); + selection.validate_num_groups(self.len())?; if self.group_values.is_empty() { return Ok(self .schema @@ -2032,9 +2031,9 @@ mod tests { data_set.load_to_group_values(&mut group_values); let selection = [16, 0, 4, 0]; - let actual = group_values - .values_preserving(GroupSelection::Indices(&selection)) - .unwrap(); + let group_selection = + GroupSelection::try_from_indices(&selection, group_values.len()).unwrap(); + let actual = group_values.values_preserving(group_selection).unwrap(); let indices = UInt32Array::from_iter_values(selection.map(|index| index as u32)); let mut destructive_group_values = GroupValuesColumn::::try_new(data_set.schema()).unwrap(); @@ -2049,9 +2048,7 @@ mod tests { assert_eq!(actual, expected); // A repeated preserving read returns the same rows and leaves all groups. - let repeated = group_values - .values_preserving(GroupSelection::Indices(&selection)) - .unwrap(); + let repeated = group_values.values_preserving(group_selection).unwrap(); assert_eq!( RecordBatch::try_new(data_set.schema(), repeated).unwrap(), expected diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index d3a036ab66a42..26b8c1cbd23d3 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -280,13 +280,12 @@ where } fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.group_values.len())?; let values: Vec = selection - .iter(self.group_values.len()) + .iter() .map(|index| self.group_values[index]) .collect(); - let nulls = self - .nulls - .build_preserving(selection, self.group_values.len())?; + let nulls = self.nulls.build_preserving(selection)?; Ok(Arc::new( PrimitiveArray::::new(ScalarBuffer::from(values), nulls) .with_data_type(self.data_type.clone()), diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index 6067a18ac5b0f..198f6ba0f13bc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -296,10 +296,8 @@ impl GroupColumn for RowsGroupColumn { } fn values_preserving(&self, selection: GroupSelection<'_>) -> Result { - debug_assert!(selection.validate(self.group_values.num_rows()).is_ok()); - let rows = selection - .iter(self.group_values.num_rows()) - .map(|index| self.group_values.row(index)); + selection.validate_num_groups(self.group_values.num_rows())?; + let rows = selection.iter().map(|index| self.group_values.row(index)); Ok(self.rows_to_array(rows)) } @@ -575,7 +573,7 @@ mod tests { col.vectorized_append(&input, &[0, 1]).unwrap(); let output = col - .values_preserving(GroupSelection::Indices(&[1, 0, 1])) + .values_preserving(GroupSelection::try_from_indices(&[1, 0, 1], 2).unwrap()) .unwrap(); let output = output.as_any().downcast_ref::().unwrap(); assert_eq!( diff --git a/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs b/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs index 602017369532b..77f8984215cbe 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs @@ -78,16 +78,15 @@ impl MaybeNullBufferBuilder { pub fn build_preserving( &self, selection: GroupSelection<'_>, - total_num_values: usize, ) -> Result> { - let selected_len = selection.len(total_num_values); + let selected_len = selection.len(); if self.nulls.as_slice().is_none() { return Ok(None); } - debug_assert_eq!(self.nulls.len(), total_num_values); + debug_assert_eq!(self.nulls.len(), selection.total_num_groups()); let mut selected = NullBufferBuilder::new(selected_len); - for index in selection.iter(total_num_values) { + for index in selection.iter() { selected.append(self.nulls.is_valid(index)); } Ok(selected.finish()) diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 970fb4d6a22dc..01e9f3eaa71ee 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -266,10 +266,8 @@ impl GroupValues for GroupValuesRows { empty_rows = self.row_converter.empty_rows(0, 0); &empty_rows }; - debug_assert!(selection.validate(group_values.num_rows()).is_ok()); - let rows = selection - .iter(group_values.num_rows()) - .map(|index| group_values.row(index)); + selection.validate_num_groups(group_values.num_rows())?; + let rows = selection.iter().map(|index| group_values.row(index)); let mut output = self.row_converter.convert_rows(rows)?; // TODO: Materialize dictionaries in group keys @@ -467,7 +465,7 @@ mod tests { group_values.intern(&[input], &mut groups)?; assert_eq!(groups, vec![0, 1, 2, 0]); - let selection = GroupSelection::Indices(&[2, 0, 1, 2]); + let selection = GroupSelection::try_from_indices(&[2, 0, 1, 2], 3)?; let expected = ListArray::from_iter_primitive::(vec![ Some(vec![Some(3)]), Some(vec![Some(1), Some(2)]), diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs index 8ddb9abd85444..76fb0529488f0 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs @@ -150,10 +150,10 @@ impl GroupValues for GroupValuesBoolean { selection: GroupSelection<'_>, ) -> Result> { let num_groups = self.len(); - debug_assert!(selection.validate(num_groups).is_ok()); - let mut values = BooleanBufferBuilder::new(selection.len(num_groups)); - let mut nulls = NullBufferBuilder::new(selection.len(num_groups)); - for index in selection.iter(num_groups) { + selection.validate_num_groups(num_groups)?; + let mut values = BooleanBufferBuilder::new(selection.len()); + let mut nulls = NullBufferBuilder::new(selection.len()); + for index in selection.iter() { if self.null_group == Some(index) { values.append(false); nulls.append_null(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 33faad8dae4d4..34ec36be31d2e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -124,8 +124,8 @@ impl GroupValues for GroupValuesBytes { &mut self, selection: GroupSelection<'_>, ) -> Result> { - debug_assert!(selection.validate(self.len()).is_ok()); - Ok(vec![self.map.keys(selection.iter(self.len()))?]) + selection.validate_num_groups(self.len())?; + Ok(vec![self.map.keys(selection.iter())?]) } fn supports_values_preserving(&self) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 5814e8289d5f2..997a7ce166a71 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -126,8 +126,8 @@ impl GroupValues for GroupValuesBytesView { &mut self, selection: GroupSelection<'_>, ) -> datafusion_common::Result> { - debug_assert!(selection.validate(self.len()).is_ok()); - Ok(vec![self.map.keys(selection.iter(self.len()))?]) + selection.validate_num_groups(self.len())?; + Ok(vec![self.map.keys(selection.iter())?]) } fn supports_values_preserving(&self) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index d5fa0521c0c9f..21b62457e3831 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -243,14 +243,12 @@ where &mut self, selection: GroupSelection<'_>, ) -> Result> { - debug_assert!(selection.validate(self.values.len()).is_ok()); - let values: Vec = selection - .iter(self.values.len()) - .map(|index| self.values[index]) - .collect(); + selection.validate_num_groups(self.values.len())?; + let values: Vec = + selection.iter().map(|index| self.values[index]).collect(); let nulls = if let Some(null_group) = self.null_group { let mut nulls = NullBufferBuilder::new(values.len()); - for index in selection.iter(self.values.len()) { + for index in selection.iter() { if index == null_group { nulls.append_null(); } else { From faa7a2e0b2a8d65faf2d3e46e93d50c981273510 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:08:56 -0700 Subject: [PATCH 6/9] Require preserving reads for group columns --- .../src/aggregates/group_values/multi_group_by/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 22f28b0543aa1..554ab4971abde 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -111,9 +111,7 @@ pub trait GroupColumn: Send + Sync { /// Builds a new array from selected stored rows without changing this /// column. Rows are returned in selection order. - fn values_preserving(&self, _selection: GroupSelection<'_>) -> Result { - not_impl_err!("Preserving group column values are not implemented") - } + fn values_preserving(&self, selection: GroupSelection<'_>) -> Result; /// Builds a new array from the first `n` stored rows, shifting the /// remaining rows to the start of the builder From ccb44f4b26219de81c62a6e8f705b73ab5382b7d Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:09:12 -0700 Subject: [PATCH 7/9] Handle empty struct min-max preserving reads --- .../src/min_max/min_max_struct.rs | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs index 1fe551ef488f4..8b0132f070c48 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs @@ -20,6 +20,7 @@ use std::{cmp::Ordering, sync::Arc}; use arrow::{ array::{ Array, ArrayData, ArrayRef, AsArray, BooleanArray, MutableArrayData, StructArray, + new_empty_array, }, datatypes::DataType, }; @@ -69,6 +70,9 @@ impl MinMaxStructAccumulator { DataType::Struct(fields) => fields, _ => return internal_err!("Data type is not a struct"), }; + if num_values == 0 { + return Ok(new_empty_array(&self.inner.data_type)); + } let null_array = StructArray::new_null(fields.clone(), 1); let min_maxes_data: Vec = min_maxes .map(|value| match value { @@ -531,6 +535,70 @@ mod tests { assert_eq!(str_array.value(1), "d"); } + #[test] + fn test_min_struct_preserving_reads_and_empty_selection() -> Result<()> { + let array = create_test_struct_array( + vec![Some(3), Some(2), Some(1)], + vec![Some("c"), Some("b"), Some("a")], + ); + let data_type = array.data_type().clone(); + let mut accumulator = MinMaxStructAccumulator::new_min(data_type.clone()); + accumulator.update_batch(&[Arc::new(array)], &[0, 1, 0], None, 3)?; + + let selection = GroupSelection::try_from_indices(&[1, 0, 2, 1], 3)?; + for _ in 0..2 { + let actual = accumulator.evaluate_preserving(selection)?; + let actual = actual.as_struct(); + assert_eq!(actual.len(), 4); + assert_eq!( + actual + .column(0) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(2), Some(1), None, Some(2)] + ); + assert_eq!( + actual + .column(1) + .as_string::() + .iter() + .collect::>(), + vec![Some("b"), Some("a"), None, Some("b")] + ); + assert!(actual.is_null(2)); + } + + let empty_selection = GroupSelection::try_from_indices(&[], 3)?; + let actual = accumulator.evaluate_preserving(empty_selection)?; + assert_eq!(actual.data_type(), &data_type); + assert!(actual.is_empty()); + let state = accumulator.state_preserving(empty_selection)?; + assert_eq!(state.len(), 1); + assert_eq!(state[0].data_type(), &data_type); + assert!(state[0].is_empty()); + + let update = create_test_struct_array(vec![Some(0)], vec![Some("z")]); + accumulator.update_batch(&[Arc::new(update)], &[2], None, 3)?; + let actual = accumulator.evaluate_preserving(GroupSelection::all(3))?; + assert_eq!( + actual + .as_struct() + .column(0) + .as_primitive::() + .values(), + &[1, 2, 0] + ); + + let mut empty_accumulator = MinMaxStructAccumulator::new_min(data_type); + assert!( + empty_accumulator + .evaluate_preserving(GroupSelection::all(0))? + .is_empty() + ); + Ok(()) + } + #[test] fn test_min_max_with_filter() { let array = create_test_struct_array( From 3ada7afe14db5718a3e9b0837d4bf258b8af0e70 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:09:33 -0700 Subject: [PATCH 8/9] Check offsets for selected byte group values --- .../group_values/multi_group_by/bytes.rs | 37 ++++++++++++++++--- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index a76a06d2e814d..9976482e5c8a7 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -34,6 +34,16 @@ use std::mem::size_of; use std::sync::Arc; use std::vec; +fn checked_output_offset( + current_len: usize, + additional_len: usize, +) -> Result { + current_len + .checked_add(additional_len) + .and_then(O::from_usize) + .ok_or_else(|| exec_datafusion_err!("Offset overflow while copying group values")) +} + /// An implementation of [`GroupColumn`] for binary and utf8 types. /// /// Stores a collection of binary or utf8 group values in a single buffer @@ -383,14 +393,16 @@ where for index in selection.iter() { let is_null = self.nulls.is_null(index); + let value = if is_null { &[] } else { self.value(index) }; + let offset = checked_output_offset::(buffer.len(), value.len())?; + nulls.append(is_null); - if !is_null { - buffer.append_slice(self.value(index)); - } - offsets.push(O::usize_as(buffer.len())); + buffer.append_slice(value); + offsets.push(offset); } - // SAFETY: offsets are constructed from the length of `buffer`. + // SAFETY: every offset was checked for representability and is the + // monotonically increasing length of `buffer` after an append. let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; let values = buffer.finish(); let nulls = nulls.build(); @@ -465,7 +477,7 @@ mod tests { use datafusion_common::DataFusionError; use datafusion_physical_expr::binary_map::OutputType; - use super::GroupColumn; + use super::{GroupColumn, checked_output_offset}; fn make_true_buffer(n: usize) -> BooleanBufferBuilder { let mut buf = BooleanBufferBuilder::new(n); @@ -477,6 +489,19 @@ mod tests { (0..buf.len()).map(|i| buf.get_bit(i)).collect() } + #[test] + fn test_selected_copy_offset_overflow_is_checked() { + assert_eq!( + checked_output_offset::(i32::MAX as usize, 0).unwrap(), + i32::MAX + ); + assert!(matches!( + checked_output_offset::(i32::MAX as usize, 1), + Err(DataFusionError::Execution(e)) if e.contains("Offset overflow") + )); + assert!(checked_output_offset::(usize::MAX, 1).is_err()); + } + #[test] fn test_byte_group_value_builder_overflow() { let mut builder = ByteGroupValueBuilder::::new(OutputType::Utf8); From fbd4ac4de0839da7c17b79630a492e5b55c2d8d7 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:09:49 -0700 Subject: [PATCH 9/9] Expand preserving grouped read coverage --- .../src/aggregate/count_distinct/groups.rs | 41 ++++++++++++ .../src/aggregate/groups_accumulator.rs | 48 ++++++++++++++ .../aggregate/groups_accumulator/bool_op.rs | 47 ++++++++++++++ .../functions-aggregate/src/correlation.rs | 65 +++++++++++++++++++ datafusion/functions-aggregate/src/stddev.rs | 32 +++++++++ .../functions-aggregate/src/string_agg.rs | 42 ++++++++++++ .../functions-aggregate/src/variance.rs | 52 +++++++++++++++ 7 files changed, 327 insertions(+) diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 8acee0554a264..6e3e3b91a74f7 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -239,6 +239,47 @@ mod tests { use arrow::datatypes::Int32Type; use datafusion_common::Result; + #[test] + fn preserving_reads_keep_distinct_state() -> Result<()> { + let mut accumulator = PrimitiveDistinctCountGroupsAccumulator::::new(); + let values = Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(1), + None, + Some(3), + ])); + accumulator.update_batch(&[values], &[0, 0, 1, 2, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; + let expected = Int64Array::from(vec![1, 2, 0, 1]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + assert!( + accumulator + .evaluate_preserving(GroupSelection::try_from_indices(&[], 4)?)? + .is_empty() + ); + + let values = Arc::new(Int32Array::from(vec![2, 4, 1])); + accumulator.update_batch(&[values], &[0, 0, 1], None, 4)?; + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Int64Array::from(vec![3, 1, 1, 0]) + ); + assert!(accumulator.supports_evaluate_preserving()); + assert!(!accumulator.supports_state_preserving()); + Ok(()) + } + #[test] fn convert_to_state_roundtrips_through_merge() -> Result<()> { let values = Arc::new(Int32Array::from(vec![ diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index afe376f87cfee..6704b068acf0b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -529,3 +529,51 @@ pub(crate) fn slice_and_maybe_filter( Ok(sliced_arrays) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::min_max::MaxAccumulator; + use arrow::array::{AsArray, Int64Array}; + use arrow::datatypes::{DataType, Int64Type}; + + #[test] + fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> { + let mut accumulator = GroupsAccumulatorAdapter::new(|| { + Ok(Box::new(MaxAccumulator::try_new(&DataType::Int64)?) + as Box) + }); + let values = Arc::new(Int64Array::from(vec![Some(1), Some(5), Some(2), None])); + accumulator.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[1, 0, 3, 1], 4)?; + let expected = Int64Array::from(vec![Some(2), Some(5), None, Some(2)]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + + let empty = + accumulator.evaluate_preserving(GroupSelection::try_from_indices(&[], 4)?)?; + assert_eq!(empty.data_type(), &DataType::Int64); + assert!(empty.is_empty()); + + let values = Arc::new(Int64Array::from(vec![7, 4, 9])); + accumulator.update_batch(&[values], &[0, 2, 3], None, 4)?; + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Int64Array::from(vec![Some(7), Some(2), Some(4), Some(9)]) + ); + assert!(accumulator.supports_evaluate_preserving()); + assert!(!accumulator.supports_state_preserving()); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs index 77b5efc06acb6..d5c0e79b6420c 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs @@ -184,3 +184,50 @@ where Ok(vec![Arc::new(values_filtered)]) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn boolean_groups_preserving_reads() -> Result<()> { + let mut accumulator = + BooleanGroupsAccumulator::new(|current, value| current && value, true); + let values = Arc::new(BooleanArray::from(vec![ + Some(true), + Some(false), + None, + Some(true), + ])); + accumulator.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; + let expected = + BooleanArray::from(vec![Some(true), Some(false), None, Some(true)]); + for _ in 0..2 { + assert_eq!( + accumulator.evaluate_preserving(selection)?.as_boolean(), + &expected + ); + assert_eq!( + accumulator.state_preserving(selection)?[0].as_boolean(), + &expected + ); + } + + let empty = + accumulator.evaluate_preserving(GroupSelection::try_from_indices(&[], 4)?)?; + assert!(empty.is_empty()); + + let values = Arc::new(BooleanArray::from(vec![false, true])); + accumulator.update_batch(&[values], &[1, 3], None, 4)?; + let expected = BooleanArray::from(vec![false, false, true, true]); + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::all(4))? + .as_boolean(), + &expected + ); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index e509e0a943267..21859006ea943 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -683,6 +683,71 @@ mod tests { assert!(result.is_err()); } + #[test] + fn correlation_groups_preserving_reads() -> Result<()> { + let mut accumulator = CorrelationGroupsAccumulator::new(); + let x = Arc::new(Float64Array::from(vec![1.0, 2.0, 1.0, 1.0, 2.0])); + let y = Arc::new(Float64Array::from(vec![2.0, 4.0, 2.0, 3.0, 1.0])); + accumulator.update_batch(&[x, y], &[0, 0, 1, 2, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; + let expected = Float64Array::from(vec![Some(-1.0), Some(1.0), None, Some(-1.0)]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + let state = accumulator.state_preserving(selection)?; + assert_eq!(state.len(), 6); + assert_eq!( + state[0].as_primitive::(), + &UInt64Array::from(vec![2, 2, 0, 2]) + ); + assert_eq!( + state[1].as_primitive::(), + &Float64Array::from(vec![3.0, 3.0, 0.0, 3.0]) + ); + assert_eq!( + state[2].as_primitive::(), + &Float64Array::from(vec![4.0, 6.0, 0.0, 4.0]) + ); + assert_eq!( + state[3].as_primitive::(), + &Float64Array::from(vec![5.0, 10.0, 0.0, 5.0]) + ); + assert_eq!( + state[4].as_primitive::(), + &Float64Array::from(vec![5.0, 5.0, 0.0, 5.0]) + ); + assert_eq!( + state[5].as_primitive::(), + &Float64Array::from(vec![10.0, 20.0, 0.0, 10.0]) + ); + } + + let empty_selection = GroupSelection::try_from_indices(&[], 4)?; + assert!(accumulator.evaluate_preserving(empty_selection)?.is_empty()); + assert!( + accumulator + .state_preserving(empty_selection)? + .iter() + .all(|array| array.is_empty()) + ); + + let x = Arc::new(Float64Array::from(vec![2.0, 1.0, 4.0])); + let y = Arc::new(Float64Array::from(vec![4.0, 2.0, 8.0])); + accumulator.update_batch(&[x, y], &[1, 3, 3], None, 4)?; + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Float64Array::from(vec![1.0, 1.0, -1.0, 1.0]) + ); + Ok(()) + } + #[test] fn convert_to_state_roundtrips_through_merge() -> Result<()> { let x = Arc::new(Float64Array::from(vec![ diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs index d7b7f45d86d77..42a36a1b11991 100644 --- a/datafusion/functions-aggregate/src/stddev.rs +++ b/datafusion/functions-aggregate/src/stddev.rs @@ -386,6 +386,38 @@ mod tests { use datafusion_functions_aggregate_common::utils::get_accum_scalar_values_as_arrays; use datafusion_physical_expr::expressions::col; + #[test] + fn stddev_groups_preserving_reads() -> Result<()> { + let mut accumulator = StddevGroupsAccumulator::new(StatsType::Population); + let values = Arc::new(Float64Array::from(vec![1.0, 3.0, 2.0, 2.0, 6.0])); + accumulator.update_batch(&[values], &[0, 0, 1, 2, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; + let expected = Float64Array::from(vec![Some(2.0), Some(1.0), None, Some(2.0)]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + assert_eq!( + accumulator.state_preserving(selection)?[0].as_primitive::(), + &UInt64Array::from(vec![2, 2, 0, 2]) + ); + } + + let values = Arc::new(Float64Array::from(vec![4.0, 5.0, 7.0])); + accumulator.update_batch(&[values], &[1, 3, 3], None, 4)?; + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Float64Array::from(vec![1.0, 1.0, 2.0, 1.0]) + ); + Ok(()) + } + #[test] fn stddev_f64_merge_1() -> Result<()> { let a = Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64])); diff --git a/datafusion/functions-aggregate/src/string_agg.rs b/datafusion/functions-aggregate/src/string_agg.rs index 10af87b338e6b..de995c50d0ff8 100644 --- a/datafusion/functions-aggregate/src/string_agg.rs +++ b/datafusion/functions-aggregate/src/string_agg.rs @@ -818,6 +818,48 @@ mod tests { arr.iter().map(|v| v.map(|s| s.to_string())).collect() } + #[test] + fn groups_preserving_reads() -> Result<()> { + let mut acc = make_groups_acc(","); + let values: ArrayRef = Arc::new(LargeStringArray::from(vec![ + Some("a"), + Some("b"), + None, + Some("c"), + ])); + acc.update_batch(&[values], &[0, 1, 2, 0], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[1, 0, 3, 1], 4)?; + let expected = + LargeStringArray::from(vec![Some("b"), Some("a,c"), None, Some("b")]); + let bytes_before = acc.total_data_bytes; + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_string::(), + &expected + ); + assert_eq!( + acc.state_preserving(selection)?[0].as_string::(), + &expected + ); + assert_eq!(acc.total_data_bytes, bytes_before); + } + + assert!( + acc.evaluate_preserving(GroupSelection::try_from_indices(&[], 4)?)? + .is_empty() + ); + + let values: ArrayRef = Arc::new(LargeStringArray::from(vec!["d", "e"])); + acc.update_batch(&[values], &[2, 3], None, 4)?; + assert_eq!( + acc.evaluate_preserving(GroupSelection::all(4))? + .as_string::(), + &LargeStringArray::from(vec!["a,c", "b", "d", "e"]) + ); + Ok(()) + } + #[test] fn groups_basic() -> Result<()> { let mut acc = make_groups_acc(","); diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index 390bf61e3a9f8..34f24b7bc76d5 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -749,6 +749,8 @@ impl Accumulator for DistinctVarianceAccumulator { #[cfg(test)] mod tests { + use arrow::array::AsArray; + use arrow::datatypes::UInt64Type; use datafusion_expr::EmitTo; use super::*; @@ -822,6 +824,56 @@ mod tests { Ok(()) } + #[test] + fn variance_groups_preserving_reads() -> Result<()> { + let mut accumulator = VarianceGroupsAccumulator::new(StatsType::Population); + let values = Arc::new(Float64Array::from(vec![1.0, 3.0, 2.0, 2.0, 6.0])); + accumulator.update_batch(&[values], &[0, 0, 1, 2, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 3, 2], 4)?; + let expected = Float64Array::from(vec![Some(4.0), Some(1.0), None, Some(4.0)]); + for _ in 0..2 { + assert_eq!( + accumulator + .evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + let state = accumulator.state_preserving(selection)?; + assert_eq!( + state[0].as_primitive::(), + &UInt64Array::from(vec![2, 2, 0, 2]) + ); + assert_eq!( + state[1].as_primitive::(), + &Float64Array::from(vec![4.0, 2.0, 0.0, 4.0]) + ); + assert_eq!( + state[2].as_primitive::(), + &Float64Array::from(vec![8.0, 2.0, 0.0, 8.0]) + ); + } + + let empty_selection = GroupSelection::try_from_indices(&[], 4)?; + assert!(accumulator.evaluate_preserving(empty_selection)?.is_empty()); + assert!( + accumulator + .state_preserving(empty_selection)? + .iter() + .all(|array| array.is_empty()) + ); + + let values = Arc::new(Float64Array::from(vec![4.0, 5.0, 7.0])); + accumulator.update_batch(&[values], &[1, 3, 3], None, 4)?; + assert_eq!( + accumulator + .evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Float64Array::from(vec![1.0, 1.0, 4.0, 1.0]) + ); + Ok(()) + } + #[test] fn test_groups_accumulator_merge_empty_states() -> Result<()> { let state_1 = vec![