From cc14ae5afe06471e00c4ed4714418cd07ad015e6 Mon Sep 17 00:00:00 2001
From: Daniel Heres <danielheres@gmail.com>
Date: Mon, 14 Jun 2021 08:40:13 +0200
Subject: [PATCH 1/3] Push down filter through UNION

---
 datafusion/src/optimizer/filter_push_down.rs | 24 ++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/datafusion/src/optimizer/filter_push_down.rs b/datafusion/src/optimizer/filter_push_down.rs
index 85d1f812f41a..1747bfbc0552 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<LogicalPlan> {
             // 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,26 @@ mod tests {
         Ok(())
     }
 
+
+    /// verifies that when two limits are in place, we jump neither
+    #[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 does not just any of the limits
+        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<()> {

From 78e37426dac7e8277fcc76de1a6c5bce17e22524 Mon Sep 17 00:00:00 2001
From: Daniel Heres <danielheres@gmail.com>
Date: Mon, 14 Jun 2021 08:43:59 +0200
Subject: [PATCH 2/3] Fix comments

---
 datafusion/src/optimizer/filter_push_down.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/datafusion/src/optimizer/filter_push_down.rs b/datafusion/src/optimizer/filter_push_down.rs
index 1747bfbc0552..cb3593a47abc 100644
--- a/datafusion/src/optimizer/filter_push_down.rs
+++ b/datafusion/src/optimizer/filter_push_down.rs
@@ -771,7 +771,6 @@ mod tests {
     }
 
 
-    /// verifies that when two limits are in place, we jump neither
     #[test]
     fn union_all() -> Result<()> {
         let table_scan = test_table_scan()?;
@@ -779,7 +778,7 @@ mod tests {
             .union(LogicalPlanBuilder::from(&table_scan).build()?)?
             .filter(col("a").eq(lit(1i64)))?
             .build()?;
-        // filter does not just any of the limits
+        // filter appears below Union
         let expected = "\
             Union\
             \n  Filter: #a Eq Int64(1)\

From 138534265edb64442f60b134ddac62cf58ea1d53 Mon Sep 17 00:00:00 2001
From: Daniel Heres <danielheres@gmail.com>
Date: Mon, 14 Jun 2021 08:46:31 +0200
Subject: [PATCH 3/3] Fmt

---
 datafusion/src/optimizer/filter_push_down.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/datafusion/src/optimizer/filter_push_down.rs b/datafusion/src/optimizer/filter_push_down.rs
index cb3593a47abc..dc4d5e993a38 100644
--- a/datafusion/src/optimizer/filter_push_down.rs
+++ b/datafusion/src/optimizer/filter_push_down.rs
@@ -770,7 +770,6 @@ mod tests {
         Ok(())
     }
 
-
     #[test]
     fn union_all() -> Result<()> {
         let table_scan = test_table_scan()?;