Speed up avg

alamb · Jul 5, 2023 · 6846970 · 6846970
1 parent ffd5cbe
commit 6846970
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 13 deletions.
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;
             },
         );
 

diff --git a/datafusion/physical-expr/src/aggregate/min_max.rs b/datafusion/physical-expr/src/aggregate/min_max.rs
@@ -37,7 +37,7 @@ use arrow::{
 };
 use arrow_array::cast::AsArray;
 use arrow_array::types::{
-    ArrowPrimitiveType, Decimal128Type, Float32Type, Float64Type, UInt32Type, UInt64Type,
+    Decimal128Type, Float32Type, Float64Type, UInt32Type, UInt64Type,
 };
 use arrow_array::{ArrowNumericType, PrimitiveArray};
 use datafusion_common::ScalarValue;