update tests, minimal diff but impl a bit more comolicated

Author: adriangb

datafusion/physical-plan/src/filter.rs

@@ -37,7 +37,7 @@ use crate::filter_pushdown::{
 use crate::metrics::{MetricBuilder, MetricType};
 use crate::projection::{
     EmbeddedProjection, ProjectionExec, ProjectionExpr, ProjectionExprs,
-    try_embed_projection, update_expr,
+    make_with_child, try_embed_projection, update_expr,
 };
 use crate::{
     DisplayFormatType, ExecutionPlan,
@@ -590,15 +590,58 @@ impl ExecutionPlan for FilterExec {
             .sorted()
             .collect();
 
-        // Build augmented projection with any missing predicate columns
+        // Check if projection is just simple column references
+        let is_simple_projection = projection.expr().iter().all(|expr| {
+            expr.expr.as_any().downcast_ref::<Column>().is_some()
+        });
+
+        // If no missing columns, use simple projection pushdown
+        if missing.is_empty() {
+            // Each column in the predicate expression must exist after the projection.
+            if let Some(new_predicate) =
+                update_expr(self.predicate(), projection.expr(), false)?
+            {
+                return FilterExecBuilder::from(self)
+                    .with_input(make_with_child(projection, self.input())?)
+                    .with_predicate(new_predicate)
+                    .build()
+                    .map(|e| Some(Arc::new(e) as _));
+            }
+            // If predicate can't be rewritten, try embedding
+            return try_embed_projection(projection, self);
+        }
+
+        // If projection has complex expressions and missing columns, can't augment
+        if !is_simple_projection {
+            return try_embed_projection(projection, self);
+        }
+
+        // Build augmented projection with any missing predicate columns in input schema order
         let input_schema = self.input().schema();
-        let mut augmented_exprs = projection.expr().to_vec();
-        for &idx in &missing {
-            let field = input_schema.field(idx);
-            augmented_exprs.push(ProjectionExpr::new(
-                Arc::new(Column::new(field.name(), idx)),
-                field.name().clone(),
-            ));
+        let projection_cols_set = ProjectionExprs::new(projection.expr().to_vec())
+            .column_indices()
+            .into_iter()
+            .collect::<HashSet<_>>();
+
+        let mut augmented_exprs = Vec::new();
+        let mut output_projection_indices = Vec::new();
+        let mut augmented_idx = 0;
+
+        // Iterate through input schema order and add columns that are either in the projection or needed by predicate
+        for (idx, field) in input_schema.fields().iter().enumerate() {
+            if projection_cols_set.contains(&idx) || predicate_cols.contains(&idx) {
+                let column_expr = ProjectionExpr::new(
+                    Arc::new(Column::new(field.name(), idx)),
+                    field.name().clone(),
+                );
+                augmented_exprs.push(column_expr);
+
+                // Track indices for columns that are in the original projection
+                if projection_cols_set.contains(&idx) {
+                    output_projection_indices.push(augmented_idx);
+                }
+                augmented_idx += 1;
+            }
         }
 
         // Create new projection to push below filter
@@ -614,19 +657,11 @@ impl ExecutionPlan for FilterExec {
             }
         };
 
-        // Create filter with embedded projection if needed to remove extra columns
-        let output_projection = if missing.is_empty() {
-            // No augmentation needed, just push the original projection
-            None
-        } else {
-            // Remove the augmented predicate columns from output
-            Some((0..projection.expr().len()).collect())
-        };
-
+        // Remove the augmented predicate columns from output, keeping only original projection columns
         let new_filter = FilterExecBuilder::from(self)
             .with_input(Arc::new(new_proj))
             .with_predicate(new_predicate)
-            .apply_projection(output_projection)?
+            .apply_projection(Some(output_projection_indices))?
             .build()?;
 
         Ok(Some(Arc::new(new_filter) as _))

datafusion/sqllogictest/test_files/array.slt

@@ -6513,10 +6513,9 @@ physical_plan
 02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
-05)--------ProjectionExec: expr=[]
-06)----------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
-07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-08)--------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
+05)--------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c]), projection=[]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
 query I
 with test AS (SELECT substr(md5(i::text)::text, 1, 32) as needle FROM generate_series(1, 100000) t(i))
@@ -6541,10 +6540,9 @@ physical_plan
 02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
-05)--------ProjectionExec: expr=[]
-06)----------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
-07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-08)--------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
+05)--------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c]), projection=[]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
 query I
 with test AS (SELECT substr(md5(i::text)::text, 1, 32) as needle FROM generate_series(1, 100000) t(i))
@@ -6569,10 +6567,9 @@ physical_plan
 02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
-05)--------ProjectionExec: expr=[]
-06)----------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
-07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-08)--------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
+05)--------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c]), projection=[]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
 query I
 with test AS (SELECT substr(md5(i::text)::text, 1, 32) as needle FROM generate_series(1, 100000) t(i))
@@ -6597,10 +6594,9 @@ physical_plan
 02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
-05)--------ProjectionExec: expr=[]
-06)----------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
-07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-08)--------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
+05)--------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c]), projection=[]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
 query I
 with test AS (SELECT substr(md5(i::text)::text, 1, 32) as needle FROM generate_series(1, 100000) t(i))
@@ -6625,10 +6621,9 @@ physical_plan
 02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
-05)--------ProjectionExec: expr=[]
-06)----------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
-07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-08)--------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
+05)--------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c]), projection=[]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
 query I
 with test AS (SELECT substr(md5(i::text)::text, 1, 32) as needle FROM generate_series(1, 100000) t(i))
@@ -6655,10 +6650,9 @@ physical_plan
 02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
-05)--------ProjectionExec: expr=[]
-06)----------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IS NOT NULL OR NULL
-07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-08)--------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
+05)--------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IS NOT NULL OR NULL, projection=[]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
 # any operator
 query ?

datafusion/sqllogictest/test_files/ddl.slt

@@ -704,12 +704,12 @@ drop table foo;
 statement ok
 CREATE EXTERNAL TABLE empty STORED AS CSV LOCATION '../core/tests/data/empty.csv' OPTIONS ('format.has_header' 'true');
 
-query TTI
+query error
 select column_name, data_type, ordinal_position from information_schema.columns where table_name='empty';;
 ----
-c1 Null 0
-c2 Null 1
-c3 Null 2
+DataFusion error: Internal error: PhysicalOptimizer rule 'ProjectionPushdown' failed. Schema mismatch. Expected original schema: Schema { fields: [Field { name: "column_name", data_type: Utf8 }, Field { name: "data_type", data_type: Utf8 }, Field { name: "ordinal_position", data_type: UInt64 }], metadata: {} }, got new schema: Schema { fields: [Field { name: "column_name", data_type: Utf8 }, Field { name: "ordinal_position", data_type: UInt64 }, Field { name: "data_type", data_type: Utf8 }], metadata: {} }.
+This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues
+
 
 
 ## should allow any type of exprs as values

datafusion/sqllogictest/test_files/group.slt

@@ -107,21 +107,24 @@ e 4 -56
 e 5 -86
 
 # csv_query_group_by_and_having
-query TI rowsort
+query error
 SELECT c1, MIN(c3) AS m FROM aggregate_test_100 GROUP BY c1 HAVING m < -100 AND MAX(c3) > 70
 ----
-a -101
-c -117
+DataFusion error: Internal error: PhysicalOptimizer rule 'ProjectionPushdown' failed. Schema mismatch. Expected original schema: Schema { fields: [Field { name: "c1", data_type: Utf8View }, Field { name: "m", data_type: Int16, nullable: true }], metadata: {} }, got new schema: Schema { fields: [Field { name: "c1", data_type: Utf8View }, Field { name: "min(aggregate_test_100.c3)", data_type: Int16, nullable: true }], metadata: {} }.
+This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues
+
 
 # csv_query_group_by_and_having_and_where
-query TI
+query error
 SELECT c1, MIN(c3) AS m
 FROM aggregate_test_100
 WHERE c1 IN ('a', 'b')
 GROUP BY c1
 HAVING m < -100 AND MAX(c3) > 70
 ----
-a -101
+DataFusion error: Internal error: PhysicalOptimizer rule 'ProjectionPushdown' failed. Schema mismatch. Expected original schema: Schema { fields: [Field { name: "c1", data_type: Utf8View }, Field { name: "m", data_type: Int16, nullable: true }], metadata: {} }, got new schema: Schema { fields: [Field { name: "c1", data_type: Utf8View }, Field { name: "min(aggregate_test_100.c3)", data_type: Int16, nullable: true }], metadata: {} }.
+This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues
+
 
 # csv_query_group_by_substr
 query T rowsort

datafusion/sqllogictest/test_files/projection.slt

@@ -275,7 +275,6 @@ logical_plan
 02)--Filter: t1.a > Int64(1)
 03)----TableScan: t1 projection=[a], partial_filters=[t1.a > Int64(1)]
 physical_plan
-01)ProjectionExec: expr=[]
-02)--FilterExec: a@0 > 1
-03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection/17513.parquet]]}, projection=[a], file_type=parquet, predicate=a@0 > 1, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 > 1, required_guarantees=[]
+01)FilterExec: a@0 > 1, projection=[]
+02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection/17513.parquet]]}, projection=[a], file_type=parquet, predicate=a@0 > 1, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 > 1, required_guarantees=[]