Skip to content

Commit afff2bb

Browse files
authored
Support unparsing UNNEST plan for Snowflake dialect (apache#9)
* support table function flatten for snowflake * add test for lateral cases * implement dialect specific rule
1 parent 50edc43 commit afff2bb

7 files changed

Lines changed: 517 additions & 107 deletions

File tree

datafusion/sql/src/unparser/ast.rs

Lines changed: 112 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -398,16 +398,18 @@ impl Default for TableWithJoinsBuilder {
398398

399399
#[derive(Clone)]
400400
pub struct RelationBuilder {
401-
relation: Option<TableFactorBuilder>,
401+
pub relation: Option<TableFactorBuilder>,
402402
}
403403

404404
#[allow(dead_code)]
405405
#[derive(Clone)]
406406
#[allow(clippy::large_enum_variant)]
407-
enum TableFactorBuilder {
407+
pub enum TableFactorBuilder {
408408
Table(TableRelationBuilder),
409409
Derived(DerivedRelationBuilder),
410410
Unnest(UnnestRelationBuilder),
411+
Function(FunctionRelationBuilder),
412+
TableFunction(TableFunctionRelationBuilder),
411413
Empty,
412414
}
413415

@@ -430,6 +432,16 @@ impl RelationBuilder {
430432
self
431433
}
432434

435+
pub fn function(&mut self, value: FunctionRelationBuilder) -> &mut Self {
436+
self.relation = Some(TableFactorBuilder::Function(value));
437+
self
438+
}
439+
440+
pub fn table_function(&mut self, value: TableFunctionRelationBuilder) -> &mut Self {
441+
self.relation = Some(TableFactorBuilder::TableFunction(value));
442+
self
443+
}
444+
433445
pub fn empty(&mut self) -> &mut Self {
434446
self.relation = Some(TableFactorBuilder::Empty);
435447
self
@@ -446,6 +458,12 @@ impl RelationBuilder {
446458
Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => {
447459
rel_builder.alias = value;
448460
}
461+
Some(TableFactorBuilder::Function(ref mut rel_builder)) => {
462+
rel_builder.alias = value;
463+
}
464+
Some(TableFactorBuilder::TableFunction(ref mut rel_builder)) => {
465+
rel_builder.alias = value;
466+
}
449467
Some(TableFactorBuilder::Empty) => (),
450468
None => (),
451469
}
@@ -456,6 +474,8 @@ impl RelationBuilder {
456474
Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
457475
Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
458476
Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?),
477+
Some(TableFactorBuilder::Function(ref value)) => Some(value.build()?),
478+
Some(TableFactorBuilder::TableFunction(ref value)) => Some(value.build()?),
459479
Some(TableFactorBuilder::Empty) => None,
460480
None => return Err(Into::into(UninitializedFieldError::from("relation"))),
461481
})
@@ -662,6 +682,96 @@ impl Default for UnnestRelationBuilder {
662682
}
663683
}
664684

685+
#[derive(Clone)]
686+
pub struct FunctionRelationBuilder {
687+
lateral: bool,
688+
name: ast::ObjectName,
689+
args: Vec<ast::FunctionArg>,
690+
alias: Option<ast::TableAlias>,
691+
}
692+
693+
#[allow(dead_code)]
694+
impl FunctionRelationBuilder {
695+
pub fn lateral(&mut self, value: bool) -> &mut Self {
696+
self.lateral = value;
697+
self
698+
}
699+
700+
pub fn name(&mut self, value: ast::ObjectName) -> &mut Self {
701+
self.name = value;
702+
self
703+
}
704+
705+
pub fn args(&mut self, value: Vec<ast::FunctionArg>) -> &mut Self {
706+
self.args = value;
707+
self
708+
}
709+
710+
pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
711+
self.alias = value;
712+
self
713+
}
714+
715+
pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
716+
Ok(ast::TableFactor::Function {
717+
lateral: self.lateral,
718+
name: self.name.clone(),
719+
args: self.args.clone(),
720+
alias: self.alias.clone(),
721+
})
722+
}
723+
724+
fn create_empty() -> Self {
725+
Self {
726+
lateral: Default::default(),
727+
name: ast::ObjectName(vec![]),
728+
args: Default::default(),
729+
alias: Default::default(),
730+
}
731+
}
732+
}
733+
734+
#[derive(Clone)]
735+
pub struct TableFunctionRelationBuilder {
736+
pub expr: Option<ast::Expr>,
737+
pub alias: Option<ast::TableAlias>,
738+
}
739+
740+
impl TableFunctionRelationBuilder {
741+
pub fn expr(&mut self, value: ast::Expr) -> &mut Self {
742+
self.expr = Some(value);
743+
self
744+
}
745+
746+
pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
747+
self.alias = value;
748+
self
749+
}
750+
751+
pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
752+
Ok(ast::TableFactor::TableFunction {
753+
expr: match self.expr {
754+
Some(ref value) => value.clone(),
755+
None => return Err(Into::into(UninitializedFieldError::from("expr"))),
756+
},
757+
alias: self.alias.clone(),
758+
})
759+
}
760+
761+
fn create_empty() -> Self {
762+
Self {
763+
expr: Default::default(),
764+
alias: Default::default(),
765+
}
766+
}
767+
}
768+
769+
impl Default for TableFunctionRelationBuilder {
770+
fn default() -> Self {
771+
Self::create_empty()
772+
}
773+
}
774+
665775
/// Runtime error when a `build()` method is called and one or more required fields
666776
/// do not have a value.
667777
#[derive(Debug, Clone)]

0 commit comments

Comments
 (0)