Add the missing equivalence info for filter pushdown

Author: liamzwbao

datafusion/datasource/src/source.rs

@@ -33,8 +33,11 @@ use crate::file_scan_config::FileScanConfig;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{Constraints, Result, Statistics};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
-use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
+use datafusion_physical_expr::{
+    conjunction, EquivalenceProperties, Partitioning, PhysicalExpr,
+};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
+use datafusion_physical_plan::filter::collect_columns_from_predicate;
 use datafusion_physical_plan::filter_pushdown::{
     ChildPushdownResult, FilterPushdownPropagation,
 };
@@ -328,6 +331,9 @@ impl ExecutionPlan for DataSourceExec {
                 new_node.data_source = data_source;
                 new_node.cache =
                     Self::compute_properties(Arc::clone(&new_node.data_source));
+                // Add the missing filters' equivalence info when filters pushdown is applied
+                let filter = conjunction(res.filters.collect_supported());
+                new_node = new_node.add_filter_equivalence_info(filter)?;
                 Ok(FilterPushdownPropagation {
                     filters: res.filters,
                     updated_node: Some(Arc::new(new_node)),
@@ -374,6 +380,20 @@ impl DataSourceExec {
         self
     }
 
+    /// Add filters' equivalence info
+    fn add_filter_equivalence_info(
+        mut self,
+        filter: Arc<dyn PhysicalExpr>,
+    ) -> Result<Self> {
+        let (equal_pairs, _) = collect_columns_from_predicate(&filter);
+        for (lhs, rhs) in equal_pairs {
+            self.cache
+                .eq_properties
+                .add_equal_conditions(&Arc::clone(lhs), &Arc::clone(rhs))?
+        }
+        Ok(self)
+    }
+
     fn compute_properties(data_source: Arc<dyn DataSource>) -> PlanProperties {
         PlanProperties::new(
             data_source.eq_properties(),

datafusion/physical-plan/src/filter.rs

@@ -749,7 +749,9 @@ impl RecordBatchStream for FilterExecStream {
 }
 
 /// Return the equals Column-Pairs and Non-equals Column-Pairs
-fn collect_columns_from_predicate(predicate: &Arc<dyn PhysicalExpr>) -> EqualAndNonEqual {
+pub fn collect_columns_from_predicate(
+    predicate: &Arc<dyn PhysicalExpr>,
+) -> EqualAndNonEqual {
     let mut eq_predicate_columns = Vec::<PhysicalExprPairRef>::new();
     let mut ne_predicate_columns = Vec::<PhysicalExprPairRef>::new();
 

datafusion/physical-plan/src/filter_pushdown.rs

@@ -102,6 +102,18 @@ impl PredicateSupports {
             .collect()
     }
 
+    /// Collect supported filters into a Vec, without removing them from the original
+    /// [`PredicateSupport`].
+    pub fn collect_supported(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+        self.0
+            .iter()
+            .filter_map(|f| match f {
+                PredicateSupport::Supported(expr) => Some(Arc::clone(expr)),
+                PredicateSupport::Unsupported(_) => None,
+            })
+            .collect()
+    }
+
     /// Collect all filters into a Vec, without removing them from the original
     /// FilterPushdowns.
     pub fn collect_all(self) -> Vec<Arc<dyn PhysicalExpr>> {