feat: impl UDTF Args replace

Author: KKould

src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs

@@ -36,9 +36,13 @@ use databend_common_expression::FunctionKind;
 use databend_common_expression::Scalar;
 use databend_common_functions::BUILTIN_FUNCTIONS;
 use databend_common_meta_app::principal::UDFDefinition;
+use databend_common_meta_app::principal::UDTF;
 use databend_common_storages_result_cache::ResultCacheMetaManager;
 use databend_common_storages_result_cache::ResultScan;
 use databend_common_users::UserApiProvider;
+use derive_visitor::DriveMut;
+use derive_visitor::VisitorMut;
+use itertools::Itertools;
 
 use crate::binder::scalar::ScalarBinder;
 use crate::binder::table_args::bind_table_args;
@@ -70,6 +74,37 @@ impl Binder {
         alias: &Option<TableAlias>,
         sample: &Option<SampleConfig>,
     ) -> Result<(SExpr, BindContext)> {
+        #[derive(VisitorMut)]
+        #[visitor(Expr(enter))]
+        struct UDTFArgVisitor<'a> {
+            udtf: &'a UDTF,
+            table_args: &'a TableArgs,
+        }
+
+        impl UDTFArgVisitor<'_> {
+            fn enter_expr(&mut self, expr: &mut Expr) {
+                if let Expr::ColumnRef { span, column } = expr {
+                    if column.database.is_some() || column.table.is_some() {
+                        return;
+                    }
+                    assert_eq!(self.udtf.arg_types.len(), self.table_args.positioned.len());
+                    let Some((pos, (_, _ty))) = self
+                        .udtf
+                        .arg_types
+                        .iter()
+                        .find_position(|(name, _)| name == column.column.name())
+                    else {
+                        return;
+                    };
+
+                    *expr = Expr::Literal {
+                        span: *span,
+                        value: Literal::String(self.table_args.positioned[pos].to_string()),
+                    }
+                }
+            }
+        }
+
         let func_name = normalize_identifier(name, &self.name_resolution_ctx);
 
         if BUILTIN_FUNCTIONS
@@ -142,30 +177,32 @@ impl Binder {
                 .await?
                 .map(|udf| udf.definition)
             {
-                let mut sql = udtf.sql;
-
-                for (name, arg) in table_args.named.iter() {
-                    // FIXME: Parameter substitution
-                    let Some((_, _ty)) =
-                        udtf.arg_types.iter().find(|(arg_name, _)| arg_name == name)
-                    else {
-                        return Err(ErrorCode::InvalidArgument(format!(
-                            "Function '{func_name}' does not have a parameter named '{name}'"
-                        )));
-                    };
+                let mut stmt = Planner::new(self.ctx.clone())
+                    .parse_sql(&udtf.sql)?
+                    .statement;
 
-                    sql = sql.replace(name, &arg.to_string());
+                if udtf.arg_types.len() != table_args.positioned.len() {
+                    return Err(ErrorCode::UDFSchemaMismatch(format!(
+                        "UDTF '{}' argument types length {} does not match input arguments length {}",
+                        func_name,
+                        udtf.arg_types.len(),
+                        table_args.positioned.len()
+                    )));
                 }
-                let mut planner = Planner::new(self.ctx.clone());
-                let (_, extras) = planner.plan_sql(&sql).await?;
+                let mut visitor = UDTFArgVisitor {
+                    udtf: &udtf,
+                    table_args: &table_args,
+                };
+                stmt.drive_mut(&mut visitor);
+
                 let binder = Binder::new(
                     self.ctx.clone(),
                     CatalogManager::instance(),
                     self.name_resolution_ctx.clone(),
                     self.metadata.clone(),
                 )
                 .with_subquery_executor(self.subquery_executor.clone());
-                let plan = binder.bind(&extras.statement).await?;
+                let plan = binder.bind(&stmt).await?;
 
                 let Plan::Query {
                     s_expr,
@@ -182,7 +219,8 @@ impl Binder {
 
                 if udtf.return_types.len() != bind_context.columns.len() {
                     return Err(ErrorCode::UDFSchemaMismatch(format!(
-                        "return types length {} does not match output columns length {}",
+                        "UDTF '{}' return types length {} does not match output columns length {}",
+                        func_name,
                         udtf.return_types.len(),
                         bind_context.columns.len()
                     )));

tests/sqllogictests/suites/base/03_common/03_0013_select_udf.test

@@ -213,7 +213,7 @@ statement ok
 insert into t1 values(0), (1);
 
 statement ok
-create function scan_t1 () RETURNS TABLE (c1_string string) as $$ select * from t1 $$;
+create or replace function scan_t1 () RETURNS TABLE (c1_string string) as $$ select * from t1 $$;
 
 query T
 select * from scan_t1();
@@ -222,13 +222,23 @@ select * from scan_t1();
 1
 
 statement ok
-create function invalid_udtf_0 () RETURNS TABLE () as $$ select * from t1 $$;
+create or replace function filter_t1 (arg0 int) RETURNS TABLE (c1_string string) as $$ select * from t1 where c1 = arg0 $$;
+
+query T
+select * from filter_t1(0);
+----
+0
+
+query T
+select * from filter_t1(1);
+----
+1
 
 statement error
-select * from invalid_udtf_0();
+select * from filter_t1();
 
 statement ok
-create function invalid_udtf_1 () RETURNS TABLE (c1_string string, c2 int) as $$ select * from t1 $$;
+create or replace function invalid_udtf_0 () RETURNS TABLE (c1_string string, c2 int) as $$ select * from t1 $$;
 
 statement error
-select * from invalid_udtf_1();
+select * from invalid_udtf_0();