[arrow-datafusion] branch master updated: bugfix: fix propagating empty_relation generates an illegal plan (#5219)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new e16650073 bugfix: fix propagating empty_relation generates an illegal 
plan (#5219)
e16650073 is described below

commit e16650073320dec3b0517aba12c6ce738eb4a503
Author: yujie.zhang <yukkit.zh...@cnosdb.com>
AuthorDate: Thu Feb 9 02:43:15 2023 +0800

    bugfix: fix propagating empty_relation generates an illegal plan (#5219)
---
 .../optimizer/src/propagate_empty_relation.rs      | 40 ++++++++++++++++++++--
 datafusion/optimizer/src/test/mod.rs               | 12 ++++---
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/datafusion/optimizer/src/propagate_empty_relation.rs 
b/datafusion/optimizer/src/propagate_empty_relation.rs
index d533803e5..8214298a0 100644
--- a/datafusion/optimizer/src/propagate_empty_relation.rs
+++ b/datafusion/optimizer/src/propagate_empty_relation.rs
@@ -104,7 +104,7 @@ impl OptimizerRule for PropagateEmptyRelation {
                         schema: plan.schema().clone(),
                     })));
                 } else if new_inputs.len() == 1 {
-                    let child = (**(union.inputs.get(0).unwrap())).clone();
+                    let child = (*new_inputs[0]).clone();
                     if child.schema().eq(plan.schema()) {
                         return Ok(Some(child));
                     } else {
@@ -200,11 +200,12 @@ mod tests {
     use crate::eliminate_filter::EliminateFilter;
     use crate::optimizer::Optimizer;
     use crate::test::{
-        assert_optimized_plan_eq, test_table_scan, test_table_scan_with_name,
+        assert_optimized_plan_eq, test_table_scan, test_table_scan_fields,
+        test_table_scan_with_name,
     };
     use crate::OptimizerContext;
     use arrow::datatypes::{DataType, Field, Schema};
-    use datafusion_common::{Column, ScalarValue};
+    use datafusion_common::{Column, DFField, DFSchema, ScalarValue};
     use datafusion_expr::logical_plan::table_scan;
     use datafusion_expr::{
         binary_expr, col, lit, logical_plan::builder::LogicalPlanBuilder, 
Expr, JoinType,
@@ -392,4 +393,37 @@ mod tests {
         let expected = "EmptyRelation";
         assert_together_optimized_plan_eq(&plan, expected)
     }
+
+    #[test]
+    fn test_empty_with_non_empty() -> Result<()> {
+        let table_scan = test_table_scan()?;
+
+        let fields = test_table_scan_fields()
+            .into_iter()
+            .map(DFField::from)
+            .collect();
+
+        let empty = LogicalPlan::EmptyRelation(EmptyRelation {
+            produce_one_row: false,
+            schema: Arc::new(DFSchema::new_with_metadata(fields, 
Default::default())?),
+        });
+
+        let one = LogicalPlanBuilder::from(empty.clone()).build()?;
+        let two = LogicalPlanBuilder::from(table_scan).build()?;
+        let three = LogicalPlanBuilder::from(empty).build()?;
+
+        // Union
+        //  EmptyRelation
+        //  TableScan: test
+        //  EmptyRelation
+        let plan = LogicalPlanBuilder::from(one)
+            .union(two)?
+            .union(three)?
+            .build()?;
+
+        let expected = "Projection: a, b, c\
+        \n  TableScan: test";
+
+        assert_together_optimized_plan_eq(&plan, expected)
+    }
 }
diff --git a/datafusion/optimizer/src/test/mod.rs 
b/datafusion/optimizer/src/test/mod.rs
index bd8c1f98a..cec21da31 100644
--- a/datafusion/optimizer/src/test/mod.rs
+++ b/datafusion/optimizer/src/test/mod.rs
@@ -24,13 +24,17 @@ use std::sync::Arc;
 
 pub mod user_defined;
 
-/// some tests share a common table with different names
-pub fn test_table_scan_with_name(name: &str) -> Result<LogicalPlan> {
-    let schema = Schema::new(vec![
+pub fn test_table_scan_fields() -> Vec<Field> {
+    vec![
         Field::new("a", DataType::UInt32, false),
         Field::new("b", DataType::UInt32, false),
         Field::new("c", DataType::UInt32, false),
-    ]);
+    ]
+}
+
+/// some tests share a common table with different names
+pub fn test_table_scan_with_name(name: &str) -> Result<LogicalPlan> {
+    let schema = Schema::new(test_table_scan_fields());
     table_scan(Some(name), &schema, None)?.build()
 }