From 0e99c3c8a9db6135a88065fe1c1ce8912b5cafbf Mon Sep 17 00:00:00 2001 From: Jon Mease Date: Sat, 27 Aug 2022 10:07:01 -0400 Subject: [PATCH 1/3] Add failing union_all_on_projection test --- datafusion/optimizer/src/filter_push_down.rs | 25 ++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/datafusion/optimizer/src/filter_push_down.rs b/datafusion/optimizer/src/filter_push_down.rs index 2ac5b6e3b0ff..22abf6cceb8e 100644 --- a/datafusion/optimizer/src/filter_push_down.rs +++ b/datafusion/optimizer/src/filter_push_down.rs @@ -1009,6 +1009,31 @@ mod tests { Ok(()) } + #[test] + fn union_all_on_projection() -> Result<()> { + let table_scan = test_table_scan()?; + let table = LogicalPlanBuilder::from(table_scan.clone()) + .project_with_alias(vec![col("a").alias("b")], Some("test2".to_string()))?; + + let plan = table + .union(table.build()?)? + .filter(col("b").eq(lit(1i64)))? + .build()?; + + println!("Input plan:\n{:?}", plan); + + // filter appears below Union + let expected = "\ + Filter: #b = Int64(1)\ + \n Union\ + \n Projection: #test.a AS b, alias=test2\ + \n TableScan: test\ + \n Projection: #test.a AS b, alias=test2\ + \n TableScan: test"; + assert_optimized_plan_eq(&plan, expected); + Ok(()) + } + /// verifies that filters with the same columns are correctly placed #[test] fn filter_2_breaks_limits() -> Result<()> { From 6e8c897600ae16d78a2f69db2571aeca786b8cb8 Mon Sep 17 00:00:00 2001 From: Jon Mease Date: Sat, 27 Aug 2022 11:41:43 -0400 Subject: [PATCH 2/3] Update projection step to replace both qualified and unqualified columns --- datafusion/optimizer/src/filter_push_down.rs | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/datafusion/optimizer/src/filter_push_down.rs b/datafusion/optimizer/src/filter_push_down.rs index 22abf6cceb8e..dcae790aa528 100644 --- a/datafusion/optimizer/src/filter_push_down.rs +++ b/datafusion/optimizer/src/filter_push_down.rs @@ -379,14 +379,18 @@ fn optimize(plan: &LogicalPlan, mut state: State) -> Result { .fields() .iter() .enumerate() - .map(|(i, field)| { + .flat_map(|(i, field)| { // strip alias, as they should not be part of filters let expr = match &expr[i] { Expr::Alias(expr, _) => expr.as_ref().clone(), expr => expr.clone(), }; - (field.qualified_name(), expr) + // Convert both qualified and unqualified fields + [ + (field.name().clone(), expr.clone()), + (field.qualified_name(), expr), + ] }) .collect::>(); @@ -1020,15 +1024,14 @@ mod tests { .filter(col("b").eq(lit(1i64)))? .build()?; - println!("Input plan:\n{:?}", plan); - // filter appears below Union let expected = "\ - Filter: #b = Int64(1)\ - \n Union\ - \n Projection: #test.a AS b, alias=test2\ + Union\ + \n Projection: #test.a AS b, alias=test2\ + \n Filter: #test.a = Int64(1)\ \n TableScan: test\ - \n Projection: #test.a AS b, alias=test2\ + \n Projection: #test.a AS b, alias=test2\ + \n Filter: #test.a = Int64(1)\ \n TableScan: test"; assert_optimized_plan_eq(&plan, expected); Ok(()) From 2cecc9519321ac8278790faed4aae1c45ec84a05 Mon Sep 17 00:00:00 2001 From: Jon Mease Date: Sat, 27 Aug 2022 13:09:01 -0400 Subject: [PATCH 3/3] Remove redundant clone --- datafusion/optimizer/src/filter_push_down.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/optimizer/src/filter_push_down.rs b/datafusion/optimizer/src/filter_push_down.rs index dcae790aa528..96d33e972e76 100644 --- a/datafusion/optimizer/src/filter_push_down.rs +++ b/datafusion/optimizer/src/filter_push_down.rs @@ -1016,7 +1016,7 @@ mod tests { #[test] fn union_all_on_projection() -> Result<()> { let table_scan = test_table_scan()?; - let table = LogicalPlanBuilder::from(table_scan.clone()) + let table = LogicalPlanBuilder::from(table_scan) .project_with_alias(vec![col("a").alias("b")], Some("test2".to_string()))?; let plan = table