fix(sql-udf): correctly handle udf_binding_flag & udf_global_count

e2e_test/udf/sql_udf.slt

@@ -26,8 +26,8 @@ statement ok
 create function add_sub_mix(INT, a INT, INT) returns int language sql as 'select $1 - a + $3';
 
 # Mixed parameter with calling inner sql udfs
-# statement ok
-# create function add_sub_mix_wrapper(INT, a INT, INT) returns int language sql as 'select add($1, a) + a + sub(a, $3)';
+statement ok
+create function add_sub_mix_wrapper(INT, a INT, INT) returns int language sql as 'select add($1, a) + a + sub(a, $3)';
 
 # Named sql udf with corner case
 statement ok
@@ -150,6 +150,11 @@ select add_sub_mix(1, 2, 3);
 ----
 2
 
+query I
+select add_sub_mix_wrapper(1, 2, 3);
+----
+4
+
 query T
 select corner_case(1, 2, 3);
 ----
@@ -392,6 +397,9 @@ drop function add_named_wrapper;
 statement ok
 drop function type_match;
 
+statement ok
+drop function add_sub_mix_wrapper;
+
 # Drop the mock table
 statement ok
 drop table t1;

src/frontend/src/binder/expr/function.rs

@@ -241,10 +241,24 @@ impl Binder {
                 if let Ok(expr) = UdfContext::extract_udf_expression(ast) {
                     self.set_udf_binding_flag();
                     let bind_result = self.bind_expr(expr);
-                    self.unset_udf_binding_flag();
+
+                    // We should properly decrement global count after a successful binding
+                    // Since the subsequent `unset flag` operation relies on global counting
+                    self.udf_context.decr_global_count();
+
+                    // Only the top-most sql udf binding should unset the flag
+                    // Otherwise the subsequent binding may not be able to
+                    // find the corresponding context, consider the following example:
+                    // e.g., `select add_wrapper(a INT, b INT) returns int language sql as 'select add(a, b) + a';`
+                    // The inner `add` should not unset the flag, otherwise the `a` will be treated as
+                    // a normal column, which is then invalid in this context.
+                    if self.udf_context.global_count() == 0 {
+                        self.unset_udf_binding_flag();
+                    }
 
                     // Restore context information for subsequent binding
                     self.udf_context.update_context(stashed_udf_context);
+
                     return bind_result;
                 } else {
                     return Err(ErrorCode::InvalidInputSyntax(

src/frontend/src/binder/mod.rs

@@ -155,6 +155,10 @@ impl UdfContext {
         self.udf_global_counter += 1;
     }
 
+    pub fn decr_global_count(&mut self) {
+        self.udf_global_counter -= 1;
+    }
+
     pub fn _is_empty(&self) -> bool {
         self.udf_param_context.is_empty()
     }
@@ -219,6 +223,8 @@ impl UdfContext {
         Ok(expr)
     }
 
+    /// Create the sql udf context
+    /// used per `bind_function` for sql udf & semantic check at definition time
     pub fn create_udf_context(
         args: &[FunctionArg],
         catalog: &Arc<FunctionCatalog>,
@@ -228,9 +234,10 @@ impl UdfContext {
             match current_arg {
                 FunctionArg::Unnamed(arg) => {
                     let FunctionArgExpr::Expr(e) = arg else {
-                        return Err(
-                            ErrorCode::InvalidInputSyntax("invalid syntax".to_string()).into()
-                        );
+                        return Err(ErrorCode::InvalidInputSyntax(
+                            "expect `FunctionArgExpr` for unnamed argument".to_string(),
+                        )
+                        .into());
                     };
                     if catalog.arg_names[i].is_empty() {
                         ret.insert(format!("${}", i + 1), e.clone());
@@ -240,7 +247,12 @@ impl UdfContext {
                         ret.insert(catalog.arg_names[i].clone(), e.clone());
                     }
                 }
-                _ => return Err(ErrorCode::InvalidInputSyntax("invalid syntax".to_string()).into()),
+                _ => {
+                    return Err(ErrorCode::InvalidInputSyntax(
+                        "expect unnamed argument when creating sql udf context".to_string(),
+                    )
+                    .into())
+                }
             }
         }
         Ok(ret)

src/frontend/src/handler/create_sql_function.rs

@@ -190,6 +190,10 @@ pub async fn handle_create_sql_function(
 
         binder.set_udf_binding_flag();
 
+        // Need to set the initial global count to 1
+        // otherwise the flag will be unset during the semantic check
+        binder.udf_context_mut().incr_global_count();
+
         if let Ok(expr) = UdfContext::extract_udf_expression(ast) {
             match binder.bind_expr(expr) {
                 Ok(expr) => {
@@ -204,7 +208,7 @@ pub async fn handle_create_sql_function(
                     }
                 }
                 Err(e) => return Err(ErrorCode::InvalidInputSyntax(format!(
-                    "failed to conduct semantic check, please see if you are calling non-existence functions: {}",
+                    "failed to conduct semantic check, please see if you are calling non-existence functions or parameters\ndetailed error message: {}",
                     e.as_report()
                 ))
                 .into()),