Refactor push_down_filter for simplicity and efficiency

Eliminate JoinInputShape/PredicateDestination layers and integrate
scalar-subquery handling into the main join-pushdown flow. Revert
join-inference replacement logic to an explicit loop. Remove
unnecessary helper functions and simplify the null-restriction
evaluator for cleaner code.

Author: kosiew

datafusion/optimizer/src/push_down_filter.rs

@@ -285,81 +285,30 @@ fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
     Ok(is_evaluate)
 }
 
-#[derive(Clone, Copy)]
-struct JoinInputShape<'a> {
-    base_plan: &'a LogicalPlan,
-    is_derived_relation: bool,
-}
-
-fn classify_join_input(plan: &LogicalPlan) -> JoinInputShape<'_> {
+fn classify_join_input(plan: &LogicalPlan) -> (bool, bool) {
     match plan {
         LogicalPlan::SubqueryAlias(subquery_alias) => {
-            let JoinInputShape { base_plan, .. } =
+            let (is_scalar_aggregate, _) =
                 classify_join_input(subquery_alias.input.as_ref());
-            JoinInputShape {
-                base_plan,
-                is_derived_relation: true,
-            }
+            (is_scalar_aggregate, true)
         }
         LogicalPlan::Projection(projection) => {
-            let shape = classify_join_input(projection.input.as_ref());
-            JoinInputShape {
-                is_derived_relation: shape.is_derived_relation,
-                ..shape
-            }
+            classify_join_input(projection.input.as_ref())
         }
-        _ => JoinInputShape {
-            base_plan: plan,
-            is_derived_relation: false,
-        },
+        LogicalPlan::Aggregate(aggregate) => (aggregate.group_expr.is_empty(), false),
+        _ => (false, false),
     }
 }
 
-fn is_scalar_aggregate_subquery(shape: JoinInputShape<'_>) -> bool {
-    matches!(
-        shape.base_plan,
-        LogicalPlan::Aggregate(aggregate) if aggregate.group_expr.is_empty()
-    )
-}
-
 fn is_scalar_subquery_cross_join(join: &Join) -> bool {
-    let left_shape = classify_join_input(join.left.as_ref());
-    let right_shape = classify_join_input(join.right.as_ref());
+    let (left_scalar_aggregate, left_is_derived_relation) =
+        classify_join_input(join.left.as_ref());
+    let (right_scalar_aggregate, right_is_derived_relation) =
+        classify_join_input(join.right.as_ref());
     join.on.is_empty()
         && join.filter.is_none()
-        && ((is_scalar_aggregate_subquery(left_shape) && right_shape.is_derived_relation)
-            || (is_scalar_aggregate_subquery(right_shape)
-                && left_shape.is_derived_relation))
-}
-
-// Keep post-join filters above certain scalar-subquery cross joins to preserve
-// behavior for the window-over-scalar-subquery regression shape.
-fn should_keep_filter_above_scalar_subquery_cross_join(
-    mut checker: ColumnChecker<'_>,
-    predicate: &Expr,
-) -> bool {
-    !checker.is_left_only(predicate) && !checker.is_right_only(predicate)
-}
-
-enum PredicateDestination {
-    Left,
-    Right,
-    Keep,
-}
-
-fn classify_predicate_destination(
-    checker: &mut ColumnChecker<'_>,
-    predicate: &Expr,
-    allow_left: bool,
-    allow_right: bool,
-) -> PredicateDestination {
-    if allow_left && checker.is_left_only(predicate) {
-        PredicateDestination::Left
-    } else if allow_right && checker.is_right_only(predicate) {
-        PredicateDestination::Right
-    } else {
-        PredicateDestination::Keep
-    }
+        && ((left_scalar_aggregate && right_is_derived_relation)
+            || (right_scalar_aggregate && left_is_derived_relation))
 }
 
 /// examine OR clause to see if any useful clauses can be extracted and push down.
@@ -506,41 +455,28 @@ fn push_down_all_join(
     let keep_mixed_scalar_subquery_filters =
         is_inner_join && is_scalar_subquery_cross_join(&join);
     for predicate in predicates {
-        match classify_predicate_destination(
-            &mut checker,
-            &predicate,
-            left_preserved,
-            right_preserved,
-        ) {
-            PredicateDestination::Left => left_push.push(predicate),
-            PredicateDestination::Right => right_push.push(predicate),
-            PredicateDestination::Keep => {
-                let should_keep_above_join = keep_mixed_scalar_subquery_filters
-                    && should_keep_filter_above_scalar_subquery_cross_join(
-                        ColumnChecker::new(left_schema, right_schema),
-                        &predicate,
-                    );
-
-                if is_inner_join
-                    && !should_keep_above_join
-                    && can_evaluate_as_join_condition(&predicate)?
-                {
-                    // Here we do not differ it is eq or non-eq predicate, ExtractEquijoinPredicate will extract the eq predicate
-                    // and convert to the join on condition
-                    join_conditions.push(predicate);
-                } else {
-                    keep_predicates.push(predicate);
-                }
-            }
+        if left_preserved && checker.is_left_only(&predicate) {
+            left_push.push(predicate);
+        } else if right_preserved && checker.is_right_only(&predicate) {
+            right_push.push(predicate);
+        } else if is_inner_join
+            && !keep_mixed_scalar_subquery_filters
+            && can_evaluate_as_join_condition(&predicate)?
+        {
+            // Here we do not differ it is eq or non-eq predicate, ExtractEquijoinPredicate will extract the eq predicate
+            // and convert to the join on condition
+            join_conditions.push(predicate);
+        } else {
+            keep_predicates.push(predicate);
         }
     }
 
     // Push predicates inferred from the join expression
     for predicate in inferred_join_predicates {
-        match classify_predicate_destination(&mut checker, &predicate, true, true) {
-            PredicateDestination::Left => left_push.push(predicate),
-            PredicateDestination::Right => right_push.push(predicate),
-            PredicateDestination::Keep => {}
+        if checker.is_left_only(&predicate) {
+            left_push.push(predicate);
+        } else if checker.is_right_only(&predicate) {
+            right_push.push(predicate);
         }
     }
 
@@ -549,15 +485,12 @@ fn push_down_all_join(
 
     if !on_filter.is_empty() {
         for on in on_filter {
-            match classify_predicate_destination(
-                &mut checker,
-                &on,
-                on_left_preserved,
-                on_right_preserved,
-            ) {
-                PredicateDestination::Left => left_push.push(on),
-                PredicateDestination::Right => right_push.push(on),
-                PredicateDestination::Keep => on_filter_join_conditions.push(on),
+            if on_left_preserved && checker.is_left_only(&on) {
+                left_push.push(on)
+            } else if on_right_preserved && checker.is_right_only(&on) {
+                right_push.push(on)
+            } else {
+                on_filter_join_conditions.push(on)
             }
         }
     }
@@ -821,24 +754,26 @@ fn infer_join_predicates_impl<
     inferred_predicates: &mut InferredPredicates,
 ) -> Result<()> {
     for predicate in input_predicates {
-        let column_refs = predicate.column_refs();
+        let mut join_cols_to_replace = HashMap::new();
         let mut saw_non_replaceable_ref = false;
-        let join_cols_to_replace = column_refs
-            .iter()
-            .filter_map(|&col| {
-                let replacement = join_col_keys.iter().find_map(|(l, r)| {
-                    if ENABLE_LEFT_TO_RIGHT && col == *l {
-                        Some((col, *r))
-                    } else if ENABLE_RIGHT_TO_LEFT && col == *r {
-                        Some((col, *l))
-                    } else {
-                        None
-                    }
-                });
-                saw_non_replaceable_ref |= replacement.is_none();
-                replacement
-            })
-            .collect::<HashMap<_, _>>();
+
+        for &col in &predicate.column_refs() {
+            let replacement = join_col_keys.iter().find_map(|(l, r)| {
+                if ENABLE_LEFT_TO_RIGHT && col == *l {
+                    Some((col, *r))
+                } else if ENABLE_RIGHT_TO_LEFT && col == *r {
+                    Some((col, *l))
+                } else {
+                    None
+                }
+            });
+
+            if let Some((source, target)) = replacement {
+                join_cols_to_replace.insert(source, target);
+            } else {
+                saw_non_replaceable_ref = true;
+            }
+        }
 
         if join_cols_to_replace.is_empty()
             || (!inferred_predicates.is_inner_join && saw_non_replaceable_ref)

datafusion/optimizer/src/utils.rs

@@ -46,7 +46,9 @@ use self::test_eval_mode::{
 
 /// Returns true if `expr` contains all columns in `schema_cols`
 pub(crate) fn has_all_column_refs(expr: &Expr, schema_cols: &HashSet<Column>) -> bool {
-    column_refs_all_in(&expr.column_refs(), |column| schema_cols.contains(column))
+    expr.column_refs()
+        .iter()
+        .all(|column| schema_cols.contains(*column))
 }
 
 pub(crate) fn replace_qualified_name(
@@ -70,13 +72,6 @@ pub fn log_plan(description: &str, plan: &LogicalPlan) {
     trace!("{description}::\n{}\n", plan.display_indent_schema());
 }
 
-pub(super) fn column_refs_all_in<'a>(
-    column_refs: &HashSet<&'a Column>,
-    mut contains: impl FnMut(&Column) -> bool,
-) -> bool {
-    column_refs.iter().all(|column| contains(column))
-}
-
 /// Determine whether a predicate can restrict NULLs. e.g.
 /// `c0 > 8` return true;
 /// `c0 IS NULL` return false.
@@ -98,7 +93,7 @@ pub fn is_restrict_null_predicate<'a>(
     // contains a placeholder for the join key columns. Callers treat such errors as
     // non-restricting (false) via `matches!(_, Ok(true))`, so we return false early
     // and avoid the expensive physical-expression compilation pipeline entirely.
-    if !column_refs_all_in(&column_refs, |column| join_cols.contains(&column)) {
+    if !column_refs.iter().all(|column| join_cols.contains(*column)) {
         return Ok(false);
     }
 
@@ -181,31 +176,28 @@ fn authoritative_restrict_null_predicate<'a>(
     predicate: Expr,
     join_cols_of_predicate: impl IntoIterator<Item = &'a Column>,
 ) -> Result<bool> {
-    evaluate_expr_with_null_column(predicate, join_cols_of_predicate)
-        .and_then(is_false_or_null_boolean_result)
+    Ok(
+        match evaluate_expr_with_null_column(predicate, join_cols_of_predicate)? {
+            ColumnarValue::Array(array) if array.len() == 1 => {
+                let boolean_array = as_boolean_array(&array)?;
+                boolean_array.is_null(0) || !boolean_array.value(0)
+            }
+            ColumnarValue::Array(_) => false,
+            ColumnarValue::Scalar(scalar) => matches!(
+                scalar,
+                ScalarValue::Boolean(None)
+                    | ScalarValue::Boolean(Some(false))
+                    | ScalarValue::Null
+            ),
+        },
+    )
 }
 
 fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> {
     let mut expr_rewrite = TypeCoercionRewriter { schema };
     expr.rewrite(&mut expr_rewrite).data()
 }
 
-fn is_false_or_null_boolean_result(result: ColumnarValue) -> Result<bool> {
-    Ok(match result {
-        ColumnarValue::Array(array) if array.len() == 1 => {
-            let boolean_array = as_boolean_array(&array)?;
-            boolean_array.is_null(0) || !boolean_array.value(0)
-        }
-        ColumnarValue::Array(_) => false,
-        ColumnarValue::Scalar(scalar) => matches!(
-            scalar,
-            ScalarValue::Boolean(None)
-                | ScalarValue::Boolean(Some(false))
-                | ScalarValue::Null
-        ),
-    })
-}
-
 #[cfg(test)]
 mod test_eval_mode {
     use std::cell::Cell;

datafusion/optimizer/src/utils/null_restriction.rs

@@ -32,21 +32,13 @@ enum NullSubstitutionValue {
     Boolean(bool),
 }
 
-impl NullSubstitutionValue {
-    fn is_null(self) -> bool {
-        matches!(self, Self::Null)
-    }
-}
-
 pub(super) fn syntactic_restrict_null_predicate(
     predicate: &Expr,
     join_cols: &HashSet<&Column>,
 ) -> Option<bool> {
     match syntactic_null_substitution_value(predicate, join_cols) {
-        Some(NullSubstitutionValue::Boolean(true)) => Some(false),
-        Some(NullSubstitutionValue::Boolean(false) | NullSubstitutionValue::Null) => {
-            Some(true)
-        }
+        Some(NullSubstitutionValue::Boolean(value)) => Some(!value),
+        Some(NullSubstitutionValue::Null) => Some(true),
         Some(NullSubstitutionValue::NonNull) | None => None,
     }
 }
@@ -117,7 +109,7 @@ fn null_if_contains_null(
 fn strict_null_only(
     value: Option<NullSubstitutionValue>,
 ) -> Option<NullSubstitutionValue> {
-    value.filter(|value| value.is_null())
+    value.filter(|value| matches!(value, NullSubstitutionValue::Null))
 }
 
 fn syntactic_null_substitution_value(