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,
@@ -3207,7 +3288,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")])?
@@ -3218,7 +3299,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)]
         "
         )
     }
@@ -3228,7 +3309,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")])?
@@ -3238,7 +3319,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)]
         "
         )
     }
@@ -4294,4 +4375,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]"
+        )
+    }
 }