Fix bug in simplify expressions (apache#214) (apache#7699)

Co-authored-by: Dan Harris <[email protected]>

Author: 2 people

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -65,10 +65,13 @@ impl SimplifyExpressions {
     ) -> Result<LogicalPlan> {
         let schema = if !plan.inputs().is_empty() {
             DFSchemaRef::new(merge_schema(plan.inputs()))
-        } else if let LogicalPlan::TableScan(_) = plan {
+        } else if let LogicalPlan::TableScan(scan) = plan {
             // When predicates are pushed into a table scan, there needs to be
             // a schema to resolve the fields against.
-            Arc::clone(plan.schema())
+            Arc::new(DFSchema::try_from_qualified_schema(
+                &scan.table_name,
+                &scan.source.schema(),
+            )?)
         } else {
             Arc::new(DFSchema::empty())
         };
@@ -111,7 +114,7 @@ mod tests {
     use crate::simplify_expressions::utils::for_test::{
         cast_to_int64_expr, now_expr, to_timestamp_expr,
     };
-    use crate::test::test_table_scan_with_name;
+    use crate::test::{assert_fields_eq, test_table_scan_with_name};
 
     use super::*;
     use arrow::datatypes::{DataType, Field, Schema};
@@ -174,6 +177,48 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_simplify_table_full_filter_in_scan() -> Result<()> {
+        let fields = vec![
+            Field::new("a", DataType::UInt32, false),
+            Field::new("b", DataType::UInt32, false),
+            Field::new("c", DataType::UInt32, false),
+        ];
+
+        let schema = Schema::new(fields);
+
+        let table_scan = table_scan_with_filters(
+            Some("test"),
+            &schema,
+            Some(vec![0]),
+            vec![col("b").is_not_null()],
+        )?
+        .build()?;
+        assert_eq!(1, table_scan.schema().fields().len());
+        assert_fields_eq(&table_scan, vec!["a"]);
+
+        let expected = "TableScan: test projection=[a], full_filters=[Boolean(true) AS b IS NOT NULL]";
+
+        assert_optimized_plan_eq(&table_scan, expected)
+    }
+
+    #[test]
+    fn test_simplify_filter_pushdown() -> Result<()> {
+        let table_scan = test_table_scan();
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .project(vec![col("a")])?
+            .filter(and(col("b").gt(lit(1)), col("b").gt(lit(1))))?
+            .build()?;
+
+        assert_optimized_plan_eq(
+            &plan,
+            "\
+	        Filter: test.b > Int32(1)\
+            \n  Projection: test.a\
+            \n    TableScan: test",
+        )
+    }
+
     #[test]
     fn test_simplify_optimized_plan() -> Result<()> {
         let table_scan = test_table_scan();