Add INFERRED_PREDICATE_ALIAS to manage inferred join predicates and ensure they remain as filters during optimization

Author: kosiew

datafusion/expr/src/expr.rs

@@ -38,6 +38,7 @@ use datafusion_common::tree_node::{
 use datafusion_common::{
     Column, DFSchema, HashMap, Result, ScalarValue, Spans, TableReference,
 };
+use crate::utils::INFERRED_PREDICATE_ALIAS;
 use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use sqlparser::ast::{
     display_comma_separated, ExceptSelectItem, ExcludeSelectItem, IlikeSelectItem,
@@ -3183,7 +3184,13 @@ pub const UNNEST_COLUMN_PREFIX: &str = "UNNEST";
 impl Display for Expr {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         match self {
-            Expr::Alias(Alias { expr, name, .. }) => write!(f, "{expr} AS {name}"),
+            Expr::Alias(Alias { expr, name, .. }) => {
+                if name == INFERRED_PREDICATE_ALIAS {
+                    write!(f, "{expr}")
+                } else {
+                    write!(f, "{expr} AS {name}")
+                }
+            }
             Expr::Column(c) => write!(f, "{c}"),
             Expr::OuterReferenceColumn(_, c) => {
                 write!(f, "{OUTER_REFERENCE_COLUMN_PREFIX}({c})")

datafusion/expr/src/utils.rs

@@ -28,6 +28,10 @@ use crate::{
 };
 use datafusion_expr_common::signature::{Signature, TypeSignature};
 
+/// Alias used to mark inferred join predicates that should remain join filters
+/// and not be converted into additional equijoin keys during optimization.
+pub const INFERRED_PREDICATE_ALIAS: &str = "__datafusion_inferred_join_predicate";
+
 use arrow::datatypes::{DataType, Field, Schema};
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
@@ -1063,7 +1067,7 @@ fn split_binary_owned_impl(
             let exprs = split_binary_owned_impl(*left, operator, exprs);
             split_binary_owned_impl(*right, operator, exprs)
         }
-        Expr::Alias(Alias { expr, .. }) => {
+        Expr::Alias(Alias { expr, name, .. }) if name != INFERRED_PREDICATE_ALIAS => {
             split_binary_owned_impl(*expr, operator, exprs)
         }
         other => {

datafusion/optimizer/src/push_down_filter.rs

@@ -35,6 +35,7 @@ use datafusion_expr::expr_rewriter::replace_col;
 use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan, TableScan, Union};
 use datafusion_expr::utils::{
     conjunction, expr_to_columns, split_conjunction, split_conjunction_owned,
+    INFERRED_PREDICATE_ALIAS,
 };
 use datafusion_expr::{
     and, or, BinaryExpr, Expr, Filter, Operator, Projection, TableProviderFilterPushDown,
@@ -463,7 +464,11 @@ fn push_down_all_join(
         } else if right_preserved && checker.is_right_only(&predicate) {
             right_push.push(predicate);
         } else {
-            join_conditions.push(predicate);
+            // Mark inferred predicates so subsequent optimization passes do not
+            // treat them as additional equijoin keys. They should remain as
+            // residual join filters to enable dynamic filtering without
+            // widening the join key set.
+            join_conditions.push(predicate.alias(INFERRED_PREDICATE_ALIAS));
         }
     }
 
@@ -1487,6 +1492,53 @@ mod tests {
         }};
     }
 
+    #[test]
+    fn inferred_predicate_stays_in_filter() -> Result<()> {
+        use datafusion_common::NullEquality;
+        use datafusion_expr::logical_plan::JoinConstraint;
+
+        let left = test_table_scan_with_name("l")?;
+        let right = test_table_scan_with_name("r")?;
+
+        let join = Join::try_new(
+            Arc::new(left),
+            Arc::new(right),
+            vec![],
+            None,
+            JoinType::Left,
+            JoinConstraint::On,
+            NullEquality::NullEqualsNull,
+        )?;
+
+        let inferred = col("l.a").eq(col("r.a"));
+        let Transformed { data: plan, .. } =
+            push_down_all_join(vec![], vec![inferred], join, vec![])?;
+
+        let join = match plan {
+            LogicalPlan::Join(j) => j,
+            _ => panic!("expected join"),
+        };
+
+        assert!(join.on.is_empty());
+        let filter = join.filter.clone().expect("expected filter");
+        match filter {
+            Expr::Alias(alias) => assert_eq!(alias.name, INFERRED_PREDICATE_ALIAS),
+            _ => panic!("expected aliased filter"),
+        }
+
+        let optimizer = Optimizer::with_rules(vec![
+            Arc::new(SimplifyExpressions::new()),
+            Arc::new(crate::extract_equijoin_predicate::ExtractEquijoinPredicate::new()),
+        ]);
+        let optimized = optimizer.optimize(LogicalPlan::Join(join), &OptimizerContext::new(), observe)?;
+        if let LogicalPlan::Join(j) = optimized {
+            assert!(j.on.is_empty());
+        } else {
+            panic!("expected join");
+        }
+        Ok(())
+    }
+
     #[test]
     fn filter_before_projection() -> Result<()> {
         let table_scan = test_table_scan()?;