Refine the size() calculation of accumulator

Author: kyotoYaho

datafusion/core/src/physical_plan/aggregates/row_hash.rs

@@ -318,6 +318,7 @@ impl GroupedHashAggregateStream {
             ..
         } = &mut self.row_aggr_state;
 
+        let mut accumulator_set_init_size = None;
         for (row, hash) in batch_hashes.into_iter().enumerate() {
             let entry = row_map.get_mut(hash, |(_hash, group_idx)| {
                 // verify that a group that we are inserting with hash is
@@ -364,13 +365,15 @@ impl GroupedHashAggregateStream {
                         + (std::mem::size_of::<u32>() * group_state.indices.capacity());
 
                     // Allocation done by normal accumulators
-                    *allocated += (std::mem::size_of::<Box<dyn Accumulator>>()
-                        * group_state.accumulator_set.capacity())
-                        + group_state
-                            .accumulator_set
-                            .iter()
-                            .map(|accu| accu.size())
-                            .sum::<usize>();
+                    *allocated += *accumulator_set_init_size.get_or_insert_with(|| {
+                        std::mem::size_of::<Box<dyn Accumulator>>()
+                            * group_state.accumulator_set.capacity()
+                            + group_state
+                                .accumulator_set
+                                .iter()
+                                .map(|accu| accu.size())
+                                .sum::<usize>()
+                    });
 
                     // for hasher function, use precomputed hash value
                     row_map.insert_accounted(
@@ -389,14 +392,23 @@ impl GroupedHashAggregateStream {
     }
 
     // Update the accumulator results, according to row_aggr_state.
-    fn update_accumulators(
+    #[allow(clippy::too_many_arguments)]
+    fn update_accumulators<F1, F2>(
         &mut self,
         groups_with_rows: &[usize],
         offsets: &[usize],
         row_values: &[Vec<ArrayRef>],
         normal_values: &[Vec<ArrayRef>],
+        func_row: F1,
+        func_normal: F2,
         allocated: &mut usize,
-    ) -> Result<()> {
+    ) -> Result<()>
+    where
+        F1: Fn(&mut RowAccumulatorItem, &mut RowAccessor, &[ArrayRef]) -> Result<()>,
+        F2: Fn(&mut AccumulatorItem, &[ArrayRef]) -> Result<()>,
+    {
+        let accumulator_set_pre =
+            get_accumulator_set_size(groups_with_rows, &self.row_aggr_state.group_states);
         // 2.1 for each key in this batch
         // 2.2 for each aggregation
         // 2.3 `slice` from each of its arrays the keys' values
@@ -428,15 +440,7 @@ impl GroupedHashAggregateStream {
                             RowAccessor::new_from_layout(self.row_aggr_layout.clone());
                         state_accessor
                             .point_to(0, group_state.aggregation_buffer.as_mut_slice());
-                        match self.mode {
-                            AggregateMode::Partial => {
-                                accumulator.update_batch(&values, &mut state_accessor)
-                            }
-                            AggregateMode::FinalPartitioned | AggregateMode::Final => {
-                                // note: the aggregation here is over states, not values, thus the merge
-                                accumulator.merge_batch(&values, &mut state_accessor)
-                            }
-                        }
+                        func_row(accumulator, &mut state_accessor, &values)
                     })
                     // 2.5
                     .and(Ok(()))?;
@@ -458,24 +462,17 @@ impl GroupedHashAggregateStream {
                         )
                     })
                     .try_for_each(|(accumulator, values)| {
-                        let size_pre = accumulator.size();
-                        let res = match self.mode {
-                            AggregateMode::Partial => accumulator.update_batch(&values),
-                            AggregateMode::FinalPartitioned | AggregateMode::Final => {
-                                // note: the aggregation here is over states, not values, thus the merge
-                                accumulator.merge_batch(&values)
-                            }
-                        };
-                        let size_post = accumulator.size();
-                        *allocated += size_post.saturating_sub(size_pre);
-                        res
+                        func_normal(accumulator, &values)
                     })
                     // 2.5
                     .and({
                         group_state.indices.clear();
                         Ok(())
                     })
             })?;
+        let accumulator_set_post =
+            get_accumulator_set_size(groups_with_rows, &self.row_aggr_state.group_states);
+        *allocated += accumulator_set_post.saturating_sub(accumulator_set_pre);
         Ok(())
     }
 
@@ -517,13 +514,39 @@ impl GroupedHashAggregateStream {
             let row_values = get_at_indices(&row_aggr_input_values, &batch_indices)?;
             let normal_values =
                 get_at_indices(&normal_aggr_input_values, &batch_indices)?;
-            self.update_accumulators(
-                &groups_with_rows,
-                &offsets,
-                &row_values,
-                &normal_values,
-                &mut allocated,
-            )?;
+            match self.mode {
+                AggregateMode::Partial => self.update_accumulators(
+                    &groups_with_rows,
+                    &offsets,
+                    &row_values,
+                    &normal_values,
+                    |accumulator: &mut RowAccumulatorItem,
+                     state_accessor: &mut RowAccessor,
+                     values: &[ArrayRef]| {
+                        accumulator.update_batch(values, state_accessor)
+                    },
+                    |accumulator: &mut AccumulatorItem, values: &[ArrayRef]| {
+                        accumulator.update_batch(values)
+                    },
+                    &mut allocated,
+                )?,
+                AggregateMode::FinalPartitioned | AggregateMode::Final => self
+                    .update_accumulators(
+                        &groups_with_rows,
+                        &offsets,
+                        &row_values,
+                        &normal_values,
+                        |accumulator: &mut RowAccumulatorItem,
+                         state_accessor: &mut RowAccessor,
+                         values: &[ArrayRef]| {
+                            accumulator.merge_batch(values, state_accessor)
+                        },
+                        |accumulator: &mut AccumulatorItem, values: &[ArrayRef]| {
+                            accumulator.merge_batch(values)
+                        },
+                        &mut allocated,
+                    )?,
+            };
         }
         allocated += self
             .row_converter
@@ -533,6 +556,19 @@ impl GroupedHashAggregateStream {
     }
 }
 
+fn get_accumulator_set_size(
+    groups_with_rows: &[usize],
+    row_group_states: &[RowGroupState],
+) -> usize {
+    groups_with_rows.iter().fold(0usize, |acc, group_idx| {
+        let group_state = &row_group_states[*group_idx];
+        group_state
+            .accumulator_set
+            .iter()
+            .fold(acc, |acc, accumulator| acc + accumulator.size())
+    })
+}
+
 /// The state that is built for each output group.
 #[derive(Debug)]
 pub struct RowGroupState {

datafusion/physical-expr/src/aggregate/average.rs

@@ -242,7 +242,7 @@ impl Accumulator for AvgAccumulator {
     }
 
     fn size(&self) -> usize {
-        std::mem::size_of_val(self) - std::mem::size_of_val(&self.sum) + self.sum.size()
+        std::mem::size_of_val(self)
     }
 }
 

datafusion/physical-expr/src/aggregate/sum.rs

@@ -289,7 +289,7 @@ impl Accumulator for SumAccumulator {
     }
 
     fn size(&self) -> usize {
-        std::mem::size_of_val(self) - std::mem::size_of_val(&self.sum) + self.sum.size()
+        std::mem::size_of_val(self)
     }
 }