ee

Author: forsaken628

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

@@ -29,6 +29,7 @@ use crate::ProjectedBlock;
 use crate::Scalar;
 use crate::ScalarRef;
 use crate::StateSerdeItem;
+use crate::StateSerdeType;
 
 pub type AggregateFunctionRef = Arc<dyn AggregateFunction>;
 
@@ -71,6 +72,11 @@ pub trait AggregateFunction: fmt::Display + Sync + Send {
 
     fn serialize_type(&self) -> Vec<StateSerdeItem>;
 
+    fn serialize_data_type(&self) -> DataType {
+        let serde_type = StateSerdeType::new(self.serialize_type());
+        serde_type.data_type()
+    }
+
     fn serialize(&self, place: AggrState, builders: &mut [ColumnBuilder]) -> Result<()> {
         let binary_builder = builders[0].as_binary_mut().unwrap();
         self.serialize_binary(place, &mut binary_builder.data)?;

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

@@ -203,6 +203,10 @@ pub enum StateSerdeItem {
 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

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

@@ -22,6 +22,7 @@ use databend_common_expression::AggrStateRegistry;
 use databend_common_expression::ColumnBuilder;
 use databend_common_expression::ProjectedBlock;
 use databend_common_expression::Scalar;
+use databend_common_expression::ScalarRef;
 use databend_common_expression::StateSerdeItem;
 
 use super::AggregateFunctionFactory;
@@ -71,7 +72,7 @@ impl AggregateFunction for AggregateStateCombinator {
     }
 
     fn return_type(&self) -> Result<DataType> {
-        Ok(DataType::Binary)
+        Ok(self.nested.serialize_data_type())
     }
 
     fn init_state(&self, place: AggrState) {
@@ -108,27 +109,32 @@ impl AggregateFunction for AggregateStateCombinator {
     }
 
     fn serialize_type(&self) -> Vec<StateSerdeItem> {
-        vec![StateSerdeItem::Binary(None)]
+        self.nested.serialize_type()
     }
 
-    fn serialize_binary(&self, place: AggrState, writer: &mut Vec<u8>) -> Result<()> {
-        self.nested.serialize_binary(place, writer)
+    fn serialize(&self, place: AggrState, builders: &mut [ColumnBuilder]) -> Result<()> {
+        self.nested.serialize(place, builders)
     }
 
-    #[inline]
-    fn merge_binary(&self, place: AggrState, reader: &mut &[u8]) -> Result<()> {
-        self.nested.merge_binary(place, reader)
+    fn serialize_binary(&self, _: AggrState, _: &mut Vec<u8>) -> Result<()> {
+        unreachable!()
+    }
+
+    fn merge(&self, place: AggrState, data: &[ScalarRef]) -> Result<()> {
+        self.nested.merge(place, data)
+    }
+
+    fn merge_binary(&self, _: AggrState, _: &mut &[u8]) -> Result<()> {
+        unreachable!()
     }
 
     fn merge_states(&self, place: AggrState, rhs: AggrState) -> Result<()> {
         self.nested.merge_states(place, rhs)
     }
 
     fn merge_result(&self, place: AggrState, builder: &mut ColumnBuilder) -> Result<()> {
-        let builder = builder.as_binary_mut().unwrap();
-        self.nested.serialize_binary(place, &mut builder.data)?;
-        builder.commit_row();
-        Ok(())
+        let builders = builder.as_tuple_mut().unwrap().as_mut_slice();
+        self.nested.serialize(place, builders)
     }
 
     fn need_manual_drop_state(&self) -> bool {

src/query/sql/src/executor/physical_plans/physical_aggregate_partial.rs

@@ -17,7 +17,6 @@ use databend_common_expression::types::DataType;
 use databend_common_expression::DataField;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::DataSchemaRefExt;
-use databend_common_expression::StateSerdeItem;
 use databend_common_functions::aggregates::AggregateFunctionFactory;
 
 use super::SortDesc;
@@ -69,16 +68,7 @@ impl AggregatePartial {
                 )
                 .unwrap();
 
-            let tuple = func
-                .serialize_type()
-                .iter()
-                .map(|serde_type| match serde_type {
-                    StateSerdeItem::DataType(data_type) => data_type.clone(),
-                    StateSerdeItem::Binary(_) => DataType::Binary,
-                })
-                .collect();
-
-            fields.push(DataField::new(&name, DataType::Tuple(tuple)))
+            fields.push(DataField::new(&name, func.serialize_data_type()))
         }
 
         for (idx, field) in self.group_by.iter().zip(

src/query/sql/src/planner/optimizer/optimizers/operator/aggregate/stats_aggregate.rs

@@ -129,6 +129,11 @@ impl RuleStatsAggregateOptimizer {
                     for (need_rewrite_agg, agg) in
                         need_rewrite_aggs.iter().zip(agg.aggregate_functions.iter())
                     {
+                        if matches!(agg.scalar, ScalarExpr::UDAFCall(_)) {
+                            agg_results.push(agg.clone());
+                            continue;
+                        }
+
                         let agg_func = AggregateFunction::try_from(agg.scalar.clone())?;
 
                         if let Some((col_id, name)) = need_rewrite_agg {

src/query/sql/src/planner/optimizer/optimizers/rule/agg_rules/agg_index/query_rewrite.rs

@@ -25,6 +25,7 @@ use databend_common_expression::types::DataType;
 use databend_common_expression::Scalar;
 use databend_common_expression::TableField;
 use databend_common_expression::TableSchemaRefExt;
+use databend_common_functions::aggregates::AggregateFunctionFactory;
 use itertools::Itertools;
 use log::info;
 
@@ -306,6 +307,12 @@ impl QueryInfo {
                 }
                 return Ok(None);
             }
+            ScalarExpr::UDAFCall(udaf) => {
+                for arg in &udaf.arguments {
+                    self.check_output_cols(arg, index_output_cols, new_selection_set)?;
+                }
+                return Ok(None);
+            }
             ScalarExpr::UDFCall(udf) => {
                 let mut valid = true;
                 let mut new_args = Vec::with_capacity(udf.arguments.len());
@@ -361,23 +368,38 @@ impl ViewInfo {
         // query can use those columns to compute expressions.
         let mut index_fields = Vec::with_capacity(query_info.output_cols.len());
         let mut index_output_cols = HashMap::with_capacity(query_info.output_cols.len());
+        let factory = AggregateFunctionFactory::instance();
         for (index, item) in query_info.output_cols.iter().enumerate() {
             let display_name = format_scalar(&item.scalar, &query_info.column_map);
 
-            let mut is_agg = false;
-            if let Some(ref aggregate) = query_info.aggregate {
-                for agg_func in &aggregate.aggregate_functions {
-                    if item.index == agg_func.index {
-                        is_agg = true;
-                        break;
-                    }
+            let aggr_scalar_item = query_info.aggregate.as_ref().and_then(|aggregate| {
+                aggregate
+                    .aggregate_functions
+                    .iter()
+                    .find(|agg_func| agg_func.index == item.index)
+            });
+
+            let (data_type, is_agg) = match aggr_scalar_item {
+                Some(item) => {
+                    let func = match &item.scalar {
+                        ScalarExpr::AggregateFunction(func) => func,
+                        _ => unreachable!(),
+                    };
+                    let func = factory.get(
+                        &func.func_name,
+                        func.params.clone(),
+                        func.args
+                            .iter()
+                            .map(|arg| arg.data_type())
+                            .collect::<Result<_>>()?,
+                        func.sort_descs
+                            .iter()
+                            .map(|desc| desc.try_into())
+                            .collect::<Result<_>>()?,
+                    )?;
+                    (func.serialize_data_type(), true)
                 }
-            }
-            // we store the value of aggregate function as binary data.
-            let data_type = if is_agg {
-                DataType::Binary
-            } else {
-                item.scalar.data_type().unwrap()
+                None => (item.scalar.data_type().unwrap(), false),
             };
 
             let name = format!("{index}");
@@ -1299,6 +1321,15 @@ fn format_scalar(scalar: &ScalarExpr, column_map: &HashMap<IndexType, ScalarExpr
             }
             scalar
         }
+        ScalarExpr::UDAFCall(udaf) => {
+            let args = udaf
+                .arguments
+                .iter()
+                .map(|arg| format_scalar(arg, column_map))
+                .collect::<Vec<_>>()
+                .join(", ");
+            format!("{}({})", &udaf.name, args)
+        }
         ScalarExpr::UDFCall(udf) => format!(
             "{}({})",
             &udf.handler,

tests/sqllogictests/suites/ee/02_ee_aggregating_index/02_0001_agg_index_projected_scan.test

@@ -75,7 +75,42 @@ SELECT b, SUM(a) from t WHERE c > 1 GROUP BY b ORDER BY b
 1 1
 2 3
 
-query IIT
+query IIR
 SELECT MAX(a), MIN(b), AVG(c) from t
 ----
 2 1 3.5
+
+statement ok
+CREATE or REPLACE FUNCTION weighted_avg (INT, INT) STATE {sum INT, weight INT} RETURNS FLOAT
+LANGUAGE javascript AS $$
+export function create_state() {
+    return {sum: 0, weight: 0};
+}
+export function accumulate(state, value, weight) {
+    state.sum += value * weight;
+    state.weight += weight;
+    return state;
+}
+export function retract(state, value, weight) {
+    state.sum -= value * weight;
+    state.weight -= weight;
+    return state;
+}
+export function merge(state1, state2) {
+    state1.sum += state2.sum;
+    state1.weight += state2.weight;
+    return state1;
+}
+export function finish(state) {
+    return state.sum / state.weight;
+}
+$$;
+
+query IIR
+SELECT MAX(a), MIN(b), weighted_avg(a,b) from t group by b;
+----
+2 2 1.3333334
+1 1 1.0
+
+# fix me
+# SELECT MAX(a), MIN(b), weighted_avg(a,b) from t;