feat(query): array function support variant type as argument (#17428)

Author: b41sh

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

@@ -1857,43 +1857,6 @@ impl<'a> TypeChecker<'a> {
         args: &[&Expr],
         lambda: &Lambda,
     ) -> Result<Box<(ScalarExpr, DataType)>> {
-        if func_name.starts_with("json_") && !args.is_empty() {
-            let target_type = if func_name.starts_with("json_array") {
-                TypeName::Array(Box::new(TypeName::Nullable(Box::new(TypeName::Variant))))
-            } else {
-                TypeName::Map {
-                    key_type: Box::new(TypeName::String),
-                    val_type: Box::new(TypeName::Nullable(Box::new(TypeName::Variant))),
-                }
-            };
-            let func_name = &func_name[5..];
-            let mut new_args: Vec<Expr> = args.iter().map(|v| (*v).to_owned()).collect();
-            new_args[0] = Expr::Cast {
-                span: new_args[0].span(),
-                expr: Box::new(new_args[0].clone()),
-                target_type,
-                pg_style: false,
-            };
-
-            let args: Vec<&Expr> = new_args.iter().collect();
-            let result = self.resolve_lambda_function(span, func_name, &args, lambda)?;
-
-            let target_type = if result.1.is_nullable() {
-                DataType::Variant.wrap_nullable()
-            } else {
-                DataType::Variant
-            };
-
-            let result_expr = ScalarExpr::CastExpr(CastExpr {
-                span: new_args[0].span(),
-                is_try: false,
-                argument: Box::new(result.0.clone()),
-                target_type: Box::new(target_type.clone()),
-            });
-
-            return Ok(Box::new((result_expr, target_type)));
-        }
-
         if matches!(
             self.bind_context.expr_context,
             ExprContext::InLambdaFunction
@@ -1903,13 +1866,6 @@ impl<'a> TypeChecker<'a> {
             )
             .set_span(span));
         }
-        let params = lambda
-            .params
-            .iter()
-            .map(|param| param.name.to_lowercase())
-            .collect::<Vec<_>>();
-
-        self.check_lambda_param_count(func_name, params.len(), span)?;
 
         if args.len() != 1 {
             return Err(ErrorCode::SemanticError(format!(
@@ -1919,7 +1875,46 @@ impl<'a> TypeChecker<'a> {
             ))
             .set_span(span));
         }
-        let box (mut arg, arg_type) = self.resolve(args[0])?;
+        let box (mut arg, mut arg_type) = self.resolve(args[0])?;
+
+        let mut func_name = func_name;
+        let mut is_cast_variant = false;
+        if arg_type.remove_nullable() == DataType::Variant {
+            if func_name.starts_with("json_") {
+                func_name = &func_name[5..];
+            }
+            // Try auto cast the Variant type to Array(Variant) or Map(String, Variant),
+            // so that the lambda functions support variant type as argument.
+            let mut target_type = if func_name.starts_with("array") {
+                DataType::Array(Box::new(DataType::Nullable(Box::new(DataType::Variant))))
+            } else {
+                DataType::Map(Box::new(DataType::Tuple(vec![
+                    DataType::String,
+                    DataType::Nullable(Box::new(DataType::Variant)),
+                ])))
+            };
+            if arg_type.is_nullable() {
+                target_type = target_type.wrap_nullable();
+            }
+
+            arg = ScalarExpr::CastExpr(CastExpr {
+                span: None,
+                is_try: false,
+                argument: Box::new(arg.clone()),
+                target_type: Box::new(target_type.clone()),
+            });
+            arg_type = target_type;
+
+            is_cast_variant = true;
+        }
+
+        let params = lambda
+            .params
+            .iter()
+            .map(|param| param.name.to_lowercase())
+            .collect::<Vec<_>>();
+
+        self.check_lambda_param_count(func_name, params.len(), span)?;
 
         let inner_ty = match arg_type.remove_nullable() {
             DataType::Array(box inner_ty) => inner_ty.clone(),
@@ -2134,7 +2129,22 @@ impl<'a> TypeChecker<'a> {
             }
         };
 
-        Ok(Box::new((lambda_func, data_type)))
+        if is_cast_variant {
+            let result_target_type = if data_type.is_nullable() {
+                DataType::Nullable(Box::new(DataType::Variant))
+            } else {
+                DataType::Variant
+            };
+            let result_target_scalar = ScalarExpr::CastExpr(CastExpr {
+                span: None,
+                is_try: false,
+                argument: Box::new(lambda_func),
+                target_type: Box::new(result_target_type.clone()),
+            });
+            Ok(Box::new((result_target_scalar, result_target_type)))
+        } else {
+            Ok(Box::new((lambda_func, data_type)))
+        }
     }
 
     fn check_lambda_param_count(
@@ -2768,6 +2778,12 @@ impl<'a> TypeChecker<'a> {
             )));
         }
 
+        if let Some(rewritten_func_func) =
+            self.try_rewrite_array_function(span, func_name, &params, &mut args, &mut arg_types)
+        {
+            return rewritten_func_func;
+        }
+
         self.resolve_scalar_function_call(span, func_name, params, args)
     }
 
@@ -3641,6 +3657,91 @@ impl<'a> TypeChecker<'a> {
         }
     }
 
+    fn array_functions() -> &'static [Ascii<&'static str>] {
+        static ARRAY_FUNCTIONS: &[Ascii<&'static str>] = &[
+            Ascii::new("array_count"),
+            Ascii::new("array_max"),
+            Ascii::new("array_min"),
+            Ascii::new("array_any"),
+            Ascii::new("array_approx_count_distinct"),
+            Ascii::new("array_unique"),
+            Ascii::new("array_sort_asc_null_first"),
+            Ascii::new("array_sort_desc_null_first"),
+            Ascii::new("array_sort_asc_null_last"),
+            Ascii::new("array_sort_desc_null_last"),
+            Ascii::new("array_remove_first"),
+            Ascii::new("array_remove_last"),
+            Ascii::new("array_distinct"),
+        ];
+        ARRAY_FUNCTIONS
+    }
+
+    fn try_rewrite_array_function(
+        &mut self,
+        span: Span,
+        func_name: &str,
+        params: &[Scalar],
+        args: &mut [ScalarExpr],
+        arg_types: &mut [DataType],
+    ) -> Option<Result<Box<(ScalarExpr, DataType)>>> {
+        // Try auto cast the Variant type to Array(Variant),
+        // so that the array functions support Variant type as argument.
+        let uni_case_func_name = Ascii::new(func_name);
+        if Self::array_functions().contains(&uni_case_func_name)
+            && !arg_types.is_empty()
+            && arg_types[0].remove_nullable() == DataType::Variant
+        {
+            let target_type = if arg_types[0].is_nullable() {
+                DataType::Nullable(Box::new(DataType::Array(Box::new(DataType::Nullable(
+                    Box::new(DataType::Variant),
+                )))))
+            } else {
+                DataType::Array(Box::new(DataType::Nullable(Box::new(DataType::Variant))))
+            };
+            let arg = args[0].clone();
+            args[0] = ScalarExpr::CastExpr(CastExpr {
+                span: None,
+                is_try: false,
+                argument: Box::new(arg),
+                target_type: Box::new(target_type.clone()),
+            });
+            arg_types[0] = target_type;
+
+            let result =
+                self.resolve_scalar_function_call(span, func_name, params.to_vec(), args.to_vec());
+            if func_name == "array_remove_first"
+                || func_name == "array_remove_last"
+                || func_name == "array_distinct"
+                || func_name == "array_sort_asc_null_first"
+                || func_name == "array_sort_desc_null_first"
+                || func_name == "array_sort_asc_null_last"
+                || func_name == "array_sort_desc_null_last"
+            {
+                if result.is_err() {
+                    return Some(result);
+                }
+                let box (result_scalar, result_type) = result.unwrap();
+
+                let result_target_type = if result_type.is_nullable() {
+                    DataType::Nullable(Box::new(DataType::Variant))
+                } else {
+                    DataType::Variant
+                };
+                let result_target_scalar = ScalarExpr::CastExpr(CastExpr {
+                    span: None,
+                    is_try: false,
+                    argument: Box::new(result_scalar),
+                    target_type: Box::new(result_target_type.clone()),
+                });
+                Some(Ok(Box::new((result_target_scalar, result_target_type))))
+            } else {
+                Some(result)
+            }
+        } else {
+            None
+        }
+    }
+
     fn resolve_trim_function(
         &mut self,
         span: Span,

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

@@ -393,11 +393,46 @@ select json_array_transform(try_cast(col1 as Variant), a -> a::Int + col2), json
 NULL NULL
 [12,13] []
 
+query TT
+select array_transform(try_cast(col1 as Variant), a -> a::Int + col2), array_filter(try_cast(col1 as Variant), a -> a::Int = col2) from t3;
+----
+[3,4,5] [2]
+[null,null] []
+NULL NULL
+[12,13] []
+
 query TT
 select json_array_reduce([1,2,3,4]::Variant, (x, y) -> 3 + x + y), json_array_transform(parse_json('"aa"'), data -> CONCAT(data::String, 'bend'));
 ----
 19 []
 
+query TT
+select array_reduce([1,2,3,4]::Variant, (x, y) -> 3 + x + y), array_transform(parse_json('"aa"'), data -> CONCAT(data::String, 'bend'));
+----
+19 []
+
+statement ok
+create or replace table t4(col1 Variant Null)
+
+statement ok
+insert into t4 values('[3,2,1,1]'),('[4,5,null,true]'),(null),('[7,"c","d"]')
+
+query TTTT
+select array_count(col1), array_max(col1), array_min(col1), array_approx_count_distinct(col1) from t4;
+----
+4 3 1 3
+4 null true 4
+NULL NULL NULL NULL
+3 "d" 7 3
+
+query TTTTT
+select array_sort(col1), array_unique(col1), array_remove_first(col1), array_remove_last(col1), array_distinct(col1) from t4;
+----
+[1,1,2,3] 3 [2,1,1] [3,2,1] [3,2,1]
+[true,4,5,null] 4 [5,null,true] [4,5,null] [4,5,null,true]
+NULL NULL NULL NULL NULL
+[7,"c","d"] 3 ["c","d"] [7,"c"] [7,"c","d"]
+
 query T
 SELECT arrays_zip(1, 'a', null);
 ----