Fix pushdown filter with projection

Dandandan · Dec 3, 2023 · 5c2c8cc · 5c2c8cc
1 parent cbb7e34
commit 5c2c8cc
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/datafusion/core/src/dataframe/parquet.rs b/datafusion/core/src/dataframe/parquet.rs
@@ -103,9 +103,10 @@ mod tests {
             .filter(col("id").eq(lit(1)))?
             .select_columns(&["bool_col", "int_col"])?;
 
-        let plan = df.explain(false, false)?.collect().await?;
+        let plan = df.explain(true, false)?.collect().await?;
         // Filters all the way to Parquet
         let formatted = pretty::pretty_format_batches(&plan)?.to_string();
+        println!("{}", formatted);
         assert!(formatted.contains("FilterExec: id@0 = 1"));
 
         Ok(())

diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs
@@ -23,14 +23,14 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::Alias;
 use datafusion_expr::utils::{conjunction, split_conjunction, split_conjunction_owned};
-use datafusion_expr::Volatility;
 use datafusion_expr::{
     and,
     expr_rewriter::replace_col,
     logical_plan::{CrossJoin, Join, JoinType, LogicalPlan, TableScan, Union},
     or, BinaryExpr, Expr, Filter, Operator, ScalarFunctionDefinition,
     TableProviderFilterPushDown,
 };
+use datafusion_expr::{Projection, Volatility};
 use itertools::Itertools;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
@@ -907,9 +907,17 @@ impl OptimizerRule for PushDownFilter {
                     Some(predicate) => LogicalPlan::Filter(Filter::try_new(
                         predicate,
                         Arc::new(new_scan),
-                        None,
+                        filter.projected_schema.clone(),
                     )?),
-                    None => new_scan,
+                    None => match &filter.projected_schema {
+                        Some(projected_schema) => {
+                            LogicalPlan::Projection(Projection::new_from_schema(
+                                Arc::new(new_scan),
+                                projected_schema.clone(),
+                            ))
+                        }
+                        None => new_scan,
+                    },
                 }
             }
             LogicalPlan::Extension(extension_plan) => {