comphead commented on code in PR #26213:
URL: https://github.com/apache/datafusion/pull/26213#discussion_r4319046693


##########
datafusion/physical-plan/src/joins/hash_join/stream.rs:
##########
@@ -930,31 +940,55 @@ impl HashJoinStream {
             // Must precede the empty-build-map return below: an all-NULL-key
             // build side has an empty full-key map but still needs UNKNOWN 
marks.
             if mode.is_correlated() && state.is_first_chunk() {
+                // `RightAnti` marks the rows of this probe batch, the others
+                // the build rows shared by all probe partitions.
+                let probe_marks = (!mode.marks_build_rows()).then(|| {

Review Comment:
   For a correlated `RightAnti` this allocates an `Arc<Mutex<_>>` and a zeroed 
bitmap for every probe batch, plus a `Box<BooleanBuffer>` when something is 
marked. That also happens for batches where neither side has a NULL-valued row, 
so nothing can be marked. Q10 has no NULLs and is flat, so the cost is probably 
small. If it shows up, allocating on the first mark, or skipping when 
`null_value_build_rows` is `None` and the probe batch has no NULL-valued row, 
would avoid it. The bitmap is local to the stream, so a plain 
`BooleanBufferBuilder` reused across batches like `null_mark_*_buffer` would 
also drop the mutex. Unmeasured.



##########
datafusion/core/tests/physical_optimizer/join_selection.rs:
##########
@@ -641,16 +642,70 @@ async fn 
test_null_aware_left_anti_with_filter_does_not_swap() -> Result<()> {
 
     let optimized_join =
         JoinSelection::new().optimize(Arc::new(join), &ConfigOptions::new())?;
-    let unswapped_join = optimized_join
+    let swapped_join = optimized_join
         .downcast_ref::<HashJoinExec>()
-        .expect("filtered null-aware anti join should not swap");
+        .expect("anti join swap should not require a projection");
 
-    assert_eq!(*unswapped_join.join_type(), JoinType::LeftAnti);
-    assert_eq!(*unswapped_join.partition_mode(), PartitionMode::CollectLeft);
-    assert!(unswapped_join.null_aware);
-    assert!(unswapped_join.filter().is_some());
-    assert_eq!(unswapped_join.left().schema().field(0).name(), "big_col");
-    assert_eq!(unswapped_join.right().schema().field(0).name(), "small_col");
+    assert_eq!(*swapped_join.join_type(), JoinType::RightAnti);
+    assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft);
+    assert!(swapped_join.null_aware);
+    let filter = swapped_join.filter().expect("filter should be kept");
+    assert_eq!(
+        filter
+            .column_indices()
+            .iter()
+            .map(|c| c.side)
+            .collect::<Vec<_>>(),
+        vec![JoinSide::Right, JoinSide::Left]
+    );
+    assert_eq!(swapped_join.left().schema().field(0).name(), "small_col");
+    assert_eq!(swapped_join.right().schema().field(0).name(), "big_col");
+
+    Ok(())
+}
+
+/// A multi-key null-aware `LeftAnti` (a multi-column `NOT IN`, correlation
+/// scope keys, or both) swaps to a `RightAnti` that keeps the number of
+/// `NOT IN` value keys and the key order.
+#[rstest]
+#[case(1)]
+#[case(2)]
+#[tokio::test]
+async fn test_null_aware_multi_key_left_anti_swaps_to_right_anti(

Review Comment:
   This test and `test_null_aware_left_anti_with_filter_swaps` repeat the setup 
and most assertions of `test_null_aware_left_anti_swaps_to_right_anti`. 
`null_aware_anti_join.slt` already shows the swapped plans too (`RightAnti`, 
key order, `null_aware(value_keys=2)`, the filter). One more `rstest` case or a 
shared helper taking the keys, `value_keys` and filter would cover it. The 
swapped filter sides are the only assertion the slt plans do not show.



##########
datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs:
##########
@@ -68,6 +78,43 @@ async fn null_aware_join_matches_reference() {
             }
         }
     }
+    // A fuzz run that never reaches one of the plan shapes would be vacuous.
+    assert!(plans.left_anti > 0, "{plans:?}");
+    assert!(plans.right_anti > 0, "{plans:?}");
+    assert!(plans.correlated_right_anti > 0, "{plans:?}");
+}
+
+/// Counts the null-aware anti join plans the fuzz queries ran.
+#[derive(Debug, Default)]
+struct NullAwarePlans {
+    left_anti: usize,
+    right_anti: usize,
+    /// `RightAnti` joins with more than one key or a join filter, which
+    /// decide `NOT IN` per probe row.
+    correlated_right_anti: usize,
+}
+
+impl NullAwarePlans {
+    async fn record(&mut self, ctx: &SessionContext, sql: &str) {
+        let plan = ctx
+            .sql(sql)
+            .await
+            .unwrap()
+            .create_physical_plan()
+            .await
+            .unwrap();
+        let plan = displayable(plan.as_ref()).indent(true).to_string();
+        for line in plan.lines().filter(|line| line.contains("null_aware")) {
+            if line.contains("join_type=LeftAnti") {
+                self.left_anti += 1;
+            } else if line.contains("join_type=RightAnti") {
+                self.right_anti += 1;
+                if line.contains("), (") || line.contains("filter=") {

Review Comment:
   The plan shape is detected by matching substrings of the displayed plan. A 
change to the display format would make the final asserts fail with a 
misleading message. Walking the plan for `HashJoinExec` and checking 
`join_type()`, `on().len()` and `filter()` would be sturdier. `record` also 
plans every query a second time because `run` plans it again, and the reference 
queries are recorded too. Recording only the `NOT IN` query, or reusing the 
physical plan, would avoid that.



##########
datafusion/core/tests/memory_limit/null_aware_not_in.rs:
##########
@@ -0,0 +1,103 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A multi-column `NOT IN` over a small subquery must build its hash join on
+//! the subquery, so a large outer table runs under a memory limit that only
+//! fits the subquery side.
+//!
+//! The null-aware join of a multi-column `NOT IN` used to stay a `LeftAnti`
+//! that builds on the outer side, so it ran out of memory where the scalar
+//! `NOT IN`, swapped to a `RightAnti` that builds on the subquery, ran.
+
+use std::sync::Arc;
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::MemTable;
+use datafusion::prelude::*;
+use datafusion_common::{Result, assert_contains};
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+
+const OUTER_ROWS: i64 = 200_000;
+/// Far below the outer side's build state, far above the subquery's.
+const MEMORY_LIMIT: usize = 2 * 1024 * 1024;
+
+/// A table `(a, b)` of `a = v, b = v % 1000` for `v` in `1..=rows`, followed 
by
+/// `extra` rows.
+fn table(rows: i64, extra: &[(Option<i64>, Option<i64>)]) -> Arc<MemTable> {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Int64, true),
+        Field::new("b", DataType::Int64, true),
+    ]));
+    let a = (1..=rows).map(Some).chain(extra.iter().map(|row| row.0));
+    let b = (1..=rows)
+        .map(|v| Some(v % 1000))
+        .chain(extra.iter().map(|row| row.1));
+    let batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![
+            Arc::new(a.collect::<Int64Array>()),
+            Arc::new(b.collect::<Int64Array>()),
+        ],
+    )
+    .unwrap();
+    // Several batches, so the outer side streams through the probe.
+    let batches = (0..batch.num_rows())
+        .step_by(8192)
+        .map(|offset| batch.slice(offset, 8192.min(batch.num_rows() - offset)))
+        .collect();
+    Arc::new(MemTable::try_new(schema, vec![batches]).unwrap())
+}
+
+async fn count_not_in(sql: &str, join_reordering: bool) -> Result<i64> {
+    let config = SessionConfig::new()
+        .with_target_partitions(1)
+        .set_bool("datafusion.optimizer.join_reordering", join_reordering);
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_limit(MEMORY_LIMIT, 1.0)
+        .build_arc()?;
+    let ctx = SessionContext::new_with_config_rt(config, runtime);
+    ctx.register_table("big", table(OUTER_ROWS, &[]))?;
+    // A NULL-valued subquery row makes `NOT IN` UNKNOWN for every outer row
+    // with `b = 7` (and no definite mismatch in `a`).
+    ctx.register_table("small", table(100, &[(None, Some(7))]))?;
+
+    let batches = ctx.sql(sql).await?.collect().await?;
+    Ok(batches[0]
+        .column(0)
+        .as_any()
+        .downcast_ref::<Int64Array>()
+        .unwrap()
+        .value(0))
+}
+
+#[tokio::test]
+async fn multi_column_not_in_builds_on_small_subquery() -> Result<()> {

Review Comment:
   This could probably be an `.slt` instead of 100 lines of Rust: 
`generate_series` tables, `SET datafusion.runtime.memory_limit` after they are 
created, then `query error Resources exhausted` with `join_reordering` off 
(`ordered_aggregate_spill.slt` has the same pattern). I did not check that 200k 
rows fit the slt runner defaults.



##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -11307,6 +11301,141 @@ mod tests {
         Ok(())
     }
 
+    /// Swaps a null-aware `LeftAnti` into the `RightAnti` that `JoinSelection`
+    /// plans, with its own build state so that both can be executed.
+    fn swap_for_separate_execution(
+        left_anti: &HashJoinExec,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let swapped = left_anti.swap_inputs(PartitionMode::CollectLeft)?;
+        let swapped = swapped.downcast_ref::<HashJoinExec>().unwrap();
+        Ok(Arc::new(swapped.builder().reset_state().build()?))
+    }
+
+    /// Uncorrelated two-column `(a, b) NOT IN (SELECT a, b ...)`, run both as
+    /// `LeftAnti` (built on the outer side) and as the swapped `RightAnti`
+    /// (built on the subquery side), which decides each outer row per probe
+    /// batch.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_multi_column_right_anti(batch_size: usize) -> 
Result<()> {

Review Comment:
   This test overlaps with `null_aware_multi_column.slt`, which now runs every 
query with join reordering on and off across the batch sizes. The 
NULL-in-subquery and empty-subquery cases here match `mc_inner_null` and 
`mc_empty` there. The only new case is the subquery where every row is 
NULL-valued (empty full-key map). Adding a table like that to the slt and 
dropping this test (about 120 lines) would keep the coverage. 
`test_null_aware_validation_multi_column_right_anti` and 
`test_null_aware_filter_accepted_for_anti_joins` also only check construction, 
and the mode is asserted here already. One parametrized construction test would 
be enough.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to