bugfix: fix propagating empty_relation generates an illegal plan (#5219)

apache · Feb 8, 2023 · e166500 · e166500
1 parent b4cf60a
commit e166500
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 7 deletions.
diff --git 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
@@ -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()
 }