Support inferring new predicates to push down

Author: xudong963

datafusion/expr/src/expr_rewriter/mod.rs

@@ -148,6 +148,25 @@ pub fn replace_col(expr: Expr, replace_map: &HashMap<&Column, &Column>) -> Resul
     .data()
 }
 
+pub fn replace_col_with_expr(
+    expr: Expr,
+    replace_map: &HashMap<Column, &Expr>,
+) -> Result<Expr> {
+    expr.transform(|expr| {
+        Ok({
+            if let Expr::Column(c) = &expr {
+                match replace_map.get(c) {
+                    Some(new_expr) => Transformed::yes((**new_expr).to_owned()),
+                    None => Transformed::no(expr),
+                }
+            } else {
+                Transformed::no(expr)
+            }
+        })
+    })
+    .data()
+}
+
 /// Recursively 'unnormalize' (remove all qualifiers) from an
 /// expression tree.
 ///

datafusion/optimizer/src/push_down_filter.rs

@@ -30,7 +30,7 @@ use datafusion_common::{
     internal_err, plan_err, qualified_name, Column, DFSchema, Result,
 };
 use datafusion_expr::expr::WindowFunction;
-use datafusion_expr::expr_rewriter::replace_col;
+use datafusion_expr::expr_rewriter::{replace_col, replace_col_with_expr};
 use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan, TableScan, Union};
 use datafusion_expr::utils::{
     conjunction, expr_to_columns, split_conjunction, split_conjunction_owned,
@@ -784,14 +784,16 @@ impl OptimizerRule for PushDownFilter {
 
                 // remove duplicated filters
                 let child_predicates = split_conjunction_owned(child_filter.predicate);
-                let new_predicates = parents_predicates
+                let mut new_predicates = parents_predicates
                     .into_iter()
                     .chain(child_predicates)
                     // use IndexSet to remove dupes while preserving predicate order
                     .collect::<IndexSet<_>>()
                     .into_iter()
                     .collect::<Vec<_>>();
 
+                new_predicates = infer_predicates_from_equalities(new_predicates)?;
+
                 let Some(new_predicate) = conjunction(new_predicates) else {
                     return plan_err!("at least one expression exists");
                 };
@@ -1382,6 +1384,73 @@ fn contain(e: &Expr, check_map: &HashMap<String, Expr>) -> bool {
     is_contain
 }
 
+/// Infers new predicates by substituting equalities.
+/// For example, with predicates `t2.b = 3` and `t1.b > t2.b`,
+/// we can infer `t1.b > 3`.
+fn infer_predicates_from_equalities(predicates: Vec<Expr>) -> Result<Vec<Expr>> {
+    // Map from column names to their literal values (from equality predicates)
+    let mut equality_map: HashMap<Column, Expr> =
+        HashMap::with_capacity(predicates.len());
+    let mut final_predicates = Vec::with_capacity(predicates.len());
+    // First pass: collect column=literal equalities
+    for predicate in predicates.iter() {
+        if let Expr::BinaryExpr(BinaryExpr {
+            left,
+            op: Operator::Eq,
+            right,
+        }) = predicate
+        {
+            if let Expr::Column(col) = left.as_ref() {
+                // Only add to map if right side is a literal
+                if matches!(right.as_ref(), Expr::Literal(_)) {
+                    equality_map.insert(col.clone(), *right.clone());
+                    final_predicates.push(predicate.clone());
+                }
+            } else if let Expr::Column(col) = right.as_ref() {
+                // Only add to map if left side is a literal
+                if matches!(left.as_ref(), Expr::Literal(_)) {
+                    equality_map.insert(col.clone(), *right.clone());
+                    final_predicates.push(predicate.clone());
+                }
+            }
+        }
+    }
+
+    // If no equality mappings found, nothing to infer
+    if equality_map.is_empty() {
+        return Ok(predicates);
+    }
+
+    // Second pass: apply substitutions to create new predicates
+    for predicate in predicates {
+        // Skip equality predicates we already used for mapping
+        if final_predicates.contains(&predicate) {
+            continue;
+        }
+
+        // Try to replace columns with their literal values
+        let mut columns_in_expr = HashSet::new();
+        expr_to_columns(&predicate, &mut columns_in_expr)?;
+
+        // Create a combined replacement map for all columns in this predicate
+        let replace_map: HashMap<_, _> = columns_in_expr
+            .into_iter()
+            .filter_map(|col| equality_map.get(&col).map(|lit| (col, lit)))
+            .collect();
+
+        if replace_map.is_empty() {
+            final_predicates.push(predicate);
+            continue;
+        }
+        // Apply all substitutions at once to get the fully substituted predicate
+        let new_pred = replace_col_with_expr(predicate, &replace_map)?;
+
+        final_predicates.push(new_pred);
+    }
+
+    Ok(final_predicates)
+}
+
 #[cfg(test)]
 mod tests {
     use std::any::Any;

datafusion/physical-optimizer/src/enforce_sorting/replace_with_order_preserving_variants.rs

@@ -197,7 +197,9 @@ fn plan_with_order_breaking_variants(
         // Replace `SortPreservingMergeExec` with a `CoalescePartitionsExec`
         // SPM may have `fetch`, so pass it to the `CoalescePartitionsExec`
         let child = Arc::clone(&sort_input.children[0].plan);
-        let coalesce = CoalescePartitionsExec::new(child).with_fetch(plan.fetch()).unwrap();
+        let coalesce = CoalescePartitionsExec::new(child)
+            .with_fetch(plan.fetch())
+            .unwrap();
         sort_input.plan = coalesce;
     } else {
         return sort_input.update_plan_from_children();