[arrow-datafusion] branch main updated: Fix bug in simplify expressions (#214) (#7699)

[dheres]

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

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


The following commit(s) were added to refs/heads/main by this push:
     new 2abacf4a07 Fix bug in simplify expressions (#214) (#7699)
2abacf4a07 is described below

commit 2abacf4a070a8cb48ac08da73d5e77331652903d
Author: Daniël Heres <danielhe...@gmail.com>
AuthorDate: Fri Sep 29 15:24:57 2023 +0200

    Fix bug in simplify expressions (#214) (#7699)
    
    Co-authored-by: Dan Harris <1327726+thinkharder...@users.noreply.github.com>
---
 .../src/simplify_expressions/simplify_exprs.rs     | 51 ++++++++++++++++++++--
 1 file changed, 48 insertions(+), 3 deletions(-)

diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs 
b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
index 35a698b709..355d556d5d 100644
--- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
+++ b/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();