fix: run logical optimizer rules for TableScan expressions (#4614)

apache · Dec 14, 2022 · 40e6a67 · 40e6a67
1 parent 508ba80
commit 40e6a67
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 5 deletions.
diff --git a/datafusion/core/tests/sql/subqueries.rs b/datafusion/core/tests/sql/subqueries.rs
@@ -398,7 +398,7 @@ order by cntrycode;"#;
     \n                Projection: AVG(customer.c_acctbal) AS __value\
     \n                  Aggregate: groupBy=[[]], aggr=[[AVG(customer.c_acctbal)]]\
     \n                    Filter: customer.c_acctbal > Decimal128(Some(0),15,2) AND substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8(\"13\"), Utf8(\"31\"), Utf8(\"23\"), Utf8(\"29\"), Utf8(\"30\"), Utf8(\"18\"), Utf8(\"17\")])\
-    \n                      TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[CAST(customer.c_acctbal AS Decimal128(30, 15)) > Decimal128(Some(0),30,15), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8(\"13\"), Utf8(\"31\"), Utf8(\"23\"), Utf8(\"29\"), Utf8(\"30\"), Utf8(\"18\"), Utf8(\"17\")]), customer.c_acctbal > Decimal128(Some(0),15,2)]";
+    \n                      TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[customer.c_acctbal > Decimal128(Some(0),15,2) AS customer.c_acctbal > Decimal128(Some(0),30,15), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8(\"13\"), Utf8(\"31\"), Utf8(\"23\"), Utf8(\"29\"), Utf8(\"30\"), Utf8(\"18\"), Utf8(\"17\")]), customer.c_acctbal > Decimal128(Some(0),15,2)]";
     assert_eq!(expected, actual);
 
     // assert data

diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
@@ -262,9 +262,9 @@ impl LogicalPlan {
                 .collect(),
             LogicalPlan::Sort(Sort { expr, .. }) => expr.clone(),
             LogicalPlan::Extension(extension) => extension.node.expressions(),
+            LogicalPlan::TableScan(TableScan { filters, .. }) => filters.clone(),
             // plans without expressions
-            LogicalPlan::TableScan(_)
-            | LogicalPlan::EmptyRelation(_)
+            LogicalPlan::EmptyRelation(_)
             | LogicalPlan::Subquery(_)
             | LogicalPlan::SubqueryAlias(_)
             | LogicalPlan::Limit(_)

diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs
@@ -25,7 +25,7 @@ use crate::logical_plan::{
     Limit, Partitioning, Prepare, Projection, Repartition, Sort, Subquery, SubqueryAlias,
     Union, Values, Window,
 };
-use crate::{Cast, Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder};
+use crate::{Cast, Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder, TableScan};
 use arrow::datatypes::{DataType, TimeUnit};
 use datafusion_common::{
     Column, DFField, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue,
@@ -587,8 +587,14 @@ pub fn from_plan(
             data_types: data_types.clone(),
             input: Arc::new(inputs[0].clone()),
         })),
+        LogicalPlan::TableScan(ts) => {
+            assert!(inputs.is_empty(), "{:?}  should have no inputs", plan);
+            Ok(LogicalPlan::TableScan(TableScan {
+                filters: expr.to_vec(),
+                ..ts.clone()
+            }))
+        }
         LogicalPlan::EmptyRelation(_)
-        | LogicalPlan::TableScan { .. }
         | LogicalPlan::CreateExternalTable(_)
         | LogicalPlan::DropTable(_)
         | LogicalPlan::DropView(_)

diff --git a/datafusion/optimizer/src/type_coercion.rs b/datafusion/optimizer/src/type_coercion.rs
@@ -95,6 +95,12 @@ fn optimize_internal(
         },
     );
 
+    if let LogicalPlan::TableScan(ts) = plan {
+        let source_schema =
+            DFSchema::try_from_qualified_schema(&ts.table_name, &ts.source.schema())?;
+        schema.merge(&source_schema);
+    }
+
     // merge the outer schema for correlated subqueries
     // like case:
     // select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where t2.c2=t1.c3)