alamb commented on code in PR #6904:
URL: https://github.com/apache/arrow-datafusion/pull/6904#discussion_r1258634453
##########
datafusion/physical-expr/src/aggregate/sum.rs:
##########
@@ -465,6 +518,114 @@ impl RowAccumulator for SumRowAccumulator {
}
}
+/// An accumulator to compute the sum of values in [`PrimitiveArray<T>`]
+#[derive(Debug)]
+struct SumGroupsAccumulator<T>
+where
+ T: ArrowNumericType + Send,
+{
+ /// The type of the computed sum
+ sum_data_type: DataType,
+
+ /// The type of the returned sum
+ return_data_type: DataType,
+
+ /// Sums per group, stored as the native type
+ sums: Vec<T::Native>,
+
+ /// Track nulls in the input / filters
+ null_state: NullState,
+}
+
+impl<T> SumGroupsAccumulator<T>
+where
+ T: ArrowNumericType + Send,
+{
+ pub fn new(sum_data_type: &DataType, return_data_type: &DataType) -> Self {
+ debug!(
+ "SumGroupsAccumulator ({}, sum type: {sum_data_type:?}) -->
{return_data_type:?}",
+ std::any::type_name::<T>()
+ );
+
+ Self {
+ return_data_type: sum_data_type.clone(),
+ sum_data_type: sum_data_type.clone(),
+ sums: vec![],
+ null_state: NullState::new(),
+ }
+ }
+}
+
+impl<T> GroupsAccumulator for SumGroupsAccumulator<T>
+where
+ T: ArrowNumericType + Send,
+{
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&arrow_array::BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ assert_eq!(values.len(), 1, "single argument to update_batch");
+ let values = values.get(0).unwrap().as_primitive::<T>();
+
+ // update sums
+ self.sums
+ .resize_with(total_num_groups, || T::default_value());
+
+ // NullState dispatches / handles tracking nulls and groups that saw
no values
Review Comment:
I am quite pleased that most of the update logic for handling nulls /
filters is encapsulated (and tested) in `NullState::accumulate`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
