chore(query): migrate v2 agg

Author: sundy-li

Cargo.lock

@@ -72,6 +72,15 @@ pub enum DataType {
     Generic(usize),
 }
 
+impl DataType {
+    pub fn wrap_nullable(&self) -> Self {
+        match self {
+            DataType::Nullable(_) => self.clone(),
+            _ => Self::Nullable(Box::new(self.clone())),
+        }
+    }
+}
+
 pub trait ValueType: Debug + Clone + PartialEq + Sized + 'static {
     type Scalar: Debug + Clone + PartialEq;
     type ScalarRef<'a>: Debug + Clone + PartialEq;

src/query/expression/src/types.rs

@@ -13,11 +13,16 @@ doctest = false
 # Workspace dependencies
 common-arrow = { path = "../../common/arrow" }
 common-base = { path = "../../common/base" }
+common-exception = { path = "../../common/exception" }
 common-expression = { path = "../expression" }
+common-hashtable = { path = "../../common/hashtable" }
+common-io = { path = "../../common/io" }
 
 # Crates.io dependencies
 base64 = "0.13.0"
 bstr = "0.2.17"
+bumpalo = "3.10.0"
+bytes = "1.1.0"
 crc32fast = "1.3.2"
 hex = "0.4.3"
 itertools = "0.10.3"
@@ -33,6 +38,7 @@ rand = { version = "0.8.5", features = ["small_rng"] }
 strength_reduce = "0.2.3"
 # TODO: Switch to jsonb. bson is used for placeholder.
 bson = "2.4"
+once_cell = "1.12.0"
 serde_json = "1.0"
 
 [dev-dependencies]

src/query/functions-v2/Cargo.toml

@@ -0,0 +1,142 @@
+// Copyright 2021 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::fmt;
+use std::sync::Arc;
+
+use common_arrow::arrow::bitmap::Bitmap;
+use common_exception::Result;
+use common_expression::Column;
+use common_expression::ColumnBuilder;
+use common_expression::Scalar;
+use common_expression::types::DataType;
+
+use crate::aggregates::AggregateFunction;
+use crate::aggregates::AggregateFunctionRef;
+use crate::aggregates::StateAddr;
+
+/// BasicAdaptor will convert all args into full column and apply the inner aggregation
+pub struct AggregateFunctionBasicAdaptor {
+    inner: AggregateFunctionRef,
+}
+
+impl AggregateFunctionBasicAdaptor {
+    pub fn create(inner: AggregateFunctionRef) -> Result<AggregateFunctionRef> {
+        Ok(Arc::new(AggregateFunctionBasicAdaptor { inner }))
+    }
+}
+
+impl AggregateFunction for AggregateFunctionBasicAdaptor {
+    fn name(&self) -> &str {
+        self.inner.name()
+    }
+
+    fn return_type(&self) -> Result<DataType> {
+        self.inner.return_type()
+    }
+
+    #[inline]
+    fn init_state(&self, place: StateAddr) {
+        self.inner.init_state(place)
+    }
+
+    #[inline]
+    fn state_layout(&self) -> std::alloc::Layout {
+        self.inner.state_layout()
+    }
+
+    #[inline]
+    fn accumulate(
+        &self,
+        place: StateAddr,
+        columns: &[Column],
+        validity: Option<&Bitmap>,
+        input_rows: usize,
+    ) -> Result<()> {
+        if self.inner.convert_const_to_full() && columns.iter().any(|c| c.is_const()) {
+            let columns: Vec<Column> = columns.iter().map(|c| c.convert_full_column()).collect();
+            self.inner.accumulate(place, &columns, validity, input_rows)
+        } else {
+            self.inner.accumulate(place, columns, validity, input_rows)
+        }
+    }
+
+    fn accumulate_keys(
+        &self,
+        places: &[StateAddr],
+        offset: usize,
+        columns: &[Column],
+        input_rows: usize,
+    ) -> Result<()> {
+        // if self.inner.convert_const_to_full() && columns.iter().any(|c| c.is_const()) {
+        //     let columns: Vec<Column> = columns.iter().map(|c| c.convert_full_column()).collect();
+        //     self.inner
+        //         .accumulate_keys(places, offset, &columns, input_rows)?;
+        // } else {
+        //     self.inner
+        //         .accumulate_keys(places, offset, columns, input_rows)?;
+        // }
+        Ok(())
+    }
+
+    #[inline]
+    fn accumulate_row(&self, place: StateAddr, columns: &[Column], row: usize) -> Result<()> {
+        self.inner.accumulate_row(place, columns, row)
+    }
+
+    #[inline]
+    fn serialize(&self, place: StateAddr, writer: &mut bytes::BytesMut) -> Result<()> {
+        self.inner.serialize(place, writer)
+    }
+
+    #[inline]
+    fn deserialize(&self, place: StateAddr, reader: &mut &[u8]) -> Result<()> {
+        self.inner.deserialize(place, reader)
+    }
+
+    fn merge(&self, place: StateAddr, rhs: StateAddr) -> Result<()> {
+        self.inner.merge(place, rhs)
+    }
+
+    fn merge_result(&self, place: StateAddr, array: &mut ColumnBuilder) -> Result<()> {
+        self.inner.merge_result(place, array)
+    }
+
+    fn get_own_null_adaptor(
+        &self,
+        nested_function: AggregateFunctionRef,
+        params: Vec<Scalar>,
+        arguments: Vec<DataType>,
+    ) -> Result<Option<AggregateFunctionRef>> {
+        self.inner
+            .get_own_null_adaptor(nested_function, params, arguments)
+    }
+
+    fn need_manual_drop_state(&self) -> bool {
+        self.inner.need_manual_drop_state()
+    }
+
+    unsafe fn drop_state(&self, place: StateAddr) {
+        self.inner.drop_state(place)
+    }
+
+    fn convert_const_to_full(&self) -> bool {
+        self.inner.convert_const_to_full()
+    }
+}
+impl fmt::Display for AggregateFunctionBasicAdaptor {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.inner)
+    }
+}

src/query/functions-v2/src/aggregates/adaptors/aggregate_base_adaptor.rs

@@ -0,0 +1,92 @@
+// Copyright 2021 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 common_expression::types::DataType;
+use common_expression::{Result, Scalar};
+
+use super::aggregate_null_variadic_adaptor::AggregateNullVariadicAdaptor;
+use super::AggregateNullUnaryAdaptor;
+use crate::aggregates::aggregate_function_factory::AggregateFunctionFeatures;
+use crate::aggregates::aggregate_null_result::AggregateNullResultFunction;
+use crate::aggregates::AggregateFunctionRef;
+
+#[derive(Clone)]
+pub struct AggregateFunctionCombinatorNull {}
+
+impl AggregateFunctionCombinatorNull {
+    pub fn transform_arguments(arguments: &[DataType]) -> Result<Vec<DataType>> {
+        let mut results = Vec::with_capacity(arguments.len());
+
+        for arg in arguments.iter() {
+            if arg.is_nullable() {
+                let ty = remove_nullable(arg.data_type());
+                results.push(DataType::new(arg.name(), ty));
+            } else {
+                results.push(arg.clone());
+            }
+        }
+        Ok(results)
+    }
+
+    pub fn transform_params(params: &[Scalar]) -> Result<Vec<Scalar>> {
+        Ok(params.to_owned())
+    }
+
+    pub fn try_create(
+        _name: &str,
+        params: Vec<Scalar>,
+        arguments: Vec<DataType>,
+        nested: AggregateFunctionRef,
+        properties: AggregateFunctionFeatures,
+    ) -> Result<AggregateFunctionRef> {
+        // has_null_types
+        if !arguments.is_empty()
+            && arguments
+                .iter()
+                .any(|f| f.data_type().data_type_id() == TypeID::Null)
+        {
+            if properties.returns_default_when_only_null {
+                return AggregateNullResultFunction::try_create(u64::to_data_type());
+            } else {
+                return AggregateNullResultFunction::try_create(NullType::new_impl());
+            }
+        }
+        let params = Self::transform_params(&params)?;
+        let arguments = Self::transform_arguments(&arguments)?;
+        let size = arguments.len();
+
+        // Some functions may have their own null adaptor
+        if let Some(null_adaptor) =
+            nested.get_own_null_adaptor(nested.clone(), params, arguments)?
+        {
+            return Ok(null_adaptor);
+        }
+
+        let return_type = nested.return_type()?;
+        let result_is_null =
+            !properties.returns_default_when_only_null && return_type.can_inside_nullable();
+
+        match size {
+            1 => match result_is_null {
+                true => Ok(AggregateNullUnaryAdaptor::<true>::create(nested)),
+                false => Ok(AggregateNullUnaryAdaptor::<false>::create(nested)),
+            },
+
+            _ => match result_is_null {
+                true => Ok(AggregateNullVariadicAdaptor::<true, true>::create(nested)),
+                false => Ok(AggregateNullVariadicAdaptor::<false, true>::create(nested)),
+            },
+        }
+    }
+}