batch_serialize

Author: forsaken628

Cargo.lock

@@ -33,7 +33,6 @@ educe = { workspace = true }
 either = { workspace = true }
 enum-as-inner = { workspace = true }
 ethnum = { workspace = true, features = ["serde", "macros"] }
-fastrace = { workspace = true }
 futures = { workspace = true }
 geo = { workspace = true }
 geozero = { workspace = true }

src/query/expression/Cargo.toml

@@ -84,6 +84,18 @@ pub trait AggregateFunction: fmt::Display + Sync + Send {
 
     fn serialize(&self, place: AggrState, builders: &mut [ColumnBuilder]) -> Result<()>;
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        for place in places {
+            self.serialize(AggrState::new(*place, loc), builders)?;
+        }
+        Ok(())
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()>;
 
     /// Batch deserialize the state and merge

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

@@ -167,7 +167,6 @@ impl AggregateHashTable {
         }
     }
 
-    #[fastrace::trace(name = "AggregateHashTable::add_groups_inner")]
     // Add new groups and combine the states
     fn add_groups_inner(
         &mut self,
@@ -382,7 +381,6 @@ impl AggregateHashTable {
         Ok(())
     }
 
-    #[fastrace::trace(name = "AggregateHashTable::combine_payload")]
     pub fn combine_payload(
         &mut self,
         payload: &Payload,

src/query/expression/src/aggregate/aggregate_hashtable.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;
@@ -141,18 +140,14 @@ 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 builders = builders[idx].as_tuple_mut().unwrap().as_mut_slice();
-                        func.serialize(AggrState::new(*place, loc), builders)?;
-                    }
-                }
+            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| builder.build().into()));

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

@@ -191,6 +191,15 @@ impl<const NULLABLE_RESULT: bool> AggregateFunction for AggregateNullUnaryAdapto
         self.0.serialize(place, builders)
     }
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        self.0.batch_serialize(places, loc, builders)
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
         self.0.merge(place, data)
     }
@@ -318,6 +327,15 @@ impl<const NULLABLE_RESULT: bool> AggregateFunction
         self.0.serialize(place, builders)
     }
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        self.0.batch_serialize(places, loc, builders)
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
         self.0.merge(place, data)
     }
@@ -530,6 +548,29 @@ impl<const NULLABLE_RESULT: bool> CommonNullAdaptor<NULLABLE_RESULT> {
             .serialize(place.remove_last_loc(), &mut builders[..(n - 1)])
     }
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        if !NULLABLE_RESULT {
+            return self.nested.batch_serialize(places, loc, builders);
+        }
+        let n = builders.len();
+        debug_assert_eq!(self.nested.serialize_type().len() + 1, n);
+        let flag_builder = builders
+            .last_mut()
+            .and_then(ColumnBuilder::as_boolean_mut)
+            .unwrap();
+        for place in places {
+            let place = AggrState::new(*place, loc);
+            flag_builder.push(get_flag(place));
+        }
+        self.nested
+            .batch_serialize(places, &loc[..loc.len() - 1], &mut builders[..(n - 1)])
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
         if !NULLABLE_RESULT {
             return self.nested.merge(place, data);

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

@@ -40,23 +40,23 @@ use super::StateAddr;
 /// Use a single additional byte of data after the nested function data:
 /// 0 means there was no input, 1 means there was some.
 pub struct AggregateFunctionOrNullAdaptor {
-    inner: AggregateFunctionRef,
+    nested: AggregateFunctionRef,
     inner_nullable: bool,
 }
 
 impl AggregateFunctionOrNullAdaptor {
     pub fn create(
-        inner: AggregateFunctionRef,
+        nested: AggregateFunctionRef,
         features: AggregateFunctionFeatures,
     ) -> Result<AggregateFunctionRef> {
         // count/count distinct should not be nullable for empty set, just return zero
-        let inner_return_type = inner.return_type()?;
+        let inner_return_type = nested.return_type()?;
         if features.returns_default_when_only_null || inner_return_type == DataType::Null {
-            return Ok(inner);
+            return Ok(nested);
         }
 
         Ok(Arc::new(AggregateFunctionOrNullAdaptor {
-            inner,
+            nested,
             inner_nullable: inner_return_type.is_nullable(),
         }))
     }
@@ -83,22 +83,22 @@ fn flag_offset(place: AggrState) -> usize {
 
 impl AggregateFunction for AggregateFunctionOrNullAdaptor {
     fn name(&self) -> &str {
-        self.inner.name()
+        self.nested.name()
     }
 
     fn return_type(&self) -> Result<DataType> {
-        Ok(self.inner.return_type()?.wrap_nullable())
+        Ok(self.nested.return_type()?.wrap_nullable())
     }
 
     #[inline]
     fn init_state(&self, place: AggrState) {
         let c = place.addr.next(flag_offset(place)).get::<u8>();
         *c = 0;
-        self.inner.init_state(place.remove_last_loc())
+        self.nested.init_state(place.remove_last_loc())
     }
 
     fn register_state(&self, registry: &mut AggrStateRegistry) {
-        self.inner.register_state(registry);
+        self.nested.register_state(registry);
         registry.register(AggrStateType::Bool);
     }
 
@@ -114,7 +114,7 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
             return Ok(());
         }
 
-        let if_cond = self.inner.get_if_condition(columns);
+        let if_cond = self.nested.get_if_condition(columns);
 
         let validity = match (if_cond, validity) {
             (None, None) => None,
@@ -129,7 +129,7 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
             .unwrap_or(true)
         {
             set_flag(place, true);
-            self.inner.accumulate(
+            self.nested.accumulate(
                 place.remove_last_loc(),
                 columns,
                 validity.as_ref(),
@@ -146,9 +146,9 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
         columns: ProjectedBlock,
         input_rows: usize,
     ) -> Result<()> {
-        self.inner
+        self.nested
             .accumulate_keys(places, &loc[..loc.len() - 1], columns, input_rows)?;
-        let if_cond = self.inner.get_if_condition(columns);
+        let if_cond = self.nested.get_if_condition(columns);
 
         match if_cond {
             Some(v) if v.null_count() > 0 => {
@@ -175,14 +175,14 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
 
     #[inline]
     fn accumulate_row(&self, place: AggrState, columns: ProjectedBlock, row: usize) -> Result<()> {
-        self.inner
+        self.nested
             .accumulate_row(place.remove_last_loc(), columns, row)?;
         set_flag(place, true);
         Ok(())
     }
 
     fn serialize_type(&self) -> Vec<StateSerdeItem> {
-        self.inner
+        self.nested
             .serialize_type()
             .into_iter()
             .chain(Some(StateSerdeItem::DataType(DataType::Boolean)))
@@ -191,7 +191,7 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
 
     fn serialize(&self, place: AggrState, builders: &mut [ColumnBuilder]) -> Result<()> {
         let n = builders.len();
-        debug_assert_eq!(self.inner.serialize_type().len() + 1, n);
+        debug_assert_eq!(self.nested.serialize_type().len() + 1, n);
 
         let flag = get_flag(place);
         builders
@@ -200,13 +200,33 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
             .unwrap()
             .push(flag);
 
-        self.inner
+        self.nested
             .serialize(place.remove_last_loc(), &mut builders[..n - 1])
     }
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        let n = builders.len();
+        debug_assert_eq!(self.nested.serialize_type().len() + 1, n);
+        let flag_builder = builders
+            .last_mut()
+            .and_then(ColumnBuilder::as_boolean_mut)
+            .unwrap();
+        for place in places {
+            let place = AggrState::new(*place, loc);
+            flag_builder.push(get_flag(place));
+        }
+        self.nested
+            .batch_serialize(places, &loc[..loc.len() - 1], &mut builders[..(n - 1)])
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
         merge_flag(place, *data.last().and_then(ScalarRef::as_boolean).unwrap());
-        self.inner
+        self.nested
             .merge(place.remove_last_loc(), &data[..data.len() - 1])?;
         Ok(())
     }
@@ -233,7 +253,7 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
                     }
                 }
                 let inner_state = Column::Tuple(tuple[0..tuple.len() - 1].to_vec()).into();
-                self.inner
+                self.nested
                     .batch_merge(places, &loc[0..loc.len() - 1], &inner_state, filter)?;
             }
             _ => {
@@ -262,7 +282,7 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
     }
 
     fn merge_states(&self, place: AggrState, rhs: AggrState) -> Result<()> {
-        self.inner
+        self.nested
             .merge_states(place.remove_last_loc(), rhs.remove_last_loc())?;
         let flag = get_flag(place) || get_flag(rhs);
         set_flag(place, flag);
@@ -275,9 +295,9 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
                 if !get_flag(place) {
                     inner_mut.push_null();
                 } else if self.inner_nullable {
-                    self.inner.merge_result(place.remove_last_loc(), builder)?;
+                    self.nested.merge_result(place.remove_last_loc(), builder)?;
                 } else {
-                    self.inner
+                    self.nested
                         .merge_result(place.remove_last_loc(), &mut inner_mut.builder)?;
                     inner_mut.validity.push(true);
                 }
@@ -293,21 +313,21 @@ impl AggregateFunction for AggregateFunctionOrNullAdaptor {
         params: Vec<Scalar>,
         arguments: Vec<DataType>,
     ) -> Result<Option<AggregateFunctionRef>> {
-        self.inner
+        self.nested
             .get_own_null_adaptor(nested_function, params, arguments)
     }
 
     fn need_manual_drop_state(&self) -> bool {
-        self.inner.need_manual_drop_state()
+        self.nested.need_manual_drop_state()
     }
 
     unsafe fn drop_state(&self, place: AggrState) {
-        self.inner.drop_state(place.remove_last_loc())
+        self.nested.drop_state(place.remove_last_loc())
     }
 }
 
 impl fmt::Display for AggregateFunctionOrNullAdaptor {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}", self.inner)
+        write!(f, "{}", self.nested)
     }
 }

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

@@ -163,10 +163,29 @@ impl AggregateFunction for AggregateIfCombinator {
         self.nested.serialize(place, builders)
     }
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        self.nested.batch_serialize(places, loc, builders)
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
         self.nested.merge(place, data)
     }
 
+    fn batch_merge(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        state: &databend_common_expression::BlockEntry,
+        filter: Option<&Bitmap>,
+    ) -> Result<()> {
+        self.nested.batch_merge(places, loc, state, filter)
+    }
+
     fn merge_states(&self, place: AggrState, rhs: AggrState) -> Result<()> {
         self.nested.merge_states(place, rhs)
     }

src/query/functions/src/aggregates/aggregate_combinator_if.rs

@@ -117,6 +117,15 @@ impl AggregateFunction for AggregateStateCombinator {
         self.nested.serialize(place, builders)
     }
 
+    fn batch_serialize(
+        &self,
+        places: &[StateAddr],
+        loc: &[AggrStateLoc],
+        builders: &mut [ColumnBuilder],
+    ) -> Result<()> {
+        self.nested.batch_serialize(places, loc, builders)
+    }
+
     fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
         self.nested.merge(place, data)
     }