feat: Optimise convert_to_state for SUM and BIT_AND_OR

Author: Mark1626

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs

@@ -52,6 +52,12 @@ where
     /// The starting value for new groups
     starting_value: T::Native,
 
+    /// When true, `starting_value` is the identity element for `prim_fn`,
+    /// i.e. `prim_fn(starting_value, x) == x` for all x. This allows
+    /// `convert_to_state` to skip allocating an initial-state array and
+    /// the element-wise arithmetic, returning the input values directly.
+    starting_value_is_identity: bool,
+
     /// Track nulls in the input / filters
     null_state: NullState,
 
@@ -70,6 +76,7 @@ where
             data_type: data_type.clone(),
             null_state: NullState::new(),
             starting_value: T::default_value(),
+            starting_value_is_identity: false,
             prim_fn,
         }
     }
@@ -79,6 +86,12 @@ where
         self.starting_value = starting_value;
         self
     }
+
+    /// Mark that `starting_value` is the identity element
+    pub fn with_starting_value_as_identity(mut self) -> Self {
+        self.starting_value_is_identity = true;
+        self
+    }
 }
 
 impl<T, F> GroupsAccumulator for PrimitiveGroupsAccumulator<T, F>
@@ -150,10 +163,6 @@ where
     ) -> Result<Vec<ArrayRef>> {
         let values = values[0].as_primitive::<T>().clone();
 
-        // Initializing state with starting values
-        let initial_state =
-            PrimitiveArray::<T>::from_value(self.starting_value, values.len());
-
         // Recalculating values in case there is filter
         let values = match opt_filter {
             None => values,
@@ -175,6 +184,20 @@ where
             }
         };
 
+        // When starting_value is the identity element for prim_fn
+        // (e.g. 0 for SUM/BIT_OR/BIT_XOR), prim_fn(starting_value, x) == x,
+        // so we can skip allocating the initial_state array and the binary_mut
+        // arithmetic entirely — just return the (filtered) input values.
+        if self.starting_value_is_identity {
+            return Ok(vec![Arc::new(
+                values.with_data_type(self.data_type.clone()),
+            )]);
+        }
+
+        // For non-identity starting values (MIN, MAX, BIT_AND, etc.),
+        // apply the operation against the starting value array.
+        let initial_state =
+            PrimitiveArray::<T>::from_value(self.starting_value, values.len());
         let state_values = compute::binary_mut(initial_state, &values, |mut x, y| {
             (self.prim_fn)(&mut x, y);
             x

datafusion/functions-aggregate/src/bit_and_or_xor.rs

@@ -54,10 +54,12 @@ macro_rules! group_accumulator_helper {
                     .with_starting_value(!0),
             )),
             BitwiseOperationType::Or => Ok(Box::new(
-                PrimitiveGroupsAccumulator::<$t, _>::new($dt, |x, y| x.bitor_assign(y)),
+                PrimitiveGroupsAccumulator::<$t, _>::new($dt, |x, y| x.bitor_assign(y))
+                    .with_starting_value_as_identity(),
             )),
             BitwiseOperationType::Xor => Ok(Box::new(
-                PrimitiveGroupsAccumulator::<$t, _>::new($dt, |x, y| x.bitxor_assign(y)),
+                PrimitiveGroupsAccumulator::<$t, _>::new($dt, |x, y| x.bitxor_assign(y))
+                    .with_starting_value_as_identity(),
             )),
         }
     };

datafusion/functions-aggregate/src/sum.rs

@@ -292,7 +292,7 @@ impl AggregateUDFImpl for Sum {
                 Ok(Box::new(PrimitiveGroupsAccumulator::<$t, _>::new(
                     &$dt,
                     |x, y| *x = x.add_wrapping(y),
-                )))
+                ).with_starting_value_as_identity()))
             };
         }
         downcast_sum!(args, helper)