[branch-56] fix: reject a multi-key null-aware join without a NOT IN value key count (backport #27111)

datafusion/core/tests/sql/joins.rs

@@ -688,3 +688,75 @@ async fn asof_join_unparser_preserves_right_preselection() -> Result<()> {
     }
     Ok(())
 }
+
+/// The `NOT IN` value key count of every decorrelated mark join in `sql`'s
+/// optimized plan, with whether the join is null-aware.
+async fn mark_join_value_keys(
+    ctx: &SessionContext,
+    sql: &str,
+) -> Result<Vec<(bool, usize)>> {
+    use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+    use datafusion::logical_expr::{JoinType, LogicalPlan};
+
+    let plan = ctx.sql(sql).await?.into_optimized_plan()?;
+    let mut joins = vec![];
+    plan.apply(|plan| {
+        if let LogicalPlan::Join(join) = plan
+            && join.join_type == JoinType::LeftMark
+        {
+            joins.push((join.null_aware, join.null_aware_value_keys));
+        }
+        Ok(TreeNodeRecursion::Continue)
+    })?;
+    Ok(joins)
+}
+
+/// Decorrelation sets the value key count only on a null-aware join, which is
+/// the only one that reads it. Any other join keeps it unset (`0`), the same as a
+/// join built through `Join::try_new` or `LogicalPlanBuilder`, so equal plans
+/// compare equal however they were built.
+#[tokio::test]
+async fn decorrelated_mark_join_sets_value_keys_only_when_null_aware() -> Result<()> {
+    let ctx = SessionContext::new();
+    for (name, nullable) in [("t", false), ("n", true)] {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, nullable),
+            Field::new("b", DataType::Int32, nullable),
+        ]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2])),
+                Arc::new(Int32Array::from(vec![3, 4])),
+            ],
+        )?;
+        ctx.register_table(
+            name,
+            Arc::new(MemTable::try_new(schema, vec![vec![batch]])?),
+        )?;
+    }
+
+    // Non-nullable keys: neither the scalar nor the multi-column `IN` is
+    // null-aware.
+    for sql in [
+        "SELECT a IN (SELECT a FROM t AS s) FROM t",
+        "SELECT (a, b) IN (SELECT a, b FROM t AS s) FROM t",
+    ] {
+        assert_eq!(
+            mark_join_value_keys(&ctx, sql).await?,
+            [(false, 0)],
+            "{sql}"
+        );
+    }
+    // Nullable keys: both are null-aware and state their count.
+    assert_eq!(
+        mark_join_value_keys(&ctx, "SELECT a IN (SELECT a FROM n AS s) FROM n").await?,
+        [(true, 1)]
+    );
+    assert_eq!(
+        mark_join_value_keys(&ctx, "SELECT (a, b) IN (SELECT a, b FROM n AS s) FROM n")
+            .await?,
+        [(true, 2)]
+    );
+    Ok(())
+}

datafusion/expr/src/logical_plan/builder.rs

@@ -1440,7 +1440,7 @@ impl LogicalPlanBuilder {
             schema: DFSchemaRef::new(join_schema),
             null_equality,
             null_aware,
-            null_aware_value_keys: 1,
+            null_aware_value_keys: 0,
         })))
     }
 

datafusion/expr/src/logical_plan/plan.rs

@@ -4650,6 +4650,11 @@ pub struct Join {
     /// after them are correlation scope keys of a correlated subquery. A NULL
     /// in a value key follows SQL three-valued logic across the whole tuple,
     /// while the scope keys only select which subquery rows are compared.
+    ///
+    /// `0` leaves the count unset, which is the default of [`Self::try_new`]
+    /// and the `LogicalPlanBuilder` join methods. Physical planning accepts an
+    /// unset count only for a single-key join, a scalar `NOT IN`, and rejects
+    /// it for more keys: no default can tell which of them are value keys.
     pub null_aware_value_keys: usize,
 }
 
@@ -4878,7 +4883,7 @@ impl Join {
             schema: Arc::new(join_schema),
             null_equality,
             null_aware,
-            null_aware_value_keys: 1,
+            null_aware_value_keys: 0,
         })
     }
 
@@ -7228,7 +7233,7 @@ mod tests {
                 schema: Arc::new(left_schema.join(&right_schema)?),
                 null_equality: NullEquality::NullEqualsNothing,
                 null_aware: false,
-                null_aware_value_keys: 1,
+                null_aware_value_keys: 0,
             }))
         }
 

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -738,14 +738,21 @@ fn subquery_output_expr(subquery: &LogicalPlan, i: usize) -> Result<Expr> {
 }
 
 /// Sets [`Join::null_aware_value_keys`](datafusion_expr::Join::null_aware_value_keys)
-/// on a join built by [`build_join`].
+/// on a null-aware join built by [`build_join`]. Only a null-aware join reads
+/// the count, so any other join keeps it unset, as [`Join::try_new`] leaves
+/// it, and equal joins compare equal however they were built.
+///
+/// [`Join::try_new`]: datafusion_expr::Join::try_new
 fn with_null_aware_value_keys(
     plan: LogicalPlan,
     null_aware_value_keys: usize,
 ) -> Result<LogicalPlan> {
     let LogicalPlan::Join(join) = plan else {
         return internal_err!("expected a join, got {}", plan.display());
     };
+    if !join.null_aware {
+        return Ok(LogicalPlan::Join(join));
+    }
     Ok(LogicalPlan::Join(
         join.with_null_aware_value_keys(null_aware_value_keys),
     ))

datafusion/optimizer/src/eliminate_cross_join.rs

@@ -345,7 +345,7 @@ fn find_inner_join(
                 schema: join_schema,
                 null_equality,
                 null_aware: false,
-                null_aware_value_keys: 1,
+                null_aware_value_keys: 0,
             }));
         }
     }
@@ -369,7 +369,7 @@ fn find_inner_join(
         join_constraint: JoinConstraint::On,
         null_equality,
         null_aware: false,
-        null_aware_value_keys: 1,
+        null_aware_value_keys: 0,
     }))
 }
 
@@ -1443,7 +1443,7 @@ mod tests {
             schema: join_schema,
             null_equality: NullEquality::NullEqualsNull, // Test preservation
             null_aware: false,
-            null_aware_value_keys: 1,
+            null_aware_value_keys: 0,
         });
 
         // Apply filter that can create join conditions