Handle EmptyRelation during SQL unparsing (apache#10803)

* handle empty relation for sql unparsing * revert the doc test change * remove dbg macro * cargo fmt * improve the error handling
comphead · Jun 5, 2024 · c5cefa8 · c5cefa8
1 parent 70256ba
commit c5cefa8
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 17 deletions.
diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs
@@ -242,7 +242,7 @@ impl SelectBuilder {
             from: self
                 .from
                 .iter()
-                .map(|b| b.build())
+                .filter_map(|b| b.build().transpose())
                 .collect::<Result<Vec<_>, BuilderError>>()?,
             lateral_views: self.lateral_views.clone(),
             selection: self.selection.clone(),
@@ -314,18 +314,17 @@ impl TableWithJoinsBuilder {
         new
     }
 
-    pub fn build(&self) -> Result<ast::TableWithJoins, BuilderError> {
-        Ok(ast::TableWithJoins {
-            relation: match self.relation {
-                Some(ref value) => value.build()?,
-                None => {
-                    return Result::Err(Into::into(UninitializedFieldError::from(
-                        "relation",
-                    )))
-                }
+    pub fn build(&self) -> Result<Option<ast::TableWithJoins>, BuilderError> {
+        match self.relation {
+            Some(ref value) => match value.build()? {
+                Some(relation) => Ok(Some(ast::TableWithJoins {
+                    relation,
+                    joins: self.joins.clone(),
+                })),
+                None => Ok(None),
             },
-            joins: self.joins.clone(),
-        })
+            None => Err(Into::into(UninitializedFieldError::from("relation"))),
+        }
     }
     fn create_empty() -> Self {
         Self {
@@ -350,6 +349,7 @@ pub(super) struct RelationBuilder {
 enum TableFactorBuilder {
     Table(TableRelationBuilder),
     Derived(DerivedRelationBuilder),
+    Empty,
 }
 
 #[allow(dead_code)]
@@ -367,6 +367,11 @@ impl RelationBuilder {
         new.relation = Option::Some(TableFactorBuilder::Derived(value));
         new
     }
+    pub fn empty(&mut self) -> &mut Self {
+        let new = self;
+        new.relation = Some(TableFactorBuilder::Empty);
+        new
+    }
     pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
         let new = self;
         match new.relation {
@@ -376,14 +381,16 @@ impl RelationBuilder {
             Some(TableFactorBuilder::Derived(ref mut rel_builder)) => {
                 rel_builder.alias = value;
             }
+            Some(TableFactorBuilder::Empty) => (),
             None => (),
         }
         new
     }
-    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
+    pub fn build(&self) -> Result<Option<ast::TableFactor>, BuilderError> {
         Ok(match self.relation {
-            Some(TableFactorBuilder::Table(ref value)) => value.build()?,
-            Some(TableFactorBuilder::Derived(ref value)) => value.build()?,
+            Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
+            Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
+            Some(TableFactorBuilder::Empty) => None,
             None => {
                 return Result::Err(Into::into(UninitializedFieldError::from("relation")))
             }

diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs
@@ -389,7 +389,10 @@ impl Unparser<'_> {
                 )?;
 
                 let ast_join = ast::Join {
-                    relation: right_relation.build()?,
+                    relation: match right_relation.build()? {
+                        Some(relation) => relation,
+                        None => return internal_err!("Failed to build right relation"),
+                    },
                     join_operator: self
                         .join_operator_to_sql(join.join_type, join_constraint),
                 };
@@ -417,7 +420,10 @@ impl Unparser<'_> {
                 )?;
 
                 let ast_join = ast::Join {
-                    relation: right_relation.build()?,
+                    relation: match right_relation.build()? {
+                        Some(relation) => relation,
+                        None => return internal_err!("Failed to build right relation"),
+                    },
                     join_operator: self.join_operator_to_sql(
                         JoinType::Inner,
                         ast::JoinConstraint::On(ast::Expr::Value(ast::Value::Boolean(
@@ -483,6 +489,10 @@ impl Unparser<'_> {
                     relation,
                 )
             }
+            LogicalPlan::EmptyRelation(_) => {
+                relation.empty();
+                Ok(())
+            }
             LogicalPlan::Extension(_) => not_impl_err!("Unsupported operator: {plan:?}"),
             _ => not_impl_err!("Unsupported operator: {plan:?}"),
         }

diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs
@@ -78,6 +78,11 @@ fn roundtrip_expr() {
 #[test]
 fn roundtrip_statement() -> Result<()> {
     let tests: Vec<&str> = vec![
+            "select 1;",
+            "select 1 limit 0;",
+            "select ta.j1_id from j1 ta join (select 1 as j1_id) tb on ta.j1_id = tb.j1_id;",
+            "select ta.j1_id from j1 ta join (select 1 as j1_id) tb on ta.j1_id = tb.j1_id where ta.j1_id > 1;",
+            "select ta.j1_id from (select 1 as j1_id) ta;",
             "select ta.j1_id from j1 ta;",
             "select ta.j1_id from j1 ta order by ta.j1_id;",
             "select * from j1 ta order by ta.j1_id, ta.j1_string desc;",