refactor: Mutil Column Aggregate Function State Serialization Interface (#18398)

Author: forsaken628

src/common/base/src/runtime/profile/profile.rs

@@ -23,7 +23,7 @@ use crate::runtime::metrics::ScopedRegistry;
 use crate::runtime::profile::ProfileStatisticsName;
 use crate::runtime::ThreadTracker;
 
-#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct ProfileLabel {
     pub name: String,
     pub value: Vec<String>,

src/common/storage/src/copy.rs

@@ -20,7 +20,7 @@ use serde::Deserialize;
 use serde::Serialize;
 use thiserror::Error;
 
-#[derive(Default, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct CopyStatus {
     /// Key is file path.
     pub files: DashMap<String, FileStatus>,
@@ -45,7 +45,7 @@ impl CopyStatus {
     }
 }
 
-#[derive(Default, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct FileStatus {
     pub num_rows_loaded: usize,
     pub error: Option<FileErrorsInfo>,
@@ -79,7 +79,7 @@ impl FileStatus {
     }
 }
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct FileErrorsInfo {
     pub num_errors: usize,
     pub first_error: FileParseErrorAtLine,
@@ -156,7 +156,7 @@ impl FileParseError {
     }
 }
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct FileParseErrorAtLine {
     pub error: FileParseError,
     pub line: usize,

src/common/storage/src/merge.rs

@@ -15,7 +15,7 @@
 use serde::Deserialize;
 use serde::Serialize;
 
-#[derive(Default, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct MutationStatus {
     pub insert_rows: u64,
     pub deleted_rows: u64,

src/query/catalog/src/plan/agg_index.rs

@@ -51,7 +51,7 @@ pub struct AggIndexInfo {
 }
 
 /// This meta just indicate the block is from aggregating index.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub struct AggIndexMeta {
     pub is_agg: bool,
     // Number of aggregation functions.
@@ -75,6 +75,6 @@ local_block_meta_serde!(AggIndexMeta);
 #[typetag::serde(name = "agg_index_meta")]
 impl BlockMetaInfo for AggIndexMeta {
     fn clone_self(&self) -> Box<dyn BlockMetaInfo> {
-        Box::new(self.clone())
+        Box::new(*self)
     }
 }

src/query/catalog/src/statistics/data_cache_statistics.rs

@@ -24,7 +24,7 @@ pub struct DataCacheMetrics {
     bytes_from_memory: AtomicUsize,
 }
 
-#[derive(Default, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct DataCacheMetricValues {
     pub bytes_from_remote_disk: usize,
     pub bytes_from_local_disk: usize,

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

@@ -22,13 +22,13 @@ use super::AggrState;
 use super::AggrStateLoc;
 use super::AggrStateRegistry;
 use super::StateAddr;
-use crate::types::BinaryType;
 use crate::types::DataType;
 use crate::BlockEntry;
 use crate::ColumnBuilder;
-use crate::ColumnView;
 use crate::ProjectedBlock;
 use crate::Scalar;
+use crate::StateSerdeItem;
+use crate::StateSerdeType;
 
 pub type AggregateFunctionRef = Arc<dyn AggregateFunction>;
 
@@ -69,35 +69,28 @@ pub trait AggregateFunction: fmt::Display + Sync + Send {
     // Used in aggregate_null_adaptor
     fn accumulate_row(&self, place: AggrState, columns: ProjectedBlock, row: usize) -> Result<()>;
 
-    fn serialize(&self, place: AggrState, writer: &mut Vec<u8>) -> Result<()>;
+    fn serialize_type(&self) -> Vec<StateSerdeItem>;
 
-    fn serialize_size_per_row(&self) -> Option<usize> {
-        None
+    fn serialize_data_type(&self) -> DataType {
+        let serde_type = StateSerdeType::new(self.serialize_type());
+        serde_type.data_type()
     }
 
-    fn merge(&self, place: AggrState, reader: &mut &[u8]) -> Result<()>;
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()>;
 
-    /// Batch merge and deserialize the state from binary array
+    /// Batch deserialize the state and merge
     fn batch_merge(
         &self,
         places: &[StateAddr],
         loc: &[AggrStateLoc],
-        state: &ColumnView<BinaryType>,
-    ) -> Result<()> {
-        for (place, mut data) in places.iter().zip(state.iter()) {
-            self.merge(AggrState::new(*place, loc), &mut data)?;
-        }
-
-        Ok(())
-    }
-
-    fn batch_merge_single(&self, place: AggrState, state: &BlockEntry) -> Result<()> {
-        let view = state.downcast::<BinaryType>().unwrap();
-        for mut data in view.iter() {
-            self.merge(place, &mut data)?;
-        }
-        Ok(())
-    }
+        state: &BlockEntry,
+        filter: Option<&Bitmap>,
+    ) -> Result<()>;
 
     fn batch_merge_states(
         &self,
@@ -149,9 +142,4 @@ pub trait AggregateFunction: fmt::Display + Sync + Send {
     fn get_if_condition(&self, _columns: ProjectedBlock) -> Option<Bitmap> {
         None
     }
-
-    // some features
-    fn convert_const_to_full(&self) -> bool {
-        true
-    }
 }

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

@@ -20,6 +20,8 @@ use enum_as_inner::EnumAsInner;
 
 use super::AggregateFunctionRef;
 use crate::types::binary::BinaryColumnBuilder;
+use crate::types::DataType;
+use crate::ColumnBuilder;
 
 #[derive(Clone, Copy, Debug)]
 pub struct StateAddr {
@@ -113,11 +115,11 @@ impl From<StateAddr> for usize {
 
 pub fn get_states_layout(funcs: &[AggregateFunctionRef]) -> Result<StatesLayout> {
     let mut registry = AggrStateRegistry::default();
-    let mut serialize_size = Vec::with_capacity(funcs.len());
+    let mut serialize_type = Vec::with_capacity(funcs.len());
     for func in funcs {
         func.register_state(&mut registry);
         registry.commit();
-        serialize_size.push(func.serialize_size_per_row());
+        serialize_type.push(StateSerdeType(func.serialize_type().into()));
     }
 
     let AggrStateRegistry { states, offsets } = registry;
@@ -132,7 +134,7 @@ pub fn get_states_layout(funcs: &[AggregateFunctionRef]) -> Result<StatesLayout>
     Ok(StatesLayout {
         layout,
         states_loc,
-        serialize_size,
+        serialize_type,
     })
 }
 
@@ -191,18 +193,66 @@ impl AggrStateLoc {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum StateSerdeItem {
+    DataType(DataType),
+    Binary(Option<usize>),
+}
+
+#[derive(Debug, Clone)]
+pub struct StateSerdeType(Box<[StateSerdeItem]>);
+
+impl StateSerdeType {
+    pub fn new(items: impl Into<Box<[StateSerdeItem]>>) -> Self {
+        StateSerdeType(items.into())
+    }
+
+    pub fn data_type(&self) -> DataType {
+        DataType::Tuple(
+            self.0
+                .iter()
+                .map(|item| match item {
+                    StateSerdeItem::DataType(data_type) => data_type.clone(),
+                    StateSerdeItem::Binary(_) => DataType::Binary,
+                })
+                .collect(),
+        )
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct StatesLayout {
     pub layout: Layout,
     pub states_loc: Vec<Box<[AggrStateLoc]>>,
-    serialize_size: Vec<Option<usize>>,
+    pub(super) serialize_type: Vec<StateSerdeType>,
 }
 
 impl StatesLayout {
-    pub fn serialize_builders(&self, num_rows: usize) -> Vec<BinaryColumnBuilder> {
-        self.serialize_size
+    pub fn num_aggr_func(&self) -> usize {
+        self.states_loc.len()
+    }
+
+    pub fn serialize_builders(&self, num_rows: usize) -> Vec<ColumnBuilder> {
+        self.serialize_type
             .iter()
-            .map(|size| BinaryColumnBuilder::with_capacity(num_rows, num_rows * size.unwrap_or(0)))
+            .map(|serde_type| {
+                let builder = serde_type
+                    .0
+                    .iter()
+                    .map(|item| match item {
+                        StateSerdeItem::DataType(data_type) => {
+                            ColumnBuilder::with_capacity(data_type, num_rows)
+                        }
+                        StateSerdeItem::Binary(size) => {
+                            ColumnBuilder::Binary(BinaryColumnBuilder::with_capacity(
+                                num_rows,
+                                num_rows * size.unwrap_or(0),
+                            ))
+                        }
+                    })
+                    .collect();
+                ColumnBuilder::Tuple(builder)
+            })
             .collect()
     }
 }

src/query/expression/src/aggregate/aggregate_hashtable.rs

@@ -27,7 +27,6 @@ use crate::aggregate::payload_row::row_match_columns;
 use crate::group_hash_columns;
 use crate::new_sel;
 use crate::read;
-use crate::types::BinaryType;
 use crate::types::DataType;
 use crate::AggregateFunctionRef;
 use crate::BlockEntry;
@@ -219,7 +218,7 @@ impl AggregateHashTable {
                     .zip(agg_states.iter())
                     .zip(states_layout.states_loc.iter())
                 {
-                    func.batch_merge(state_places, loc, &state.downcast::<BinaryType>().unwrap())?;
+                    func.batch_merge(state_places, loc, state, None)?;
                 }
             }
         }

src/query/expression/src/aggregate/payload.rs

@@ -421,11 +421,15 @@ impl Payload {
         true
     }
 
-    pub fn empty_block(&self, fake_rows: Option<usize>) -> DataBlock {
-        let fake_rows = fake_rows.unwrap_or(0);
-        let entries = (0..self.aggrs.len())
-            .map(|_| {
-                ColumnBuilder::repeat_default(&DataType::Binary, fake_rows)
+    pub fn empty_block(&self, fake_rows: usize) -> DataBlock {
+        assert_eq!(self.aggrs.is_empty(), self.states_layout.is_none());
+        let entries = self
+            .states_layout
+            .as_ref()
+            .iter()
+            .flat_map(|layout| layout.serialize_type.iter())
+            .map(|serde_type| {
+                ColumnBuilder::repeat_default(&serde_type.data_type(), fake_rows)
                     .build()
                     .into()
             })

src/query/expression/src/aggregate/payload_flush.rs

@@ -18,7 +18,6 @@ use databend_common_io::prelude::bincode_deserialize_from_slice;
 use super::partitioned_payload::PartitionedPayload;
 use super::payload::Payload;
 use super::probe_state::ProbeState;
-use super::AggrState;
 use crate::read;
 use crate::types::binary::BinaryColumn;
 use crate::types::binary::BinaryColumnBuilder;
@@ -125,7 +124,7 @@ impl Payload {
         }
 
         if blocks.is_empty() {
-            return Ok(self.empty_block(None));
+            return Ok(self.empty_block(0));
         }
         DataBlock::concat(&blocks)
     }
@@ -141,26 +140,17 @@ impl Payload {
         if let Some(state_layout) = self.states_layout.as_ref() {
             let mut builders = state_layout.serialize_builders(row_count);
 
-            for place in state.state_places.as_slice()[0..row_count].iter() {
-                for (idx, (loc, func)) in state_layout
-                    .states_loc
-                    .iter()
-                    .zip(self.aggrs.iter())
-                    .enumerate()
-                {
-                    {
-                        let builder = &mut builders[idx];
-                        func.serialize(AggrState::new(*place, loc), &mut builder.data)?;
-                        builder.commit_row();
-                    }
-                }
+            for ((loc, func), builder) in state_layout
+                .states_loc
+                .iter()
+                .zip(self.aggrs.iter())
+                .zip(builders.iter_mut())
+            {
+                let builders = builder.as_tuple_mut().unwrap().as_mut_slice();
+                func.batch_serialize(&state.state_places.as_slice()[0..row_count], loc, builders)?;
             }
 
-            entries.extend(
-                builders
-                    .into_iter()
-                    .map(|builder| Column::Binary(builder.build()).into()),
-            );
+            entries.extend(builders.into_iter().map(|builder| builder.build().into()));
         }
 
         entries.extend_from_slice(&state.take_group_columns());
@@ -177,7 +167,7 @@ impl Payload {
         }
 
         if blocks.is_empty() {
-            return Ok(self.empty_block(None));
+            return Ok(self.empty_block(0));
         }
 
         DataBlock::concat(&blocks)