viirya commented on code in PR #22038:
URL: https://github.com/apache/datafusion/pull/22038#discussion_r3887045555
##########
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:
Agreed — added
`test_nlj_memory_limited_fallback_disabled_right_join_still_spills`, covering
both `RIGHT` and `INNER` under the same multi-partition memory pressure with
`enable_nlj_coordinated_fallback = false`, asserting `spill_count > 0`. That
pins the guard's scope so a future change can't quietly disable the spill
fallback for join types that never needed left-bitmap coordination.
One thing worth recording from writing it: my first version drained the
partitions sequentially (reusing `multi_partition_memory_limited_join_collect`)
and **hung**. The coordinator seeds every chunk's probe counter with
`right_partition_count` regardless of join type — `with_visited_bitmap` only
controls whether a bitmap is allocated, not the counter — so with a multi-chunk
left side, partition 0 waits on the other three to report before chunk 0 can be
released, and collecting them one after another deadlocks. That's why all the
existing multi-chunk coordinated tests use the concurrent helper; the new test
does too, and I noted the reason in a comment so nobody reuses the sequential
helper here.
It's a test-harness constraint rather than a product issue (a real plan
always polls its partitions concurrently), but it does mean "RIGHT joins don't
need coordination" is true of the *bitmap* and not of the chunk lockstep. Happy
to make the counter join-type-aware in a follow-up if you think that lockstep
is worth removing for non-left-emitting joins — it would let those partitions
advance independently, but it's a behavioural change beyond this PR's scope.
--
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]
