physical-plan: push filters down to UnionExec children

Filters are safe to be pushed down, so we can override the default behavior
here.

Signed-off-by: Alfonso Subiotto Marques <[email protected]>

Author: asubiotto

datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs

@@ -60,6 +60,7 @@ use datafusion_physical_plan::{
     ExecutionPlan,
 };
 
+use datafusion_physical_plan::union::UnionExec;
 use futures::StreamExt;
 use object_store::{memory::InMemory, ObjectStore};
 use util::{format_plan_for_test, OptimizationTest, TestNode, TestScanBuilder};
@@ -1833,6 +1834,34 @@ STORED AS PARQUET;
     // Pushdown pruned most rows
 }
 
+#[test]
+fn test_filter_pushdown_through_union() {
+    let scan1 = TestScanBuilder::new(schema()).with_support(true).build();
+    let scan2 = TestScanBuilder::new(schema()).with_support(true).build();
+
+    let union = UnionExec::try_new(vec![scan1, scan2]).unwrap();
+
+    let predicate = col_lit_predicate("a", "foo", &schema());
+    let plan = Arc::new(FilterExec::try_new(predicate, union).unwrap());
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, FilterPushdown::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - FilterExec: a@0 = foo
+        -   UnionExec
+        -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
+        -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
+      output:
+        Ok:
+          - UnionExec
+          -   DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=a@0 = foo
+          -   DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=a@0 = foo
+    "
+    );
+}
+
 /// Schema:
 /// a: String
 /// b: String

datafusion/physical-plan/src/union.rs

@@ -36,16 +36,18 @@ use crate::execution_plan::{
     boundedness_from_children, check_default_invariants, emission_type_from_children,
     InvariantLevel,
 };
+use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase};
 use crate::metrics::BaselineMetrics;
 use crate::projection::{make_with_child, ProjectionExec};
 use crate::stream::ObservedStream;
 
 use arrow::datatypes::{Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
+use datafusion_common::config::ConfigOptions;
 use datafusion_common::stats::Precision;
 use datafusion_common::{exec_err, internal_datafusion_err, internal_err, Result};
 use datafusion_execution::TaskContext;
-use datafusion_physical_expr::{calculate_union, EquivalenceProperties};
+use datafusion_physical_expr::{calculate_union, EquivalenceProperties, PhysicalExpr};
 
 use futures::Stream;
 use itertools::Itertools;
@@ -218,10 +220,6 @@ impl ExecutionPlan for UnionExec {
         })
     }
 
-    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        self.inputs.iter().collect()
-    }
-
     fn maintains_input_order(&self) -> Vec<bool> {
         // If the Union has an output ordering, it maintains at least one
         // child's ordering (i.e. the meet).
@@ -247,6 +245,14 @@ impl ExecutionPlan for UnionExec {
         }
     }
 
+    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+        vec![false; self.children().len()]
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        self.inputs.iter().collect()
+    }
+
     fn with_new_children(
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
@@ -324,10 +330,6 @@ impl ExecutionPlan for UnionExec {
         }
     }
 
-    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
-        vec![false; self.children().len()]
-    }
-
     fn supports_limit_pushdown(&self) -> bool {
         true
     }
@@ -352,6 +354,15 @@ impl ExecutionPlan for UnionExec {
 
         Ok(Some(UnionExec::try_new(new_children.clone())?))
     }
+
+    fn gather_filters_for_pushdown(
+        &self,
+        _phase: FilterPushdownPhase,
+        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
+        _config: &ConfigOptions,
+    ) -> Result<FilterDescription> {
+        FilterDescription::from_children(parent_filters, &self.children())
+    }
 }
 
 /// Combines multiple input streams by interleaving them.