extend a scan projection if some pushed filters become unsupported

Consider the following optimizer-run scenario:

1. `supports_filters_pushdown` returns `Exact` on some filter, e.g. "a = 1",
   where the column "a" is not required by the query projection.

2. "a" is removed from the table provider projection by "optimize projection"
   rule.

3. `supports_filters_pushdown` changes a decision and returns `Inexact` on
   this filter the next time. e.g., input filters are changed and it prefers
   to use a new one.

4. "a" is not returned to the table provider projection which leads to filter
   that references a column which is not a part of the input schema.

This patch fixes issue introducing the following logic within a filter push-down rule:

1. Collect columns that are not used in the current table provider scan projection,
   but required for filter expressions. Call it `additional_projection`.

2. If `additional_projection` is empty -- leave logic as is prior the patch.

3. Otherwise extend a table provider projection and wrap a plan with
   an additional projection node to preserve schema used prior to the rule.

Author: askalt

datafusion/optimizer/src/push_down_filter.rs

@@ -38,7 +38,8 @@ use datafusion_expr::utils::{
     conjunction, expr_to_columns, split_conjunction, split_conjunction_owned,
 };
 use datafusion_expr::{
-    BinaryExpr, Expr, Filter, Operator, Projection, TableProviderFilterPushDown, and, or,
+    BinaryExpr, Expr, Filter, LogicalPlanBuilder, Operator, Projection,
+    TableProviderFilterPushDown, and, or,
 };
 
 use crate::optimizer::ApplyOrder;
@@ -1128,7 +1129,7 @@ impl OptimizerRule for PushDownFilter {
             LogicalPlan::TableScan(scan) => {
                 let filter_predicates: Vec<_> = split_conjunction(&filter.predicate)
                     .into_iter()
-                    // Add already pushed filters.
+                    // Add already pushed filters to ensure that the rule is idempotent.
                     .chain(scan.filters.iter())
                     .unique()
                     .collect();
@@ -1162,26 +1163,106 @@ impl OptimizerRule for PushDownFilter {
                 let new_scan_filters: Vec<Expr> =
                     new_scan_filters.unique().cloned().collect();
 
+                let source_schema = scan.source.schema();
+                let mut additional_projection = HashSet::new();
+
                 // Compose predicates to be of `Unsupported` or `Inexact` pushdown type, and also include volatile filters
                 let new_predicate: Vec<Expr> = zip
-                    .filter(|(_, res)| res != &TableProviderFilterPushDown::Exact)
+                    .filter(|(expr, res)| {
+                        if *res == TableProviderFilterPushDown::Exact {
+                            return false;
+                        }
+                        // For each not exactly supported filter we must ensure that all columns are projected,
+                        // so we collect all columns which are not currently projected.
+                        expr.apply(|expr| {
+                            if let Expr::Column(column) = expr
+                                && let Ok(idx) = source_schema.index_of(column.name())
+                                && scan
+                                    .projection
+                                    .as_ref()
+                                    .is_some_and(|p| !p.contains(&idx))
+                            {
+                                additional_projection.insert(idx);
+                            }
+                            Ok(TreeNodeRecursion::Continue)
+                        })
+                        .unwrap();
+                        true
+                    })
                     .map(|(pred, _)| pred)
                     .chain(volatile_filters)
                     .cloned()
                     .collect();
 
-                let new_scan = LogicalPlan::TableScan(TableScan {
-                    filters: new_scan_filters,
-                    ..scan
-                });
-
-                Transformed::yes(new_scan).transform_data(|new_scan| {
-                    if let Some(predicate) = conjunction(new_predicate) {
-                        make_filter(predicate, Arc::new(new_scan)).map(Transformed::yes)
+                // Wraps with a filter if some filters are not supported exactly.
+                let filtered = move |plan| {
+                    if let Some(new_predicate) = conjunction(new_predicate) {
+                        Filter::try_new(new_predicate, Arc::new(plan))
+                            .map(LogicalPlan::Filter)
                     } else {
-                        Ok(Transformed::no(new_scan))
+                        Ok(plan)
                     }
-                })
+                };
+
+                if additional_projection.is_empty() {
+                    // No additional projection is required.
+                    let new_scan = LogicalPlan::TableScan(TableScan {
+                        filters: new_scan_filters,
+                        ..scan
+                    });
+                    return filtered(new_scan).map(Transformed::yes);
+                }
+
+                let scan_table_name = &scan.table_name;
+                let new_scan = filtered(
+                    LogicalPlanBuilder::scan_with_filters_fetch(
+                        scan_table_name.clone(),
+                        Arc::clone(&scan.source),
+                        scan.projection.clone().map(|mut projection| {
+                            // Extend a projection.
+                            projection.extend(additional_projection);
+                            projection
+                        }),
+                        new_scan_filters,
+                        scan.fetch,
+                    )?
+                    .build()?,
+                )?;
+
+                // Project fields required by the initial projection.
+                let new_plan = LogicalPlan::Projection(Projection::try_new_with_schema(
+                    scan.projection
+                        .as_ref()
+                        .map(|projection| {
+                            projection
+                                .iter()
+                                .cloned()
+                                .map(|idx| {
+                                    Expr::Column(Column::new(
+                                        Some(scan_table_name.clone()),
+                                        source_schema.field(idx).name(),
+                                    ))
+                                })
+                                .collect()
+                        })
+                        .unwrap_or_else(|| {
+                            source_schema
+                                .fields()
+                                .iter()
+                                .map(|field| {
+                                    Expr::Column(Column::new(
+                                        Some(scan_table_name.clone()),
+                                        field.name(),
+                                    ))
+                                })
+                                .collect()
+                        }),
+                    Arc::new(new_scan),
+                    // Preserve a projected schema metadata.
+                    scan.projected_schema,
+                )?);
+
+                Ok(Transformed::yes(new_plan))
             }
             LogicalPlan::Extension(extension_plan) => {
                 // This check prevents the Filter from being removed when the extension node has no children,
@@ -3205,7 +3286,7 @@ mod tests {
         let plan = table_scan_with_pushdown_provider_builder(
             TableProviderFilterPushDown::Inexact,
             vec![col("a").eq(lit(10i64)), col("b").gt(lit(11i64))],
-            Some(vec![0]),
+            Some(vec![0, 1]),
         )?
         .filter(and(col("a").eq(lit(10i64)), col("b").gt(lit(11i64))))?
         .project(vec![col("a"), col("b")])?
@@ -3216,7 +3297,7 @@ mod tests {
             @r"
         Projection: a, b
           Filter: a = Int64(10) AND b > Int64(11)
-            TableScan: test projection=[a], partial_filters=[a = Int64(10), b > Int64(11)]
+            TableScan: test projection=[a, b], partial_filters=[a = Int64(10), b > Int64(11)]
         "
         )
     }
@@ -3226,7 +3307,7 @@ mod tests {
         let plan = table_scan_with_pushdown_provider_builder(
             TableProviderFilterPushDown::Exact,
             vec![],
-            Some(vec![0]),
+            Some(vec![0, 1]),
         )?
         .filter(and(col("a").eq(lit(10i64)), col("b").gt(lit(11i64))))?
         .project(vec![col("a"), col("b")])?
@@ -3236,7 +3317,7 @@ mod tests {
             plan,
             @r"
         Projection: a, b
-          TableScan: test projection=[a], full_filters=[a = Int64(10), b > Int64(11)]
+          TableScan: test projection=[a, b], full_filters=[a = Int64(10), b > Int64(11)]
         "
         )
     }
@@ -4292,4 +4373,34 @@ mod tests {
         "
         )
     }
+
+    #[test]
+    fn test_projection_is_updated_when_filter_becomes_unsupported() -> Result<()> {
+        let test_provider = PushDownProvider {
+            filter_support: TableProviderFilterPushDown::Unsupported,
+        };
+
+        let projected_schema = test_provider.schema().project(&[0])?;
+        let table_scan = LogicalPlan::TableScan(TableScan {
+            table_name: "test".into(),
+            // Emulate that there were pushed filters but now
+            // provider cannot support it.
+            filters: vec![col("b").eq(lit(1i64))],
+            projected_schema: Arc::new(DFSchema::try_from(projected_schema)?),
+            projection: Some(vec![0]),
+            source: Arc::new(test_provider),
+            fetch: None,
+        });
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .filter(col("a").eq(lit(1i64)))?
+            .build()?;
+
+        assert_optimized_plan_equal!(plan,
+            @r"
+        Projection: test.a
+          Filter: a = Int64(1) AND b = Int64(1)
+            TableScan: test projection=[a, b]"
+        )
+    }
 }

datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part

@@ -60,9 +60,9 @@ logical_plan
 04)------Inner Join: lineitem.l_partkey = part.p_partkey Filter: part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2) AND part.p_size <= Int32(15)
 05)--------Projection: lineitem.l_partkey, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount
 06)----------Filter: (lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2)) AND (lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG")) AND lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON")
-07)------------TableScan: lineitem projection=[l_partkey, l_quantity, l_extendedprice, l_discount, l_shipinstruct, l_shipmode], partial_filters=[lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG"), lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON"), lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2)]
+07)------------TableScan: lineitem projection=[l_partkey, l_quantity, l_extendedprice, l_discount, l_shipinstruct, l_shipmode], partial_filters=[lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2), lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG"), lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON")]
 08)--------Filter: (part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND part.p_size <= Int32(15)) AND part.p_size >= Int32(1)
-09)----------TableScan: part projection=[p_partkey, p_brand, p_size, p_container], partial_filters=[part.p_size >= Int32(1), part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND part.p_size <= Int32(15)]
+09)----------TableScan: part projection=[p_partkey, p_brand, p_size, p_container], partial_filters=[part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND part.p_size <= Int32(15), part.p_size >= Int32(1)]
 physical_plan
 01)ProjectionExec: expr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@0 as revenue]
 02)--AggregateExec: mode=Final, gby=[], aggr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]

datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part

@@ -64,7 +64,7 @@ logical_plan
 06)----------Inner Join:  Filter: CAST(customer.c_acctbal AS Decimal128(19, 6)) > __scalar_sq_2.avg(customer.c_acctbal)
 07)------------Projection: customer.c_phone, customer.c_acctbal
 08)--------------LeftAnti Join: customer.c_custkey = __correlated_sq_1.o_custkey
-09)----------------Filter: substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])
+09)----------------Filter: substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")]) AND Boolean(true)
 10)------------------TableScan: customer projection=[c_custkey, c_phone, c_acctbal], partial_filters=[substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")]), Boolean(true)]
 11)----------------SubqueryAlias: __correlated_sq_1
 12)------------------TableScan: orders projection=[o_custkey]
@@ -87,7 +87,7 @@ physical_plan
 11)--------------------CoalescePartitionsExec
 12)----------------------HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(c_custkey@0, o_custkey@0)], projection=[c_phone@1, c_acctbal@2]
 13)------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4
-14)--------------------------FilterExec: substr(c_phone@1, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17])
+14)--------------------------FilterExec: substr(c_phone@1, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17]) AND true
 15)----------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 16)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_phone, c_acctbal], file_type=csv, has_header=false
 17)------------------------RepartitionExec: partitioning=Hash([o_custkey@0], 4), input_partitions=4