feat(query): Support st_collect aggregate function (#16996)

feat(query): Support st_collect aggregate function
databendlabs · Dec 4, 2024 · 0343b57 · 0343b57
1 parent 2b4782d
commit 0343b57
Show file tree

Hide file tree

Showing 15 changed files with 1,187 additions and 20 deletions.
diff --git a/src/query/functions/src/aggregates/aggregate_array_agg.rs b/src/query/functions/src/aggregates/aggregate_array_agg.rs
@@ -15,6 +15,7 @@
 use std::alloc::Layout;
 use std::fmt;
 use std::marker::PhantomData;
+use std::mem;
 use std::sync::Arc;
 
 use borsh::BorshDeserialize;
@@ -101,8 +102,9 @@ where
         match inner_type.remove_nullable() {
             DataType::Decimal(decimal_type) => {
                 let size = decimal_type.size();
-                for value in &self.values {
-                    let val = T::upcast_scalar(value.clone());
+                let values = mem::take(&mut self.values);
+                for value in values.into_iter() {
+                    let val = T::upcast_scalar(value);
                     let decimal_val = val.as_decimal().unwrap();
                     let new_val = match decimal_val {
                         DecimalScalar::Decimal128(v, _) => {
@@ -116,8 +118,9 @@ where
                 }
             }
             _ => {
-                for value in &self.values {
-                    let val = T::upcast_scalar(value.clone());
+                let values = mem::take(&mut self.values);
+                for value in values.into_iter() {
+                    let val = T::upcast_scalar(value);
                     inner_builder.push(val.as_ref());
                 }
             }

diff --git a/src/query/functions/src/aggregates/aggregate_function_factory.rs b/src/query/functions/src/aggregates/aggregate_function_factory.rs
@@ -26,6 +26,17 @@ use super::AggregateFunctionOrNullAdaptor;
 use crate::aggregates::AggregateFunctionRef;
 use crate::aggregates::Aggregators;
 
+// The NULL value in the those function needs to be handled separately.
+const NEED_NULL_AGGREGATE_FUNCTIONS: [&str; 7] = [
+    "array_agg",
+    "list",
+    "json_array_agg",
+    "json_object_agg",
+    "group_array_moving_avg",
+    "group_array_moving_sum",
+    "st_collect",
+];
+
 const STATE_SUFFIX: &str = "_state";
 
 pub type AggregateFunctionCreator =
@@ -172,15 +183,8 @@ impl AggregateFunctionFactory {
     ) -> Result<AggregateFunctionRef> {
         let name = name.as_ref();
         let mut features = AggregateFunctionFeatures::default();
-        // The NULL value in the array_agg function needs to be added to the returned array column,
-        // so handled separately.
-        if name == "array_agg"
-            || name == "list"
-            || name == "json_array_agg"
-            || name == "json_object_agg"
-            || name == "group_array_moving_avg"
-            || name == "group_array_moving_sum"
-        {
+
+        if NEED_NULL_AGGREGATE_FUNCTIONS.contains(&name) {
             let agg = self.get_impl(name, params, arguments, &mut features)?;
             return Ok(agg);
         }

diff --git a/src/query/functions/src/aggregates/aggregate_json_array_agg.rs b/src/query/functions/src/aggregates/aggregate_json_array_agg.rs
@@ -15,6 +15,7 @@
 use std::alloc::Layout;
 use std::fmt;
 use std::marker::PhantomData;
+use std::mem;
 use std::sync::Arc;
 
 use borsh::BorshDeserialize;
@@ -102,8 +103,9 @@ where
     fn merge_result(&mut self, builder: &mut ColumnBuilder) -> Result<()> {
         let tz = TimeZone::UTC;
         let mut items = Vec::with_capacity(self.values.len());
-        for value in &self.values {
-            let v = T::upcast_scalar(value.clone());
+        let values = mem::take(&mut self.values);
+        for value in values.into_iter() {
+            let v = T::upcast_scalar(value);
             // NULL values are omitted from the output.
             if v == Scalar::Null {
                 continue;

diff --git a/src/query/functions/src/aggregates/aggregate_json_object_agg.rs b/src/query/functions/src/aggregates/aggregate_json_object_agg.rs
@@ -16,6 +16,7 @@ use std::alloc::Layout;
 use std::collections::BTreeMap;
 use std::fmt;
 use std::marker::PhantomData;
+use std::mem;
 use std::sync::Arc;
 
 use borsh::BorshDeserialize;
@@ -166,19 +167,20 @@ where
 
     fn merge_result(&mut self, builder: &mut ColumnBuilder) -> Result<()> {
         let tz = TimeZone::UTC;
-        let mut kvs = Vec::with_capacity(self.kvs.len());
-        for (key, value) in &self.kvs {
-            let v = V::upcast_scalar(value.clone());
+        let mut values = Vec::with_capacity(self.kvs.len());
+        let kvs = mem::take(&mut self.kvs);
+        for (key, value) in kvs.into_iter() {
+            let v = V::upcast_scalar(value);
             // NULL values are omitted from the output.
             if v == Scalar::Null {
                 continue;
             }
             let mut val = vec![];
             cast_scalar_to_variant(v.as_ref(), &tz, &mut val);
-            kvs.push((key, val));
+            values.push((key, val));
         }
         let mut data = vec![];
-        jsonb::build_object(kvs.iter().map(|(k, v)| (k, &v[..])), &mut data).unwrap();
+        jsonb::build_object(values.iter().map(|(k, v)| (k, &v[..])), &mut data).unwrap();
 
         let object_value = Scalar::Variant(data);
         builder.push(object_value.as_ref());