From 0e2df274459f54d0f397d52288d721c992e85720 Mon Sep 17 00:00:00 2001 From: Ran Reichman Date: Mon, 31 Aug 2026 07:20:17 -0400 Subject: [PATCH 1/4] buffer the NestedLoopJoin build side as coalesced chunks instead of one concat_batches allocation --- .../src/joins/nested_loop_join.rs | 340 ++++++++++++++---- 1 file changed, 265 insertions(+), 75 deletions(-) diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index bb91735369b9d..a9f7f5a8c59a6 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -54,9 +54,7 @@ use arrow::array::{ UInt64Array, new_null_array, }; use arrow::buffer::BooleanBuffer; -use arrow::compute::{ - BatchCoalescer, concat_batches, filter, filter_record_batch, not, take, -}; +use arrow::compute::{BatchCoalescer, filter, filter_record_batch, not, take}; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_schema::DataType; @@ -697,6 +695,7 @@ impl ExecutionPlan for NestedLoopJoinExec { need_produce_result_in_final(self.join_type), right_partition_count, left_spill_manager, + batch_size, )) })?; @@ -1014,8 +1013,16 @@ impl EmbeddedProjection for NestedLoopJoinExec { /// Left (build-side) data pub(crate) struct JoinLeftData { - /// Build-side data collected to single batch - batch: RecordBatch, + /// Build-side data as bounded chunks, in input order. Kept as chunks rather than one + /// `concat_batches` result so buffering never needs input and output to coexist, and a + /// chunk that already arrived at target size is retained without being copied at all. + chunks: Vec, + /// Row index of the first row of each chunk, i.e. prefix sums over the chunk lengths. + /// The visited-left bitmap is indexed by these global row numbers. + row_offsets: Vec, + total_rows: usize, + /// Build-side schema, kept so an empty chunk list still knows its shape + schema: SchemaRef, /// Shared bitmap builder for visited left indices bitmap: SharedBitmapBuilder, /// Counter of running probe-threads, potentially able to update `bitmap` @@ -1029,21 +1036,57 @@ pub(crate) struct JoinLeftData { impl JoinLeftData { pub(crate) fn new( - batch: RecordBatch, + chunks: Vec, + schema: SchemaRef, bitmap: SharedBitmapBuilder, probe_threads_counter: AtomicUsize, reservation: MemoryReservation, ) -> Self { + // A zero-row chunk would stall the probe and emit cursors (its range is empty, so the + // global row index never advances past it), so drop them here. + let chunks: Vec = + chunks.into_iter().filter(|c| c.num_rows() > 0).collect(); + let mut row_offsets = Vec::with_capacity(chunks.len()); + let mut total_rows = 0; + for chunk in &chunks { + row_offsets.push(total_rows); + total_rows += chunk.num_rows(); + } Self { - batch, + chunks, + row_offsets, + total_rows, + schema, bitmap, probe_threads_counter, reservation, } } - pub(crate) fn batch(&self) -> &RecordBatch { - &self.batch + pub(crate) fn chunk(&self, idx: usize) -> &RecordBatch { + &self.chunks[idx] + } + + pub(crate) fn row_offset(&self, idx: usize) -> usize { + self.row_offsets[idx] + } + + pub(crate) fn total_rows(&self) -> usize { + self.total_rows + } + + pub(crate) fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + /// Index of the chunk holding the given global row, and that row's offset inside it. + pub(crate) fn locate(&self, row: usize) -> Option<(usize, usize)> { + let idx = match self.row_offsets.binary_search(&row) { + Ok(idx) => idx, + Err(0) => return None, + Err(next) => next - 1, + }; + (row < self.total_rows).then(|| (idx, row - self.row_offsets[idx])) } pub(crate) fn bitmap(&self) -> &SharedBitmapBuilder { @@ -1071,10 +1114,14 @@ async fn collect_left_input( with_visited_left_side: bool, probe_threads_count: usize, spill_manager: Option, + target_batch_size: usize, ) -> Result { let schema = stream.schema(); let metrics = join_metrics; - let mut batches: Vec = Vec::new(); + let mut chunks: Vec = Vec::new(); + // Batches at or above half the target size pass through without being copied. + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size) + .with_biggest_coalesce_batch_size(Some(target_batch_size / 2)); while let Some(batch) = stream.next().await { let batch = batch?; @@ -1084,15 +1131,23 @@ async fn collect_left_input( metrics.build_mem_used.add(batch_size); metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); - batches.push(batch); + coalescer.push_batch(batch)?; + while let Some(chunk) = coalescer.next_completed_batch() { + chunks.push(chunk); + } } Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { let spill_manager = spill_manager.expect("checked by is_spillable_oom"); + // The batch that hit the limit is already in memory, so it joins the + // coalescer unreserved; the spill drains it from there. + metrics.build_input_batches.add(1); + metrics.build_input_rows.add(batch.num_rows()); + coalescer.push_batch(batch)?; let spilled = spill_left_input( spill_manager, Arc::clone(&schema), - batches, - Some(batch), + chunks, + coalescer, stream, metrics, &reservation, @@ -1108,23 +1163,24 @@ async fn collect_left_input( Err(e) => return Err(e), } } - - let merged_batch = concat_batches(&schema, &batches)?; + coalescer.finish_buffered_batch()?; + while let Some(chunk) = coalescer.next_completed_batch() { + chunks.push(chunk); + } // Reserve memory for visited_left_side bitmap if required by join type let visited_left_side = if with_visited_left_side { - let n_rows = merged_batch.num_rows(); + let n_rows: usize = chunks.iter().map(|c| c.num_rows()).sum(); let buffer_size = n_rows.div_ceil(8); match reservation.try_grow(buffer_size) { Ok(()) => {} Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { let spill_manager = spill_manager.expect("checked by is_spillable_oom"); - drop(batches); let spilled = spill_left_input( spill_manager, Arc::clone(&schema), - vec![merged_batch], - None, + chunks, + coalescer, stream, metrics, &reservation, @@ -1149,7 +1205,8 @@ async fn collect_left_input( }; Ok(LeftLoad::InMemory(Arc::new(JoinLeftData::new( - merged_batch, + chunks, + schema, Mutex::new(visited_left_side), AtomicUsize::new(probe_threads_count), reservation, @@ -1165,9 +1222,10 @@ fn left_load_from_spill( ) -> LeftLoad { match spilled { Some(data) => LeftLoad::Spilled(Arc::new(data)), - // No rows means no bitmap either, whatever the join type. + // No rows means no chunks and no bitmap, whatever the join type. None => LeftLoad::InMemory(Arc::new(JoinLeftData::new( - RecordBatch::new_empty(schema), + Vec::new(), + schema, Mutex::new(BooleanBufferBuilder::new(0)), AtomicUsize::new(probe_threads_count), reservation, @@ -1187,13 +1245,15 @@ fn is_spillable_oom( ) } -/// Write the already-buffered left batches plus the remainder of the same stream to one spill file. +/// Write the already-completed chunks plus the remainder of the same stream to one spill file. +/// The remainder keeps flowing through the same coalescer, so the file holds uniformly sized +/// chunks and the memory-limited replay reads them back at that granularity. /// Returns `None` when the left side carried no rows at all, which needs no spill file. async fn spill_left_input( spill_manager: SpillManager, schema: SchemaRef, - buffered: Vec, - pending: Option, + chunks: Vec, + mut coalescer: BatchCoalescer, mut stream: SendableRecordBatchStream, metrics: BuildProbeJoinMetrics, reservation: &MemoryReservation, @@ -1201,21 +1261,17 @@ async fn spill_left_input( let mut spill_file = spill_manager.create_in_progress_file("NestedLoopJoin left spill")?; - for batch in buffered { + for batch in chunks { if batch.num_rows() > 0 { spill_file.append_batch(&batch)?; } } - // The in-memory batches are spilled and dropped, so their reservation goes back to the pool - // before the rest of the stream is drained. + // The in-memory chunks are spilled and dropped, so their reservation goes back to the pool + // before the rest of the stream is drained; only the coalescer's one in-progress chunk + // stays resident past this point. reservation.free(); - - for batch in pending.into_iter() { - if batch.num_rows() > 0 { - metrics.build_input_batches.add(1); - metrics.build_input_rows.add(batch.num_rows()); - spill_file.append_batch(&batch)?; - } + while let Some(chunk) = coalescer.next_completed_batch() { + spill_file.append_batch(&chunk)?; } while let Some(batch) = stream.next().await { @@ -1223,9 +1279,16 @@ async fn spill_left_input( if batch.num_rows() > 0 { metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); - spill_file.append_batch(&batch)?; + coalescer.push_batch(batch)?; + while let Some(chunk) = coalescer.next_completed_batch() { + spill_file.append_batch(&chunk)?; + } } } + coalescer.finish_buffered_batch()?; + while let Some(chunk) = coalescer.next_completed_batch() { + spill_file.append_batch(&chunk)?; + } Ok(spill_file.finish()?.map(|file| LeftSpillData { spill_manager, @@ -1981,23 +2044,19 @@ impl NestedLoopJoinStream { return ControlFlow::Continue(()); } - let merged_batch = match concat_batches( + // The spill file was written as coalesced chunks, so the batches read back are already + // at target granularity and become the chunk list as-is, with no concatenation. + let chunks = std::mem::take(&mut active.pending_batches); + let left_schema = Arc::clone( 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 let with_visited = need_produce_result_in_final(self.join_type); - let n_rows = merged_batch.num_rows(); + let n_rows: usize = chunks.iter().map(|c| c.num_rows()).sum(); let visited_left_side = if with_visited { let buffer_size = n_rows.div_ceil(8); // Use infallible grow for bitmap — it's small @@ -2015,7 +2074,8 @@ impl NestedLoopJoinStream { let dummy_reservation = active.reservation.new_empty(); let left_data = JoinLeftData::new( - merged_batch, + chunks, + left_schema, Mutex::new(visited_left_side), // In memory-limited mode, only 1 probe thread per chunk AtomicUsize::new(1), @@ -2115,7 +2175,7 @@ impl NestedLoopJoinStream { if let (Ok(left_data), Some(right_batch)) = (self.get_left_data(), self.current_right_batch.as_ref()) { - let left_rows = left_data.batch().num_rows(); + let left_rows = left_data.total_rows(); let right_rows = right_batch.num_rows(); self.metrics.selectivity.add_total(left_rows * right_rows); } @@ -2424,10 +2484,16 @@ impl NestedLoopJoinStream { .clone(); // stop probing, the caller will go to the next state - if self.left_probe_idx >= left_data.batch().num_rows() { + if self.left_probe_idx >= left_data.total_rows() { return Ok(false); } + // Probe ranges never cross a chunk boundary, so locate the chunk once here. + let (chunk_idx, local_idx) = + left_data.locate(self.left_probe_idx).ok_or_else(|| { + internal_datafusion_err!("left_probe_idx must be within the left data") + })?; + // ======== // Join (l_row x right_batch) // and push the result into output_buffer @@ -2453,7 +2519,7 @@ impl NestedLoopJoinStream { // batch. let l_row_count = std::cmp::min( l_row_cnt_ratio, - left_data.batch().num_rows() - self.left_probe_idx, + left_data.chunk(chunk_idx).num_rows() - local_idx, ); debug_assert!( @@ -2463,7 +2529,8 @@ impl NestedLoopJoinStream { let joined_batch = self.process_left_range_join( &left_data, &right_batch, - self.left_probe_idx, + chunk_idx, + local_idx, l_row_count, )?; @@ -2476,9 +2543,12 @@ impl NestedLoopJoinStream { return Ok(true); } - let l_idx = self.left_probe_idx; - let joined_batch = - self.process_single_left_row_join(&left_data, &right_batch, l_idx)?; + let joined_batch = self.process_single_left_row_join( + &left_data, + &right_batch, + chunk_idx, + local_idx, + )?; if let Some(batch) = joined_batch { self.output_buffer.push_batch(batch)?; @@ -2493,7 +2563,8 @@ impl NestedLoopJoinStream { Ok(true) } - /// Process [l_start_index, l_start_index + l_count) JOIN right_batch + /// Process the left rows starting at `l_local_start` (local to the given chunk, never + /// crossing its end) JOIN right_batch. /// Returns a RecordBatch containing the join results (None if empty) /// /// Side Effect: If the join type requires, left or right side matched bitmap @@ -2502,7 +2573,8 @@ impl NestedLoopJoinStream { &mut self, left_data: &JoinLeftData, right_batch: &RecordBatch, - l_start_index: usize, + chunk_idx: usize, + l_local_start: usize, l_row_count: usize, ) -> Result> { // Construct the Cartesian product between the specified range of left rows @@ -2510,13 +2582,17 @@ impl NestedLoopJoinStream { // materializes the intermediate batch, and finally applies the join filter // to it. // ----------------------------------------------------------- + let left_chunk = left_data.chunk(chunk_idx); + // The visited-left bitmap is indexed by global row numbers. + let l_global_start = left_data.row_offset(chunk_idx) + l_local_start; let right_rows = right_batch.num_rows(); let total_rows = l_row_count * right_rows; - // Build index arrays for cartesian product: left_range X right_batch + // Build index arrays for cartesian product: left_range X right_batch. + // The indices are local to the chunk, since they address its columns. let left_indices: UInt32Array = UInt32Array::from_iter_values((0..l_row_count).flat_map(|i| { - std::iter::repeat_n((l_start_index + i) as u32, right_rows) + std::iter::repeat_n((l_local_start + i) as u32, right_rows) })); let right_indices: UInt32Array = UInt32Array::from_iter_values( (0..l_row_count).flat_map(|_| 0..right_rows as u32), @@ -2543,7 +2619,7 @@ impl NestedLoopJoinStream { Vec::with_capacity(filter.column_indices().len()); for column_index in filter.column_indices() { let array = if column_index.side == JoinSide::Left { - let col = left_data.batch().column(column_index.index); + let col = left_chunk.column(column_index.index); take(col.as_ref(), &left_indices, None)? } else { let col = right_batch.column(column_index.index); @@ -2596,7 +2672,7 @@ impl NestedLoopJoinStream { internal_datafusion_err!("Must be Some after the previous combining step") })?; - let l_index = l_start_index + i / right_rows; + let l_index = l_global_start + i / right_rows; let r_index = i % right_rows; if let Some(bitmap) = left_bitmap.as_mut() @@ -2664,7 +2740,7 @@ impl NestedLoopJoinStream { Vec::with_capacity(self.output_schema.fields().len()); for column_index in &self.column_indices { let array = if column_index.side == JoinSide::Left { - let col = left_data.batch().column(column_index.index); + let col = left_chunk.column(column_index.index); take(col.as_ref(), &left_indices, None)? } else { let col = right_batch.column(column_index.index); @@ -2687,17 +2763,22 @@ impl NestedLoopJoinStream { &mut self, left_data: &JoinLeftData, right_batch: &RecordBatch, - l_index: usize, + chunk_idx: usize, + l_local_index: usize, ) -> Result> { let right_row_count = right_batch.num_rows(); if right_row_count == 0 { return Ok(None); } + let left_chunk = left_data.chunk(chunk_idx); + // The visited-left bitmap is indexed by global row numbers. + let l_global_index = left_data.row_offset(chunk_idx) + l_local_index; + let cur_right_bitmap = if let Some(filter) = &self.join_filter { apply_filter_to_row_join_batch( - left_data.batch(), - l_index, + left_chunk, + l_local_index, right_batch, filter, )? @@ -2705,7 +2786,7 @@ impl NestedLoopJoinStream { BooleanArray::from(vec![true; right_row_count]) }; - self.update_matched_bitmap(l_index, &cur_right_bitmap)?; + self.update_matched_bitmap(l_global_index, &cur_right_bitmap)?; // For the following join types: here we only have to set the left/right // bitmap, and no need to output result @@ -2728,8 +2809,8 @@ impl NestedLoopJoinStream { // Use the optimized approach similar to build_intermediate_batch_for_single_left_row let join_batch = build_row_join_batch( &self.output_schema, - left_data.batch(), - l_index, + left_chunk, + l_local_index, right_batch, Some(cur_right_bitmap), &self.column_indices, @@ -2744,7 +2825,6 @@ impl NestedLoopJoinStream { /// false -> next state (Done) fn process_left_unmatched(&mut self) -> Result { let left_data = self.get_left_data()?; - let left_batch = left_data.batch(); // ======== // Check early return conditions @@ -2753,7 +2833,7 @@ impl NestedLoopJoinStream { // Early return if join type can't have unmatched rows let join_type_no_produce_left = !need_produce_result_in_final(self.join_type); // Stop processing unmatched rows, the caller will go to the next state - let finished = self.left_emit_idx >= left_batch.num_rows(); + let finished = self.left_emit_idx >= left_data.total_rows(); // `ProbeEnd` already recorded whether this stream emits unmatched-left // rows. Every probe partition passes through this state, but only the @@ -2768,7 +2848,14 @@ impl NestedLoopJoinStream { // Each time, the number to process is up to batch size // ======== let start_idx = self.left_emit_idx; - let end_idx = std::cmp::min(start_idx + self.batch_size, left_batch.num_rows()); + // Emission ranges never cross a chunk boundary; the output buffer re-coalesces the + // possibly smaller batch emitted at a chunk's tail. + let (chunk_idx, _) = left_data.locate(start_idx).ok_or_else(|| { + internal_datafusion_err!("left_emit_idx must be within the left data") + })?; + let chunk_end = + left_data.row_offset(chunk_idx) + left_data.chunk(chunk_idx).num_rows(); + let end_idx = std::cmp::min(start_idx + self.batch_size, chunk_end); if let Some(batch) = self.process_left_unmatched_range(left_data, start_idx, end_idx)? @@ -2805,10 +2892,17 @@ impl NestedLoopJoinStream { return Ok(None); } - // Slice both left batch, and bitmap to range [start_idx, end_idx) - // The range is bit index (not byte) - let left_batch = left_data.batch(); - let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx); + // Slice both left chunk, and bitmap to range [start_idx, end_idx) + // The range is bit index (not byte). The caller never lets a range cross a + // chunk boundary, so the whole range lives in one chunk. + let (chunk_idx, local_start) = left_data.locate(start_idx).ok_or_else(|| { + internal_datafusion_err!( + "unmatched-left range must start within the left data" + ) + })?; + let left_batch_sliced = left_data + .chunk(chunk_idx) + .slice(local_start, end_idx - start_idx); // Can this be more efficient? let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx); @@ -2852,7 +2946,7 @@ impl NestedLoopJoinStream { let cur_right_batch = unwrap_or_internal_err!(right_batch); let left_data = self.get_left_data()?; - let left_schema = left_data.batch().schema(); + let left_schema = left_data.schema(); let res = build_unmatched_batch( &self.output_schema, @@ -3440,6 +3534,102 @@ pub(crate) mod tests { Arc::new(TestMemoryExec::update_cache(&source)) } + /// A build side that already arrives in target-sized batches is retained as-is: the chunks + /// share their buffers with the input, so nothing is copied. Concatenating the build side + /// into one batch, as this operator used to, copies every byte and holds both copies. + #[tokio::test] + async fn build_side_chunks_reuse_the_input_buffers() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batches: Vec = (0..4) + .map(|b| { + let values: Vec = (0..1024).map(|i| b * 1024 + i).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(arrow::array::Int32Array::from(values))], + ) + .unwrap() + }) + .collect(); + let input_ptrs: Vec<*const u8> = batches + .iter() + .map(|b| b.column(0).to_data().buffers()[0].as_ptr()) + .collect(); + + let stream: SendableRecordBatchStream = + Box::pin(crate::stream::RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(batches.into_iter().map(Ok)), + )); + let task_ctx = Arc::new(TaskContext::default()); + let reservation = MemoryConsumer::new("test").register(task_ctx.memory_pool()); + let metrics = ExecutionPlanMetricsSet::new(); + + let load = collect_left_input( + stream, + BuildProbeJoinMetrics::new(0, &metrics), + reservation, + false, + 1, + None, + // Target below the input batch size, so every batch takes the large-batch bypass. + 512, + ) + .await?; + + let LeftLoad::InMemory(data) = load else { + panic!("the build side fit in memory"); + }; + assert_eq!(data.chunks.len(), 4, "each input batch is its own chunk"); + assert_eq!(data.total_rows(), 4096); + let chunk_ptrs: Vec<*const u8> = data + .chunks + .iter() + .map(|c| c.column(0).to_data().buffers()[0].as_ptr()) + .collect(); + assert_eq!( + input_ptrs, chunk_ptrs, + "chunks must reuse the input buffers instead of copying them" + ); + Ok(()) + } + + /// Zero-row chunks are dropped at construction: an empty chunk creates duplicate row + /// offsets, and the probe and emit cursors clamped to such a chunk's end would never + /// advance. Both chunk producers already normalize empties away (the load coalescer emits + /// none; the spill read loop skips them), so this guards future producers. + #[test] + fn join_left_data_drops_zero_row_chunks() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let empty = RecordBatch::new_empty(Arc::clone(&schema)); + let data = |from: i32, n: i32| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(arrow::array::Int32Array::from( + (from..from + n).collect::>(), + ))], + ) + .unwrap() + }; + let task_ctx = Arc::new(TaskContext::default()); + let reservation = MemoryConsumer::new("test").register(task_ctx.memory_pool()); + let left_data = JoinLeftData::new( + vec![empty.clone(), data(0, 5), empty.clone(), data(5, 7), empty], + Arc::clone(&schema), + Mutex::new(BooleanBufferBuilder::new(0)), + AtomicUsize::new(1), + reservation, + ); + assert_eq!(left_data.chunks.len(), 2); + assert_eq!(left_data.total_rows(), 12); + // Offsets are strictly increasing, so locate() is unambiguous at every row. + assert_eq!(left_data.locate(0), Some((0, 0))); + assert_eq!(left_data.locate(4), Some((0, 4))); + assert_eq!(left_data.locate(5), Some((1, 0))); + assert_eq!(left_data.locate(11), Some((1, 6))); + assert_eq!(left_data.locate(12), None); + } + /// An input that can be executed only once: later executions yield no batches, the way a /// stream backed by an external one-shot iterator behaves. #[derive(Debug)] From a5bbd22a0a1e20dbeffdc289b06479083af0186d Mon Sep 17 00:00:00 2001 From: Ran Reichman Date: Tue, 8 Sep 2026 23:09:24 -0400 Subject: [PATCH 2/4] test the NestedLoopJoin probe and unmatched-left emission across build chunk boundaries --- .../src/joins/nested_loop_join.rs | 171 ++++++++++++++++++ 1 file changed, 171 insertions(+) diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index a9f7f5a8c59a6..56c3887bf5a7a 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -3482,6 +3482,7 @@ pub(crate) mod tests { use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field}; use datafusion_common::assert_contains; + use datafusion_common::cast::as_int32_array; use datafusion_common::test_util::batches_to_sort_string; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; @@ -3630,6 +3631,176 @@ pub(crate) mod tests { assert_eq!(left_data.locate(12), None); } + const CHUNKED_LEFT_ROWS: i32 = 36; + /// Left rows that [`chunked_right_table`] matches: the first and last row, and the rows on + /// both sides of a chunk boundary in every layout the test builds (11|12 for 4- and 12-row + /// chunks, 27|28 for 4- and 7-row chunks). Every other boundary has unmatched rows on both + /// sides. + const CHUNKED_LEFT_MATCHED: [i32; 6] = [0, 11, 12, 27, 28, 35]; + + fn chunked_left_table(batch_rows: Option) -> Arc { + let ids: Vec = (0..CHUNKED_LEFT_ROWS).collect(); + build_table( + ("a1", &ids), + ("b1", &ids), + ("c1", &ids), + batch_rows, + Vec::new(), + ) + } + + /// One row per batch, so a batch size of 12 takes the multi-row probe path (12 / 1 > 10) + /// and a batch size of 4 the single-row one. + fn chunked_right_table() -> Arc { + let mut ids = CHUNKED_LEFT_MATCHED.to_vec(); + ids.push(99); + build_table( + ("a2", &ids), + ("b2", &ids), + ("c2", &ids), + Some(1), + Vec::new(), + ) + } + + /// left.a1 = right.a2 + fn equality_join_filter() -> JoinFilter { + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let intermediate_schema = Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("a2", DataType::Int32, true), + ]); + let expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a1", 0)), + Operator::Eq, + Arc::new(Column::new("a2", 1)), + )) as Arc; + JoinFilter::new(expression, column_indices, Arc::new(intermediate_schema)) + } + + /// Number of chunks the build-side load produces for `left` at this target batch size. + async fn build_side_chunk_count( + left: &Arc, + target_batch_size: usize, + ) -> Result { + let task_ctx = Arc::new(TaskContext::default()); + let reservation = MemoryConsumer::new("test").register(task_ctx.memory_pool()); + let metrics = ExecutionPlanMetricsSet::new(); + let load = collect_left_input( + left.execute(0, task_ctx)?, + BuildProbeJoinMetrics::new(0, &metrics), + reservation, + false, + 1, + None, + target_batch_size, + ) + .await?; + let LeftLoad::InMemory(data) = load else { + panic!("the build side fit in memory"); + }; + Ok(data.chunks.len()) + } + + /// The build side arrives either as single-row batches, which the load coalesces into chunks + /// of exactly `batch_size` rows, or as 7-row batches, which bypass the coalescer and leave + /// chunk edges off the output batch size, so probe and unmatched-left ranges must be clamped + /// at them. The same rows delivered as one batch form a single chunk, so the two runs differ + /// only in chunk layout and must agree. + #[rstest] + #[tokio::test] + async fn join_across_build_chunk_boundaries( + #[values(4, 12)] batch_size: usize, + #[values(1, 7)] left_batch_rows: usize, + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightMark + )] + join_type: JoinType, + ) -> Result<()> { + let rows = CHUNKED_LEFT_ROWS as usize; + // Batches above half the target bypass the coalescer and stay whole. + let expected_chunks = if left_batch_rows > batch_size / 2 { + rows.div_ceil(left_batch_rows) + } else { + rows.div_ceil(batch_size) + }; + assert!(expected_chunks > 1); + let chunked_left = chunked_left_table(Some(left_batch_rows)); + assert_eq!( + build_side_chunk_count(&chunked_left, batch_size).await?, + expected_chunks + ); + let single_chunk_left = chunked_left_table(None); + assert_eq!( + build_side_chunk_count(&single_chunk_left, batch_size).await?, + 1 + ); + + let (columns, chunked, _) = multi_partitioned_join_collect( + chunked_left, + chunked_right_table(), + &join_type, + Some(equality_join_filter()), + new_task_ctx(batch_size), + ) + .await?; + let (_, single_chunk, _) = multi_partitioned_join_collect( + single_chunk_left, + chunked_right_table(), + &join_type, + Some(equality_join_filter()), + new_task_ctx(batch_size), + ) + .await?; + assert_eq!( + batches_to_sort_string(&chunked), + batches_to_sort_string(&single_chunk) + ); + + if join_type == JoinType::Left { + let a2 = columns.iter().position(|c| c == "a2").unwrap(); + let (mut matched, mut unmatched) = (Vec::new(), Vec::new()); + for batch in &chunked { + let a1 = as_int32_array(batch.column(0))?; + for row in 0..batch.num_rows() { + if batch.column(a2).is_null(row) { + unmatched.push(a1.value(row)); + } else { + matched.push(a1.value(row)); + } + } + } + matched.sort_unstable(); + unmatched.sort_unstable(); + assert_eq!(matched, CHUNKED_LEFT_MATCHED); + assert_eq!( + unmatched, + (0..CHUNKED_LEFT_ROWS) + .filter(|i| !CHUNKED_LEFT_MATCHED.contains(i)) + .collect::>() + ); + } + Ok(()) + } + /// An input that can be executed only once: later executions yield no batches, the way a /// stream backed by an external one-shot iterator behaves. #[derive(Debug)] From cc8badfff3066b49d47464ca938e9591c07eba70 Mon Sep 17 00:00:00 2001 From: Ran Reichman Date: Fri, 11 Sep 2026 01:47:41 -0400 Subject: [PATCH 3/4] Flush the partial chunk when the build side spills and write the rest at input granularity, coalesce each replay pass, and move chunk range clamping into JoinLeftData::range --- .../src/joins/nested_loop_join.rs | 290 ++++++++---------- 1 file changed, 122 insertions(+), 168 deletions(-) diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 56c3887bf5a7a..7a3fe729617c4 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -1034,6 +1034,16 @@ pub(crate) struct JoinLeftData { reservation: MemoryReservation, } +/// A run of build-side rows inside one chunk. Ranges never cross a chunk boundary, so the +/// chunk's columns are addressed with `local_start`, while `global_start` indexes the +/// visited-left bitmap. +pub(crate) struct ChunkRange<'a> { + pub(crate) chunk: &'a RecordBatch, + pub(crate) local_start: usize, + pub(crate) global_start: usize, + pub(crate) len: usize, +} + impl JoinLeftData { pub(crate) fn new( chunks: Vec, @@ -1063,14 +1073,6 @@ impl JoinLeftData { } } - pub(crate) fn chunk(&self, idx: usize) -> &RecordBatch { - &self.chunks[idx] - } - - pub(crate) fn row_offset(&self, idx: usize) -> usize { - self.row_offsets[idx] - } - pub(crate) fn total_rows(&self) -> usize { self.total_rows } @@ -1079,14 +1081,25 @@ impl JoinLeftData { Arc::clone(&self.schema) } - /// Index of the chunk holding the given global row, and that row's offset inside it. - pub(crate) fn locate(&self, row: usize) -> Option<(usize, usize)> { - let idx = match self.row_offsets.binary_search(&row) { + /// Up to `max_len` rows starting at global row `start`, clamped at the end of the chunk + /// holding `start`. `None` once `start` is past the last row. + pub(crate) fn range(&self, start: usize, max_len: usize) -> Option> { + let idx = match self.row_offsets.binary_search(&start) { Ok(idx) => idx, Err(0) => return None, Err(next) => next - 1, }; - (row < self.total_rows).then(|| (idx, row - self.row_offsets[idx])) + if start >= self.total_rows { + return None; + } + let chunk = &self.chunks[idx]; + let local_start = start - self.row_offsets[idx]; + Some(ChunkRange { + chunk, + local_start, + global_start: start, + len: max_len.min(chunk.num_rows() - local_start), + }) } pub(crate) fn bitmap(&self) -> &SharedBitmapBuilder { @@ -1119,9 +1132,7 @@ async fn collect_left_input( let schema = stream.schema(); let metrics = join_metrics; let mut chunks: Vec = Vec::new(); - // Batches at or above half the target size pass through without being copied. - let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size) - .with_biggest_coalesce_batch_size(Some(target_batch_size / 2)); + let mut coalescer = chunk_coalescer(&schema, target_batch_size); while let Some(batch) = stream.next().await { let batch = batch?; @@ -1138,16 +1149,17 @@ async fn collect_left_input( } Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { let spill_manager = spill_manager.expect("checked by is_spillable_oom"); - // The batch that hit the limit is already in memory, so it joins the - // coalescer unreserved; the spill drains it from there. metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); - coalescer.push_batch(batch)?; + coalescer.finish_buffered_batch()?; + while let Some(chunk) = coalescer.next_completed_batch() { + chunks.push(chunk); + } let spilled = spill_left_input( spill_manager, Arc::clone(&schema), chunks, - coalescer, + Some(batch), stream, metrics, &reservation, @@ -1180,7 +1192,7 @@ async fn collect_left_input( spill_manager, Arc::clone(&schema), chunks, - coalescer, + None, stream, metrics, &reservation, @@ -1245,15 +1257,46 @@ fn is_spillable_oom( ) } -/// Write the already-completed chunks plus the remainder of the same stream to one spill file. -/// The remainder keeps flowing through the same coalescer, so the file holds uniformly sized -/// chunks and the memory-limited replay reads them back at that granularity. +/// Chunks are `target_batch_size` rows; a batch already at or above half that passes through +/// without being copied. +fn chunk_coalescer(schema: &SchemaRef, target_batch_size: usize) -> BatchCoalescer { + BatchCoalescer::new(Arc::clone(schema), target_batch_size) + .with_biggest_coalesce_batch_size(Some(target_batch_size / 2)) +} + +/// Compacts already-buffered batches into chunks. Every input must be reserved by the caller: +/// the copy this makes is bounded by the input, so the pass that reserved it bounds the copy. +fn coalesce_chunks( + batches: Vec, + schema: &SchemaRef, + target_batch_size: usize, +) -> Result> { + let mut coalescer = chunk_coalescer(schema, target_batch_size); + let mut chunks = Vec::with_capacity(batches.len()); + for batch in batches { + coalescer.push_batch(batch)?; + while let Some(chunk) = coalescer.next_completed_batch() { + chunks.push(chunk); + } + } + coalescer.finish_buffered_batch()?; + while let Some(chunk) = coalescer.next_completed_batch() { + chunks.push(chunk); + } + Ok(chunks) +} + +/// Write the already-completed chunks, the batch that hit the limit, and the remainder of the +/// same stream to one spill file. Nothing is reserved past this point, so the remainder is +/// written batch by batch as it arrives instead of being coalesced, which would hold up to a +/// chunk's worth of unreserved rows; the memory-limited replay coalesces each pass it reads +/// back, after reserving it. /// Returns `None` when the left side carried no rows at all, which needs no spill file. async fn spill_left_input( spill_manager: SpillManager, schema: SchemaRef, chunks: Vec, - mut coalescer: BatchCoalescer, + pending: Option, mut stream: SendableRecordBatchStream, metrics: BuildProbeJoinMetrics, reservation: &MemoryReservation, @@ -1262,16 +1305,13 @@ async fn spill_left_input( spill_manager.create_in_progress_file("NestedLoopJoin left spill")?; for batch in chunks { - if batch.num_rows() > 0 { - spill_file.append_batch(&batch)?; - } + spill_file.append_batch(&batch)?; } // The in-memory chunks are spilled and dropped, so their reservation goes back to the pool - // before the rest of the stream is drained; only the coalescer's one in-progress chunk - // stays resident past this point. + // before the rest of the stream is drained. reservation.free(); - while let Some(chunk) = coalescer.next_completed_batch() { - spill_file.append_batch(&chunk)?; + if let Some(batch) = pending.filter(|b| b.num_rows() > 0) { + spill_file.append_batch(&batch)?; } while let Some(batch) = stream.next().await { @@ -1279,16 +1319,9 @@ async fn spill_left_input( if batch.num_rows() > 0 { metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); - coalescer.push_batch(batch)?; - while let Some(chunk) = coalescer.next_completed_batch() { - spill_file.append_batch(&chunk)?; - } + spill_file.append_batch(&batch)?; } } - coalescer.finish_buffered_batch()?; - while let Some(chunk) = coalescer.next_completed_batch() { - spill_file.append_batch(&chunk)?; - } Ok(spill_file.finish()?.map(|file| LeftSpillData { spill_manager, @@ -2044,15 +2077,21 @@ impl NestedLoopJoinStream { return ControlFlow::Continue(()); } - // The spill file was written as coalesced chunks, so the batches read back are already - // at target granularity and become the chunk list as-is, with no concatenation. - let chunks = std::mem::take(&mut active.pending_batches); let left_schema = Arc::clone( active .left_schema .as_ref() .expect("left_schema must be set"), ); + // Every batch of the pass is reserved above, so compacting it here stays within budget. + let chunks = match coalesce_chunks( + std::mem::take(&mut active.pending_batches), + &left_schema, + self.batch_size, + ) { + Ok(chunks) => chunks, + Err(e) => return ControlFlow::Break(Poll::Ready(Some(Err(e)))), + }; // Build visited bitmap if needed for this join type let with_visited = need_produce_result_in_final(self.join_type); @@ -2488,12 +2527,6 @@ impl NestedLoopJoinStream { return Ok(false); } - // Probe ranges never cross a chunk boundary, so locate the chunk once here. - let (chunk_idx, local_idx) = - left_data.locate(self.left_probe_idx).ok_or_else(|| { - internal_datafusion_err!("left_probe_idx must be within the left data") - })?; - // ======== // Join (l_row x right_batch) // and push the result into output_buffer @@ -2514,41 +2547,31 @@ impl NestedLoopJoinStream { let l_row_cnt_ratio = self.batch_size / right_batch.num_rows(); if l_row_cnt_ratio > 10 { - // Calculate max left rows to handle at once. This operator tries to handle - // up to `datafusion.execution.batch_size` rows at once in the intermediate - // batch. - let l_row_count = std::cmp::min( - l_row_cnt_ratio, - left_data.chunk(chunk_idx).num_rows() - local_idx, - ); - - debug_assert!( - l_row_count != 0, - "This function should only be entered when there are remaining left rows to process" - ); - let joined_batch = self.process_left_range_join( - &left_data, - &right_batch, - chunk_idx, - local_idx, - l_row_count, - )?; + // Handle up to `datafusion.execution.batch_size` rows at once in the intermediate + // batch, clamped at the current chunk's end. + let range = left_data + .range(self.left_probe_idx, l_row_cnt_ratio) + .ok_or_else(|| { + internal_datafusion_err!( + "left_probe_idx must be within the left data" + ) + })?; + let joined_batch = + self.process_left_range_join(&left_data, &range, &right_batch)?; if let Some(batch) = joined_batch { self.output_buffer.push_batch(batch)?; } - self.left_probe_idx += l_row_count; + self.left_probe_idx += range.len; return Ok(true); } - let joined_batch = self.process_single_left_row_join( - &left_data, - &right_batch, - chunk_idx, - local_idx, - )?; + let range = left_data.range(self.left_probe_idx, 1).ok_or_else(|| { + internal_datafusion_err!("left_probe_idx must be within the left data") + })?; + let joined_batch = self.process_single_left_row_join(&range, &right_batch)?; if let Some(batch) = joined_batch { self.output_buffer.push_batch(batch)?; @@ -2563,8 +2586,7 @@ impl NestedLoopJoinStream { Ok(true) } - /// Process the left rows starting at `l_local_start` (local to the given chunk, never - /// crossing its end) JOIN right_batch. + /// Process the left rows of `range` JOIN right_batch. /// Returns a RecordBatch containing the join results (None if empty) /// /// Side Effect: If the join type requires, left or right side matched bitmap @@ -2572,19 +2594,18 @@ impl NestedLoopJoinStream { fn process_left_range_join( &mut self, left_data: &JoinLeftData, + range: &ChunkRange<'_>, right_batch: &RecordBatch, - chunk_idx: usize, - l_local_start: usize, - l_row_count: usize, ) -> Result> { // Construct the Cartesian product between the specified range of left rows // and the entire right_batch. First, it calculates the index vectors, then // materializes the intermediate batch, and finally applies the join filter // to it. // ----------------------------------------------------------- - let left_chunk = left_data.chunk(chunk_idx); - // The visited-left bitmap is indexed by global row numbers. - let l_global_start = left_data.row_offset(chunk_idx) + l_local_start; + let left_chunk = range.chunk; + let l_local_start = range.local_start; + let l_global_start = range.global_start; + let l_row_count = range.len; let right_rows = right_batch.num_rows(); let total_rows = l_row_count * right_rows; @@ -2761,19 +2782,17 @@ impl NestedLoopJoinStream { /// will be set for matched indices. fn process_single_left_row_join( &mut self, - left_data: &JoinLeftData, + range: &ChunkRange<'_>, right_batch: &RecordBatch, - chunk_idx: usize, - l_local_index: usize, ) -> Result> { let right_row_count = right_batch.num_rows(); if right_row_count == 0 { return Ok(None); } - let left_chunk = left_data.chunk(chunk_idx); - // The visited-left bitmap is indexed by global row numbers. - let l_global_index = left_data.row_offset(chunk_idx) + l_local_index; + let left_chunk = range.chunk; + let l_local_index = range.local_start; + let l_global_index = range.global_start; let cur_right_bitmap = if let Some(filter) = &self.join_filter { apply_filter_to_row_join_batch( @@ -2847,19 +2866,15 @@ impl NestedLoopJoinStream { // Process unmatched rows and push the result into output_buffer // Each time, the number to process is up to batch size // ======== - let start_idx = self.left_emit_idx; - // Emission ranges never cross a chunk boundary; the output buffer re-coalesces the - // possibly smaller batch emitted at a chunk's tail. - let (chunk_idx, _) = left_data.locate(start_idx).ok_or_else(|| { - internal_datafusion_err!("left_emit_idx must be within the left data") - })?; - let chunk_end = - left_data.row_offset(chunk_idx) + left_data.chunk(chunk_idx).num_rows(); - let end_idx = std::cmp::min(start_idx + self.batch_size, chunk_end); + // The output buffer re-coalesces the possibly smaller batch emitted at a chunk's tail. + let range = left_data + .range(self.left_emit_idx, self.batch_size) + .ok_or_else(|| { + internal_datafusion_err!("left_emit_idx must be within the left data") + })?; + let end_idx = range.global_start + range.len; - if let Some(batch) = - self.process_left_unmatched_range(left_data, start_idx, end_idx)? - { + if let Some(batch) = self.process_left_unmatched_range(left_data, &range)? { self.output_buffer.push_batch(batch)?; } @@ -2870,39 +2885,16 @@ impl NestedLoopJoinStream { Ok(true) } - /// Process unmatched rows from the left data within the specified range. + /// Process unmatched rows from the left data within `range`. /// Returns a RecordBatch containing the unmatched rows (None if empty). - /// - /// # Arguments - /// * `left_data` - The left side data containing the batch and bitmap - /// * `start_idx` - Start index (inclusive) of the range to process - /// * `end_idx` - End index (exclusive) of the range to process - /// - /// # Safety - /// The caller is responsible for ensuring that `start_idx` and `end_idx` are - /// within valid bounds of the left batch. This function does not perform - /// bounds checking. fn process_left_unmatched_range( &self, left_data: &JoinLeftData, - start_idx: usize, - end_idx: usize, + range: &ChunkRange<'_>, ) -> Result> { - if start_idx == end_idx { - return Ok(None); - } - - // Slice both left chunk, and bitmap to range [start_idx, end_idx) - // The range is bit index (not byte). The caller never lets a range cross a - // chunk boundary, so the whole range lives in one chunk. - let (chunk_idx, local_start) = left_data.locate(start_idx).ok_or_else(|| { - internal_datafusion_err!( - "unmatched-left range must start within the left data" - ) - })?; - let left_batch_sliced = left_data - .chunk(chunk_idx) - .slice(local_start, end_idx - start_idx); + let start_idx = range.global_start; + let end_idx = start_idx + range.len; + let left_batch_sliced = range.chunk.slice(range.local_start, range.len); // Can this be more efficient? let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx); @@ -3535,9 +3527,8 @@ pub(crate) mod tests { Arc::new(TestMemoryExec::update_cache(&source)) } - /// A build side that already arrives in target-sized batches is retained as-is: the chunks - /// share their buffers with the input, so nothing is copied. Concatenating the build side - /// into one batch, as this operator used to, copies every byte and holds both copies. + /// The zero-copy retention the chunked layout exists for, which no result-level test can + /// observe: a build side arriving at target size shares its buffers with the input. #[tokio::test] async fn build_side_chunks_reuse_the_input_buffers() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -3594,43 +3585,6 @@ pub(crate) mod tests { Ok(()) } - /// Zero-row chunks are dropped at construction: an empty chunk creates duplicate row - /// offsets, and the probe and emit cursors clamped to such a chunk's end would never - /// advance. Both chunk producers already normalize empties away (the load coalescer emits - /// none; the spill read loop skips them), so this guards future producers. - #[test] - fn join_left_data_drops_zero_row_chunks() { - let schema: SchemaRef = - Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let empty = RecordBatch::new_empty(Arc::clone(&schema)); - let data = |from: i32, n: i32| { - RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(arrow::array::Int32Array::from( - (from..from + n).collect::>(), - ))], - ) - .unwrap() - }; - let task_ctx = Arc::new(TaskContext::default()); - let reservation = MemoryConsumer::new("test").register(task_ctx.memory_pool()); - let left_data = JoinLeftData::new( - vec![empty.clone(), data(0, 5), empty.clone(), data(5, 7), empty], - Arc::clone(&schema), - Mutex::new(BooleanBufferBuilder::new(0)), - AtomicUsize::new(1), - reservation, - ); - assert_eq!(left_data.chunks.len(), 2); - assert_eq!(left_data.total_rows(), 12); - // Offsets are strictly increasing, so locate() is unambiguous at every row. - assert_eq!(left_data.locate(0), Some((0, 0))); - assert_eq!(left_data.locate(4), Some((0, 4))); - assert_eq!(left_data.locate(5), Some((1, 0))); - assert_eq!(left_data.locate(11), Some((1, 6))); - assert_eq!(left_data.locate(12), None); - } - const CHUNKED_LEFT_ROWS: i32 = 36; /// Left rows that [`chunked_right_table`] matches: the first and last row, and the rows on /// both sides of a chunk boundary in every layout the test builds (11|12 for 4- and 12-row From 16593136d381a1ab535ef3e4b45175917fbe3d29 Mon Sep 17 00:00:00 2001 From: Ran Reichman Date: Tue, 15 Sep 2026 03:15:07 -0400 Subject: [PATCH 4/4] Compact the build side only after it is fully reserved and keep spilled and replayed batches as they arrived so the memory-limited path never copies a pass --- .../src/joins/nested_loop_join.rs | 59 ++--- .../tests/nested_loop_join_memory.rs | 228 ++++++++++++++++++ 2 files changed, 249 insertions(+), 38 deletions(-) create mode 100644 datafusion/physical-plan/tests/nested_loop_join_memory.rs diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index b2f3b99ac8e97..ff941af4a16e1 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -1204,8 +1204,7 @@ async fn collect_left_input( ) -> Result { let schema = stream.schema(); let metrics = join_metrics; - let mut chunks: Vec = Vec::new(); - let mut coalescer = chunk_coalescer(&schema, target_batch_size); + let mut batches: Vec = Vec::new(); while let Some(batch) = stream.next().await { let batch = batch?; @@ -1216,10 +1215,7 @@ async fn collect_left_input( metrics.build_mem_used.add(batch_size); metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); - coalescer.push_batch(batch)?; - while let Some(chunk) = coalescer.next_completed_batch() { - chunks.push(chunk); - } + batches.push(batch); } Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { // Do not keep the operator timer running while the spill path @@ -1228,14 +1224,10 @@ async fn collect_left_input( let spill_manager = spill_manager.expect("checked by is_spillable_oom"); metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); - coalescer.finish_buffered_batch()?; - while let Some(chunk) = coalescer.next_completed_batch() { - chunks.push(chunk); - } let spilled = spill_left_input( spill_manager, Arc::clone(&schema), - chunks, + batches, Some(batch), stream, metrics, @@ -1257,10 +1249,9 @@ async fn collect_left_input( // polling the child stream above. let build_timer = metrics.build_time.timer(); - coalescer.finish_buffered_batch()?; - while let Some(chunk) = coalescer.next_completed_batch() { - chunks.push(chunk); - } + // Compacted only once the whole side is reserved, so a load that spills never has a + // partially built chunk to materialize while the pool is exhausted. + let chunks = coalesce_chunks(batches, &schema, target_batch_size)?; // Reserve memory for visited_left_side bitmap if required by join type let visited_left_side = if with_visited_left_side { @@ -1342,21 +1333,16 @@ fn is_spillable_oom( ) } -/// Chunks are `target_batch_size` rows; a batch already at or above half that passes through -/// without being copied. -fn chunk_coalescer(schema: &SchemaRef, target_batch_size: usize) -> BatchCoalescer { - BatchCoalescer::new(Arc::clone(schema), target_batch_size) - .with_biggest_coalesce_batch_size(Some(target_batch_size / 2)) -} - -/// Compacts already-buffered batches into chunks. Every input must be reserved by the caller: -/// the copy this makes is bounded by the input, so the pass that reserved it bounds the copy. +/// Compacts a fully buffered build side into `target_batch_size`-row chunks; a batch already at +/// or above half that passes through without being copied. Concatenation completes one chunk at +/// a time, so the unreserved copy in flight is at most one chunk, like any operator's output batch. fn coalesce_chunks( batches: Vec, schema: &SchemaRef, target_batch_size: usize, ) -> Result> { - let mut coalescer = chunk_coalescer(schema, target_batch_size); + let mut coalescer = BatchCoalescer::new(Arc::clone(schema), target_batch_size) + .with_biggest_coalesce_batch_size(Some(target_batch_size / 2)); let mut chunks = Vec::with_capacity(batches.len()); for batch in batches { coalescer.push_batch(batch)?; @@ -1371,16 +1357,14 @@ fn coalesce_chunks( Ok(chunks) } -/// Write the already-completed chunks, the batch that hit the limit, and the remainder of the -/// same stream to one spill file. Nothing is reserved past this point, so the remainder is -/// written batch by batch as it arrives instead of being coalesced, which would hold up to a -/// chunk's worth of unreserved rows; the memory-limited replay coalesces each pass it reads -/// back, after reserving it. +/// Write the batches buffered so far, the batch that hit the limit, and the remainder of the +/// same stream to one spill file, each as it arrived. Nothing is copied or coalesced here, so the +/// spill allocates nothing beyond the file writer while the pool is exhausted. /// Returns `None` when the left side carried no rows at all, which needs no spill file. async fn spill_left_input( spill_manager: SpillManager, schema: SchemaRef, - chunks: Vec, + batches: Vec, pending: Option, mut stream: SendableRecordBatchStream, metrics: BuildProbeJoinMetrics, @@ -1390,10 +1374,10 @@ async fn spill_left_input( let mut spill_file = spill_manager.create_in_progress_file("NestedLoopJoin left spill")?; - for batch in chunks { + for batch in batches { spill_file.append_batch(&batch)?; } - // The in-memory chunks are spilled and dropped, so their reservation goes back to the pool + // The buffered batches are spilled and dropped, so their reservation goes back to the pool // before the rest of the stream is drained. reservation.free(); if let Some(batch) = pending.filter(|b| b.num_rows() > 0) { @@ -1447,7 +1431,7 @@ enum NLJState { } /// Outcome of the single pass over the left (build) input. pub(crate) enum LeftLoad { - /// The left side fit the memory budget and is buffered as one batch. + /// The left side fit the memory budget and is buffered as chunks. InMemory(Arc), /// The budget ran out, so the left side was spilled during that same pass. Every partition /// shares this handle, and each left chunk pass re-opens the file. @@ -1887,7 +1871,6 @@ impl FallbackCoordinator { &mut reservation, carryover, Arc::clone(&left_schema), - task_context.session_config().batch_size(), build_time.clone(), ); let load_result = { @@ -1985,7 +1968,6 @@ impl FallbackCoordinator { reservation: &mut MemoryReservation, carryover: Option, left_schema: SchemaRef, - target_batch_size: usize, build_time: Time, ) -> Result { // The previous chunk's bytes were moved into its `JoinLeftData`, so @@ -2042,8 +2024,9 @@ impl FallbackCoordinator { } let _build_timer = build_time.timer(); - // Every batch of the pass is reserved above, so compacting it here stays within budget. - let chunks = coalesce_chunks(pending_batches, &left_schema, target_batch_size)?; + // Kept as read back: compacting the pass would copy it while its reserved inputs are + // still live, on top of a pool that is already full. + let chunks = pending_batches; let n_rows: usize = chunks.iter().map(|c| c.num_rows()).sum(); let visited_left_side = if self.with_visited_bitmap { let buffer_size = n_rows.div_ceil(8); diff --git a/datafusion/physical-plan/tests/nested_loop_join_memory.rs b/datafusion/physical-plan/tests/nested_loop_join_memory.rs new file mode 100644 index 0000000000000..c1f6cae033bfa --- /dev/null +++ b/datafusion/physical-plan/tests/nested_loop_join_memory.rs @@ -0,0 +1,228 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Peak-allocation check for the NestedLoopJoin build side under a memory limit. +//! +//! The memory pool only sees what the operator reserves, so a copy of the build side that is +//! never reserved (the `concat_batches` this operator used to make, or coalescing a pass while +//! its inputs are still live) is invisible to pool-based assertions. This binary counts live +//! bytes at the allocator instead, and must stay the only test in it so nothing else runs +//! alongside the measurement. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::any::Any; +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::array::{Int32Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{JoinSide, JoinType, Result}; +use datafusion_execution::TaskContext; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column}; +use datafusion_physical_plan::common::collect; +use datafusion_physical_plan::joins::NestedLoopJoinExec; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; +use parking_lot::RwLock; + +struct PeakTrackingAllocator; + +static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0); +static PEAK_BYTES: AtomicUsize = AtomicUsize::new(0); + +fn record_alloc(bytes: usize) { + let live = LIVE_BYTES.fetch_add(bytes, Ordering::Relaxed) + bytes; + PEAK_BYTES.fetch_max(live, Ordering::Relaxed); +} + +unsafe impl GlobalAlloc for PeakTrackingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + record_alloc(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed); + record_alloc(new_size); + } + new_ptr + } +} + +#[global_allocator] +static GLOBAL: PeakTrackingAllocator = PeakTrackingAllocator; + +const ROW_BYTES: usize = 64 * 1024; +const BUILD_ROWS: usize = 512; +const POOL_BYTES: usize = 8 * 1024 * 1024; + +/// One row per batch, each with its own 64 KiB string allocation, so a target-row coalescer +/// would merge a whole pass into a single chunk and copy every byte of it. The count is shared +/// across `reset_state` so the test can read it after the run. +#[derive(Debug)] +struct WideRows { + schema: SchemaRef, + emitted: Arc, +} + +impl fmt::Display for WideRows { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "WideRows(emitted={})", + self.emitted.load(Ordering::Relaxed) + ) + } +} + +impl LazyBatchGenerator for WideRows { + fn as_any(&self) -> &dyn Any { + self + } + + fn generate_next_batch(&mut self) -> Result> { + if self.emitted.load(Ordering::Relaxed) == BUILD_ROWS { + return Ok(None); + } + self.emitted.fetch_add(1, Ordering::Relaxed); + let batch = RecordBatch::try_new( + Arc::clone(&self.schema), + vec![ + Arc::new(Int32Array::from(vec![0])), + Arc::new(StringArray::from(vec!["x".repeat(ROW_BYTES)])), + ], + )?; + Ok(Some(batch)) + } + + fn reset_state(&self) -> Arc> { + Arc::new(RwLock::new(WideRows { + schema: Arc::clone(&self.schema), + emitted: Arc::clone(&self.emitted), + })) + } +} + +/// `left.k > right.k` with every left `k` at 0 and every right `k` at 1, so the join keeps its +/// full probe shape while producing no rows to hold on to. +fn never_matching_filter() -> JoinFilter { + let expression: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("k", 0)), + Operator::Gt, + Arc::new(Column::new("k", 1)), + )); + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let schema = Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("k", DataType::Int32, false), + ]); + JoinFilter::new(expression, column_indices, Arc::new(schema)) +} + +#[tokio::test] +async fn build_side_spill_and_replay_stay_within_the_pool() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let emitted = Arc::new(AtomicUsize::new(0)); + let left: Arc = Arc::new(LazyMemoryExec::try_new( + Arc::clone(&left_schema), + vec![Arc::new(RwLock::new(WideRows { + schema: left_schema, + emitted: Arc::clone(&emitted), + }))], + )?); + + let right_schema = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![Arc::new(Int32Array::from(vec![1, 1, 1, 1]))], + )?; + let right: Arc = + TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None)?; + + let join = NestedLoopJoinExec::try_new( + left, + right, + Some(never_matching_filter()), + &JoinType::Inner, + None, + )?; + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(POOL_BYTES, 1.0) + .build_arc()?; + let ctx = Arc::new(TaskContext::default().with_runtime(runtime)); + + let baseline = LIVE_BYTES.load(Ordering::Relaxed); + PEAK_BYTES.store(baseline, Ordering::Relaxed); + let output = collect(join.execute(0, ctx)?).await?; + let peak = PEAK_BYTES.load(Ordering::Relaxed) - baseline; + + let output_rows: usize = output.iter().map(|b| b.num_rows()).sum(); + assert_eq!(output_rows, 0); + let metrics = join.metrics().expect("metrics"); + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "the build side must spill" + ); + assert_eq!( + emitted.load(Ordering::Relaxed), + BUILD_ROWS, + "every build row must be consumed" + ); + + // The build side is four times the pool, so it spills and is replayed pass by pass, each pass + // filling the pool. Any unreserved copy of a pass (a coalesced spill chunk at the spill + // transition, or a coalesced replay pass) doubles that peak; ordinary bookkeeping (the + // carried-over batch, spill reader buffers, the probe side) stays far below half a pool. + let limit = POOL_BYTES + POOL_BYTES / 2; + assert!( + peak < limit, + "peak live allocation {:.2} MiB exceeds {:.2} MiB with an {:.0} MiB pool", + peak as f64 / (1024.0 * 1024.0), + limit as f64 / (1024.0 * 1024.0), + POOL_BYTES as f64 / (1024.0 * 1024.0), + ); + Ok(()) +}