Consolidate plan-wrapper traversal and optimizations

Combine plan-wrapper traversal and cross-join shape detection.
Shorten join-column replacement scan and share authoritative
null-result decoding. Remove unused helpers and reorganize
strict-null operator list behind a classifier helper.
Public interfaces remain unchanged.

Author: kosiew

datafusion/core/benches/sql_planner_extended.rs

@@ -357,8 +357,12 @@ fn build_non_case_left_join_df_with_push_down_filter(
 fn find_filter_predicates(plan: &LogicalPlan) -> Vec<datafusion_expr::Expr> {
     match plan {
         LogicalPlan::Filter(filter) => split_conjunction_owned(filter.predicate.clone()),
-        LogicalPlan::Projection(projection) => find_filter_predicates(projection.input.as_ref()),
-        other => panic!("expected benchmark query plan to contain a Filter, found {other:?}"),
+        LogicalPlan::Projection(projection) => {
+            find_filter_predicates(projection.input.as_ref())
+        }
+        other => {
+            panic!("expected benchmark query plan to contain a Filter, found {other:?}")
+        }
     }
 }
 
@@ -375,7 +379,8 @@ fn assert_case_heavy_left_join_inference_candidates(
     for predicate in predicates {
         let column_refs = predicate.column_refs();
         assert!(
-            column_refs.contains(&&left_join_key) || column_refs.contains(&&right_join_key),
+            column_refs.contains(&&left_join_key)
+                || column_refs.contains(&&right_join_key),
             "benchmark predicate should reference a join key: {predicate}"
         );
         assert!(

datafusion/optimizer/src/push_down_filter.rs

@@ -285,33 +285,39 @@ fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
     Ok(is_evaluate)
 }
 
-fn strip_aliases_and_projections(plan: &LogicalPlan) -> &LogicalPlan {
+fn strip_plan_wrappers(plan: &LogicalPlan) -> (&LogicalPlan, bool) {
     match plan {
         LogicalPlan::SubqueryAlias(subquery_alias) => {
-            strip_aliases_and_projections(subquery_alias.input.as_ref())
+            let (plan, _) = strip_plan_wrappers(subquery_alias.input.as_ref());
+            (plan, true)
         }
         LogicalPlan::Projection(projection) => {
-            strip_aliases_and_projections(projection.input.as_ref())
+            let (plan, is_derived_relation) =
+                strip_plan_wrappers(projection.input.as_ref());
+            (plan, is_derived_relation)
         }
-        _ => plan,
+        _ => (plan, false),
     }
 }
 
 fn is_scalar_aggregate_subquery(plan: &LogicalPlan) -> bool {
     matches!(
-        strip_aliases_and_projections(plan),
+        strip_plan_wrappers(plan).0,
         LogicalPlan::Aggregate(aggregate) if aggregate.group_expr.is_empty()
     )
 }
 
 fn is_derived_relation(plan: &LogicalPlan) -> bool {
-    match plan {
-        LogicalPlan::SubqueryAlias(_) => true,
-        LogicalPlan::Projection(projection) => {
-            is_derived_relation(projection.input.as_ref())
-        }
-        _ => false,
-    }
+    strip_plan_wrappers(plan).1
+}
+
+fn is_scalar_subquery_cross_join(join: &Join) -> bool {
+    join.on.is_empty()
+        && join.filter.is_none()
+        && ((is_scalar_aggregate_subquery(join.left.as_ref())
+            && is_derived_relation(join.right.as_ref()))
+            || (is_scalar_aggregate_subquery(join.right.as_ref())
+                && is_derived_relation(join.left.as_ref())))
 }
 
 // Keep post-join filters above certain scalar-subquery cross joins to preserve
@@ -320,19 +326,12 @@ fn should_keep_filter_above_scalar_subquery_cross_join(
     join: &Join,
     predicate: &Expr,
 ) -> bool {
-    if !join.on.is_empty() || join.filter.is_some() {
+    if !is_scalar_subquery_cross_join(join) {
         return false;
     }
 
     let mut checker = ColumnChecker::new(join.left.schema(), join.right.schema());
-    let references_both_sides =
-        !checker.is_left_only(predicate) && !checker.is_right_only(predicate);
-
-    references_both_sides
-        && ((is_scalar_aggregate_subquery(join.left.as_ref())
-            && is_derived_relation(join.right.as_ref()))
-            || (is_scalar_aggregate_subquery(join.right.as_ref())
-                && is_derived_relation(join.left.as_ref())))
+    !checker.is_left_only(predicate) && !checker.is_right_only(predicate)
 }
 
 /// examine OR clause to see if any useful clauses can be extracted and push down.
@@ -777,37 +776,31 @@ fn infer_join_predicates_impl<
 ) -> Result<()> {
     for predicate in input_predicates {
         let column_refs = predicate.column_refs();
-        let mut join_col_replacements = Vec::new();
-        let mut has_non_replaceable_refs = false;
-
-        for &col in &column_refs {
-            let mut replacement = None;
-
-            for (l, r) in join_col_keys.iter() {
-                if ENABLE_LEFT_TO_RIGHT && col == *l {
-                    replacement = Some((col, *r));
-                    break;
-                }
-                if ENABLE_RIGHT_TO_LEFT && col == *r {
-                    replacement = Some((col, *l));
-                    break;
-                }
-            }
+        let join_col_replacements: Vec<_> = column_refs
+            .iter()
+            .filter_map(|&col| {
+                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
+                    }
+                })
+            })
+            .collect();
 
-            if let Some(replacement) = replacement {
-                join_col_replacements.push(replacement);
-            } else {
-                has_non_replaceable_refs = true;
-            }
-        }
         if join_col_replacements.is_empty() {
             continue;
         }
 
         // For non-inner joins, predicates that reference any non-replaceable
         // columns cannot be inferred on the other side. Skip the null-restriction
         // helper entirely in that common mixed-reference case.
-        if !inferred_predicates.is_inner_join && has_non_replaceable_refs {
+        if !inferred_predicates.is_inner_join
+            && join_col_replacements.len() != column_refs.len()
+        {
             continue;
         }
 

datafusion/optimizer/src/utils.rs

@@ -130,13 +130,14 @@ pub fn is_restrict_null_predicate<'a>(
     // Collect join columns so they can be used in both the fast-path check and the
     // fallback evaluation path below.
     let join_cols: HashSet<&Column> = join_cols_of_predicate.into_iter().collect();
+    let column_refs = predicate.column_refs();
 
     // Fast path: if the predicate references columns outside the join key set,
     // `evaluate_expr_with_null_column` would fail because the null schema only
     // 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 !null_restriction::predicate_uses_only_columns(&predicate, &join_cols) {
+    if !null_restriction::all_columns_allowed(&column_refs, &join_cols) {
         return Ok(false);
     }
 
@@ -180,12 +181,10 @@ pub fn evaluates_to_null<'a>(
         return Ok(true);
     }
 
-    Ok(
-        match evaluate_expr_with_null_column(predicate, null_columns)? {
-            ColumnarValue::Array(_) => false,
-            ColumnarValue::Scalar(scalar) => scalar.is_null(),
-        },
-    )
+    Ok(authoritative_null_result(evaluate_expr_with_null_column(
+        predicate,
+        null_columns,
+    )?)? == AuthoritativeNullResult::AlwaysNull)
 }
 
 fn evaluate_expr_with_null_column<'a>(
@@ -219,22 +218,41 @@ fn authoritative_restrict_null_predicate<'a>(
     predicate: Expr,
     join_cols_of_predicate: impl IntoIterator<Item = &'a Column>,
 ) -> Result<bool> {
-    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)
-                } else {
-                    false
-                }
+    Ok(authoritative_null_result(evaluate_expr_with_null_column(
+        predicate,
+        join_cols_of_predicate,
+    )?)? == AuthoritativeNullResult::NullRestricting)
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum AuthoritativeNullResult {
+    AlwaysNull,
+    NullRestricting,
+    Other,
+}
+
+fn authoritative_null_result(value: ColumnarValue) -> Result<AuthoritativeNullResult> {
+    Ok(match value {
+        ColumnarValue::Array(array) => {
+            if array.len() != 1 {
+                return Ok(AuthoritativeNullResult::Other);
             }
-            ColumnarValue::Scalar(scalar) => matches!(
-                scalar,
-                ScalarValue::Boolean(None) | ScalarValue::Boolean(Some(false))
-            ),
-        },
-    )
+
+            let boolean_array = as_boolean_array(&array)?;
+            if boolean_array.is_null(0) || !boolean_array.value(0) {
+                AuthoritativeNullResult::NullRestricting
+            } else {
+                AuthoritativeNullResult::Other
+            }
+        }
+        ColumnarValue::Scalar(scalar) if scalar.is_null() => {
+            AuthoritativeNullResult::AlwaysNull
+        }
+        ColumnarValue::Scalar(
+            ScalarValue::Boolean(None) | ScalarValue::Boolean(Some(false)),
+        ) => AuthoritativeNullResult::NullRestricting,
+        ColumnarValue::Scalar(_) => AuthoritativeNullResult::Other,
+    })
 }
 
 fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> {