feat: add selective non-destructive grouped state reads

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, not_impl_err, utils::split_vec_min_alloc};
 
 /// Describes how many rows should be emitted during grouping.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -50,6 +50,91 @@ impl EmitTo {
     }
 }
 
+/// Selects groups for a non-destructive grouped aggregation read.
+///
+/// Unlike [`EmitTo`], this selection does not remove groups or change their
+/// indices. Selections created by [`Self::try_from_indices`] preserve the
+/// requested order and support duplicate indices.
+///
+/// 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 struct GroupSelection<'a> {
+    total_num_groups: usize,
+    indices: Option<&'a [usize]>,
+}
+
+impl<'a> GroupSelection<'a> {
+    /// 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`.
+    ///
+    /// 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<Self> {
+        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(Self {
+            total_num_groups,
+            indices: Some(indices),
+        })
+    }
+
+    /// 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 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<Item = usize> + '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())
+    }
+}
+
 /// `GroupsAccumulator` implements a single aggregate (e.g. AVG) and
 /// stores the state for *all* groups internally.
 ///
@@ -154,6 +239,27 @@ pub trait GroupsAccumulator: Send + std::any::Any {
     /// `n`. See [`EmitTo::First`] for more details.
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef>;
 
+    /// Returns final aggregate values without changing the logical state or
+    /// group indices.
+    ///
+    /// Rows are returned in the order specified by `selection`. An empty
+    /// selection returns a correctly typed array with no rows.
+    ///
+    /// 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<'_>,
+    ) -> Result<ArrayRef> {
+        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.
     ///
@@ -172,6 +278,28 @@ pub trait GroupsAccumulator: Send + std::any::Any {
     /// [`Accumulator::state`]: crate::accumulator::Accumulator::state
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>;
 
+    /// 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`. An empty selection returns the normal number
+    /// of correctly typed state arrays, each with no rows.
+    ///
+    /// 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<'_>,
+    ) -> Result<Vec<ArrayRef>> {
+        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.
     ///
@@ -247,7 +375,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 +421,36 @@ mod tests {
             original_capacity,
         );
     }
+
+    #[test]
+    fn group_selection_is_validated_once_and_reusable() {
+        let values = [10, 20, 30, 40];
+        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<_>>(),
+            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::<Vec<_>>(),
+            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"));
+    }
 }

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::{

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<ArrayRef> {
+        selection.validate_num_groups(self.counts.len())?;
+        let counts = selection
+            .iter()
+            .map(|index| self.counts[index])
+            .collect::<Vec<_>>();
+        Ok(Arc::new(Int64Array::from(counts)))
+    }
+
+    fn supports_evaluate_preserving(&self) -> bool {
+        true
+    }
+
     fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
         let num_emitted = match emit_to {
             EmitTo::All => self.counts.len(),
@@ -221,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::<Int32Type>::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::<arrow::datatypes::Int64Type>(),
+                &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::<arrow::datatypes::Int64Type>(),
+            &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![

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<ArrayRef> {
+        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.
+            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() {
+            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<Vec<ArrayRef>> {
         let vec_size_pre = self.states.allocated_size();
@@ -499,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<dyn Accumulator>)
+        });
+        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::<Int64Type>(),
+                &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::<Int64Type>(),
+            &Int64Array::from(vec![Some(7), Some(2), Some(4), Some(9)])
+        );
+        assert!(accumulator.supports_evaluate_preserving());
+        assert!(!accumulator.supports_state_preserving());
+        Ok(())
+    }
+}