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,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<Item = usize> + '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.
 ///
@@ -154,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<ArrayRef>;
 
+    /// 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<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 +245,27 @@ 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`.
+    ///
+    /// 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<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 +341,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 +387,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::<Vec<_>>();
+        assert_eq!(selected, vec![40, 20, 40]);
+
+        let selected = GroupSelection::All
+            .iter(values.len())
+            .map(|index| values[index])
+            .collect::<Vec<_>>();
+        assert_eq!(selected, values);
+
+        let invalid = GroupSelection::Indices(&[4]);
+        assert_eq!(invalid.len(values.len()), 1);
+        assert_eq!(invalid.iter(values.len()).collect::<Vec<_>>(), vec![4]);
+        let error = invalid.validate(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> {
+        debug_assert!(selection.validate(self.counts.len()).is_ok());
+        let counts = selection
+            .iter(self.counts.len())
+            .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(),

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> {
+        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<Vec<ArrayRef>> {
         let vec_size_pre = self.states.allocated_size();

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<Option<NullBuffer>> {
+        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.

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<ArrayRef> {
+        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<Vec<ArrayRef>> {
         self.evaluate(emit_to).map(|arr| vec![arr])
     }
 
+    fn state_preserving(
+        &mut self,
+        selection: GroupSelection<'_>,
+    ) -> Result<Vec<ArrayRef>> {
+        self.evaluate_preserving(selection).map(|arr| vec![arr])
+    }
+
+    fn supports_state_preserving(&self) -> bool {
+        true
+    }
+
     fn merge_batch(
         &mut self,
         values: &[ArrayRef],