feat: Implement UDTF

KKould · KKould · commit 61d285ddb8f8 · 2025-08-10T16:33:44.000+08:00
diff --git a/src/meta/app/src/principal/mod.rs b/src/meta/app/src/principal/mod.rs
@@ -96,6 +96,7 @@ pub use task::Status;
 pub use task::Task;
 pub use task::TaskMessage;
 pub use task::TaskRun;
+pub use user_defined_function::UDTF;
 pub use task::WarehouseOptions;
 pub use task_ident::TaskIdent;
 pub use task_ident::TaskIdentRaw;
diff --git a/src/meta/app/src/principal/user_defined_function.rs b/src/meta/app/src/principal/user_defined_function.rs
@@ -51,6 +51,13 @@ pub struct UDFScript {
     pub immutable: Option<bool>,
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct UDTF {
+    pub arg_types: BTreeMap<String, DataType>,
+    pub return_types: BTreeMap<String, DataType>,
+    pub sql: String,
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct UDAFScript {
     pub code: String,
@@ -71,6 +78,7 @@ pub enum UDFDefinition {
     UDFServer(UDFServer),
     UDFScript(UDFScript),
     UDAFScript(UDAFScript),
+    UDTF(UDTF),
 }
 
 impl UDFDefinition {
@@ -80,6 +88,7 @@ impl UDFDefinition {
             Self::UDFServer(_) => "UDFServer",
             Self::UDFScript(_) => "UDFScript",
             Self::UDAFScript(_) => "UDAFScript",
+            Self::UDTF(_) => "UDTF",
         }
     }
 
@@ -88,13 +97,15 @@ impl UDFDefinition {
             Self::LambdaUDF(_) => false,
             Self::UDFServer(_) => false,
             Self::UDFScript(_) => false,
+            Self::UDTF(_) => false,
             Self::UDAFScript(_) => true,
         }
     }
 
     pub fn language(&self) -> &str {
         match self {
             Self::LambdaUDF(_) => "SQL",
+            Self::UDTF(_) => "SQL",
             Self::UDFServer(x) => x.language.as_str(),
             Self::UDFScript(x) => x.language.as_str(),
             Self::UDAFScript(x) => x.language.as_str(),
@@ -292,6 +303,26 @@ impl Display for UDFDefinition {
                 }
                 write!(f, " }} RETURNS {return_type} LANGUAGE {language} IMPORTS = {imports:?} PACKAGES = {packages:?} RUNTIME_VERSION = {runtime_version} AS $${code}$$")?;
             }
+            UDFDefinition::UDTF(UDTF {
+                                    arg_types,
+                return_types,
+                sql,
+            }) => {
+                for (i, (name, ty)) in arg_types.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "{name} {ty}")?;
+                }
+                write!(f, ") RETURNS (")?;
+                for (i, (name, ty)) in return_types.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "{name} {ty}")?;
+                }
+                write!(f, ") AS $${sql}$$")?;
+            }
         }
         Ok(())
     }
diff --git a/src/meta/proto-conv/src/udf_from_to_protobuf_impl.rs b/src/meta/proto-conv/src/udf_from_to_protobuf_impl.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use chrono::DateTime;
 use chrono::Utc;
 use databend_common_expression::infer_schema_type;
@@ -275,6 +276,87 @@ impl FromToProto for mt::UDAFScript {
     }
 }
 
+impl FromToProto for mt::UDTF {
+    type PB = pb::Udtf;
+
+    fn get_pb_ver(p: &Self::PB) -> u64 {
+        p.ver
+    }
+
+    fn from_pb(p: Self::PB) -> Result<Self, Incompatible>
+    where
+        Self: Sized
+    {
+        reader_check_msg(p.ver, p.min_reader_ver)?;
+
+        let mut arg_types = BTreeMap::new();
+        for (arg_name, arg_ty) in p
+            .arg_names
+            .into_iter().zip(p.arg_types.into_iter()) {
+            let ty = (&TableDataType::from_pb(arg_ty)?).into();
+
+            arg_types.insert(arg_name, ty);
+        }
+
+        let mut return_types = BTreeMap::new();
+        for (return_name, arg_ty) in p
+            .return_names
+            .into_iter().zip(p.return_types.into_iter()) {
+            let ty = (&TableDataType::from_pb(arg_ty)?).into();
+
+            return_types.insert(return_name, ty);
+        }
+
+        Ok(Self {
+            arg_types,
+            return_types,
+            sql: p.sql,
+        })
+    }
+
+    fn to_pb(&self) -> Result<Self::PB, Incompatible> {
+        let mut arg_names = Vec::with_capacity(self.arg_types.len());
+        let mut arg_types = Vec::with_capacity(self.arg_types.len());
+        for (arg_name, arg_type) in self.arg_types.iter() {
+            let arg_type = infer_schema_type(arg_type)
+                .map_err(|e| {
+                    Incompatible::new(format!(
+                        "Convert DataType to TableDataType failed: {}",
+                        e.message()
+                    ))
+                })?
+                .to_pb()?;
+            arg_names.push(arg_name.clone());
+            arg_types.push(arg_type);
+        }
+
+        let mut return_names = Vec::with_capacity(self.return_types.len());
+        let mut return_types = Vec::with_capacity(self.return_types.len());
+        for (return_name, return_type) in self.return_types.iter() {
+            let return_type = infer_schema_type(return_type)
+                .map_err(|e| {
+                    Incompatible::new(format!(
+                        "Convert DataType to TableDataType failed: {}",
+                        e.message()
+                    ))
+                })?
+                .to_pb()?;
+            return_names.push(return_name.clone());
+            return_types.push(return_type);
+        }
+
+        Ok(pb::Udtf {
+            ver: VER,
+            min_reader_ver: MIN_READER_VER,
+            arg_names,
+            arg_types,
+            return_names,
+            return_types,
+            sql: self.sql.clone(),
+        })
+    }
+}
+
 impl FromToProto for mt::UserDefinedFunction {
     type PB = pb::UserDefinedFunction;
     fn get_pb_ver(p: &Self::PB) -> u64 {
@@ -295,6 +377,9 @@ impl FromToProto for mt::UserDefinedFunction {
             Some(pb::user_defined_function::Definition::UdafScript(udaf_script)) => {
                 mt::UDFDefinition::UDAFScript(mt::UDAFScript::from_pb(udaf_script)?)
             }
+            Some(pb::user_defined_function::Definition::Udtf(udtf)) => {
+                mt::UDFDefinition::UDTF(mt::UDTF::from_pb(udtf)?)
+            }
             None => {
                 return Err(Incompatible::new(
                     "UserDefinedFunction.definition cannot be None".to_string(),
@@ -327,6 +412,9 @@ impl FromToProto for mt::UserDefinedFunction {
             mt::UDFDefinition::UDAFScript(udaf_script) => {
                 pb::user_defined_function::Definition::UdafScript(udaf_script.to_pb()?)
             }
+            mt::UDFDefinition::UDTF(udtf) => {
+                pb::user_defined_function::Definition::Udtf(udtf.to_pb()?)
+            }
         };
 
         Ok(pb::UserDefinedFunction {
diff --git a/src/meta/protos/proto/udf.proto b/src/meta/protos/proto/udf.proto
@@ -69,6 +69,17 @@ message UDAFScript {
   repeated string packages = 8;
 }
 
+message UDTF {
+  uint64 ver = 100;
+  uint64 min_reader_ver = 101;
+
+  repeated string arg_names = 1;
+  repeated DataType arg_types = 2;
+  repeated string return_names = 3;
+  repeated DataType return_types = 4;
+  string sql = 5;
+}
+
 message UserDefinedFunction {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;
@@ -80,6 +91,7 @@ message UserDefinedFunction {
     UDFServer udf_server = 4;
     UDFScript udf_script = 6;
     UDAFScript udaf_script = 7;
+    UDTF udtf = 8;
   }
   // The time udf created.
   optional string created_on = 5;
diff --git a/src/query/ast/src/ast/common.rs b/src/query/ast/src/ast/common.rs
@@ -25,7 +25,7 @@ use crate::ast::WithOptions;
 use crate::Span;
 
 // Identifier of table name or column name.
-#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
+#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Drive, DriveMut)]
 pub struct Identifier {
     pub span: Span,
     pub name: String,
@@ -34,7 +34,7 @@ pub struct Identifier {
     pub ident_type: IdentifierType,
 }
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Default)]
 pub enum IdentifierType {
     #[default]
     None,
diff --git a/src/query/ast/src/ast/statements/udf.rs b/src/query/ast/src/ast/statements/udf.rs
@@ -71,6 +71,11 @@ pub enum UDFDefinition {
         language: String,
         runtime_version: String,
     },
+    UDTFSql {
+        arg_types: BTreeMap<Identifier, TypeName>,
+        return_types: BTreeMap<Identifier, TypeName>,
+        sql: String,
+    }
 }
 
 impl Display for UDFDefinition {
@@ -176,6 +181,17 @@ impl Display for UDFDefinition {
                 }
                 write!(f, " ADDRESS = '{address}'")?;
             }
+            UDFDefinition::UDTFSql {
+                arg_types,
+                return_types,
+                sql
+            } => {
+                write!(f, "(")?;
+                write_comma_separated_list(f, arg_types.iter().map(|(name, ty)| format!("{name} {ty}")))?;
+                write!(f, ") RETURNS TABLE (")?;
+                write_comma_separated_list(f, return_types.iter().map(|(name, ty)| format!("{name} {ty}")))?;
+                write!(f, ") AS $$\n{sql}\n$$")?;
+            }
             UDFDefinition::UDAFScript {
                 arg_types,
                 state_fields: state_types,
diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -4969,6 +4969,21 @@ pub fn udf_definition(i: Input) -> IResult<UDFDefinition> {
         },
     );
 
+    let udtf = map(
+        rule! {
+            "(" ~ #comma_separated_list0(udtf_arg) ~ ")"
+            ~ RETURNS ~ TABLE ~ "(" ~ #comma_separated_list0(udtf_arg) ~ ")"
+            ~ AS ~ ^#code_string
+        },
+        |(_, arg_types, _, _, _, _, return_types, _, _, sql)| {
+            UDFDefinition::UDTFSql {
+                arg_types: BTreeMap::from_iter(arg_types),
+                return_types: BTreeMap::from_iter(return_types),
+                sql,
+            }
+        }
+    );
+
     let udaf = map(
         rule! {
             "(" ~ #comma_separated_list0(type_name) ~ ")"
@@ -5033,7 +5048,14 @@ pub fn udf_definition(i: Input) -> IResult<UDFDefinition> {
         #lambda_udf: "AS (<parameter>, ...) -> <definition expr>"
         | #udaf: "(<arg_type>, ...) STATE {<state_field>, ...} RETURNS <return_type> LANGUAGE <language> { ADDRESS=<udf_server_address> | AS <language_codes> } "
         | #udf: "(<arg_type>, ...) RETURNS <return_type> LANGUAGE <language> HANDLER=<handler> { ADDRESS=<udf_server_address> | AS <language_codes> } "
+        | #udtf: "(<arg_type>, ...) RETURNS TABLE (<return_type>, ...) AS <sql> }"
+    )(i)
+}
 
+fn udtf_arg(i:Input) -> IResult<(Identifier, TypeName)> {
+    map(
+        rule! { #ident ~ ^#type_name },
+        |(name, ty)| (name, ty),
     )(i)
 }
 
diff --git a/src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs b/src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs
@@ -25,6 +25,7 @@ use databend_common_ast::ast::SelectTarget;
 use databend_common_ast::ast::TableAlias;
 use databend_common_ast::ast::TableReference;
 use databend_common_ast::Span;
+use databend_common_catalog::catalog::CatalogManager;
 use databend_common_catalog::catalog_kind::CATALOG_DEFAULT;
 use databend_common_catalog::table_args::TableArgs;
 use databend_common_catalog::table_function::TableFunction;
@@ -34,6 +35,7 @@ use databend_common_expression::types::NumberScalar;
 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_storages_result_cache::ResultCacheMetaManager;
 use databend_common_storages_result_cache::ResultScan;
 use databend_common_users::UserApiProvider;
@@ -45,11 +47,11 @@ use crate::binder::ColumnBindingBuilder;
 use crate::binder::Visibility;
 use crate::optimizer::ir::SExpr;
 use crate::planner::semantic::normalize_identifier;
-use crate::plans::EvalScalar;
+use crate::plans::{EvalScalar, Plan};
 use crate::plans::FunctionCall;
 use crate::plans::RelOperator;
 use crate::plans::ScalarItem;
-use crate::BindContext;
+use crate::{BindContext, Planner};
 use crate::ScalarExpr;
 
 impl Binder {
@@ -129,6 +131,42 @@ impl Binder {
         );
         let table_args = bind_table_args(&mut scalar_binder, params, named_params)?;
 
+        let tenant = self.ctx.get_tenant();
+        let udtf_result = databend_common_base::runtime::block_on(async {
+            if let Some(UDFDefinition::UDTF(udtf)) = UserApiProvider::instance().get_udf(&tenant, &func_name.name).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.get(name) else { return Err(ErrorCode::InvalidArgument(format!("Function '{func_name}' does not have a parameter named '{name}'"))) };
+
+                    sql = sql.replace(name, &arg.to_string());
+                }
+                let mut planner = Planner::new(self.ctx.clone());
+                let (_, extras) = planner.plan_sql(&sql).await?;
+                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::Query {
+                    s_expr, bind_context, ..
+                } = plan else {
+                    return Err(ErrorCode::UDFRuntimeError("Query in UDTF returned no result set"))
+                };
+
+                return Ok(Some((*s_expr, *bind_context)))
+            }
+            Ok(None)
+        });
+        if let Some(result) = udtf_result? {
+            return Ok(result);
+        }
+
         if func_name.name.eq_ignore_ascii_case("result_scan") {
             self.bind_result_scan(bind_context, span, alias, &table_args)
         } else {
diff --git a/src/query/sql/src/planner/binder/udf.rs b/src/query/sql/src/planner/binder/udf.rs
diff --git a/src/query/sql/src/planner/semantic/type_check.rs b/src/query/sql/src/planner/semantic/type_check.rs
diff --git a/src/query/storages/system/src/user_functions_table.rs b/src/query/storages/system/src/user_functions_table.rs
