Refactor is_restrict_null_predicate in utils.rs

Restore is_restrict_null_predicate to its original behavior by
removing the broader syntactic null-restriction path and its
associated test scaffolding. Maintain early rejection for the
push_down_filter caller while eliminating extra tree-walk work
on common paths. Confirmed changes by running tests and
formatting.

Author: kosiew

datafusion/optimizer/src/utils.rs

@@ -17,8 +17,6 @@
 
 //! Utility functions leveraged by the query optimizer rules
 
-mod null_restriction;
-
 use std::collections::{BTreeSet, HashMap, HashSet};
 use std::sync::Arc;
 
@@ -38,12 +36,6 @@ use log::{debug, trace};
 /// as it was initially placed here and then moved elsewhere.
 pub use datafusion_expr::expr_rewriter::NamePreserver;
 
-#[cfg(test)]
-use self::test_eval_mode::{
-    NullRestrictionEvalMode, null_restriction_eval_mode,
-    set_null_restriction_eval_mode_for_test, with_null_restriction_eval_mode_for_test,
-};
-
 /// Returns true if `expr` contains all columns in `schema_cols`
 pub(crate) fn has_all_column_refs(expr: &Expr, schema_cols: &HashSet<Column>) -> bool {
     expr.column_refs()
@@ -86,40 +78,17 @@ 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 !column_refs.iter().all(|column| join_cols.contains(*column)) {
-        return Ok(false);
-    }
-
-    #[cfg(test)]
-    if matches!(
-        null_restriction_eval_mode(),
-        NullRestrictionEvalMode::AuthoritativeOnly
-    ) {
-        return authoritative_restrict_null_predicate(predicate, join_cols);
-    }
-
-    if let Some(is_restricting) =
-        null_restriction::syntactic_restrict_null_predicate(&predicate, &join_cols)
+    if !predicate
+        .column_refs()
+        .iter()
+        .all(|column| join_cols.contains(*column))
     {
-        #[cfg(debug_assertions)]
-        {
-            let authoritative = authoritative_restrict_null_predicate(
-                predicate.clone(),
-                join_cols.iter().copied(),
-            )?;
-            debug_assert_eq!(
-                is_restricting, authoritative,
-                "syntactic fast path disagrees with authoritative null-restriction evaluation for predicate: {predicate}"
-            );
-        }
-        return Ok(is_restricting);
+        return Ok(false);
     }
 
     authoritative_restrict_null_predicate(predicate, join_cols)
@@ -198,75 +167,13 @@ fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> {
     expr.rewrite(&mut expr_rewrite).data()
 }
 
-#[cfg(test)]
-mod test_eval_mode {
-    use std::cell::Cell;
-
-    /// Null restriction evaluation mode for optimizer tests.
-    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
-    pub(crate) enum NullRestrictionEvalMode {
-        Auto,
-        AuthoritativeOnly,
-    }
-
-    thread_local! {
-        static NULL_RESTRICTION_EVAL_MODE: Cell<NullRestrictionEvalMode> =
-            const { Cell::new(NullRestrictionEvalMode::Auto) };
-    }
-
-    pub(crate) fn set_null_restriction_eval_mode_for_test(mode: NullRestrictionEvalMode) {
-        NULL_RESTRICTION_EVAL_MODE.with(|eval_mode| eval_mode.set(mode));
-    }
-
-    pub(crate) fn null_restriction_eval_mode() -> NullRestrictionEvalMode {
-        NULL_RESTRICTION_EVAL_MODE.with(Cell::get)
-    }
-
-    pub(crate) fn with_null_restriction_eval_mode_for_test<T>(
-        mode: NullRestrictionEvalMode,
-        f: impl FnOnce() -> T,
-    ) -> T {
-        struct NullRestrictionEvalModeReset(NullRestrictionEvalMode);
-
-        impl Drop for NullRestrictionEvalModeReset {
-            fn drop(&mut self) {
-                set_null_restriction_eval_mode_for_test(self.0);
-            }
-        }
-
-        let previous_mode = null_restriction_eval_mode();
-        set_null_restriction_eval_mode_for_test(mode);
-        let _reset = NullRestrictionEvalModeReset(previous_mode);
-        f()
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::panic::{AssertUnwindSafe, catch_unwind};
-
     use datafusion_expr::{
         Operator, binary_expr, case, col, in_list, is_null, lit, when,
     };
 
-    fn restrict_null_predicate_in_modes(
-        predicate: Expr,
-        join_cols: &[Column],
-    ) -> Result<(bool, bool)> {
-        let auto_result = with_null_restriction_eval_mode_for_test(
-            NullRestrictionEvalMode::Auto,
-            || is_restrict_null_predicate(predicate.clone(), join_cols.iter()),
-        )?;
-
-        let authoritative_result = with_null_restriction_eval_mode_for_test(
-            NullRestrictionEvalMode::AuthoritativeOnly,
-            || is_restrict_null_predicate(predicate.clone(), join_cols.iter()),
-        )?;
-
-        Ok((auto_result, authoritative_result))
-    }
-
     #[test]
     fn expr_is_restrict_null_predicate() -> Result<()> {
         let test_cases = vec![
@@ -409,118 +316,4 @@ mod tests {
 
         Ok(())
     }
-
-    #[test]
-    fn syntactic_fast_path_matches_authoritative_evaluator() -> Result<()> {
-        let test_cases = vec![
-            is_null(col("a")),
-            Expr::IsNotNull(Box::new(col("a"))),
-            binary_expr(col("a"), Operator::Gt, lit(8i64)),
-            binary_expr(col("a"), Operator::Eq, lit(ScalarValue::Null)),
-            binary_expr(col("a"), Operator::And, lit(true)),
-            binary_expr(col("a"), Operator::Or, lit(false)),
-            Expr::Not(Box::new(col("a").is_true())),
-            col("a").is_true(),
-            col("a").is_false(),
-            col("a").is_unknown(),
-            col("a").is_not_true(),
-            col("a").is_not_false(),
-            col("a").is_not_unknown(),
-            col("a").between(lit(1i64), lit(10i64)),
-            binary_expr(
-                when(Expr::IsNotNull(Box::new(col("a"))), col("a"))
-                    .otherwise(col("b"))?,
-                Operator::Gt,
-                lit(2i64),
-            ),
-            case(col("a"))
-                .when(lit(1i64), lit(true))
-                .otherwise(lit(false))?,
-            case(col("a"))
-                .when(lit(0i64), lit(false))
-                .otherwise(lit(true))?,
-            binary_expr(
-                case(col("a"))
-                    .when(lit(0i64), lit(true))
-                    .otherwise(lit(false))?,
-                Operator::Or,
-                lit(false),
-            ),
-            binary_expr(
-                case(lit(1i64))
-                    .when(lit(1i64), lit(ScalarValue::Null))
-                    .otherwise(lit(false))?,
-                Operator::IsNotDistinctFrom,
-                lit(true),
-            ),
-        ];
-
-        for predicate in test_cases {
-            let join_cols = predicate.column_refs();
-            if let Some(syntactic) = null_restriction::syntactic_restrict_null_predicate(
-                &predicate, &join_cols,
-            ) {
-                let authoritative = authoritative_restrict_null_predicate(
-                    predicate.clone(),
-                    join_cols.iter().copied(),
-                )
-                .unwrap_or_else(|error| {
-                    panic!(
-                        "authoritative evaluator failed for predicate `{predicate}`: {error}"
-                    )
-                });
-                assert_eq!(
-                    syntactic, authoritative,
-                    "syntactic fast path disagrees with authoritative evaluator for predicate: {predicate}",
-                );
-            }
-        }
-
-        Ok(())
-    }
-
-    #[test]
-    fn null_restriction_eval_mode_auto_vs_authoritative_only() -> Result<()> {
-        let predicate = binary_expr(col("a"), Operator::Gt, lit(8i64));
-        let join_cols_of_predicate = predicate
-            .column_refs()
-            .into_iter()
-            .cloned()
-            .collect::<Vec<_>>();
-        let (auto_result, authoritative_result) =
-            restrict_null_predicate_in_modes(predicate, &join_cols_of_predicate)?;
-
-        assert_eq!(auto_result, authoritative_result);
-
-        Ok(())
-    }
-
-    #[test]
-    fn mixed_reference_predicate_remains_fast_pathed_in_authoritative_mode() -> Result<()>
-    {
-        let predicate = binary_expr(col("a"), Operator::Gt, col("b"));
-        let join_cols = vec![Column::from_name("a")];
-        let (auto_result, authoritative_only_result) =
-            restrict_null_predicate_in_modes(predicate.clone(), &join_cols)?;
-
-        assert!(!auto_result, "{predicate}");
-        assert!(!authoritative_only_result, "{predicate}");
-
-        Ok(())
-    }
-
-    #[test]
-    fn null_restriction_eval_mode_guard_restores_on_panic() {
-        set_null_restriction_eval_mode_for_test(NullRestrictionEvalMode::Auto);
-
-        let result = catch_unwind(AssertUnwindSafe(|| {
-            with_null_restriction_eval_mode_for_test(
-                NullRestrictionEvalMode::AuthoritativeOnly,
-                || panic!("intentional panic to verify test mode reset"),
-            )
-        }));
-
-        assert!(result.is_err());
-        assert_eq!(null_restriction_eval_mode(), NullRestrictionEvalMode::Auto);
-    }
 }