feat: impl UDTF Server

Author: KKould

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

@@ -116,6 +116,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/protos/proto/udf.proto

@@ -86,6 +86,20 @@ 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;
+    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 +126,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;

src/query/ast/src/ast/statements/udf.rs

@@ -88,6 +88,15 @@ pub enum UDFDefinition {
         return_types: Vec<(Identifier, TypeName)>,
         sql: String,
     },
+    UDTFServer {
+        arg_types: Vec<(Identifier, TypeName)>,
+        return_types: Vec<(Identifier, TypeName)>,
+        address: String,
+        handler: String,
+        headers: BTreeMap<String, String>,
+        language: String,
+        immutable: Option<bool>,
+    },
     ScalarUDF {
         arg_types: Vec<(Identifier, TypeName)>,
         definition: String,
@@ -282,6 +291,46 @@ impl Display for UDFDefinition {
                 )?;
                 write!(f, ") AS $$\n{sql}\n$$")?;
             }
+            UDFDefinition::UDTFServer {
+                arg_types,
+                return_types,
+                address,
+                handler,
+                headers,
+                language,
+                immutable,
+            } => {
+                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, ") 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 {
                 arg_types,
                 definition,