feat: impl UDTF Server

Author: KKould

src/meta/app/src/principal/mod.rs

@@ -123,6 +123,7 @@ pub use user_defined_function::UDAFScript;
 pub use user_defined_function::UDFDefinition;
 pub use user_defined_function::UDFScript;
 pub use user_defined_function::UDFServer;
+pub use user_defined_function::UDTFServer;
 pub use user_defined_function::UserDefinedFunction;
 pub use user_defined_function::UDTF;
 pub use user_grant::GrantEntry;

src/meta/app/src/principal/user_defined_function.rs

@@ -52,6 +52,18 @@ pub struct UDFScript {
     pub immutable: Option<bool>,
 }
 
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub struct UDTFServer {
+    pub address: String,
+    pub handler: String,
+    pub headers: BTreeMap<String, String>,
+    pub language: String,
+    pub arg_names: Vec<String>,
+    pub arg_types: Vec<DataType>,
+    pub return_types: Vec<(String, DataType)>,
+    pub immutable: Option<bool>,
+}
+
 /// User Defined Table Function (UDTF)
 ///
 /// # Fields
@@ -98,6 +110,7 @@ pub enum UDFDefinition {
     UDFServer(UDFServer),
     UDFScript(UDFScript),
     UDAFScript(UDAFScript),
+    UDTFServer(UDTFServer),
     UDTF(UDTF),
     ScalarUDF(ScalarUDF),
 }
@@ -110,7 +123,8 @@ impl UDFDefinition {
             Self::UDFScript(_) => "UDFScript",
             Self::UDAFScript(_) => "UDAFScript",
             Self::UDTF(_) => "UDTF",
-            UDFDefinition::ScalarUDF(_) => "ScalarUDF",
+            Self::UDTFServer(_) => "UDTFServer",
+            Self::ScalarUDF(_) => "ScalarUDF",
         }
     }
 
@@ -120,6 +134,7 @@ impl UDFDefinition {
             Self::UDFServer(_) => false,
             Self::UDFScript(_) => false,
             Self::UDTF(_) => false,
+            Self::UDTFServer(_) => false,
             Self::ScalarUDF(_) => false,
             Self::UDAFScript(_) => true,
         }
@@ -130,6 +145,7 @@ impl UDFDefinition {
             Self::LambdaUDF(_) => "SQL",
             Self::UDTF(_) => "SQL",
             Self::ScalarUDF(_) => "SQL",
+            Self::UDTFServer(x) => x.language.as_str(),
             Self::UDFServer(x) => x.language.as_str(),
             Self::UDFScript(x) => x.language.as_str(),
             Self::UDAFScript(x) => x.language.as_str(),
@@ -220,6 +236,13 @@ impl UserDefinedFunction {
             created_on: Utc::now(),
         }
     }
+
+    pub fn as_udtf_server(self) -> Option<UDTFServer> {
+        if let UDFDefinition::UDTFServer(udtf_server) = self.definition {
+            return Some(udtf_server);
+        }
+        None
+    }
 }
 
 impl Display for UDFDefinition {
@@ -353,6 +376,50 @@ impl Display for UDFDefinition {
                 }
                 write!(f, ") AS $${sql}$$")?;
             }
+            UDFDefinition::UDTFServer(UDTFServer {
+                address,
+                handler,
+                headers,
+                language,
+                arg_names,
+                arg_types,
+                return_types,
+                immutable,
+            }) => {
+                for (i, (name, ty)) in arg_names.iter().zip(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, ") LANGUAGE {language}")?;
+                if let Some(immutable) = immutable {
+                    if *immutable {
+                        write!(f, " IMMUTABLE")?;
+                    } else {
+                        write!(f, " VOLATILE")?;
+                    }
+                }
+                write!(f, " HANDLER = {handler}")?;
+                if !headers.is_empty() {
+                    write!(f, " HEADERS = (")?;
+                    for (i, (key, value)) in headers.iter().enumerate() {
+                        if i > 0 {
+                            write!(f, ", ")?;
+                        }
+                        write!(f, "{key} = {value}")?;
+                    }
+                    write!(f, ")")?;
+                }
+                write!(f, " ADDRESS = {address}")?;
+            }
             UDFDefinition::ScalarUDF(ScalarUDF {
                 arg_types,
                 return_type,

src/meta/proto-conv/src/udf_from_to_protobuf_impl.rs

@@ -359,6 +359,88 @@ impl FromToProto for mt::UDTF {
     }
 }
 
+impl FromToProto for mt::UDTFServer {
+    type PB = pb::UdtfServer;
+
+    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 = Vec::with_capacity(p.arg_types.len());
+        for arg_type in p.arg_types {
+            let arg_type = DataType::from(&TableDataType::from_pb(arg_type)?);
+            arg_types.push(arg_type);
+        }
+        let mut return_types = Vec::new();
+        for return_ty in p.return_types {
+            let ty_pb = return_ty.ty.ok_or_else(|| {
+                Incompatible::new("UDTF.arg_types.ty can not be None".to_string())
+            })?;
+            let ty = TableDataType::from_pb(ty_pb)?;
+
+            return_types.push((return_ty.name, (&ty).into()));
+        }
+
+        Ok(mt::UDTFServer {
+            address: p.address,
+            arg_types,
+            return_types,
+            handler: p.handler,
+            headers: p.headers,
+            language: p.language,
+            immutable: p.immutable,
+            arg_names: p.arg_names,
+        })
+    }
+
+    fn to_pb(&self) -> Result<Self::PB, Incompatible> {
+        let mut arg_types = Vec::with_capacity(self.arg_types.len());
+        for 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_types.push(arg_type);
+        }
+        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_types.push(UdtfArg {
+                name: return_name.clone(),
+                ty: Some(return_type),
+            });
+        }
+
+        Ok(pb::UdtfServer {
+            ver: VER,
+            min_reader_ver: MIN_READER_VER,
+            address: self.address.clone(),
+            handler: self.handler.clone(),
+            headers: self.headers.clone(),
+            language: self.language.clone(),
+            arg_types,
+            return_types,
+            immutable: self.immutable,
+            arg_names: self.arg_names.clone(),
+        })
+    }
+}
+
 impl FromToProto for mt::ScalarUDF {
     type PB = pb::ScalarUdf;
 
@@ -454,6 +536,9 @@ impl FromToProto for mt::UserDefinedFunction {
             Some(pb::user_defined_function::Definition::ScalarUdf(scalar_udf)) => {
                 mt::UDFDefinition::ScalarUDF(mt::ScalarUDF::from_pb(scalar_udf)?)
             }
+            Some(pb::user_defined_function::Definition::UdtfServer(udtf_server)) => {
+                mt::UDFDefinition::UDTFServer(mt::UDTFServer::from_pb(udtf_server)?)
+            }
             None => {
                 return Err(Incompatible::new(
                     "UserDefinedFunction.definition cannot be None".to_string(),
@@ -492,6 +577,9 @@ impl FromToProto for mt::UserDefinedFunction {
             mt::UDFDefinition::ScalarUDF(scalar_udf) => {
                 pb::user_defined_function::Definition::ScalarUdf(scalar_udf.to_pb()?)
             }
+            mt::UDFDefinition::UDTFServer(udtf_server) => {
+                pb::user_defined_function::Definition::UdtfServer(udtf_server.to_pb()?)
+            }
         };
 
         Ok(pb::UserDefinedFunction {

src/meta/proto-conv/src/util.rs

@@ -187,6 +187,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (155, "2025-10-24: Add: RowAccessPolicyMeta::RowAccessPolicyArg"),
     (156, "2025-10-22: Add: DataMaskMeta add DataMaskArg"),
     (157, "2025-10-22: Add: TableDataType TimestampTz"),
+    (158, "2025-10-22: Add: Server UDTF"),
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)

src/meta/proto-conv/tests/it/main.rs

@@ -149,3 +149,4 @@ mod v154_vacuum_watermark;
 mod v155_row_access_policy_args;
 mod v156_data_mask_args;
 mod v157_type_timestamp_tz;
+mod v158_udtf_server;

src/meta/proto-conv/tests/it/v158_udtf_server.rs

@@ -0,0 +1,81 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// 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::types::DataType;
+use databend_common_expression::types::NumberDataType;
+use databend_common_meta_app::principal::UDFDefinition;
+use databend_common_meta_app::principal::UDTFServer;
+use databend_common_meta_app::principal::UserDefinedFunction;
+use fastrace::func_name;
+
+use crate::common;
+
+// These bytes are built when a new version in introduced,
+// and are kept for backward compatibility test.
+//
+// *************************************************************
+// * These messages should never be updated,                   *
+// * only be added when a new version is added,                *
+// * or be removed when an old version is no longer supported. *
+// *************************************************************
+//
+// The message bytes are built from the output of `test_pb_from_to()`
+#[test]
+fn test_decode_v158_server_udtf() -> anyhow::Result<()> {
+    let bytes = vec![
+        10, 15, 116, 101, 115, 116, 95, 115, 99, 97, 108, 97, 114, 95, 117, 100, 102, 18, 21, 84,
+        104, 105, 115, 32, 105, 115, 32, 97, 32, 100, 101, 115, 99, 114, 105, 112, 116, 105, 111,
+        110, 82, 144, 1, 10, 21, 104, 116, 116, 112, 58, 47, 47, 108, 111, 99, 97, 108, 104, 111,
+        115, 116, 58, 56, 56, 56, 56, 18, 11, 112, 108, 117, 115, 95, 105, 110, 116, 95, 112, 121,
+        26, 6, 112, 121, 116, 104, 111, 110, 34, 10, 146, 2, 0, 160, 6, 158, 1, 168, 6, 24, 34, 10,
+        138, 2, 0, 160, 6, 158, 1, 168, 6, 24, 42, 16, 10, 2, 99, 49, 18, 10, 146, 2, 0, 160, 6,
+        158, 1, 168, 6, 24, 42, 25, 10, 2, 99, 50, 18, 19, 154, 2, 9, 42, 0, 160, 6, 158, 1, 168,
+        6, 24, 160, 6, 158, 1, 168, 6, 24, 50, 14, 10, 4, 107, 101, 121, 49, 18, 6, 118, 97, 108,
+        117, 101, 49, 66, 2, 99, 49, 66, 2, 99, 50, 160, 6, 158, 1, 168, 6, 24, 42, 23, 50, 48, 50,
+        51, 45, 49, 50, 45, 49, 53, 32, 48, 49, 58, 50, 54, 58, 48, 57, 32, 85, 84, 67, 160, 6,
+        158, 1, 168, 6, 24,
+    ];
+
+    let want = || UserDefinedFunction {
+        name: "test_scalar_udf".to_string(),
+        description: "This is a description".to_string(),
+        definition: UDFDefinition::UDTFServer(UDTFServer {
+            address: "http://localhost:8888".to_string(),
+            handler: "plus_int_py".to_string(),
+            headers: vec![("key1".to_string(), "value1".to_string())]
+                .into_iter()
+                .collect(),
+            language: "python".to_string(),
+            arg_names: vec![s("c1"), s("c2")],
+            arg_types: vec![DataType::String, DataType::Boolean],
+            return_types: vec![
+                (s("c1"), DataType::String),
+                (s("c2"), DataType::Number(NumberDataType::Int8)),
+            ],
+            immutable: None,
+        }),
+        created_on: DateTime::<Utc>::from_timestamp(1702603569, 0).unwrap(),
+    };
+
+    common::test_pb_from_to(func_name!(), want())?;
+    common::test_load_old(func_name!(), bytes.as_slice(), 158, want())
+}
+
+fn s(ss: impl ToString) -> String {
+    ss.to_string()
+}

src/meta/protos/proto/udf.proto

@@ -86,6 +86,21 @@ message UDTF {
   string sql = 3;
 }
 
+message UDTFServer {
+  uint64 ver = 100;
+    uint64 min_reader_ver = 101;
+
+    string address = 1;
+    string handler = 2;
+    string language = 3;
+    repeated DataType arg_types = 4;
+    // return column name with data type
+    repeated UDTFArg return_types = 5;
+    map<string, string> headers = 6;
+    optional bool immutable = 7;
+    repeated string arg_names = 8;
+}
+
 message ScalarUDF {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;
@@ -112,6 +127,7 @@ message UserDefinedFunction {
     UDAFScript udaf_script = 7;
     UDTF udtf = 8;
     ScalarUDF scalar_udf = 9;
+    UDTFServer udtf_server = 10;
   }
   // The time udf created.
   optional string created_on = 5;