refactor: track NestedLoopJoin fallback matches in one left bitmap and drop the cancel protocol

datafusion/core/tests/fuzz_cases/join_fuzz.rs

@@ -21,9 +21,9 @@ use std::time::SystemTime;
 use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj};
 
 use arrow::array::{
-    Array, ArrayRef, BinaryArray, BooleanArray, Float64Array, Int32Array,
+    Array, ArrayRef, BinaryArray, BooleanArray, Float64Array, Int32Array, UInt32Array,
 };
-use arrow::compute::SortOptions;
+use arrow::compute::{SortOptions, take_record_batch};
 use arrow::datatypes::{DataType, Schema};
 use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::pretty_format_batches;
@@ -867,7 +867,30 @@ impl JoinFuzzTestCase {
     }
 
     fn nested_loop_join(&self) -> Arc<NestedLoopJoinExec> {
-        let (left, right) = self.left_right();
+        let (_, right) = self.left_right();
+        self.nested_loop_join_with_right(right)
+    }
+
+    /// A nested loop join whose right side is spread over `partitions` partitions
+    fn nested_loop_join_partitioned_right(
+        &self,
+        partitions: usize,
+    ) -> Arc<NestedLoopJoinExec> {
+        let mut right = vec![vec![]; partitions];
+        for (i, batch) in self.input2.iter().enumerate() {
+            right[i % partitions].push(batch.clone());
+        }
+        let right =
+            MemorySourceConfig::try_new_exec(&right, self.input2[0].schema(), None)
+                .unwrap();
+        self.nested_loop_join_with_right(right)
+    }
+
+    fn nested_loop_join_with_right(
+        &self,
+        right: Arc<dyn ExecutionPlan>,
+    ) -> Arc<NestedLoopJoinExec> {
+        let (left, _) = self.left_right();
 
         let column_indices = self.column_indices();
         let intermediate_schema = self.intermediate_schema();
@@ -1274,6 +1297,97 @@ async fn test_filtered_join_spill_fuzz() {
     }
 }
 
+/// Fuzz test: compare NLJ under memory pressure against the same NLJ without a
+/// memory limit, for every join type. Under pressure the left side is spilled
+/// and read back in chunks that several right partitions probe, so this
+/// exercises the unmatched rows that have to be tracked across chunks and
+/// across partitions, with chunk boundaries that random batch sizes make
+/// different on every run.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_nested_loop_join_spill_fuzz() {
+    let join_types = [
+        JoinType::Inner,
+        JoinType::Left,
+        JoinType::Right,
+        JoinType::Full,
+        JoinType::LeftSemi,
+        JoinType::LeftAnti,
+        JoinType::LeftMark,
+        JoinType::RightSemi,
+        JoinType::RightAnti,
+        JoinType::RightMark,
+    ];
+
+    let runtime_spill = RuntimeEnvBuilder::new()
+        .with_memory_limit(4096, 1.0)
+        .with_disk_manager_builder(
+            DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
+        )
+        .build_arc()
+        .unwrap();
+
+    for join_type in join_types {
+        // The staggered batches are slices that each report the size of the
+        // buffers they share, which no batch fits the memory limit with. Copy
+        // them, so that the limit is exceeded part way through the left side
+        // and chunks span several batches.
+        let left = make_staggered_batches_i32(500, true)
+            .iter()
+            .map(|batch| {
+                let all_rows = UInt32Array::from_iter_values(0..batch.num_rows() as u32);
+                take_record_batch(batch, &all_rows).unwrap()
+            })
+            .collect();
+        let test_case = JoinFuzzTestCase::new(
+            left,
+            make_staggered_batches_i32(500, true),
+            join_type,
+            Some(Box::new(col_lt_col_filter)),
+        );
+
+        for batch_size in [7, 100] {
+            let session_config = SessionConfig::new().with_batch_size(batch_size);
+
+            // Baseline (no memory limit)
+            let ctx = SessionContext::new_with_config(session_config.clone());
+            let expected = collect(
+                test_case.nested_loop_join_partitioned_right(4),
+                ctx.task_ctx(),
+            )
+            .await
+            .unwrap();
+
+            // With spilling
+            let nlj = test_case.nested_loop_join_partitioned_right(4);
+            let task_ctx_spill = Arc::new(
+                TaskContext::default()
+                    .with_session_config(session_config)
+                    .with_runtime(Arc::clone(&runtime_spill)),
+            );
+            let spilled =
+                collect(Arc::clone(&nlj) as Arc<dyn ExecutionPlan>, task_ctx_spill)
+                    .await
+                    .unwrap();
+            assert!(
+                nlj.metrics().unwrap().spill_count().unwrap_or(0) > 0,
+                "{join_type:?} batch_size={batch_size}: expected the join to spill"
+            );
+
+            let expected_fmt = pretty_format_batches(&expected).unwrap().to_string();
+            let spilled_fmt = pretty_format_batches(&spilled).unwrap().to_string();
+            let mut expected_sorted: Vec<&str> = expected_fmt.trim().lines().collect();
+            expected_sorted.sort_unstable();
+            let mut spilled_sorted: Vec<&str> = spilled_fmt.trim().lines().collect();
+            spilled_sorted.sort_unstable();
+
+            assert_eq!(
+                expected_sorted, spilled_sorted,
+                "Content mismatch for {join_type:?} batch_size={batch_size}"
+            );
+        }
+    }
+}
+
 /// Return randomly sized record batches with:
 /// two sorted int32 columns 'a', 'b' ranged from 0..99 as join columns
 /// two random int32 columns 'x', 'y' as other columns

datafusion/core/tests/memory_limit/nlj_spill_unmatched.rs

@@ -27,6 +27,7 @@ use std::sync::Arc;
 
 use arrow::array::{Int32Array, RecordBatch};
 use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::physical_plan::{ExecutionPlan, collect};
 use datafusion::prelude::*;
 use datafusion_execution::runtime_env::RuntimeEnvBuilder;
 
@@ -57,7 +58,19 @@ fn table(rows: usize, value_column: &str, seed: u64) -> RecordBatch {
     .unwrap()
 }
 
+/// Total `spill_count` of `plan` and everything below it.
+fn spill_count(plan: &Arc<dyn ExecutionPlan>) -> usize {
+    let own = plan
+        .metrics()
+        .and_then(|metrics| metrics.spill_count())
+        .unwrap_or(0);
+    own + plan.children().into_iter().map(spill_count).sum::<usize>()
+}
+
 /// Run `sql`, optionally under a memory limit, returning the sorted rows.
+///
+/// The limit is there to make the join spill, so a run under it that does not
+/// spill would compare the in-memory path with itself.
 async fn run(sql: &str, memory_limit: Option<usize>) -> Vec<String> {
     let config = SessionConfig::new()
         .with_target_partitions(1)
@@ -73,7 +86,20 @@ async fn run(sql: &str, memory_limit: Option<usize>) -> Vec<String> {
     ctx.register_batch("l", table(200, "v", 11)).unwrap();
     ctx.register_batch("r", table(90, "w", 16)).unwrap();
 
-    let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+    let plan = ctx
+        .sql(sql)
+        .await
+        .unwrap()
+        .create_physical_plan()
+        .await
+        .unwrap();
+    let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
+    assert_eq!(
+        spill_count(&plan) > 0,
+        memory_limit.is_some(),
+        "unexpected spilling with memory limit {memory_limit:?}"
+    );
+
     let mut rows = vec![];
     for batch in &batches {
         for row in 0..batch.num_rows() {