feat: support agg state combinator

src/query/functions/src/aggregates/aggregate_combinator_state.rs

@@ -0,0 +1,141 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::alloc::Layout;
+use std::fmt;
+use std::sync::Arc;
+
+use common_arrow::arrow::bitmap::Bitmap;
+use common_exception::Result;
+use common_expression::types::DataType;
+use common_expression::Column;
+use common_expression::ColumnBuilder;
+use common_expression::Scalar;
+
+use super::StateAddr;
+use crate::aggregates::aggregate_function_factory::AggregateFunctionCreator;
+use crate::aggregates::aggregate_function_factory::CombinatorDescription;
+use crate::aggregates::AggregateFunction;
+use crate::aggregates::AggregateFunctionRef;
+
+#[derive(Clone)]
+pub struct AggregateStateCombinator {
+    name: String,
+    nested: AggregateFunctionRef,
+}
+
+impl AggregateStateCombinator {
+    pub fn try_create(
+        nested_name: &str,
+        params: Vec<Scalar>,
+        arguments: Vec<DataType>,
+        nested_creator: &AggregateFunctionCreator,
+    ) -> Result<AggregateFunctionRef> {
+        let arg_name = arguments
+            .iter()
+            .map(|x| x.to_string())
+            .collect::<Vec<_>>()
+            .join(", ");
+
+        let name = format!("StateCombinator({nested_name}, {arg_name})");
+
+        let nested = nested_creator(nested_name, params, arguments)?;
+
+        Ok(Arc::new(AggregateStateCombinator { name, nested }))
+    }
+
+    pub fn combinator_desc() -> CombinatorDescription {
+        CombinatorDescription::creator(Box::new(Self::try_create))
+    }
+}
+
+impl AggregateFunction for AggregateStateCombinator {
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn return_type(&self) -> Result<DataType> {
+        Ok(DataType::String)
+    }
+
+    fn init_state(&self, place: StateAddr) {
+        self.nested.init_state(place);
+    }
+
+    fn is_state(&self) -> bool {
+        true
+    }
+
+    fn state_layout(&self) -> Layout {
+        self.nested.state_layout()
+    }
+
+    fn accumulate(
+        &self,
+        place: StateAddr,
+        columns: &[Column],
+        validity: Option<&Bitmap>,
+        input_rows: usize,
+    ) -> Result<()> {
+        self.nested.accumulate(place, columns, validity, input_rows)
+    }
+
+    fn accumulate_keys(
+        &self,
+        places: &[StateAddr],
+        offset: usize,
+        columns: &[Column],
+        input_rows: usize,
+    ) -> Result<()> {
+        self.nested
+            .accumulate_keys(places, offset, columns, input_rows)
+    }
+
+    fn accumulate_row(&self, place: StateAddr, columns: &[Column], row: usize) -> Result<()> {
+        self.nested.accumulate_row(place, columns, row)
+    }
+
+    fn serialize(&self, place: StateAddr, writer: &mut Vec<u8>) -> Result<()> {
+        self.nested.serialize(place, writer)
+    }
+
+    fn deserialize(&self, place: StateAddr, reader: &mut &[u8]) -> Result<()> {
+        self.nested.deserialize(place, reader)
+    }
+
+    fn merge(&self, place: StateAddr, rhs: StateAddr) -> Result<()> {
+        self.nested.merge(place, rhs)
+    }
+
+    fn merge_result(&self, place: StateAddr, builder: &mut ColumnBuilder) -> Result<()> {
+        let str_builder = builder.as_string_mut().unwrap();
+        self.serialize(place, &mut str_builder.data)?;
+        str_builder.commit_row();
+        Ok(())
+    }
+
+    fn need_manual_drop_state(&self) -> bool {
+        self.nested.need_manual_drop_state()
+    }
+
+    unsafe fn drop_state(&self, place: StateAddr) {
+        self.nested.drop_state(place);
+    }
+}
+
+impl fmt::Display for AggregateStateCombinator {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.name)
+    }
+}

src/query/functions/src/aggregates/aggregate_function.rs

@@ -34,6 +34,11 @@ pub trait AggregateFunction: fmt::Display + Sync + Send {
     fn return_type(&self) -> Result<DataType>;
 
     fn init_state(&self, place: StateAddr);
+
+    fn is_state(&self) -> bool {
+        false
+    }
+
     fn state_layout(&self) -> Layout;
 
     // accumulate is to accumulate the arrays in batch mode

src/query/functions/src/aggregates/aggregator.rs

@@ -18,6 +18,7 @@ use super::aggregate_arg_min_max::aggregate_arg_min_function_desc;
 use super::aggregate_avg::aggregate_avg_function_desc;
 use super::aggregate_combinator_distinct::aggregate_combinator_distinct_desc;
 use super::aggregate_combinator_distinct::aggregate_combinator_uniq_desc;
+use super::aggregate_combinator_state::AggregateStateCombinator;
 use super::aggregate_covariance::aggregate_covariance_population_desc;
 use super::aggregate_covariance::aggregate_covariance_sample_desc;
 use super::aggregate_min_max_any::aggregate_any_function_desc;
@@ -81,5 +82,6 @@ impl Aggregators {
     pub fn register_combinator(factory: &mut AggregateFunctionFactory) {
         factory.register_combinator("_if", AggregateIfCombinator::combinator_desc());
         factory.register_combinator("_distinct", aggregate_combinator_distinct_desc());
+        factory.register_combinator("_state", AggregateStateCombinator::combinator_desc());
     }
 }

src/query/functions/src/aggregates/mod.rs

@@ -25,6 +25,7 @@ mod aggregate_array_agg;
 mod aggregate_avg;
 mod aggregate_combinator_distinct;
 mod aggregate_combinator_if;
+mod aggregate_combinator_state;
 mod aggregate_covariance;
 mod aggregate_distinct_state;
 mod aggregate_kurtosis;

tests/sqllogictests/suites/query/02_function/02_0000_function_aggregate_state

@@ -0,0 +1,10 @@
+query T
+select length(max_state(number)) from numbers(100);
+----
+2
+
+query I
+select length(sum_state(number)) from numbers(10000);
+----
+5
+