feat: string_agg is compatible with multiple types as arguments instead of just String (#17570)

* feat: `string_agg` is compatible with multiple types as arguments instead of just `String`

Signed-off-by: Kould <[email protected]>

* chore: add type check for `string_agg`

Signed-off-by: Kould <[email protected]>

---------

Signed-off-by: Kould <[email protected]>

Author: KKould

src/query/functions/src/aggregates/aggregate_string_agg.rs

@@ -27,8 +27,13 @@ use databend_common_expression::types::ValueType;
 use databend_common_expression::AggrStateRegistry;
 use databend_common_expression::AggrStateType;
 use databend_common_expression::ColumnBuilder;
+use databend_common_expression::DataBlock;
+use databend_common_expression::EvaluateOptions;
+use databend_common_expression::Evaluator;
+use databend_common_expression::FunctionContext;
 use databend_common_expression::InputColumns;
 use databend_common_expression::Scalar;
+use databend_common_expression::Value;
 
 use super::aggregate_function_factory::AggregateFunctionDescription;
 use super::borsh_deserialize_state;
@@ -39,6 +44,7 @@ use crate::aggregates::assert_variadic_arguments;
 use crate::aggregates::AggrState;
 use crate::aggregates::AggrStateLoc;
 use crate::aggregates::AggregateFunction;
+use crate::BUILTIN_FUNCTIONS;
 
 #[derive(BorshSerialize, BorshDeserialize, Debug)]
 pub struct StringAggState {
@@ -49,6 +55,7 @@ pub struct StringAggState {
 pub struct AggregateStringAggFunction {
     display_name: String,
     delimiter: String,
+    value_type: DataType,
 }
 
 impl AggregateFunction for AggregateStringAggFunction {
@@ -77,7 +84,22 @@ impl AggregateFunction for AggregateStringAggFunction {
         validity: Option<&Bitmap>,
         _input_rows: usize,
     ) -> Result<()> {
-        let column = StringType::try_downcast_column(&columns[0]).unwrap();
+        let column = if self.value_type != DataType::String {
+            let block = DataBlock::new_from_columns(vec![columns[0].clone()]);
+            let func_ctx = &FunctionContext::default();
+            let evaluator = Evaluator::new(&block, func_ctx, &BUILTIN_FUNCTIONS);
+            let value = evaluator.run_cast(
+                None,
+                &self.value_type,
+                &DataType::String,
+                Value::Column(columns[0].clone()),
+                None,
+                &mut EvaluateOptions::default(),
+            )?;
+            StringType::try_downcast_column(value.as_column().unwrap()).unwrap()
+        } else {
+            StringType::try_downcast_column(&columns[0]).unwrap()
+        };
         let state = place.get::<StringAggState>();
         match validity {
             Some(validity) => {
@@ -175,10 +197,15 @@ impl fmt::Display for AggregateStringAggFunction {
 }
 
 impl AggregateStringAggFunction {
-    fn try_create(display_name: &str, delimiter: String) -> Result<Arc<dyn AggregateFunction>> {
+    fn try_create(
+        display_name: &str,
+        delimiter: String,
+        value_type: DataType,
+    ) -> Result<Arc<dyn AggregateFunction>> {
         let func = AggregateStringAggFunction {
             display_name: display_name.to_string(),
             delimiter,
+            value_type,
         };
         Ok(Arc::new(func))
     }
@@ -191,19 +218,29 @@ pub fn try_create_aggregate_string_agg_function(
     _sort_descs: Vec<AggregateFunctionSortDesc>,
 ) -> Result<Arc<dyn AggregateFunction>> {
     assert_variadic_arguments(display_name, argument_types.len(), (1, 2))?;
-    // TODO:(b41sh) support other data types
-    if argument_types[0].remove_nullable() != DataType::String {
+    let value_type = argument_types[0].remove_nullable();
+    if !matches!(
+        value_type,
+        DataType::Boolean
+            | DataType::String
+            | DataType::Number(_)
+            | DataType::Decimal(_)
+            | DataType::Timestamp
+            | DataType::Date
+            | DataType::Variant
+            | DataType::Interval
+    ) {
         return Err(ErrorCode::BadDataValueType(format!(
-            "The argument of aggregate function {} must be string",
-            display_name
+            "{} does not support type '{:?}'",
+            display_name, value_type
         )));
     }
     let delimiter = if params.len() == 1 {
         params[0].as_string().unwrap().clone()
     } else {
         String::new()
     };
-    AggregateStringAggFunction::try_create(display_name, delimiter)
+    AggregateStringAggFunction::try_create(display_name, delimiter, value_type)
 }
 
 pub fn aggregate_string_agg_function_desc() -> AggregateFunctionDescription {

src/query/sql/src/planner/semantic/type_check.rs

@@ -1785,6 +1785,7 @@ impl<'a> TypeChecker<'a> {
                 )));
             }
             let _ = arguments.pop();
+            let _ = arg_types.pop();
             let delimiter = delimiter_value.unwrap();
             vec![delimiter.value]
         } else {

tests/sqllogictests/suites/query/functions/02_0000_function_aggregate_mix.test

@@ -354,14 +354,14 @@ select group_array_moving_sum(k), group_array_moving_sum(2)(v) from aggr;
 [1,3,5,7,9,11,13,15,17,19,21] [10,20,20,20,30,40,45,55,60,60,60]
 
 statement ok
-create table t3(s string null, x int null, y int null, z int null);
+create table t3(s string null, x int null, y int null, z int null, b boolean null, arr Array(Int32));
 
 statement ok
 insert into t3 values
-    ('abc', 3, 1, 2),
-    ('def', 2, 1, null),
-    (null, 1, 2, 1),
-    ('xyz', 0, 2, 0);
+    ('abc', 3, 1, 2, true, [1,2,3]),
+    ('def', 2, 1, null, false, [4,5,6]),
+    (null, 1, 2, 1, null, [7,8,9]),
+    ('xyz', 0, 2, 0, true, [10,11,12]);
 
 query T
 select string_agg(s) from t3;
@@ -378,6 +378,19 @@ select string_agg(s, '|') from t3;
 ----
 abc|def|xyz
 
+query T
+select string_agg(x + 1, '|') from t3;
+----
+4|3|2|1
+
+query T
+select string_agg(b, '|') from t3;
+----
+true|false|true
+
+statement error
+select listagg(arr) from t3;
+
 query T
 select listagg(s) from t3;
 ----