fix(query): Disabled the shrink_scalar optimization for cast expressions

databendlabs · Dec 26, 2024 · a808187 · a808187
1 parent b72f3e1
commit a808187
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 7 deletions.
diff --git a/src/query/sql/src/planner/semantic/type_check.rs b/src/query/sql/src/planner/semantic/type_check.rs
@@ -594,7 +594,7 @@ impl<'a> TypeChecker<'a> {
                 let registry = &BUILTIN_FUNCTIONS;
                 let checked_expr = type_check::check(&raw_expr, registry)?;
 
-                if let Some(constant) = self.try_fold_constant(&checked_expr) {
+                if let Some(constant) = self.try_fold_constant(&checked_expr, false) {
                     return Ok(constant);
                 }
                 // if the source type is nullable, cast target type should also be nullable.
@@ -637,7 +637,7 @@ impl<'a> TypeChecker<'a> {
                 let registry = &BUILTIN_FUNCTIONS;
                 let checked_expr = type_check::check(&raw_expr, registry)?;
 
-                if let Some(constant) = self.try_fold_constant(&checked_expr) {
+                if let Some(constant) = self.try_fold_constant(&checked_expr, false) {
                     return Ok(constant);
                 }
 
@@ -1086,7 +1086,6 @@ impl<'a> TypeChecker<'a> {
 
             Expr::Hole { .. } => unreachable!("hole is impossible in trivial query"),
         };
-
         Ok(Box::new((scalar, data_type)))
     }
 
@@ -2749,7 +2748,7 @@ impl<'a> TypeChecker<'a> {
             } => {
                 let mut folded_args = Vec::with_capacity(args.len());
                 for (checked_arg, arg) in checked_args.iter().zip(args.iter()) {
-                    match self.try_fold_constant(checked_arg) {
+                    match self.try_fold_constant(checked_arg, true) {
                         Some(constant) if arg.evaluable() => {
                             folded_args.push(constant.0);
                         }
@@ -2767,7 +2766,7 @@ impl<'a> TypeChecker<'a> {
             self.ctx.set_cacheable(false);
         }
 
-        if let Some(constant) = self.try_fold_constant(&expr) {
+        if let Some(constant) = self.try_fold_constant(&expr, true) {
             return Ok(constant);
         }
 
@@ -5067,12 +5066,17 @@ impl<'a> TypeChecker<'a> {
     fn try_fold_constant<Index: ColumnIndex>(
         &self,
         expr: &EExpr<Index>,
+        enable_shrink: bool,
     ) -> Option<Box<(ScalarExpr, DataType)>> {
-        if expr.is_deterministic(&BUILTIN_FUNCTIONS) {
+        if expr.is_deterministic(&BUILTIN_FUNCTIONS) && enable_shrink {
             if let (EExpr::Constant { scalar, .. }, _) =
                 ConstantFolder::fold(expr, &self.func_ctx, &BUILTIN_FUNCTIONS)
             {
-                let scalar = shrink_scalar(scalar);
+                let scalar = if enable_shrink {
+                    shrink_scalar(scalar)
+                } else {
+                    scalar
+                };
                 let ty = scalar.as_ref().infer_data_type();
                 return Some(Box::new((
                     ConstantExpr {

diff --git a/tests/sqllogictests/suites/base/15_procedure/15_0002_procedure.test b/tests/sqllogictests/suites/base/15_procedure/15_0002_procedure.test
@@ -124,5 +124,36 @@ drop procedure if exists not_exists_p();
 statement error 3130
 drop procedure not_exists_p();
 
+statement ok
+drop procedure if exists sum_even_numbers(Int, Int);
+
+statement ok
+CREATE PROCEDURE sum_even_numbers(start_val Int, end_val Int)
+RETURNS UInt8 NOT NULL
+LANGUAGE SQL
+COMMENT='Calculate the sum of all even numbers'
+AS $$
+BEGIN
+    LET sum := 0;
+    FOR i IN start_val TO end_val DO
+        IF i % 2 = 0 THEN
+            sum := sum + i;
+        END IF;
+    END FOR;
+    RETURN sum;
+END;
+$$;
+
+statement error 3130
+call procedure sum_even_numbers(1, 2)
+
+query T
+call procedure sum_even_numbers(1::INT, 2::INT)
+----
+2
+
+statement ok
+drop procedure sum_even_numbers(Int, Int);
+
 statement ok
 unset global enable_experimental_procedure;