Support unparsing UNNEST plan for Snowflake dialect

datafusion/sql/src/unparser/ast.rs

@@ -398,16 +398,18 @@ impl Default for TableWithJoinsBuilder {
 
 #[derive(Clone)]
 pub struct RelationBuilder {
-    relation: Option<TableFactorBuilder>,
+    pub relation: Option<TableFactorBuilder>,
 }
 
 #[allow(dead_code)]
 #[derive(Clone)]
 #[allow(clippy::large_enum_variant)]
-enum TableFactorBuilder {
+pub enum TableFactorBuilder {
     Table(TableRelationBuilder),
     Derived(DerivedRelationBuilder),
     Unnest(UnnestRelationBuilder),
+    Function(FunctionRelationBuilder),
+    TableFunction(TableFunctionRelationBuilder),
     Empty,
 }
 
@@ -430,6 +432,16 @@ impl RelationBuilder {
         self
     }
 
+    pub fn function(&mut self, value: FunctionRelationBuilder) -> &mut Self {
+        self.relation = Some(TableFactorBuilder::Function(value));
+        self
+    }
+
+    pub fn table_function(&mut self, value: TableFunctionRelationBuilder) -> &mut Self {
+        self.relation = Some(TableFactorBuilder::TableFunction(value));
+        self
+    }
+
     pub fn empty(&mut self) -> &mut Self {
         self.relation = Some(TableFactorBuilder::Empty);
         self
@@ -446,6 +458,12 @@ impl RelationBuilder {
             Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => {
                 rel_builder.alias = value;
             }
+            Some(TableFactorBuilder::Function(ref mut rel_builder)) => {
+                rel_builder.alias = value;
+            }
+            Some(TableFactorBuilder::TableFunction(ref mut rel_builder)) => {
+                rel_builder.alias = value;
+            }
             Some(TableFactorBuilder::Empty) => (),
             None => (),
         }
@@ -456,6 +474,8 @@ impl RelationBuilder {
             Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
             Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
             Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?),
+            Some(TableFactorBuilder::Function(ref value)) => Some(value.build()?),
+            Some(TableFactorBuilder::TableFunction(ref value)) => Some(value.build()?),
             Some(TableFactorBuilder::Empty) => None,
             None => return Err(Into::into(UninitializedFieldError::from("relation"))),
         })
@@ -662,6 +682,96 @@ impl Default for UnnestRelationBuilder {
     }
 }
 
+#[derive(Clone)]
+pub struct FunctionRelationBuilder {
+    lateral: bool,
+    name: ast::ObjectName,
+    args: Vec<ast::FunctionArg>,
+    alias: Option<ast::TableAlias>,
+}
+
+#[allow(dead_code)]
+impl FunctionRelationBuilder {
+    pub fn lateral(&mut self, value: bool) -> &mut Self {
+        self.lateral = value;
+        self
+    }
+
+    pub fn name(&mut self, value: ast::ObjectName) -> &mut Self {
+        self.name = value;
+        self
+    }
+
+    pub fn args(&mut self, value: Vec<ast::FunctionArg>) -> &mut Self {
+        self.args = value;
+        self
+    }
+
+    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
+        self.alias = value;
+        self
+    }
+
+    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
+        Ok(ast::TableFactor::Function {
+            lateral: self.lateral,
+            name: self.name.clone(),
+            args: self.args.clone(),
+            alias: self.alias.clone(),
+        })
+    }
+
+    fn create_empty() -> Self {
+        Self {
+            lateral: Default::default(),
+            name: ast::ObjectName(vec![]),
+            args: Default::default(),
+            alias: Default::default(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct TableFunctionRelationBuilder {
+    pub expr: Option<ast::Expr>,
+    pub alias: Option<ast::TableAlias>,
+}
+
+impl TableFunctionRelationBuilder {
+    pub fn expr(&mut self, value: ast::Expr) -> &mut Self {
+        self.expr = Some(value);
+        self
+    }
+
+    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
+        self.alias = value;
+        self
+    }
+
+    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
+        Ok(ast::TableFactor::TableFunction {
+            expr: match self.expr {
+                Some(ref value) => value.clone(),
+                None => return Err(Into::into(UninitializedFieldError::from("expr"))),
+            },
+            alias: self.alias.clone(),
+        })
+    }
+
+    fn create_empty() -> Self {
+        Self {
+            expr: Default::default(),
+            alias: Default::default(),
+        }
+    }
+}
+
+impl Default for TableFunctionRelationBuilder {
+    fn default() -> Self {
+        Self::create_empty()
+    }
+}
+
 /// Runtime error when a `build()` method is called and one or more required fields
 /// do not have a value.
 #[derive(Debug, Clone)]