diff --git a/datafusion/src/optimizer/filter_push_down.rs b/datafusion/src/optimizer/filter_push_down.rs index 85d1f812f41a..dc4d5e993a38 100644 --- a/datafusion/src/optimizer/filter_push_down.rs +++ b/datafusion/src/optimizer/filter_push_down.rs @@ -322,6 +322,10 @@ fn optimize(plan: &LogicalPlan, mut state: State) -> Result { // sort is filter-commutable push_down(&state, plan) } + LogicalPlan::Union { .. } => { + // union all is filter-commutable + push_down(&state, plan) + } LogicalPlan::Limit { input, .. } => { // limit is _not_ filter-commutable => collect all columns from its input let used_columns = input @@ -766,6 +770,24 @@ mod tests { Ok(()) } + #[test] + fn union_all() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(&table_scan) + .union(LogicalPlanBuilder::from(&table_scan).build()?)? + .filter(col("a").eq(lit(1i64)))? + .build()?; + // filter appears below Union + let expected = "\ + Union\ + \n Filter: #a Eq Int64(1)\ + \n TableScan: test projection=None\ + \n Filter: #a Eq Int64(1)\ + \n TableScan: test projection=None"; + assert_optimized_plan_eq(&plan, expected); + Ok(()) + } + /// verifies that filters with the same columns are correctly placed #[test] fn filter_2_breaks_limits() -> Result<()> {