kosiew commented on code in PR #22038:
URL: https://github.com/apache/datafusion/pull/22038#discussion_r3886243194
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -2252,29 +2636,49 @@ impl NestedLoopJoinStream {
return ControlFlow::Break(poll);
}
- if !self.left_exhausted && self.is_memory_limited() {
- // More left data to process — free current chunk and
- // go back to BufferingLeft for the next chunk
- if let SpillState::Active(ref active) =
self.spill_state {
- active.reservation.resize(0);
+ // Drop our reference to the current chunk's
+ // `JoinLeftData` before releasing the slot. Once the
+ // last partition does this, the `Arc` reaches zero
+ // refcount and the per-chunk reservation is freed.
+ self.buffered_left_data = None;
+
+ if self.is_memory_limited() {
+ let is_emitter = self.is_unmatched_left_emitter;
+ if let SpillState::Active(active) = &mut
self.spill_state {
+ // The last partition for this chunk (the
+ // unmatched-left emitter elected in `ProbeEnd`)
+ // releases the coordinator slot so the next
+ // leader can load the following chunk.
+ if is_emitter {
+ let coordinator =
Arc::clone(&active.coordinator);
+ let released_index = active.next_chunk_index;
+ active.chunk_release_in_flight = Some(
+ async move {
+
coordinator.release_chunk(released_index).await
+ }
+ .boxed(),
+ );
+ }
+ active.next_chunk_index += 1;
}
- self.buffered_left_data = None;
+ // `is_unmatched_left_emitter` is recomputed when
+ // `ProbeEnd` is re-entered for the next chunk, so it
+ // does not need to be reset here.
+ }
+
+ if !self.left_exhausted && self.is_memory_limited() {
+ // More left data to process — go back to
+ // BufferingLeft for the next chunk.
self.left_probe_idx = 0;
self.left_emit_idx = 0;
- // Each memory-limited chunk gets a fresh per-chunk
- // `JoinLeftData`/counter; `is_unmatched_left_emitter`
is
- // recomputed when `ProbeEnd` is re-entered for the
next
- // chunk, so it does not need to be reset here.
self.state = NLJState::BufferingLeft;
} else if self.is_memory_limited()
Review Comment:
I think the final chunk is never actually released here.
The release future is created above, but `chunk_release_in_flight` is only
polled when we return to `BufferingLeft` or re-enter
`handle_emit_left_unmatched`. That works for non-final chunks because they
transition back to `BufferingLeft`. For the final chunk, though, a LEFT join
goes directly to `Done`, while a FULL join goes to `EmitGlobalRightUnmatched`.
As a result, `FallbackCoordinatorInner.current` can keep its
`Arc<JoinLeftData>` along with the coordinator reservation until the physical
plan itself is dropped. If a completed query's plan stays alive, the final
chunk's batch and bitmap remain accounted against the memory pool and could
cause later queries to hit the memory limit.
Could we make sure the final chunk release is driven to completion before
transitioning to `Done` or `EmitGlobalRightUnmatched`? It would also be good to
add a regression that keeps the plan alive after collection and verifies that
the coordinator reservation has been released.
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -4528,4 +4903,527 @@ pub(crate) mod tests {
"));
Ok(())
}
+
+ // ========================================================================
+ // Multi-partition memory-limited correctness tests
+ //
+ // These tests reproduce the cross-partition coordination bug in the
+ // memory-limited fallback path: each output partition independently
+ // constructs a per-chunk `JoinLeftData` with `AtomicUsize::new(1)`,
+ // so left-side visited state is not shared across right partitions.
+ // For join types that emit unmatched left rows in the final output
+ // (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL), this leads to a
+ // left row being emitted as unmatched by partitions whose right
+ // input did not match it — even when another partition did match.
+ // ========================================================================
+
+ /// Build the right table as one batch per row, so RepartitionExec can
+ /// distribute rows across multiple output partitions.
+ fn build_right_table_one_batch_per_row() -> Arc<dyn ExecutionPlan> {
+ build_table(
+ ("a2", &vec![12, 2, 10]),
+ ("b2", &vec![10, 2, 10]),
+ ("c2", &vec![40, 80, 100]),
+ Some(1),
+ Vec::new(),
+ )
+ }
+
+ /// Run a NLJ across 4 right partitions under a tight memory limit, so
+ /// every output partition takes the memory-limited fallback path. The
+ /// right side is shuffled via `RepartitionExec(RoundRobinBatch(4))`.
+ async fn multi_partition_memory_limited_join_collect(
+ left: Arc<dyn ExecutionPlan>,
+ right: Arc<dyn ExecutionPlan>,
+ join_type: &JoinType,
+ join_filter: Option<JoinFilter>,
+ context: Arc<TaskContext>,
+ ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> {
+ let partition_count = 4;
+ let right = Arc::new(RepartitionExec::try_new(
+ right,
+ Partitioning::RoundRobinBatch(partition_count),
+ )?) as Arc<dyn ExecutionPlan>;
+
+ let nested_loop_join =
+ NestedLoopJoinExec::try_new(left, right, join_filter, join_type,
None)?;
+ let columns = columns(&nested_loop_join.schema());
+
+ let mut batches = vec![];
+ for i in 0..partition_count {
+ let stream = nested_loop_join.execute(i, Arc::clone(&context))?;
+ let more = common::collect(stream).await?;
+ batches.extend(more.into_iter().filter(|b| b.num_rows() > 0));
+ }
+
+ let metrics = nested_loop_join.metrics().unwrap();
+ Ok((columns, batches, metrics))
+ }
+
+ #[tokio::test]
+ async fn test_nlj_memory_limited_multi_partition_left_join() -> Result<()>
{
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table();
+ let right = build_right_table_one_batch_per_row();
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
multi_partition_memory_limited_join_collect(
+ left,
+ right,
+ &JoinType::Left,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // Expected output is identical to the single-partition spill path
+ // and the multi-partition non-spill path. Each left row appears
+ // exactly once: the matched (5,5,50)+(2,2,80) row, plus the two
+ // left rows filtered out by `b1 != 8` as unmatched.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+-----+----+----+----+
+ | a1 | b1 | c1 | a2 | b2 | c2 |
+ +----+----+-----+----+----+----+
+ | 11 | 8 | 110 | | | |
+ | 5 | 5 | 50 | 2 | 2 | 80 |
+ | 9 | 8 | 90 | | | |
+ +----+----+-----+----+----+----+
+ "));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_nlj_memory_limited_multi_partition_full_join() -> Result<()>
{
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table();
+ let right = build_right_table_one_batch_per_row();
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
multi_partition_memory_limited_join_collect(
+ left,
+ right,
+ &JoinType::Full,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // Expected: 1 matched + 2 left-unmatched + 2 right-unmatched.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+-----+----+----+-----+
+ | a1 | b1 | c1 | a2 | b2 | c2 |
+ +----+----+-----+----+----+-----+
+ | | | | 10 | 10 | 100 |
+ | | | | 12 | 10 | 40 |
+ | 11 | 8 | 110 | | | |
+ | 5 | 5 | 50 | 2 | 2 | 80 |
+ | 9 | 8 | 90 | | | |
+ +----+----+-----+----+----+-----+
+ "));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_nlj_memory_limited_multi_partition_left_semi_join() ->
Result<()> {
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table();
+ // Two right rows that both match the same left row (5,5,50). Without
+ // shared left-visited state across partitions, each matching partition
+ // emits the left row once, producing duplicates.
+ let right = build_table(
+ ("a2", &vec![2, 3, 10]),
+ ("b2", &vec![2, 2, 10]),
+ ("c2", &vec![80, 70, 100]),
+ Some(1),
+ Vec::new(),
+ );
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
multi_partition_memory_limited_join_collect(
+ left,
+ right,
+ &JoinType::LeftSemi,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // Left semi: each left row appears at most once, even if it matches
+ // multiple right rows distributed across partitions.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+----+
+ | a1 | b1 | c1 |
+ +----+----+----+
+ | 5 | 5 | 50 |
+ +----+----+----+
+ "));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_nlj_memory_limited_multi_partition_left_anti_join() ->
Result<()> {
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table();
+ let right = build_right_table_one_batch_per_row();
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
multi_partition_memory_limited_join_collect(
+ left,
+ right,
+ &JoinType::LeftAnti,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // Left anti: only left rows with no matching right row.
+ // (5,5,50) matches (2,2,80) under the filter, so it must NOT appear.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+-----+
+ | a1 | b1 | c1 |
+ +----+----+-----+
+ | 11 | 8 | 110 |
+ | 9 | 8 | 90 |
+ +----+----+-----+
+ "));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_nlj_memory_limited_multi_partition_left_mark_join() ->
Result<()> {
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table();
+ let right = build_right_table_one_batch_per_row();
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
multi_partition_memory_limited_join_collect(
+ left,
+ right,
+ &JoinType::LeftMark,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1", "mark"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // Left mark: every left row appears exactly once with a bool
+ // indicating whether it matched at least one right row.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+-----+-------+
+ | a1 | b1 | c1 | mark |
+ +----+----+-----+-------+
+ | 11 | 8 | 110 | false |
+ | 5 | 5 | 50 | true |
+ | 9 | 8 | 90 | false |
+ +----+----+-----+-------+
+ "));
+ Ok(())
+ }
+
+ /// A left table with one batch per row, spanning enough rows that the
+ /// memory-limited fallback must split it across MULTIPLE chunks. Only
+ /// `(5,5,50)` matches the right side under `prepare_join_filter`
+ /// (`b1 != 8`); every other row has `b1 = 8` and is therefore an
+ /// unmatched left row.
+ fn build_left_table_multi_chunk() -> Arc<dyn ExecutionPlan> {
+ build_table(
+ ("a1", &vec![5, 9, 11, 13, 15, 17, 19, 21]),
+ ("b1", &vec![5, 8, 8, 8, 8, 8, 8, 8]),
+ ("c1", &vec![50, 90, 110, 130, 150, 170, 190, 210]),
+ // One row per batch, so the tight per-chunk memory budget forces
+ // the coordinator to load each row as a separate chunk.
+ Some(1),
+ Vec::new(),
+ )
+ }
+
+ /// Run a NLJ across 4 right partitions, collecting every output
+ /// partition CONCURRENTLY. This is required for the multi-chunk
+ /// coordinator path: a chunk is not released until all partitions
+ /// finish probing it, and a partition cannot advance to the next chunk
+ /// until the current one is released. Collecting partitions
+ /// sequentially would therefore deadlock; concurrent collection mirrors
+ /// how partitions actually run under the runtime.
+ async fn multi_partition_memory_limited_join_collect_concurrent(
+ left: Arc<dyn ExecutionPlan>,
+ right: Arc<dyn ExecutionPlan>,
+ join_type: &JoinType,
+ join_filter: Option<JoinFilter>,
+ context: Arc<TaskContext>,
+ ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> {
+ let partition_count = 4;
+ let right = Arc::new(RepartitionExec::try_new(
+ right,
+ Partitioning::RoundRobinBatch(partition_count),
+ )?) as Arc<dyn ExecutionPlan>;
+
+ let nested_loop_join = Arc::new(NestedLoopJoinExec::try_new(
+ left,
+ right,
+ join_filter,
+ join_type,
+ None,
+ )?);
+ let columns = columns(&nested_loop_join.schema());
+
+ let mut handles = vec![];
+ for i in 0..partition_count {
+ let stream = nested_loop_join.execute(i, Arc::clone(&context))?;
+ handles.push(SpawnedTask::spawn(
+ async move { common::collect(stream).await },
+ ));
+ }
+
+ let mut batches = vec![];
+ for handle in handles {
+ let more = handle.join().await.expect("partition task panicked")?;
+ batches.extend(more.into_iter().filter(|b| b.num_rows() > 0));
+ }
+
+ let metrics = nested_loop_join.metrics().unwrap();
+ Ok((columns, batches, metrics))
+ }
+
+ /// Regression test for the multi-chunk coordinator path of a LEFT join.
+ ///
+ /// Unlike the other multi-partition tests, the left side here is split
+ /// into multiple chunks (one row per batch under a tight memory limit),
+ /// so this exercises `carryover` between chunks, `release_chunk`
+ /// advancing `next_chunk_index`, waiter notification, and re-entering
+ /// `BufferingLeft` for each subsequent chunk. Every left row must appear
+ /// exactly once: duplicates would indicate the per-chunk `JoinLeftData`
+ /// (visited bitmap + probe-thread counter) was not shared across right
+ /// partitions.
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_nlj_memory_limited_multi_partition_multi_chunk_left_join()
-> Result<()>
+ {
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table_multi_chunk();
+ let right = build_right_table_one_batch_per_row();
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
+ multi_partition_memory_limited_join_collect_concurrent(
+ left,
+ right,
+ &JoinType::Left,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // (5,5,50) matches (2,2,80); all other left rows (b1 = 8) are
+ // unmatched and appear exactly once.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+-----+----+----+----+
+ | a1 | b1 | c1 | a2 | b2 | c2 |
+ +----+----+-----+----+----+----+
+ | 11 | 8 | 110 | | | |
+ | 13 | 8 | 130 | | | |
+ | 15 | 8 | 150 | | | |
+ | 17 | 8 | 170 | | | |
+ | 19 | 8 | 190 | | | |
+ | 21 | 8 | 210 | | | |
+ | 5 | 5 | 50 | 2 | 2 | 80 |
+ | 9 | 8 | 90 | | | |
+ +----+----+-----+----+----+----+
+ "));
+ Ok(())
+ }
+
+ /// Regression test for the multi-chunk coordinator path of a FULL join.
+ ///
+ /// In addition to the per-chunk left sharing exercised by the LEFT case,
+ /// this covers the global right-unmatched bitmap accumulated ACROSS all
+ /// left chunks and emitted once in `EmitGlobalRightUnmatched`. The two
+ /// right rows with `b2 = 10` are filtered out of every match and must
+ /// appear exactly once as unmatched-right rows, regardless of how many
+ /// left chunks were processed.
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_nlj_memory_limited_multi_partition_multi_chunk_full_join()
-> Result<()>
+ {
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+ let left = build_left_table_multi_chunk();
+ let right = build_right_table_one_batch_per_row();
+ let filter = prepare_join_filter();
+
+ let (columns, batches, metrics) =
+ multi_partition_memory_limited_join_collect_concurrent(
+ left,
+ right,
+ &JoinType::Full,
+ Some(filter),
+ task_ctx,
+ )
+ .await?;
+
+ assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "Expected spilling under tight memory limit"
+ );
+
+ // Matched: (5,5,50)+(2,2,80). Unmatched left: the seven b1 = 8 rows.
+ // Unmatched right: (12,10,40) and (10,10,100), each emitted once from
+ // the global right bitmap accumulated across all left chunks.
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches),
@r"
+ +----+----+-----+----+----+-----+
+ | a1 | b1 | c1 | a2 | b2 | c2 |
+ +----+----+-----+----+----+-----+
+ | | | | 10 | 10 | 100 |
+ | | | | 12 | 10 | 40 |
+ | 11 | 8 | 110 | | | |
+ | 13 | 8 | 130 | | | |
+ | 15 | 8 | 150 | | | |
+ | 17 | 8 | 170 | | | |
+ | 19 | 8 | 190 | | | |
+ | 21 | 8 | 210 | | | |
+ | 5 | 5 | 50 | 2 | 2 | 80 |
+ | 9 | 8 | 90 | | | |
+ +----+----+-----+----+----+-----+
+ "));
+ Ok(())
+ }
+
+ /// Like `task_ctx_with_memory_limit`, but also disables the coordinated
+ /// memory-limited fallback via `enable_nlj_coordinated_fallback = false`
+ /// (the opt-out a distributed engine would use).
+ fn task_ctx_with_memory_limit_no_coordinated_fallback(
+ memory_limit: usize,
+ batch_size: usize,
+ ) -> Result<Arc<TaskContext>> {
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(memory_limit, 1.0)
+ .build_arc()?;
+ let mut cfg = TaskContext::default()
+ .session_config()
+ .clone()
+ .with_batch_size(batch_size);
+ cfg.options_mut().execution.enable_nlj_coordinated_fallback = false;
+ let task_ctx = TaskContext::default()
+ .with_runtime(runtime)
+ .with_session_config(cfg);
+ Ok(Arc::new(task_ctx))
+ }
+
+ /// Collect a multi-partition NLJ under a tight memory limit and return the
+ /// first error, if any. Used to assert that disabling the coordinated
+ /// fallback makes a left-emitting multi-partition join fail with resource
+ /// exhaustion (rather than spill, or — in a distributed setting — hang).
+ async fn multi_partition_join_collect_err(
+ left: Arc<dyn ExecutionPlan>,
+ right: Arc<dyn ExecutionPlan>,
+ join_type: &JoinType,
+ join_filter: Option<JoinFilter>,
+ context: Arc<TaskContext>,
+ ) -> Result<()> {
+ let partition_count = 4;
+ let right = Arc::new(RepartitionExec::try_new(
+ right,
+ Partitioning::RoundRobinBatch(partition_count),
+ )?) as Arc<dyn ExecutionPlan>;
+ let nested_loop_join = Arc::new(NestedLoopJoinExec::try_new(
+ left,
+ right,
+ join_filter,
+ join_type,
+ None,
+ )?);
+
+ let mut handles = vec![];
+ for i in 0..partition_count {
+ let stream = nested_loop_join.execute(i, Arc::clone(&context))?;
+ handles.push(SpawnedTask::spawn(
+ async move { common::collect(stream).await },
+ ));
+ }
+ for handle in handles {
+ handle.join().await.expect("partition task panicked")?;
+ }
+ Ok(())
+ }
+
+ /// When `enable_nlj_coordinated_fallback` is disabled, a LEFT join with
Review Comment:
Could we also add an opt-out regression for one non-left-emitting join, such
as `RIGHT` or `INNER`, under the same multi-partition memory pressure?
The documented contract says `enable_nlj_coordinated_fallback = false` only
disables the coordinated fallback for left-emitting multi-partition joins. A
small regression here would help ensure that a future change to this guard does
not accidentally disable spill fallback for unaffected join types.
--
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]
