Make GroupsAccumulatorFlatAdapter accumulators not intermingle with state

Author: srh

datafusion/src/physical_plan/groups_accumulator_flat_adapter.rs

@@ -91,8 +91,14 @@ use smallvec::SmallVec;
 pub struct GroupsAccumulatorFlatAdapter<AccumulatorType: Accumulator> {
     factory: Box<dyn Fn() -> Result<AccumulatorType> + Send>,
 
-    /// state for each group, stored in group_index order
-    states: Vec<AccumulatorState<AccumulatorType>>,
+    /// [`Accumulators`] that store the per-group state
+    accumulators: Vec<AccumulatorType>,
+
+    /// scratch space: indexes in the input array that will be fed to this accumulator (this is a
+    /// parallel array to accumulators). Stores indexes as `u32` to match the arrow `take` kernel
+    /// input.  Note that we have indices.len() <= accumulators.len(); it only gets extended when
+    /// used.
+    indices: Vec<SmallVec<[u32; 4]>>,
 
     // TODO: Code maintaining this is commented.
     /// Current memory usage, in bytes.
@@ -104,31 +110,15 @@ pub struct GroupsAccumulatorFlatAdapter<AccumulatorType: Accumulator> {
     allocation_bytes: usize,
 }
 
-struct AccumulatorState<AccumulatorType> {
-    /// [`Accumulator`] that stores the per-group state
-    accumulator: AccumulatorType,
-
-    /// scratch space: indexes in the input array that will be fed to
-    /// this accumulator. Stores indexes as `u32` to match the arrow
-    /// `take` kernel input.
-    indices: SmallVec<[u32; 4]>,
-}
-
-impl<AccumulatorType> AccumulatorState<AccumulatorType> {
-    fn new(accumulator: AccumulatorType) -> Self {
-        Self {
-            accumulator,
-            indices: smallvec![],
-        }
-    }
-
-    /* TODO: Add Accumulator::size?
-    /// Returns the amount of memory taken by this structure and its accumulator
-    fn size(&self) -> usize {
-        self.accumulator.size() + size_of_val(self) + self.indices.allocated_size()
-    }
-    */
+// TODO: Remove this or remember to account for accumulators and indices in a size() calculation.
+// impl<AccumulatorType> AccumulatorState<AccumulatorType> {
+/* TODO: Add Accumulator::size?
+/// Returns the amount of memory taken by this structure and its accumulator
+fn size(&self) -> usize {
+    self.accumulator.size() + size_of_val(self) + self.indices.allocated_size()
 }
+*/
+// }
 
 impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType> {
     /// Create a new adapter that will create a new [`Accumulator`]
@@ -139,28 +129,40 @@ impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType>
     {
         Self {
             factory: Box::new(factory),
-            states: vec![],
+            accumulators: vec![],
+            indices: vec![],
             allocation_bytes: 0,
         }
     }
 
     /// Ensure that self.accumulators has total_num_groups
     fn make_accumulators_if_needed(&mut self, total_num_groups: usize) -> Result<()> {
         // can't shrink
-        assert!(total_num_groups >= self.states.len());
-        let vec_size_pre = self.states.allocated_size();
+        assert!(total_num_groups >= self.accumulators.len());
+        let vec_size_pre = self.accumulators.allocated_size();
 
         // instantiate new accumulators
-        let new_accumulators = total_num_groups - self.states.len();
+        let new_accumulators = total_num_groups - self.accumulators.len();
         for _ in 0..new_accumulators {
             let accumulator = (self.factory)()?;
-            let state = AccumulatorState::new(accumulator);
             // TODO: Add Accumulator::size()?
             // self.add_allocation(state.size());
-            self.states.push(state);
+            self.accumulators.push(accumulator);
         }
 
-        self.adjust_allocation(vec_size_pre, self.states.allocated_size());
+        self.adjust_allocation(vec_size_pre, self.accumulators.allocated_size());
+        Ok(())
+    }
+
+    fn make_indices_if_needed(&mut self, total_num_groups: usize) -> Result<()> {
+        // can't shrink
+        assert!(total_num_groups >= self.indices.len());
+        let vec_size_pre = self.indices.allocated_size();
+
+        // instantiate new indices
+        self.indices.resize(total_num_groups, smallvec![]);
+
+        self.adjust_allocation(vec_size_pre, self.indices.allocated_size());
         Ok(())
     }
 
@@ -191,12 +193,12 @@ impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType>
         let mut sizes_post = 0;
         for offsets in offsets_param.windows(2) {
             let group_idx = group_indices[offsets[0]];
-            let state = &mut self.states[group_idx];
+            let accumulator: &mut AccumulatorType = &mut self.accumulators[group_idx];
             // sizes_pre += state.size();  // TODO: Add Accumulator::size?
 
             let values_to_accumulate =
                 slice_and_maybe_filter(values, opt_filter, offsets)?;
-            f(&mut state.accumulator, &values_to_accumulate)?;
+            f(accumulator, &values_to_accumulate)?;
 
             // sizes_post += state.size();  // TODO: Add Accumulator::size?
         }
@@ -217,6 +219,7 @@ impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType>
         F: Fn(&mut dyn Accumulator, &[ArrayRef]) -> Result<()>,
     {
         self.make_accumulators_if_needed(total_num_groups)?;
+        self.make_indices_if_needed(total_num_groups)?;
 
         assert_eq!(
             values[0].len(),
@@ -228,7 +231,7 @@ impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType>
         // Note that self.state.indices starts empty for all groups
         // (it is cleared out below)
         for (idx, group_index) in group_indices.iter().enumerate() {
-            self.states[*group_index].indices.push(idx as u32);
+            self.indices[*group_index].push(idx as u32);
         }
 
         // groups_with_rows holds a list of group indexes that have
@@ -245,8 +248,7 @@ impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType>
         let mut offsets = vec![0];
 
         let mut offset_so_far = 0;
-        for (group_index, state) in self.states.iter_mut().enumerate() {
-            let indices = &state.indices;
+        for (group_index, indices) in self.indices.iter_mut().enumerate() {
             if indices.is_empty() {
                 continue;
             }
@@ -272,19 +274,18 @@ impl<AccumulatorType: Accumulator> GroupsAccumulatorFlatAdapter<AccumulatorType>
         let mut sizes_pre = 0;
         let mut sizes_post = 0;
         for (&group_idx, offsets) in iter {
-            let state = &mut self.states[group_idx];
             // sizes_pre += state.size();  // TODO: Add Accumulator::size?
 
             let values_to_accumulate = slice_and_maybe_filter(
                 &values,
                 opt_filter.as_ref().map(|f| as_boolean_array(f)),
                 offsets,
             )?;
-            f(&mut state.accumulator, &values_to_accumulate)?;
+            f(&mut self.accumulators[group_idx], &values_to_accumulate)?;
 
             // clear out the state so they are empty for next
             // iteration
-            state.indices.clear();
+            self.indices[group_idx].clear();
             // sizes_post += state.size();  // TODO: Add Accumulator::size?
         }
 
@@ -365,41 +366,43 @@ impl<AccumulatorType: Accumulator> GroupsAccumulator
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let vec_size_pre = self.states.allocated_size();
+        let vec_size_pre = self.accumulators.allocated_size();
 
-        let states = emit_to.take_needed(&mut self.states);
+        let accumulators = emit_to.take_needed(&mut self.accumulators);
+        self.indices.truncate(self.accumulators.len());
 
-        let results: Vec<ScalarValue> = states
+        let results: Vec<ScalarValue> = accumulators
             .into_iter()
-            .map(|mut state| {
+            .map(|mut accumulator| {
                 // self.free_allocation(state.size());  // TODO: Add Accumulator::size?
-                state.accumulator.evaluate()
+                accumulator.evaluate()
             })
             .collect::<Result<_>>()?;
 
         let result = ScalarValue::iter_to_array(results);
 
-        self.adjust_allocation(vec_size_pre, self.states.allocated_size());
+        self.adjust_allocation(vec_size_pre, self.accumulators.allocated_size());
 
         result
     }
 
     fn peek_evaluate(&self, group_index: usize) -> Result<ScalarValue> {
-        self.states[group_index].accumulator.evaluate()
+        self.accumulators[group_index].evaluate()
     }
 
     // filtered_null_mask(opt_filter, &values);
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        let vec_size_pre = self.states.allocated_size();
-        let states = emit_to.take_needed(&mut self.states);
+        let vec_size_pre = self.accumulators.allocated_size();
+        let accumulators = emit_to.take_needed(&mut self.accumulators);
+        self.indices.truncate(self.accumulators.len());
 
         // each accumulator produces a potential vector of values
         // which we need to form into columns
         let mut results: Vec<Vec<ScalarValue>> = vec![];
 
-        for mut state in states {
+        for mut accumulator in accumulators {
             // self.free_allocation(state.size());  // TODO: Add Accumulator::size?
-            let accumulator_state = state.accumulator.state()?;
+            let accumulator_state = accumulator.state()?;
             results.resize_with(accumulator_state.len(), Vec::new);
             for (idx, state_val) in accumulator_state.into_iter().enumerate() {
                 results[idx].push(state_val);
@@ -419,13 +422,13 @@ impl<AccumulatorType: Accumulator> GroupsAccumulator
                 assert_eq!(arr.len(), first_col.len())
             }
         }
-        self.adjust_allocation(vec_size_pre, self.states.allocated_size());
+        self.adjust_allocation(vec_size_pre, self.accumulators.allocated_size());
 
         Ok(arrays)
     }
 
     fn peek_state(&self, group_index: usize) -> Result<SmallVec<[ScalarValue; 2]>> {
-        self.states[group_index].accumulator.state()
+        self.accumulators[group_index].state()
     }
 
     fn merge_batch(