From 4084ad17c113477cdc02843bd166f2cc0e204da6 Mon Sep 17 00:00:00 2001
From: Ran Reichman <ran@flarion.io>
Date: Tue, 25 Aug 2026 16:54:10 -0400
Subject: [PATCH] fix: refuse memory-limited NestedLoopJoin fallback for
 left-emission joins with a multi-partition probe side

---
 .../src/joins/nested_loop_join.rs             | 83 ++++++++++---------
 1 file changed, 46 insertions(+), 37 deletions(-)

diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index 548a0cf1db9b..e215c7dff106 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -681,22 +681,23 @@ impl ExecutionPlan for NestedLoopJoinExec {
         // Determine if OOM fallback to memory-limited mode is possible.
         // Conditions:
         // 1. Disk manager supports temp files (needed for spilling).
-        // 2. FULL join with multiple right partitions is not yet supported
-        //    in the fallback path. FULL join needs to track BOTH left-side
-        //    matches (for unmatched left rows) AND right-side matches (for
-        //    unmatched right rows). The fallback path builds a per-partition
-        //    `JoinLeftData` with `probe_threads_counter == 1`, so each
-        //    partition emits unmatched left rows based only on its own
-        //    right-side matches, producing incorrect duplicate output for
-        //    left rows that match in another partition. Other join types
-        //    that need only one-sided final emission (LEFT, LEFT SEMI,
-        //    LEFT ANTI, LEFT MARK) have a similar latent issue in the
-        //    fallback path which predates this change; tracking is out of
-        //    scope for this PR.
-        let full_join_multi_partition =
-            matches!(self.join_type, JoinType::Full) && right_partition_count > 1;
+        // 2. Join types whose final emission reads the visited-left bitmap
+        //    (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) need that bitmap
+        //    complete across every probe partition. The fallback path builds a
+        //    per-partition `JoinLeftData` with `probe_threads_counter == 1`, so
+        //    with more than one right partition each partition emits from a
+        //    bitmap that only saw its own right rows: unmatched rows are
+        //    emitted once per partition, and rows matched only in another
+        //    partition are emitted as unmatched. Refusing the fallback turns
+        //    those wrong results into a plain ResourcesExhausted error.
+        //    Right-side emission types are safe because each partition owns
+        //    its right rows exclusively. Cross-partition coordination of the
+        //    left bitmap is tracked in
+        //    https://github.com/apache/datafusion/issues/22038.
+        let left_emission_multi_partition =
+            need_produce_result_in_final(self.join_type) && right_partition_count > 1;
         let spill_state = if context.runtime_env().disk_manager.tmp_files_enabled()
-            && !full_join_multi_partition
+            && !left_emission_multi_partition
         {
             SpillState::Pending {
                 left_plan: Arc::clone(&self.left),
@@ -3842,12 +3843,10 @@ pub(crate) mod tests {
 
         // Join types that support memory-limited fallback should succeed
         // even under tight memory limits (they spill to disk instead of OOM).
+        // Right-side emission types are safe with multiple right partitions
+        // because each partition owns its right rows exclusively.
         let fallback_join_types = vec![
             JoinType::Inner,
-            JoinType::Left,
-            JoinType::LeftSemi,
-            JoinType::LeftAnti,
-            JoinType::LeftMark,
             JoinType::Right,
             JoinType::RightSemi,
             JoinType::RightAnti,
@@ -3872,24 +3871,34 @@ pub(crate) mod tests {
             .await?;
         }
 
-        // FULL JOIN with multiple right partitions is intentionally not
-        // supported in the fallback path yet (cross-partition left-bitmap
-        // coordination is missing). It should still OOM under tight memory.
-        let runtime = RuntimeEnvBuilder::new()
-            .with_memory_limit(100, 1.0)
-            .build_arc()?;
-        let task_ctx = TaskContext::default().with_runtime(runtime);
-        let task_ctx = Arc::new(task_ctx);
-        let err = multi_partitioned_join_collect(
-            Arc::clone(&left),
-            Arc::clone(&right),
-            &JoinType::Full,
-            Some(filter.clone()),
-            task_ctx,
-        )
-        .await
-        .unwrap_err();
-        assert_contains!(err.to_string(), "Resources exhausted");
+        // Join types whose final emission reads the visited-left bitmap are
+        // intentionally not supported in the fallback path with multiple right
+        // partitions (cross-partition left-bitmap coordination is missing;
+        // per-partition bitmaps emit wrong rows). They must OOM cleanly instead.
+        let gated_join_types = vec![
+            JoinType::Left,
+            JoinType::LeftSemi,
+            JoinType::LeftAnti,
+            JoinType::LeftMark,
+            JoinType::Full,
+        ];
+        for join_type in &gated_join_types {
+            let runtime = RuntimeEnvBuilder::new()
+                .with_memory_limit(100, 1.0)
+                .build_arc()?;
+            let task_ctx = TaskContext::default().with_runtime(runtime);
+            let task_ctx = Arc::new(task_ctx);
+            let err = multi_partitioned_join_collect(
+                Arc::clone(&left),
+                Arc::clone(&right),
+                join_type,
+                Some(filter.clone()),
+                task_ctx,
+            )
+            .await
+            .unwrap_err();
+            assert_contains!(err.to_string(), "Resources exhausted");
+        }
 
         Ok(())
     }