fix: correlated NOT IN with a non-equality correlation returns wrong results

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark

@@ -3,6 +3,20 @@ group null_aware_join
 
 load sql_benchmarks/null_aware_join/init/load.sql
 
+# Correctness canary: the NOT IN result must match a reference count
+# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS.
+# As Q04, with NULL outer keys excluded unless the subquery is empty.
+assert I
+SELECT count(*) = (
+  SELECT count(*) FROM small_outer o
+  WHERE o.z <= (SELECT min(z) FROM small_inner)
+     OR (o.id_n1 IS NOT NULL AND NOT (o.id % 2 = 0 AND (o.id / 2) % 1000 < o.z))
+)
+FROM small_outer o
+WHERE o.id_n1 NOT IN (SELECT i.id_n0 FROM small_inner i WHERE i.z < o.z);
+----
+true
+
 expect_plan HashJoinExec
 expect_plan null_aware: true
 

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark

@@ -3,6 +3,20 @@ group null_aware_join
 
 load sql_benchmarks/null_aware_join/init/load.sql
 
+# Correctness canary: the NOT IN result must match a reference count
+# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS.
+# As Q04, with NULL outer keys excluded unless the subquery is empty.
+assert I
+SELECT count(*) = (
+  SELECT count(*) FROM small_outer o
+  WHERE o.z <= (SELECT min(z) FROM small_inner)
+     OR (o.id_n50 IS NOT NULL AND NOT (o.id % 2 = 0 AND (o.id / 2) % 1000 < o.z))
+)
+FROM small_outer o
+WHERE o.id_n50 NOT IN (SELECT i.id_n0 FROM small_inner i WHERE i.z < o.z);
+----
+true
+
 expect_plan HashJoinExec
 expect_plan null_aware: true
 

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark

@@ -3,6 +3,22 @@ group null_aware_join
 
 load sql_benchmarks/null_aware_join/init/load.sql
 
+# Correctness canary: the NOT IN result must match a reference count
+# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS.
+# A row is TRUE when the subquery is empty, or when the subquery holds no NULL
+# and the key is not in it. A NULL is in scope when its z is below o.z.
+assert I
+SELECT count(*) = (
+  SELECT count(*) FROM small_outer o
+  WHERE o.z <= (SELECT min(z) FROM small_inner)
+     OR (o.z <= (SELECT min(z) FROM small_inner WHERE id_n50 IS NULL)
+         AND NOT (o.id % 2 = 0 AND (o.id / 2) % 1000 < o.z))
+)
+FROM small_outer o
+WHERE o.id_n0 NOT IN (SELECT i.id_n50 FROM small_inner i WHERE i.z < o.z);
+----
+true
+
 expect_plan HashJoinExec
 expect_plan null_aware: true
 

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q08.benchmark

@@ -3,7 +3,30 @@ group null_aware_join
 
 load sql_benchmarks/null_aware_join/init/load.sql
 
+# Correctness canary: the NOT IN result must match a reference count
+# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS.
+# A row is TRUE when o.z > 900, when the subquery for its k is empty, or when
+# its key is not NULL and not in that subquery.
+assert I
+SELECT count(*) = (
+  SELECT count(*)
+  FROM small_outer o
+  JOIN (SELECT k, min(z) AS min_z FROM small_inner GROUP BY k) m ON m.k = o.k
+  WHERE o.z > 900
+     OR o.z <= m.min_z
+     OR (o.id_n50 IS NOT NULL
+         AND NOT (o.id % 2 = 0 AND (o.id / 2) % 16 = o.k AND (o.id / 2) % 1000 < o.z))
+)
+FROM small_outer o
+WHERE o.z > 900
+   OR o.id_n50 NOT IN (
+        SELECT i.id_n0 FROM small_inner i WHERE i.k = o.k AND i.z < o.z
+      );
+----
+true
+
 expect_plan HashJoinExec
+expect_plan null_aware: true
 
 run
 -- Q8: NOT IN correlated by both an equality and a non-equality, 50% NULL on

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -21,7 +21,6 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use crate::decorrelate::PullUpCorrelatedExpr;
-use crate::extract_equijoin_predicate::split_eq_and_noneq_join_predicate;
 use crate::optimizer::ApplyOrder;
 use crate::utils::replace_qualified_name;
 use crate::{OptimizerConfig, OptimizerRule};
@@ -535,25 +534,12 @@ fn build_join(
             sub_query_alias.clone()
         };
 
-        let mark_filter_is_hashable_only =
-            if join_type == JoinType::LeftMark && in_predicate_opt.is_some() {
-                let (_, residual_filter) = split_eq_and_noneq_join_predicate(
-                    join_filter.clone(),
-                    left.schema(),
-                    right_projected.schema(),
-                )?;
-                residual_filter.is_none()
-            } else {
-                false
-            };
-
         // For scalar NOT IN mark joins, propagate null-aware semantics into the
-        // nullable mark column when the predicate can be implemented by hash keys.
-        // Non-equality correlated filters stay on the legacy path because hash join
-        // execution cannot mark UNKNOWN candidates for residual predicates.
+        // nullable mark column. A non-equality correlation stays behind as a
+        // join filter, which the hash join also applies when it decides
+        // whether a NULL makes the mark UNKNOWN.
         let null_aware = join_type == JoinType::LeftMark
             && in_predicate_opt.is_some()
-            && mark_filter_is_hashable_only
             && join_keys_may_be_null(
                 &join_filter,
                 left.schema(),
@@ -707,18 +693,6 @@ mod tests {
         plan.inputs().into_iter().any(has_null_aware_left_mark_join)
     }
 
-    fn has_non_null_aware_left_mark_join(plan: &LogicalPlan) -> bool {
-        if let LogicalPlan::Join(join) = plan
-            && join.join_type == JoinType::LeftMark
-        {
-            return !join.null_aware;
-        }
-
-        plan.inputs()
-            .into_iter()
-            .any(has_non_null_aware_left_mark_join)
-    }
-
     fn optimize_with_decorrelate(plan: LogicalPlan) -> Result<LogicalPlan> {
         let optimizer = crate::Optimizer::with_rules(vec![Arc::new(
             DecorrelatePredicateSubquery::new(),
@@ -1441,8 +1415,10 @@ mod tests {
         Ok(())
     }
 
+    /// A non-equality correlation stays behind as a join filter, which the hash
+    /// join applies when it marks UNKNOWN rows, so the mark is still null-aware.
     #[test]
-    fn correlated_not_in_mark_join_is_not_null_aware_for_residual_filter() -> Result<()> {
+    fn correlated_not_in_mark_join_is_null_aware_for_residual_filter() -> Result<()> {
         let outer_scan = nullable_scalar_mark_scan("outer_t")?;
         let inner_scan = nullable_scalar_mark_scan("inner_t")?;
 
@@ -1461,7 +1437,7 @@ mod tests {
 
         let optimized = optimize_with_decorrelate(plan)?;
         assert!(
-            has_non_null_aware_left_mark_join(&optimized),
+            has_null_aware_left_mark_join(&optimized),
             "{}",
             optimized.display_indent_schema()
         );

datafusion/physical-optimizer/src/join_selection.rs

@@ -172,11 +172,13 @@ impl PhysicalOptimizerRule for JoinSelection {
     }
 }
 
-/// Determines whether it is possible to swap inputs of a hash join - for null-aware joins, we can only swap `LeftAnti` with no filters
+/// Determines whether it is possible to swap inputs of a hash join - for null-aware joins, we can only swap an uncorrelated `LeftAnti`
+/// (a single join key and no filter), because the swapped `RightAnti` has no per-row NULL handling
 fn can_swap_hash_join(hash_join: &HashJoinExec) -> bool {
     hash_join.join_type().supports_swap()
         && (!hash_join.null_aware
             || (*hash_join.join_type() == JoinType::LeftAnti
+                && hash_join.on().len() == 1
                 && hash_join.filter().is_none()))
 }