You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/19 06:36:08 UTC
(datafusion) branch gh-readonly-queue/main/pr-21699-1fbbba5c63183d377b42bf534289718c32493ea4 created (now 6aa5a7ee98)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-21699-1fbbba5c63183d377b42bf534289718c32493ea4
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 6aa5a7ee98 refactor: Share left-side spill file across partitions on OOM fallback (#21699)
This branch includes the following new commits:
new 6aa5a7ee98 refactor: Share left-side spill file across partitions on OOM fallback (#21699)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: refactor: Share left-side spill file across partitions on OOM fallback (#21699)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-21699-1fbbba5c63183d377b42bf534289718c32493ea4
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 6aa5a7ee98cf3de5a5ddd406d99d5859ebeb8d1d
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Sun Apr 19 14:35:45 2026 +0800
refactor: Share left-side spill file across partitions on OOM fallback (#21699)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes #.
## Rationale for this change
To reduce the redundant re-execution of the left side during OOM
fallback.
## What changes are included in this PR?
Previously when OnceFut failed with OOM, each partition independently
re-executed the left child to get its own stream. This was redundant
since all partitions need the same left data.
Now the first partition to initiate fallback spills the entire left side
to disk via a shared OnceAsync<LeftSpillData>. Other partitions wait on
the same future and read from the shared spill file, avoiding redundant
re-execution of the left child.
Co-authored-by: Claude Code
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
Unit test
## Are there any user-facing changes?
No
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
.../physical-plan/src/joins/nested_loop_join.rs | 160 +++++++++++++++++----
.../test_files/nested_loop_join_spill.slt | 7 +-
2 files changed, 135 insertions(+), 32 deletions(-)
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index 93d7e72587..030fa5bbf9 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -207,6 +207,13 @@ pub struct NestedLoopJoinExec {
/// Each output stream waits on the `OnceAsync` to signal the completion of
/// the build(left) side data, and buffer them all for later joining.
build_side_data: OnceAsync<JoinLeftData>,
+ /// Shared left-side spill data for OOM fallback.
+ ///
+ /// When `build_side_data` fails with OOM, the first partition to
+ /// initiate fallback spills the entire left side to disk. Other
+ /// partitions share the same spill file via this `OnceAsync`,
+ /// avoiding redundant re-execution of the left child.
+ left_spill_data: Arc<OnceAsync<LeftSpillData>>,
/// Information of index and left / right placement of columns
column_indices: Vec<ColumnIndex>,
/// Projection to apply to the output of the join
@@ -290,6 +297,7 @@ impl NestedLoopJoinExecBuilder {
join_type,
join_schema,
build_side_data: Default::default(),
+ left_spill_data: Arc::new(OnceAsync::default()),
column_indices,
projection,
metrics: Default::default(),
@@ -492,6 +500,7 @@ impl NestedLoopJoinExec {
right,
metrics: ExecutionPlanMetricsSet::new(),
build_side_data: Default::default(),
+ left_spill_data: Arc::new(OnceAsync::default()),
cache: Arc::clone(&self.cache),
filter: self.filter.clone(),
join_type: self.join_type,
@@ -655,6 +664,7 @@ impl ExecutionPlan for NestedLoopJoinExec {
SpillState::Pending {
left_plan: Arc::clone(&self.left),
task_context: Arc::clone(&context),
+ left_spill_data: Arc::clone(&self.left_spill_data),
}
} else {
SpillState::Disabled
@@ -863,6 +873,20 @@ enum NLJState {
EmitLeftUnmatched,
Done,
}
+/// Shared data for the left-side spill fallback.
+///
+/// When the in-memory `OnceFut` path fails with OOM, the first partition
+/// spills the entire left side to disk. This struct holds the spill file
+/// reference so other partitions can read from the same file.
+pub(crate) struct LeftSpillData {
+ /// SpillManager used to read the spill file (has the left schema)
+ spill_manager: SpillManager,
+ /// The spill file containing all left-side batches
+ spill_file: RefCountedTempFile,
+ /// Left-side schema
+ schema: SchemaRef,
+}
+
/// Tracks the state of the memory-limited spill fallback for NLJ.
///
/// The NLJ always starts with the standard OnceFut path. If the in-memory
@@ -882,6 +906,9 @@ pub(crate) enum SpillState {
left_plan: Arc<dyn ExecutionPlan>,
/// TaskContext for re-execution and SpillManager creation
task_context: Arc<TaskContext>,
+ /// Shared OnceAsync for left-side spill data. The first partition
+ /// to initiate fallback spills the left side; others share the file.
+ left_spill_data: Arc<OnceAsync<LeftSpillData>>,
},
/// Fallback has been triggered. Left data is being loaded in chunks
@@ -892,16 +919,20 @@ pub(crate) enum SpillState {
/// State for active memory-limited spill execution.
/// Boxed inside [`SpillState::Active`] to reduce enum size.
pub(crate) struct SpillStateActive {
- /// Left input stream for incremental buffering
- left_stream: SendableRecordBatchStream,
+ /// Shared future for left-side spill data. All partitions wait on
+ /// the same future — the first to poll triggers the actual spill.
+ left_spill_fut: OnceFut<LeftSpillData>,
+ /// Left input stream for incremental chunk reading (from spill file).
+ /// None until `left_spill_fut` resolves.
+ left_stream: Option<SendableRecordBatchStream>,
+ /// Left-side schema (set once `left_spill_fut` resolves)
+ left_schema: Option<SchemaRef>,
/// Memory reservation for left-side buffering
reservation: MemoryReservation,
/// Accumulated left batches for the current chunk
pending_batches: Vec<RecordBatch>,
- /// Left-side schema (for concat_batches)
- left_schema: SchemaRef,
/// SpillManager for right-side spilling
- spill_manager: SpillManager,
+ right_spill_manager: SpillManager,
/// In-progress spill file for writing right batches during first pass
right_spill_in_progress: Option<InProgressSpillFile>,
/// Completed right-side spill file (available after first pass)
@@ -1263,19 +1294,18 @@ impl NestedLoopJoinStream {
/// Switch from the standard OnceFut path to memory-limited mode.
///
- /// Re-executes the left child to get a fresh stream, creates a
- /// SpillManager for right-side spilling, and transitions the spill
- /// state from `Pending` to `Active`. The next call to
- /// `handle_buffering_left` will dispatch to
- /// `handle_buffering_left_memory_limited`.
+ /// Uses the shared `left_spill_data` OnceAsync so that only the first
+ /// partition to reach this point re-executes the left child and spills
+ /// it to disk. Other partitions share the same spill file.
fn initiate_fallback(&mut self) -> Result<()> {
// Take ownership of Pending state
- let (left_plan, context) =
+ let (left_plan, context, left_spill_data) =
match std::mem::replace(&mut self.spill_state, SpillState::Disabled) {
SpillState::Pending {
left_plan,
task_context,
- } => (left_plan, task_context),
+ left_spill_data,
+ } => (left_plan, task_context, left_spill_data),
_ => {
return internal_err!(
"initiate_fallback called in non-Pending spill state"
@@ -1283,9 +1313,42 @@ impl NestedLoopJoinStream {
}
};
- // Re-execute left child to get a fresh stream
- let left_stream = left_plan.execute(0, Arc::clone(&context))?;
- let left_schema = left_stream.schema();
+ // Use OnceAsync to ensure only the first partition spills the left
+ // side. Other partitions will get the same OnceFut that resolves
+ // to the shared spill file.
+ let left_spill_fut = left_spill_data.try_once(|| {
+ let plan = Arc::clone(&left_plan);
+ let ctx = Arc::clone(&context);
+ let spill_metrics = self.metrics.spill_metrics.clone();
+ Ok(async move {
+ let mut stream = plan.execute(0, Arc::clone(&ctx))?;
+ let schema = stream.schema();
+ let left_spill_manager = SpillManager::new(
+ ctx.runtime_env(),
+ spill_metrics,
+ Arc::clone(&schema),
+ )
+ .with_compression_type(ctx.session_config().spill_compression());
+
+ let result = left_spill_manager
+ .spill_record_batch_stream_and_return_max_batch_memory(
+ &mut stream,
+ "NestedLoopJoin left spill",
+ )
+ .await?;
+
+ match result {
+ Some((file, _max_batch_memory)) => Ok(LeftSpillData {
+ spill_manager: left_spill_manager,
+ spill_file: file,
+ schema,
+ }),
+ None => {
+ internal_err!("Left side produced no data to spill")
+ }
+ }
+ })
+ })?;
// Create reservation with can_spill for fair memory allocation
let reservation = MemoryConsumer::new("NestedLoopJoinLoad[fallback]".to_string())
@@ -1294,7 +1357,7 @@ impl NestedLoopJoinStream {
// Create SpillManager for right-side spilling
let right_schema = self.right_data.schema();
- let spill_manager = SpillManager::new(
+ let right_spill_manager = SpillManager::new(
context.runtime_env(),
self.metrics.spill_metrics.clone(),
right_schema,
@@ -1302,11 +1365,12 @@ impl NestedLoopJoinStream {
.with_compression_type(context.session_config().spill_compression());
self.spill_state = SpillState::Active(Box::new(SpillStateActive {
- left_stream,
+ left_spill_fut,
+ left_stream: None,
+ left_schema: None,
reservation,
pending_batches: Vec::new(),
- left_schema,
- spill_manager,
+ right_spill_manager,
right_spill_in_progress: None,
right_spill_file: None,
right_max_batch_memory: 0,
@@ -1378,11 +1442,44 @@ impl NestedLoopJoinStream {
);
};
+ // On first entry (or after re-entry for a new chunk pass when
+ // left_stream was consumed), wait for the shared left spill
+ // future to resolve and then open a stream from the spill file.
+ if active.left_stream.is_none() {
+ match active.left_spill_fut.get_shared(cx) {
+ Poll::Ready(Ok(spill_data)) => {
+ match spill_data
+ .spill_manager
+ .read_spill_as_stream(spill_data.spill_file.clone(), None)
+ {
+ Ok(stream) => {
+ active.left_schema = Some(Arc::clone(&spill_data.schema));
+ active.left_stream = Some(stream);
+ }
+ Err(e) => {
+ return ControlFlow::Break(Poll::Ready(Some(Err(e))));
+ }
+ }
+ }
+ Poll::Ready(Err(e)) => {
+ return ControlFlow::Break(Poll::Ready(Some(Err(e))));
+ }
+ Poll::Pending => {
+ return ControlFlow::Break(Poll::Pending);
+ }
+ }
+ }
+
+ let left_stream = active
+ .left_stream
+ .as_mut()
+ .expect("left_stream must be set after spill future resolves");
+
// Poll left stream for more batches.
// Note: pending_batches may already contain a batch from the
// previous chunk iteration (the batch that triggered the memory limit).
loop {
- match active.left_stream.poll_next_unpin(cx) {
+ match left_stream.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(batch))) => {
if batch.num_rows() == 0 {
continue;
@@ -1431,13 +1528,18 @@ impl NestedLoopJoinStream {
return ControlFlow::Continue(());
}
- let merged_batch =
- match concat_batches(&active.left_schema, &active.pending_batches) {
- Ok(batch) => batch,
- Err(e) => {
- return ControlFlow::Break(Poll::Ready(Some(Err(e.into()))));
- }
- };
+ let merged_batch = match concat_batches(
+ active
+ .left_schema
+ .as_ref()
+ .expect("left_schema must be set"),
+ &active.pending_batches,
+ ) {
+ Ok(batch) => batch,
+ Err(e) => {
+ return ControlFlow::Break(Poll::Ready(Some(Err(e.into()))));
+ }
+ };
active.pending_batches.clear();
// Build visited bitmap if needed for this join type
@@ -1472,7 +1574,7 @@ impl NestedLoopJoinStream {
// Set up right-side stream for this pass
if !active.is_first_right_pass {
if let Some(file) = active.right_spill_file.as_ref() {
- match active.spill_manager.read_spill_as_stream(
+ match active.right_spill_manager.read_spill_as_stream(
file.clone(),
Some(active.right_max_batch_memory),
) {
@@ -1487,7 +1589,7 @@ impl NestedLoopJoinStream {
} else {
// First pass: create InProgressSpillFile for right side
match active
- .spill_manager
+ .right_spill_manager
.create_in_progress_file("NestedLoopJoin right spill")
{
Ok(file) => {
diff --git a/datafusion/sqllogictest/test_files/nested_loop_join_spill.slt b/datafusion/sqllogictest/test_files/nested_loop_join_spill.slt
index 5b383f3edf..7b5da1d4b8 100644
--- a/datafusion/sqllogictest/test_files/nested_loop_join_spill.slt
+++ b/datafusion/sqllogictest/test_files/nested_loop_join_spill.slt
@@ -39,8 +39,9 @@ INNER JOIN generate_series(1, 1) AS t2(v2)
100000 1 100000
# --- Verify spill metrics via EXPLAIN ANALYZE ---
-# The NestedLoopJoinExec line should show spill_count=1, confirming
-# the memory-limited fallback path was taken and right side was spilled.
+# The NestedLoopJoinExec line should show spill_count=2, confirming
+# the memory-limited fallback path was taken (left side spilled once,
+# right side spilled once).
query TT
EXPLAIN ANALYZE SELECT count(*)
FROM generate_series(1, 100000) AS t1(v1)
@@ -50,7 +51,7 @@ INNER JOIN generate_series(1, 1) AS t2(v2)
Plan with Metrics
01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)], metrics=[<slt:ignore>]
02)--AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1))], metrics=[<slt:ignore>]
-03)----NestedLoopJoinExec: join_type=Inner, filter=v1@0 + v2@1 > 0, projection=[], metrics=[output_rows=100.0 K, <slt:ignore> spill_count=1, <slt:ignore>]
+03)----NestedLoopJoinExec: join_type=Inner, filter=v1@0 + v2@1 > 0, projection=[], metrics=[output_rows=100.0 K, <slt:ignore> spill_count=2, <slt:ignore>]
04)------ProjectionExec: expr=[value@0 as v1], metrics=[<slt:ignore>]
05)--------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192], metrics=[<slt:ignore>]
06)------ProjectionExec: expr=[value@0 as v2], metrics=[<slt:ignore>]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org