Dandandan commented on code in PR #6904:
URL: https://github.com/apache/arrow-datafusion/pull/6904#discussion_r1258634310
##########
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());
Review Comment:
`.resize` as well
--
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]
