Skip to content

Commit 8f2fd93

Browse files
viiryaviirya
committed
feat: Add config opt-out for coordinated NLJ memory-limited fallback
The memory-limited NestedLoopJoin fallback coordinates per-chunk left state (visited bitmap + probe-thread counter) across all right-side partitions via `FallbackCoordinator`. That coordination assumes every right partition runs in the same process. Distributed engines (Ballista, datafusion-distributed) execute each output partition as an independent task that builds its own coordinator and polls a single partition, so the shared probe-thread counter never reaches zero, no partition is elected to release the current chunk, and the waiters block forever — the memory-limited fallback deadlocks. Add `datafusion.execution.enable_nlj_coordinated_fallback` (default true, preserving current behavior). When set to false, the coordinated fallback is disabled for the join types that would deadlock — left-emitting joins (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) with a multi-partition right side — which then use `SpillState::Disabled` and fail with resource exhaustion under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are always safe and keep the coordinated fallback regardless of the flag. Raised by @andygrove on apache#22038. Co-authored-by: Claude Code
1 parent b439ffa commit 8f2fd93

3 files changed

Lines changed: 148 additions & 1 deletion

File tree

‎datafusion/common/src/config.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -860,6 +860,25 @@ config_namespace! {
860860
/// Default: 128 MB
861861
pub max_spill_file_size_bytes: ConfigNonZeroUsize, default = non_zero_usize_default(128 * 1024 * 1024)
862862

863+
/// Enables the memory-limited fallback for `NestedLoopJoinExec` join
864+
/// types that emit unmatched left rows in the final output (LEFT, LEFT
865+
/// SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple
866+
/// partitions.
867+
///
868+
/// This fallback coordinates per-chunk left state (visited bitmap and
869+
/// probe-thread counter) across all right-side partitions, which
870+
/// assumes every partition runs in the same process. Distributed
871+
/// engines that execute each output partition as an independent task
872+
/// (e.g. Ballista, datafusion-distributed) build a separate coordinator
873+
/// per task and poll only one partition, so the cross-partition
874+
/// counter never reaches zero and the fallback would stall. Such
875+
/// engines should set this to `false`: the coordinated fallback is then
876+
/// disabled for left-emitting multi-partition joins, which instead fail
877+
/// with a resource-exhaustion error under memory pressure rather than
878+
/// deadlocking. Single-partition and non-left-emitting joins are
879+
/// unaffected and always keep the fallback.
880+
pub enable_nlj_coordinated_fallback: bool, default = true
881+
863882
/// Number of files to read in parallel when inferring schema and statistics
864883
pub meta_fetch_concurrency: ConfigNonZeroUsize, default = non_zero_usize_default(32)
865884

‎datafusion/physical-plan/src/joins/nested_loop_join.rs‎

Lines changed: 128 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -677,7 +677,27 @@ impl ExecutionPlan for NestedLoopJoinExec {
677677
// counter) across all right-side partitions via
678678
// [`FallbackCoordinator`], so left-side tracking is coordinated
679679
// exactly as in the single-pass path.
680-
let spill_state = if context.runtime_env().disk_manager.tmp_files_enabled() {
680+
//
681+
// That coordination assumes all right partitions run in the same
682+
// process. Distributed engines run each partition as an independent
683+
// task with its own coordinator, so the shared probe-thread counter
684+
// would never reach zero and the fallback would stall. When
685+
// `enable_nlj_coordinated_fallback` is disabled, such engines opt
686+
// out of the coordinated fallback for the affected join types
687+
// (left-emitting joins with a multi-partition right side); those cases
688+
// use `SpillState::Disabled` and fail with resource exhaustion under
689+
// memory pressure instead of deadlocking. Single-partition and
690+
// non-left-emitting joins are always safe and keep the fallback.
691+
let coordinated_fallback_disabled = !context
692+
.session_config()
693+
.options()
694+
.execution
695+
.enable_nlj_coordinated_fallback
696+
&& need_produce_result_in_final(self.join_type)
697+
&& right_partition_count > 1;
698+
let spill_state = if context.runtime_env().disk_manager.tmp_files_enabled()
699+
&& !coordinated_fallback_disabled
700+
{
681701
SpillState::Pending {
682702
left_plan: Arc::clone(&self.left),
683703
task_context: Arc::clone(&context),
@@ -4788,4 +4808,111 @@ pub(crate) mod tests {
47884808
"));
47894809
Ok(())
47904810
}
4811+
4812+
/// Like `task_ctx_with_memory_limit`, but also disables the coordinated
4813+
/// memory-limited fallback via `enable_nlj_coordinated_fallback = false`
4814+
/// (the opt-out a distributed engine would use).
4815+
fn task_ctx_with_memory_limit_no_coordinated_fallback(
4816+
memory_limit: usize,
4817+
batch_size: usize,
4818+
) -> Result<Arc<TaskContext>> {
4819+
let runtime = RuntimeEnvBuilder::new()
4820+
.with_memory_limit(memory_limit, 1.0)
4821+
.build_arc()?;
4822+
let mut cfg = TaskContext::default()
4823+
.session_config()
4824+
.clone()
4825+
.with_batch_size(batch_size);
4826+
cfg.options_mut().execution.enable_nlj_coordinated_fallback = false;
4827+
let task_ctx = TaskContext::default()
4828+
.with_runtime(runtime)
4829+
.with_session_config(cfg);
4830+
Ok(Arc::new(task_ctx))
4831+
}
4832+
4833+
/// Collect a multi-partition NLJ under a tight memory limit and return the
4834+
/// first error, if any. Used to assert that disabling the coordinated
4835+
/// fallback makes a left-emitting multi-partition join fail with resource
4836+
/// exhaustion (rather than spill, or — in a distributed setting — hang).
4837+
async fn multi_partition_join_collect_err(
4838+
left: Arc<dyn ExecutionPlan>,
4839+
right: Arc<dyn ExecutionPlan>,
4840+
join_type: &JoinType,
4841+
join_filter: Option<JoinFilter>,
4842+
context: Arc<TaskContext>,
4843+
) -> Result<()> {
4844+
let partition_count = 4;
4845+
let right = Arc::new(RepartitionExec::try_new(
4846+
right,
4847+
Partitioning::RoundRobinBatch(partition_count),
4848+
)?) as Arc<dyn ExecutionPlan>;
4849+
let nested_loop_join = Arc::new(NestedLoopJoinExec::try_new(
4850+
left,
4851+
right,
4852+
join_filter,
4853+
join_type,
4854+
None,
4855+
)?);
4856+
4857+
let mut handles = vec![];
4858+
for i in 0..partition_count {
4859+
let stream = nested_loop_join.execute(i, Arc::clone(&context))?;
4860+
handles.push(SpawnedTask::spawn(
4861+
async move { common::collect(stream).await },
4862+
));
4863+
}
4864+
for handle in handles {
4865+
handle.join().await.expect("partition task panicked")?;
4866+
}
4867+
Ok(())
4868+
}
4869+
4870+
/// When `enable_nlj_coordinated_fallback` is disabled, a LEFT join with
4871+
/// a multi-partition right side must NOT take the coordinated fallback: it
4872+
/// fails with resource exhaustion under a tight memory limit instead. This
4873+
/// is the distributed-safe opt-out (the coordinated fallback would
4874+
/// otherwise deadlock across processes).
4875+
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
4876+
async fn test_nlj_memory_limited_fallback_disabled_left_join_oom() -> Result<()> {
4877+
let task_ctx = task_ctx_with_memory_limit_no_coordinated_fallback(50, 16)?;
4878+
let left = build_left_table_multi_chunk();
4879+
let right = build_right_table_one_batch_per_row();
4880+
let filter = prepare_join_filter();
4881+
4882+
let err = multi_partition_join_collect_err(
4883+
left,
4884+
right,
4885+
&JoinType::Left,
4886+
Some(filter),
4887+
task_ctx,
4888+
)
4889+
.await
4890+
.unwrap_err();
4891+
assert_contains!(err.to_string(), "Resources exhausted");
4892+
Ok(())
4893+
}
4894+
4895+
/// FULL join counterpart of the above: the opt-out disables the coordinated
4896+
/// fallback for FULL (also a left-emitting join) with a multi-partition
4897+
/// right side, so it fails with resource exhaustion rather than
4898+
/// coordinating.
4899+
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
4900+
async fn test_nlj_memory_limited_fallback_disabled_full_join_oom() -> Result<()> {
4901+
let task_ctx = task_ctx_with_memory_limit_no_coordinated_fallback(50, 16)?;
4902+
let left = build_left_table_multi_chunk();
4903+
let right = build_right_table_one_batch_per_row();
4904+
let filter = prepare_join_filter();
4905+
4906+
let err = multi_partition_join_collect_err(
4907+
left,
4908+
right,
4909+
&JoinType::Full,
4910+
Some(filter),
4911+
task_ctx,
4912+
)
4913+
.await
4914+
.unwrap_err();
4915+
assert_contains!(err.to_string(), "Resources exhausted");
4916+
Ok(())
4917+
}
47914918
}

‎docs/source/user-guide/configs.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ The following configuration settings are available:
126126
| datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. |
127127
| datafusion.execution.sort_pushdown_buffer_capacity | 1073741824 | Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. |
128128
| datafusion.execution.max_spill_file_size_bytes | 134217728 | Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB |
129+
| datafusion.execution.enable_nlj_coordinated_fallback | true | Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback coordinates per-chunk left state (visited bitmap and probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) build a separate coordinator per task and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. |
129130
| datafusion.execution.meta_fetch_concurrency | 32 | Number of files to read in parallel when inferring schema and statistics |
130131
| datafusion.execution.minimum_parallel_output_files | 4 | Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file is reached. |
131132
| datafusion.execution.soft_max_rows_per_output_file | 50000000 | Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max |

0 commit comments

Comments
 (0)