sum

Author: forsaken628

src/query/functions/src/aggregates/adaptors/aggregate_null_adaptor.rs

@@ -750,7 +750,7 @@ mod tests {
 
         fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
             let state = place.get::<u64>();
-            if let Some(ScalarRef::Number(NumberScalar::UInt64(value))) = data.get(0) {
+            if let Some(ScalarRef::Number(NumberScalar::UInt64(value))) = data.first() {
                 *state += value;
             }
             Ok(())
@@ -1031,7 +1031,7 @@ mod tests {
         {
             let values = vec![10u64, 20u64, 30u64, 40u64];
             let flags = vec![true, false, true, true];
-            let filter_bits = vec![true, false, true, false]; // Only process indices 0 and 2
+            let filter_bits = [true, false, true, false]; // Only process indices 0 and 2
             let filter = Bitmap::from_iter(filter_bits.iter().copied());
 
             test_both_strategies(&values, &flags, Some(&filter))?;
@@ -1040,7 +1040,7 @@ mod tests {
         {
             let values = vec![100u64, 200u64, 300u64, 400u64];
             let flags = vec![true, true, false, true]; // Skip index 2
-            let filter_bits = vec![true, false, true, true]; // Skip index 1
+            let filter_bits = [true, false, true, true]; // Skip index 1
             let filter = Bitmap::from_iter(filter_bits.iter().copied());
 
             test_both_strategies(&values, &flags, Some(&filter))?;

src/query/functions/src/aggregates/aggregate_sum.rs

@@ -32,7 +32,9 @@ use databend_common_expression::AggregateFunctionRef;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::ColumnBuilder;
 use databend_common_expression::Scalar;
+use databend_common_expression::ScalarRef;
 use databend_common_expression::StateAddr;
+use databend_common_expression::StateSerdeItem;
 use databend_common_expression::SELECTIVITY_THRESHOLD;
 use num_traits::AsPrimitive;
 
@@ -76,14 +78,14 @@ pub trait SumState: BorshSerialize + BorshDeserialize + Send + Sync + Default +
 
 #[derive(BorshSerialize, BorshDeserialize)]
 pub struct NumberSumState<N>
-where N: ValueType
+where N: ArgType
 {
     pub value: N::Scalar,
 }
 
 impl<N> Default for NumberSumState<N>
 where
-    N: ValueType,
+    N: ArgType,
     N::Scalar: Number + AsPrimitive<f64> + BorshSerialize + BorshDeserialize + std::ops::AddAssign,
 {
     fn default() -> Self {
@@ -130,7 +132,7 @@ where
 impl<T, N> UnaryState<T, N> for NumberSumState<N>
 where
     T: ArgType + Sync + Send,
-    N: ValueType,
+    N: ArgType,
     T::Scalar: Number + AsPrimitive<N::Scalar>,
     N::Scalar: Number + AsPrimitive<f64> + BorshSerialize + BorshDeserialize + std::ops::AddAssign,
     for<'a> T::ScalarRef<'a>: Number + AsPrimitive<N::Scalar>,
@@ -169,6 +171,20 @@ where
         builder.push_item(N::to_scalar_ref(&self.value));
         Ok(())
     }
+
+    fn serialize_type() -> Vec<StateSerdeItem> {
+        std::vec![StateSerdeItem::DataType(T::data_type())]
+    }
+
+    fn serialize_state(&self, builders: &mut [ColumnBuilder]) -> Result<()> {
+        N::downcast_builder(&mut builders[0]).push_item(N::to_scalar_ref(&self.value));
+        Ok(())
+    }
+
+    fn deserialize_state(data: &[ScalarRef]) -> Result<Self> {
+        let value = N::to_owned_scalar(N::try_downcast_scalar(&data[0]).unwrap());
+        Ok(Self { value })
+    }
 }
 
 #[derive(BorshDeserialize, BorshSerialize)]

src/query/functions/src/aggregates/aggregate_unary.rs

@@ -80,6 +80,22 @@ where
         builder: R::ColumnBuilderMut<'_>,
         function_data: Option<&dyn FunctionData>,
     ) -> Result<()>;
+
+    fn serialize_type() -> Vec<StateSerdeItem> {
+        vec![StateSerdeItem::Binary(None)]
+    }
+
+    fn serialize_state(&self, builders: &mut [ColumnBuilder]) -> Result<()> {
+        let binary_builder = builders[0].as_binary_mut().unwrap();
+        self.serialize(&mut binary_builder.data)?;
+        binary_builder.commit_row();
+        Ok(())
+    }
+
+    fn deserialize_state(data: &[ScalarRef]) -> Result<Self> {
+        let mut binary = *data[0].as_binary().unwrap();
+        Ok(Self::deserialize_reader(&mut binary)?)
+    }
 }
 
 pub trait FunctionData: Send + Sync {
@@ -230,21 +246,17 @@ where
     }
 
     fn serialize_type(&self) -> Vec<StateSerdeItem> {
-        vec![StateSerdeItem::Binary(None)]
+        S::serialize_type()
     }
 
     fn serialize(&self, place: AggrState, builders: &mut [ColumnBuilder]) -> Result<()> {
-        let binary_builder = builders[0].as_binary_mut().unwrap();
         let state: &mut S = place.get::<S>();
-        state.serialize(&mut binary_builder.data)?;
-        binary_builder.commit_row();
-        Ok(())
+        state.serialize_state(builders)
     }
 
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
-        let mut binary = *data[0].as_binary().unwrap();
         let state: &mut S = place.get::<S>();
-        let rhs = S::deserialize_reader(&mut binary)?;
+        let rhs = S::deserialize_state(data)?;
         state.merge(&rhs)
     }