perf: read NOT NULL tuple elements of a multi-column NOT IN as scope keys

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q11.benchmark

@@ -25,9 +25,10 @@ expect_plan HashJoinExec
 
 run
 -- Q11: uncorrelated multi-column NOT IN with 1% NULL in one subquery element.
--- Each NULL-valued subquery row is paired with every outer row and the pair is
--- kept as UNKNOWN only when the other element matches, so the cost grows with
--- (outer rows) x (NULL-valued subquery rows). Compare with Q10 and Q12.
+-- `z` is NOT NULL on both sides, so the join reads it as a scope key: each
+-- NULL-valued subquery row is paired only with the outer rows with the same
+-- `z`, through a hash lookup, so the cost grows with (NULL-valued subquery
+-- rows) x (outer rows per `z` value). Compare with Q10 and Q12.
 SELECT count(*)
 FROM small_outer o
 WHERE (o.id_n0, o.z) NOT IN (SELECT i.id_n1, i.z FROM small_inner i);

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q12.benchmark

@@ -26,10 +26,10 @@ expect_plan HashJoinExec
 
 run
 -- Q12: uncorrelated multi-column NOT IN with 1% NULL in one outer element,
--- the mirror of Q11. Each NULL-valued outer row is paired with every subquery
--- row and is UNKNOWN as soon as the other element of one of them matches: the
--- cost grows with (NULL-valued outer rows) x (subquery rows). Compare with Q10
--- and Q11.
+-- the mirror of Q11. `z` is NOT NULL on both sides, so the join reads it as a
+-- scope key: each NULL-valued outer row is paired only with the subquery rows
+-- with the same `z`, so the cost grows with (NULL-valued outer rows) x
+-- (subquery rows per `z` value). Compare with Q10 and Q11.
 SELECT count(*)
 FROM small_outer o
 WHERE (o.id_n1, o.z) NOT IN (SELECT i.id_n0, i.z FROM small_inner i);

datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs

@@ -21,6 +21,8 @@
 //! null-aware hash joins, are compared with `NOT EXISTS` / `EXISTS`
 //! formulations of the same three-valued logic, which do not. Random tables
 //! with NULLs on both sides are swept across batch sizes and partition counts.
+//! Some columns are NOT NULL, so that tuples mix elements that can be NULL with
+//! elements that cannot, which the null-aware join reads as scope keys.
 //!
 //! The subquery table is sometimes much smaller than the outer table, so
 //! `JoinSelection` also swaps `NOT IN` into a null-aware `RightAnti` that
@@ -44,11 +46,12 @@ async fn null_aware_join_matches_reference() {
     let mut plans = NullAwarePlans::default();
     for seed in 0..20 {
         let mut rng = StdRng::seed_from_u64(seed);
-        let outer = random_table(&mut rng, &["id", "k", "a", "b", "c"], 6);
+        let outer = random_table(&mut rng, &["id", "k", "a", "b", "c"], &["n", "m"], 6);
         // Alternate between a subquery table of similar size and a much
         // smaller one that `JoinSelection` moves to the build side.
         let inner_rows = if seed % 2 == 0 { 6 } else { 2 };
-        let inner = random_table(&mut rng, &["k", "x", "y", "z"], inner_rows);
+        let inner =
+            random_table(&mut rng, &["k", "x", "y", "z"], &["v", "w"], inner_rows);
         for target_partitions in [1, 4] {
             for batch_size in [1, 3, 8192] {
                 let config = SessionConfig::new()
@@ -142,6 +145,22 @@ fn cases() -> Vec<(String, String)> {
                 "o.a = i.x AND o.b = i.y AND o.c = i.z",
             ),
             ("(o.a, 1)", "i.x, i.y", "o.a = i.x AND 1 = i.y"),
+            // `o.n`, `o.m`, `i.v` and `i.w` are NOT NULL.
+            ("(o.a, o.n)", "i.x, i.v", "o.a = i.x AND o.n = i.v"),
+            ("(o.n, o.a)", "i.v, i.x", "o.n = i.v AND o.a = i.x"),
+            (
+                "(o.n, o.a, o.b)",
+                "i.v, i.x, i.y",
+                "o.n = i.v AND o.a = i.x AND o.b = i.y",
+            ),
+            (
+                "(o.n, o.a, o.m)",
+                "i.v, i.x, i.w",
+                "o.n = i.v AND o.a = i.x AND o.m = i.w",
+            ),
+            ("(o.n, o.a)", "i.y, i.x", "o.n = i.y AND o.a = i.x"),
+            ("(1, o.a)", "i.v, i.x", "1 = i.v AND o.a = i.x"),
+            ("(o.n, o.m)", "i.v, i.w", "o.n = i.v AND o.m = i.w"),
         ] {
             let subquery = format!("SELECT {output} FROM i{correlation}");
             let matched = format!(
@@ -179,14 +198,24 @@ fn cases() -> Vec<(String, String)> {
     cases
 }
 
-/// A table of nullable `Int32` columns with values in `0..3`, split into up
-/// to three partitions of fewer than `max_rows` rows each. A column named `id`
-/// holds unique row ids instead.
-fn random_table(rng: &mut StdRng, columns: &[&str], max_rows: i32) -> Arc<MemTable> {
+/// A table of `Int32` columns with values in `0..3`, split into up to three
+/// partitions of fewer than `max_rows` rows each: the `nullable` columns, then
+/// the `non_nullable` ones. A column named `id` holds unique row ids instead.
+fn random_table(
+    rng: &mut StdRng,
+    nullable: &[&str],
+    non_nullable: &[&str],
+    max_rows: i32,
+) -> Arc<MemTable> {
+    let columns = nullable
+        .iter()
+        .map(|name| (*name, true))
+        .chain(non_nullable.iter().map(|name| (*name, false)))
+        .collect::<Vec<_>>();
     let schema = Arc::new(Schema::new(
         columns
             .iter()
-            .map(|name| Field::new(*name, DataType::Int32, true))
+            .map(|(name, nullable)| Field::new(*name, DataType::Int32, *nullable))
             .collect::<Vec<_>>(),
     ));
     let null_fraction = [0.0, 0.1, 0.3][rng.random_range(0..3)];
@@ -196,13 +225,13 @@ fn random_table(rng: &mut StdRng, columns: &[&str], max_rows: i32) -> Arc<MemTab
             let rows = rng.random_range(0..max_rows);
             let arrays = columns
                 .iter()
-                .map(|name| {
+                .map(|(name, nullable)| {
                     let values: Int32Array = if *name == "id" {
                         (next_id..next_id + rows).map(Some).collect()
                     } else {
                         (0..rows)
                             .map(|_| {
-                                (!rng.random_bool(null_fraction))
+                                (!(*nullable && rng.random_bool(null_fraction)))
                                     .then(|| rng.random_range(0..3))
                             })
                             .collect()

datafusion/expr/src/logical_plan/plan.rs

@@ -4615,7 +4615,9 @@ pub struct Join {
     /// [`Self::null_aware`] is set.
     ///
     /// A scalar `x NOT IN (SELECT y ...)` has one value key, a multi-column
-    /// `(a, b) NOT IN (SELECT x, y ...)` one per tuple element. The value keys
+    /// `(a, b) NOT IN (SELECT x, y ...)` one per tuple element that can be
+    /// NULL in scope. A tuple element that cannot be NULL on either side never
+    /// makes the comparison UNKNOWN, so it is a scope key instead. The value keys
     /// are the leading equi-join keys (`on[..n]`, or the leading equality
     /// conjuncts of `filter` before they are extracted into `on`); any keys
     /// after them are correlation scope keys of a correlated subquery. A NULL

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -569,7 +569,7 @@ struct MarkJoin {
 /// Are the `IN` equalities, in order, the first equalities of `join_filter`
 /// that the hash join can use as keys? A null-aware hash join reads its first
 /// `V` keys as the `IN` values (one per tuple element of a multi-column
-/// `IN`) and the other keys as the correlation.
+/// `IN` that can be NULL in scope) and the other keys as the correlation.
 fn values_are_first_keys(
     in_values: &[InValue],
     join_filter: &Expr,
@@ -722,9 +722,11 @@ fn strip_casts(expr: &Expr) -> &Expr {
 /// Combines the `IN`/`NOT IN` equalities, one per value, with the correlation
 /// predicates, if any.
 ///
-/// The equalities stay the leading conjuncts, in order, so that they become
-/// `on[..V]`, the key positions a null-aware hash join reads as the `NOT IN`
-/// value keys (see `HashJoinExec::null_aware`).
+/// The equalities stay the leading conjuncts, in order, so that the first `V`
+/// of them become `on[..V]`, the key positions a null-aware hash join reads as
+/// the `NOT IN` value keys (see `HashJoinExec::null_aware`). The equalities
+/// after them, if any, are on tuple elements that cannot be NULL, which the
+/// join reads as scope keys like the correlation.
 fn in_join_filter(in_values: &[InValue], correlation: Option<Expr>) -> Result<Expr> {
     conjunction(
         in_values
@@ -953,16 +955,19 @@ fn build_join(
         None => vec![],
     };
 
-    // True if an `IN` value or the subquery column it is compared with can be
-    // NULL inside the scope of an outer row.
-    let mut value_may_be_null = false;
-    for in_value in &in_values {
-        value_may_be_null |= in_value.may_be_null_in_scope(
-            left.schema(),
-            sub_query_alias.schema(),
-            &pull_up.correlated_filters,
-        )?;
-    }
+    // For each `IN` value, true if it or the subquery column it is compared
+    // with can be NULL inside the scope of an outer row.
+    let elements_may_be_null = in_values
+        .iter()
+        .map(|in_value| {
+            in_value.may_be_null_in_scope(
+                left.schema(),
+                sub_query_alias.schema(),
+                &pull_up.correlated_filters,
+            )
+        })
+        .collect::<Result<Vec<_>>>()?;
+    let value_may_be_null = elements_may_be_null.iter().any(|may_be_null| *may_be_null);
 
     // Whether this join needs null-aware (`NOT IN` three-valued) semantics.
     // Decided once: the constant projection below, the mark branch and the anti
@@ -1037,6 +1042,38 @@ fn build_join(
     }
     let left = projected_left.as_ref().unwrap_or(left);
 
+    // The first `V` keys of a null-aware join are its `NOT IN` value keys.
+    //
+    // A tuple element that cannot be NULL in scope on either side never makes
+    // the tuple comparison UNKNOWN: compared with a definite value it is either
+    // equal, which leaves the other elements to decide, or a definite mismatch,
+    // which makes the comparison FALSE whatever the other elements are. It
+    // therefore only selects which subquery rows a NULL in another element is
+    // compared with, which is what a correlation scope key does. Moving such
+    // elements after the ones that can be NULL turns them into scope keys, so
+    // the hash join finds the rows a NULL-valued row is compared with through a
+    // hash lookup instead of pairing it with every row of the other side
+    // (<https://github.com/apache/datafusion/issues/26088>).
+    //
+    // A null-aware tuple has at least one element that can be NULL, so `V` is
+    // at least one. With exactly one, the join is a scalar `NOT IN` on that
+    // element, correlated by the other elements.
+    let null_aware_value_keys = if null_aware && in_values.len() > 1 {
+        let (nullable, non_nullable): (Vec<_>, Vec<_>) = in_values
+            .into_iter()
+            .zip(elements_may_be_null)
+            .partition(|(_, may_be_null)| *may_be_null);
+        let null_aware_value_keys = nullable.len();
+        in_values = nullable
+            .into_iter()
+            .chain(non_nullable)
+            .map(|(in_value, _)| in_value)
+            .collect();
+        null_aware_value_keys
+    } else {
+        in_values.len().max(1)
+    };
+
     let join_filter = match (in_values.is_empty(), join_filter_opt) {
         (false, correlation_opt) => in_join_filter(&in_values, correlation_opt)?,
         (true, Some(correlation)) => correlation,
@@ -1046,23 +1083,22 @@ fn build_join(
     // Only a null-aware join reads the key order.
     let value_is_first_key = null_aware
         && values_are_first_keys(
-            &in_values,
+            &in_values[..null_aware_value_keys],
             &join_filter,
             left.schema(),
             sub_query_alias.schema(),
         )?;
 
     // A multi-column `IN` has no fallback that materializes the UNKNOWN rows
-    // (see below), so a tuple element that does not become its key, for
-    // example because its type is not hashable, is rejected instead of being
-    // read in the wrong key position.
+    // (see below), so a tuple element that can be NULL and does not become its
+    // key, for example because its type is not hashable, is rejected instead
+    // of being read in the wrong key position. The other elements need not be
+    // keys: like any correlation, they may stay in the join filter.
     if null_aware && !value_is_first_key && in_values.len() > 1 {
         return not_impl_err!(
             "Multi-column NOT IN subquery is not supported when a tuple element is not a hashable equi-join key: {join_filter}"
         );
     }
-    // The first `V` keys of a null-aware join are its `NOT IN` value keys.
-    let null_aware_value_keys = in_values.len().max(1);
 
     // A `NOT IN` in a filter builds a `LeftAnti` join. The null-aware hash
     // join reads its first key as the `NOT IN` value and the other keys as the

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -855,16 +855,18 @@ pub struct HashJoinExec {
     /// [`Self::null_aware`] is set.
     ///
     /// A scalar `x NOT IN (SELECT y ...)` has one value key; a multi-column
-    /// `(a, b) NOT IN (SELECT x, y ...)` has one per tuple element. The tuple
-    /// comparison is UNKNOWN rather than FALSE when no element pair is a
-    /// definite mismatch and at least one of them involves a NULL, so a NULL in
-    /// one element does not make every comparison UNKNOWN the way a scalar
-    /// NULL does.
+    /// `(a, b) NOT IN (SELECT x, y ...)` has one per tuple element that can be
+    /// NULL. The tuple comparison is UNKNOWN rather than FALSE when no element
+    /// pair is a definite mismatch and at least one of them involves a NULL,
+    /// so a NULL in one element does not make every comparison UNKNOWN the
+    /// way a scalar NULL does. A tuple element that cannot be NULL on either
+    /// side never makes the comparison UNKNOWN; the planner passes it as a
+    /// scope key after the value keys, so it narrows the rows a NULL-valued
+    /// row is compared with.
     ///
     /// Such a join decides UNKNOWN per row, like a correlated scalar `NOT IN`:
-    /// without correlation keys every NULL-valued row is paired with every row
-    /// on the other side to check its other elements. See
-    /// <https://github.com/apache/datafusion/issues/26088> for the follow-up.
+    /// without scope keys every NULL-valued row is paired with every row on
+    /// the other side to check its other elements.
     pub null_aware_value_keys: usize,
     /// Cache holding plan properties like equivalences, output partitioning etc.
     cache: Arc<PlanProperties>,

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -3907,7 +3907,7 @@ async fn roundtrip_join_null_aware() -> Result<()> {
         (
             "
             SELECT k
-            FROM (VALUES (1, 1, 2), (2, 3, 4)) AS t1(k, a, b)
+            FROM (VALUES (1, 1, 2), (2, 3, CAST(NULL AS INT))) AS t1(k, a, b)
             WHERE (a, b) NOT IN (
                 SELECT x, y
                 FROM (VALUES (1, CAST(NULL AS INT), 2)) AS t2(k, x, y)
@@ -3916,6 +3916,19 @@ async fn roundtrip_join_null_aware() -> Result<()> {
             ",
             (3, 2),
         ),
+        // `b` and `y` are NOT NULL, so `b = y` is a scope key, not a value key.
+        (
+            "
+            SELECT k
+            FROM (VALUES (1, 1, 2), (2, 3, 4)) AS t1(k, a, b)
+            WHERE (a, b) NOT IN (
+                SELECT x, y
+                FROM (VALUES (1, CAST(NULL AS INT), 2)) AS t2(k, x, y)
+                WHERE t2.k = t1.k
+            )
+            ",
+            (3, 1),
+        ),
     ];
 
     for (sql, expected_keys) in cases {