Merge branch 'alamb/hash_agg_spike' of github.com:alamb/arrow-datafus…

…ion into alamb/hash_agg_spike
alamb · Jul 6, 2023 · b326b68 · b326b68
2 parents 49fc6c1 + 2f4907a
commit b326b68
Show file tree

Hide file tree

Showing 3 changed files with 300 additions and 18 deletions.
diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
diff --git a/datafusion/physical-expr/src/aggregate/average.rs b/datafusion/physical-expr/src/aggregate/average.rs
@@ -532,19 +532,8 @@ where
         assert_eq!(values.len(), 1, "single argument to update_batch");
         let values = values.get(0).unwrap().as_primitive::<T>();
 
-        // increment counts
+        // increment counts, update sums
         self.counts.resize(total_num_groups, 0);
-        self.null_state.accumulate(
-            group_indices,
-            values,
-            opt_filter,
-            total_num_groups,
-            |group_index, _new_value| {
-                self.counts[group_index] += 1;
-            },
-        );
-
-        // update sums
         self.sums.resize(total_num_groups, T::default_value());
         self.null_state.accumulate(
             group_indices,
@@ -554,6 +543,8 @@ where
             |group_index, new_value| {
                 let sum = &mut self.sums[group_index];
                 *sum = sum.add_wrapping(new_value);
+
+                self.counts[group_index] += 1;
             },
         );
 
@@ -591,7 +582,7 @@ where
             partial_sums,
             opt_filter,
             total_num_groups,
-            |group_index, new_value| {
+            |group_index, new_value: <T as ArrowPrimitiveType>::Native| {
                 let sum = &mut self.sums[group_index];
                 *sum = sum.add_wrapping(new_value);
             },