refactor(query): adjust the serialization of each aggregate function state one by one (#18462)

Author: forsaken628

Cargo.lock

@@ -113,12 +113,17 @@ pub trait AggregateFunction: fmt::Display + Sync + Send {
         builder: &mut ColumnBuilder,
     ) -> Result<()> {
         for place in places {
-            self.merge_result(AggrState::new(*place, &loc), builder)?;
+            self.merge_result(AggrState::new(*place, &loc), false, builder)?;
         }
         Ok(())
     }
 
-    fn merge_result(&self, place: AggrState, builder: &mut ColumnBuilder) -> Result<()>;
+    fn merge_result(
+        &self,
+        place: AggrState,
+        read_only: bool,
+        builder: &mut ColumnBuilder,
+    ) -> Result<()>;
 
     // std::mem::needs_drop::<State>
     // if true will call drop_state

src/query/expression/src/aggregate/aggregate_function.rs

@@ -37,7 +37,8 @@ impl StateAddr {
     }
 
     #[inline]
-    pub fn get<'a, T>(&self) -> &'a mut T {
+    pub fn get<'a, T>(&self) -> &'a mut T
+    where T: Send + 'static {
         unsafe { &mut *(self.addr as *mut T) }
     }
 
@@ -71,15 +72,19 @@ impl StateAddr {
 
     #[inline]
     pub fn write<T, F>(&self, f: F)
-    where F: FnOnce() -> T {
+    where
+        F: FnOnce() -> T,
+        T: Send + 'static,
+    {
         unsafe {
             let ptr = self.addr as *mut T;
             std::ptr::write(ptr, f());
         }
     }
 
     #[inline]
-    pub fn write_state<T>(&self, state: T) {
+    pub fn write_state<T>(&self, state: T)
+    where T: Send + 'static {
         unsafe {
             let ptr = self.addr as *mut T;
             std::ptr::write(ptr, state);
@@ -199,6 +204,12 @@ pub enum StateSerdeItem {
     Binary(Option<usize>),
 }
 
+impl From<DataType> for StateSerdeItem {
+    fn from(value: DataType) -> Self {
+        Self::DataType(value)
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct StateSerdeType(Box<[StateSerdeItem]>);
 
@@ -268,15 +279,19 @@ impl<'a> AggrState<'a> {
         Self { addr, loc }
     }
 
-    pub fn get<'b, T>(&self) -> &'b mut T {
+    pub fn get<'b, T>(&self) -> &'b mut T
+    where T: Send + 'static {
         debug_assert_eq!(self.loc.len(), 1);
         self.addr
             .next(self.loc[0].into_custom().unwrap().1)
             .get::<T>()
     }
 
     pub fn write<T, F>(&self, f: F)
-    where F: FnOnce() -> T {
+    where
+        F: FnOnce() -> T,
+        T: Send + 'static,
+    {
         debug_assert_eq!(self.loc.len(), 1);
         self.addr
             .next(self.loc[0].into_custom().unwrap().1)

src/query/expression/src/aggregate/aggregate_function_state.rs

@@ -19,7 +19,7 @@ use std::slice::SliceIndex;
 use crate::BlockEntry;
 use crate::DataBlock;
 
-#[derive(Copy, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub struct ProjectedBlock<'a> {
     map: Option<&'a [usize]>,
     entries: &'a [BlockEntry],

src/query/expression/src/projected_block.rs

@@ -45,8 +45,6 @@ use std::ops::Deref;
 use std::ops::DerefMut;
 use std::ops::Range;
 
-use borsh::BorshDeserialize;
-use borsh::BorshSerialize;
 use databend_common_ast::ast::TypeName;
 pub use databend_common_base::base::OrderedFloat;
 pub use databend_common_io::deserialize_bitmap;
@@ -105,18 +103,7 @@ use crate::ScalarRef;
 
 pub type GenericMap = [DataType];
 
-#[derive(
-    Debug,
-    Clone,
-    PartialEq,
-    Eq,
-    Hash,
-    Serialize,
-    Deserialize,
-    BorshSerialize,
-    BorshDeserialize,
-    EnumAsInner,
-)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, EnumAsInner)]
 pub enum DataType {
     Null,
     EmptyArray,
@@ -494,7 +481,7 @@ pub fn convert_to_type_name(ty: &DataType) -> TypeName {
 
 /// [AccessType] defines a series of access methods for a data type
 pub trait AccessType: Debug + Clone + PartialEq + Sized + 'static {
-    type Scalar: Debug + Clone + PartialEq;
+    type Scalar: Debug + Clone + PartialEq + Send + 'static;
     type ScalarRef<'a>: Debug + Clone + PartialEq;
     type Column: Debug + Clone + PartialEq + Send;
     type Domain: Debug + Clone + PartialEq;

src/query/expression/src/types.rs

@@ -45,6 +45,7 @@ jiff = { workspace = true }
 jsonb = { workspace = true }
 lexical-core = { workspace = true }
 libm = { workspace = true }
+log = { workspace = true }
 match-template = { workspace = true }
 md-5 = { workspace = true }
 naive-cityhash = { workspace = true }

src/query/functions/Cargo.toml

@@ -19,11 +19,9 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_expression::types::number::NumberColumnBuilder;
-use databend_common_expression::types::BinaryType;
 use databend_common_expression::types::Bitmap;
 use databend_common_expression::types::DataType;
 use databend_common_expression::types::NumberDataType;
-use databend_common_expression::types::UnaryType;
 use databend_common_expression::with_number_mapped_type;
 use databend_common_expression::AggrStateRegistry;
 use databend_common_expression::AggrStateType;
@@ -38,16 +36,16 @@ use super::aggregate_distinct_state::AggregateDistinctState;
 use super::aggregate_distinct_state::AggregateDistinctStringState;
 use super::aggregate_distinct_state::AggregateUniqStringState;
 use super::aggregate_distinct_state::DistinctStateFunc;
-use super::aggregate_function::AggregateFunction;
-use super::aggregate_function_factory::AggregateFunctionCreator;
-use super::aggregate_function_factory::AggregateFunctionDescription;
-use super::aggregate_function_factory::AggregateFunctionSortDesc;
-use super::aggregate_function_factory::CombinatorDescription;
-use super::aggregator_common::assert_variadic_arguments;
+use super::assert_variadic_arguments;
+use super::AggrState;
+use super::AggrStateLoc;
 use super::AggregateCountFunction;
-use crate::aggregates::AggrState;
-use crate::aggregates::AggrStateLoc;
-use crate::aggregates::StateAddr;
+use super::AggregateFunction;
+use super::AggregateFunctionCreator;
+use super::AggregateFunctionDescription;
+use super::AggregateFunctionSortDesc;
+use super::CombinatorDescription;
+use super::StateAddr;
 
 #[derive(Clone)]
 pub struct AggregateDistinctCombinator<State> {
@@ -56,10 +54,12 @@ pub struct AggregateDistinctCombinator<State> {
     nested_name: String,
     arguments: Vec<DataType>,
     nested: Arc<dyn AggregateFunction>,
-    _state: PhantomData<State>,
+    _s: PhantomData<fn(State)>,
 }
 
-impl<State> AggregateDistinctCombinator<State> {
+impl<State> AggregateDistinctCombinator<State>
+where State: Send + 'static
+{
     fn get_state(place: AggrState) -> &mut State {
         place
             .addr
@@ -113,7 +113,7 @@ where State: DistinctStateFunc
     }
 
     fn serialize_type(&self) -> Vec<StateSerdeItem> {
-        vec![StateSerdeItem::Binary(None)]
+        State::serialize_type(None)
     }
 
     fn batch_serialize(
@@ -122,13 +122,7 @@ where State: DistinctStateFunc
         loc: &[AggrStateLoc],
         builders: &mut [ColumnBuilder],
     ) -> Result<()> {
-        let binary_builder = builders[0].as_binary_mut().unwrap();
-        for place in places {
-            let state = Self::get_state(AggrState::new(*place, loc));
-            state.serialize(&mut binary_builder.data)?;
-            binary_builder.commit_row();
-        }
-        Ok(())
+        State::batch_serialize(places, &loc[..1], builders)
     }
 
     fn batch_merge(
@@ -138,23 +132,7 @@ where State: DistinctStateFunc
         state: &BlockEntry,
         filter: Option<&Bitmap>,
     ) -> Result<()> {
-        let view = state.downcast::<UnaryType<BinaryType>>().unwrap();
-        let iter = places.iter().zip(view.iter());
-
-        if let Some(filter) = filter {
-            for (place, mut data) in iter.zip(filter.iter()).filter_map(|(v, b)| b.then_some(v)) {
-                let state = Self::get_state(AggrState::new(*place, loc));
-                let rhs = State::deserialize(&mut data)?;
-                state.merge(&rhs)?;
-            }
-        } else {
-            for (place, mut data) in iter {
-                let state = Self::get_state(AggrState::new(*place, loc));
-                let rhs = State::deserialize(&mut data)?;
-                state.merge(&rhs)?;
-            }
-        }
-        Ok(())
+        State::batch_merge(places, &loc[..1], state, filter)
     }
 
     fn merge_states(&self, place: AggrState, rhs: AggrState) -> Result<()> {
@@ -163,7 +141,12 @@ where State: DistinctStateFunc
         state.merge(other)
     }
 
-    fn merge_result(&self, place: AggrState, builder: &mut ColumnBuilder) -> Result<()> {
+    fn merge_result(
+        &self,
+        place: AggrState,
+        read_only: bool,
+        builder: &mut ColumnBuilder,
+    ) -> Result<()> {
         let state = Self::get_state(place);
         let nested_place = place.remove_first_loc();
 
@@ -178,13 +161,13 @@ where State: DistinctStateFunc
             Ok(())
         } else {
             if state.is_empty() {
-                return self.nested.merge_result(nested_place, builder);
+                return self.nested.merge_result(nested_place, read_only, builder);
             }
             let entries = &state.build_entries(&self.arguments).unwrap();
             self.nested
                 .accumulate(nested_place, entries.into(), None, state.len())?;
             // merge_result
-            self.nested.merge_result(nested_place, builder)
+            self.nested.merge_result(nested_place, read_only, builder)
         }
     }
 
@@ -220,7 +203,7 @@ pub fn aggregate_combinator_distinct_desc() -> CombinatorDescription {
 }
 
 pub fn aggregate_combinator_uniq_desc() -> AggregateFunctionDescription {
-    let features = super::aggregate_function_factory::AggregateFunctionFeatures {
+    let features = super::AggregateFunctionFeatures {
         returns_default_when_only_null: true,
         ..Default::default()
     };
@@ -264,7 +247,7 @@ pub fn try_create(
                         arguments,
                         nested,
                         name,
-                        _state: PhantomData,
+                        _s: PhantomData,
                     }));
                 }
             }),
@@ -277,7 +260,7 @@ pub fn try_create(
                         arguments,
                         nested,
                         nested_name: nested_name.to_owned(),
-                        _state: PhantomData,
+                        _s: PhantomData,
                     })),
                     _ => Ok(Arc::new(AggregateDistinctCombinator::<
                         AggregateDistinctStringState,
@@ -286,7 +269,7 @@ pub fn try_create(
                         arguments,
                         nested,
                         name,
-                        _state: PhantomData,
+                        _s: PhantomData,
                     })),
                 };
             }
@@ -300,6 +283,6 @@ pub fn try_create(
         arguments,
         nested,
         name,
-        _state: PhantomData,
+        _s: PhantomData,
     }))
 }

src/query/functions/src/aggregates/aggregate_combinator_distinct.rs renamed to src/query/functions/src/aggregates/adaptors/aggregate_combinator_distinct.rs

@@ -28,15 +28,15 @@ use databend_common_expression::ProjectedBlock;
 use databend_common_expression::Scalar;
 use databend_common_expression::StateSerdeItem;
 
+use super::AggrState;
+use super::AggrStateLoc;
+use super::AggregateFunction;
+use super::AggregateFunctionCreator;
+use super::AggregateFunctionRef;
+use super::AggregateFunctionSortDesc;
+use super::CombinatorDescription;
 use super::StateAddr;
-use crate::aggregates::aggregate_function_factory::AggregateFunctionCreator;
-use crate::aggregates::aggregate_function_factory::AggregateFunctionSortDesc;
-use crate::aggregates::aggregate_function_factory::CombinatorDescription;
-use crate::aggregates::AggrState;
-use crate::aggregates::AggrStateLoc;
-use crate::aggregates::AggregateFunction;
-use crate::aggregates::AggregateFunctionRef;
-use crate::aggregates::StateAddrs;
+use super::StateAddrs;
 
 #[derive(Clone)]
 pub struct AggregateIfCombinator {
@@ -181,8 +181,13 @@ impl AggregateFunction for AggregateIfCombinator {
         self.nested.merge_states(place, rhs)
     }
 
-    fn merge_result(&self, place: AggrState, builder: &mut ColumnBuilder) -> Result<()> {
-        self.nested.merge_result(place, builder)
+    fn merge_result(
+        &self,
+        place: AggrState,
+        read_only: bool,
+        builder: &mut ColumnBuilder,
+    ) -> Result<()> {
+        self.nested.merge_result(place, read_only, builder)
     }
 
     fn need_manual_drop_state(&self) -> bool {