From fd218b3b86f015b4733d669303da967dd1db4f77 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Wed, 16 Sep 2026 16:34:13 +0800 Subject: [PATCH 1/3] test: validate nested loop join build-side memory usage Add one NLJ query at the existing 10-million-row memory-validation scale. Keep the 80 MB build input on the left and cap baseline-adjusted RSS at 90 MB, allowing overhead without a second build-side copy. The test fails with concatenation and passes with segmented storage. --- .../memory_limit_validation/mod.rs | 1 + .../nlj_mem_validation.rs | 57 +++++++++++++++++++ 2 files changed, 58 insertions(+) create mode 100644 datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs index 83ebb266c8257..aeaffb50af979 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs @@ -18,6 +18,7 @@ //! Validates query's actual memory usage is consistent with the specified memory //! limit. +mod nlj_mem_validation; mod smj_mem_validation; mod sort_mem_validation; mod utils; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs new file mode 100644 index 0000000000000..2b38de04d5cda --- /dev/null +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs @@ -0,0 +1,57 @@ +// 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. + +//! Memory validation for the nested loop join's buffered build side. + +use datafusion::prelude::SessionConfig; + +use crate::memory_limit::memory_limit_validation::utils; + +#[test] +fn nlj_no_mem_limit_runner() { + utils::spawn_test_process("nlj_mem_validation", "nlj_no_mem_limit"); +} + +// NLJ buffers the entire build side and uses a small, constant amount of memory +// to stream the probe side. +// +// This test verifies that an NLJ without a memory limit follows this estimate. +#[tokio::test] +async fn nlj_no_mem_limit() { + let config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(8192) + // Keep the large input on the build side, as written in the query. + .set_bool("datafusion.optimizer.join_reordering", false); + + utils::validate_query_with_memory_limits_and_config( + 90_000_000, // 80 MB for the build side + 10 MB of extra room. + None, + "SELECT count(*), sum(l.value) + FROM generate_series(1, 10000000) AS l + JOIN generate_series(1, 1) AS r + ON (l.value + r.value) % 2 = 0", + "SELECT count(*), sum(l.value) + FROM generate_series(1, 1000000) AS l + JOIN generate_series(1, 1) AS r + ON (l.value + r.value) % 2 = 0", + config, + Some("NestedLoopJoinExec"), + Some(false), + ) + .await; +} From 2a6d0d1051e24708b28caad1fe9cd95c3a46bd26 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Wed, 16 Sep 2026 16:16:50 +0800 Subject: [PATCH 2/3] perf: Reduce NLJ build side memory usage by 2X with segmented batch layout --- .../physical-plan/src/joins/logical_batch.rs | 748 ++++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 1 + .../src/joins/nested_loop_join.rs | 104 ++- 3 files changed, 796 insertions(+), 57 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/logical_batch.rs diff --git a/datafusion/physical-plan/src/joins/logical_batch.rs b/datafusion/physical-plan/src/joins/logical_batch.rs new file mode 100644 index 0000000000000..f6bf6762f6e85 --- /dev/null +++ b/datafusion/physical-plan/src/joins/logical_batch.rs @@ -0,0 +1,748 @@ +// 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. + +//! [`LogicalBatch`]: a logically contiguous batch stored as a sequence of +//! [`RecordBatch`]es. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, UInt32Array, new_empty_array}; +use arrow::compute::{TakeOptions, concat, concat_batches, take}; +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, exec_datafusion_err, exec_err}; + +/// A logically contiguous batch backed by a sequence of [`RecordBatch`]es +/// that share one schema. Methods on this struct accept global row indices. +/// +/// # Example +/// ```text +/// +/// segment 0 (3 rows) segment 1 (2 rows) segment 2 (4 rows) +/// ┌───┬───┬───┐ ┌───┬───┐ ┌───┬───┬───┬───┐ +/// │ a │ b │ c │ │ d │ e │ │ f │ g │ h │ i │ +/// └───┴───┴───┘ └───┴───┘ └───┴───┴───┴───┘ +/// 0 1 2 3 4 5 6 7 8 ◀── global row index +/// ``` +/// +/// # Motivation +/// +/// Joins (e.g. Nested Loop Join) usually buffer all build-side input, and next concatenating +/// them into a contiguous batch, before the next step. It will 2X the memory usage +/// since fragmented batches and final contiguous batch exist at the same time. This +/// struct avoids concatenation step, and helps reduce memory usage by 2X. +/// +/// Avoiding memory concatenating overhead is not the motivation, since it's usually +/// fast and not a bottleneck in real workloads; at the same time single-batch abstraction +/// help simplify join logic. +/// +/// See issue for details: +/// - +/// +/// # TODO +/// It's named 'logical batch' because it's possible to swap the physical layout +/// and keep the same interface for other usages. For example, segments are aligned +/// at the same size, so it achieves O(1) access speed. +#[derive(Debug, Clone)] +pub(crate) struct LogicalBatch { + schema: SchemaRef, + /// The underlying batches, in row order. Empty batches are dropped on + /// construction, so every segment holds at least one row. + segments: Vec, + /// `offsets[i]` is the global index of the first row of `segments[i]`; + /// `offsets[segments.len()]` is the total number of rows. + offsets: Vec, +} + +impl LogicalBatch { + /// Creates a logical batch from `batches`, which must all have `schema`. + /// + /// # Errors + /// + /// Returns an execution error if a batch has a different schema or the + /// total row count overflows. + pub(crate) fn new(schema: SchemaRef, batches: Vec) -> Result { + if batches.iter().any(|batch| batch.schema() != schema) { + return exec_err!("LogicalBatch input batches must have the same schema"); + } + let segments: Vec = batches + .into_iter() + .filter(|batch| batch.num_rows() > 0) + .collect(); + let mut offsets = Vec::with_capacity(segments.len() + 1); + let mut num_rows: usize = 0; + offsets.push(num_rows); + for segment in &segments { + num_rows = num_rows.checked_add(segment.num_rows()).ok_or_else(|| { + exec_datafusion_err!("LogicalBatch total row count exceeds usize::MAX") + })?; + offsets.push(num_rows); + } + Ok(Self { + schema, + segments, + offsets, + }) + } + + /// Creates a logical batch with no rows. + pub(crate) fn new_empty(schema: SchemaRef) -> Self { + Self { + schema, + segments: vec![], + offsets: vec![0], + } + } + + pub(crate) fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + /// Total number of rows across all segments. + pub(crate) fn num_rows(&self) -> usize { + // `offsets` always holds at least the leading 0 + self.offsets[self.offsets.len() - 1] + } + + /// Returns the underlying batches, in row order. + pub(crate) fn into_batches(self) -> Vec { + self.segments + } + + /// The row at global index `index`. + /// + /// # Errors + /// + /// Returns an execution error if `index >= num_rows()`. + pub(crate) fn row(&self, index: usize) -> Result> { + if index >= self.num_rows() { + return exec_err!( + "row index {index} out of bounds for a batch of {} rows", + self.num_rows() + ); + } + let segment = self.segment_of(index, 0); + Ok(BatchRow { + batch: &self.segments[segment], + index: index - self.offsets[segment], + }) + } + + /// Resolves global row indices for [`Self::take_column`]. + /// + /// This plays the role of the index array given to [`take`] on a plain + /// batch: resolve the indices once, then gather as many columns as + /// needed with them. Rows may repeat or appear in any order; gathers + /// preserve that order. + /// + /// # Errors + /// + /// Returns an execution error if any index is `>= num_rows()` or its + /// index within the segment cannot be represented as a `u32`. + pub(crate) fn row_indices( + &self, + rows: impl IntoIterator, + ) -> Result { + let rows = rows.into_iter(); + let mut groups = Vec::new(); + let mut indices = Vec::with_capacity(rows.size_hint().0); + // Consecutive rows usually sit in the same segment, so retry the last + // segment before searching + let mut segment = 0; + for row in rows { + if row >= self.num_rows() { + return exec_err!( + "row index {row} out of bounds for a batch of {} rows", + self.num_rows() + ); + } + let next_segment = self.segment_of(row, segment); + // Start a new group whenever the segment changes. Keeping separate + // groups for repeated visits to a segment preserves output order. + if next_segment != segment && !indices.is_empty() { + groups.push(Single { + segment, + indices: UInt32Array::from(std::mem::take(&mut indices)), + }); + } + segment = next_segment; + let index = u32::try_from(row - self.offsets[segment]).map_err(|_| { + exec_datafusion_err!( + "row index {row} within segment {segment} exceeds u32::MAX" + ) + })?; + indices.push(index); + } + + if !indices.is_empty() { + groups.push(Single { + segment, + indices: UInt32Array::from(indices), + }); + } + + Ok(RowIndices(match groups.len() { + 0 => IndicesVariant::Empty, + 1 => IndicesVariant::Single(groups.pop().unwrap()), + _ => IndicesVariant::Multi(groups), + })) + } + + /// Gathers the rows selected by `indices` from column `column`, like + /// [`take`] on the column of a plain batch. + /// + /// # Errors + /// + /// Returns an execution error if the column or resolved row indices + /// are out of bounds for this batch. + pub(crate) fn take_column( + &self, + column: usize, + indices: &RowIndices, + ) -> Result { + let field = self.schema.fields().get(column).ok_or_else(|| { + exec_datafusion_err!( + "column index {column} out of bounds for a batch of {} columns", + self.schema.fields().len() + ) + })?; + let take_single = |single: &Single| -> Result { + let batch = self.segments.get(single.segment).ok_or_else(|| { + exec_datafusion_err!( + "segment index {} out of bounds for a batch of {} segments", + single.segment, + self.segments.len() + ) + })?; + let values = batch.columns().get(column).ok_or_else(|| { + exec_datafusion_err!( + "column index {column} out of bounds for a batch of {} columns", + batch.num_columns() + ) + })?; + take( + values.as_ref(), + &single.indices, + Some(TakeOptions { check_bounds: true }), + ) + .map_err(|error| exec_datafusion_err!("Failed to gather rows: {error}")) + }; + + Ok(match &indices.0 { + IndicesVariant::Empty => new_empty_array(field.data_type()), + IndicesVariant::Single(single) => take_single(single)?, + IndicesVariant::Multi(groups) => { + let arrays = groups + .iter() + .map(take_single) + .collect::, _>>()?; + let arrays: Vec<&dyn Array> = + arrays.iter().map(|array| array.as_ref()).collect(); + concat(&arrays)? + } + }) + } + + /// Returns rows `offset..offset + length` as one batch. + /// + /// This is zero-copy when the range lies within a single segment, and + /// copies the rows into a new batch otherwise. + /// + /// # Errors + /// + /// Returns an execution error if the range is out of bounds or + /// `offset + length` overflows. + pub(crate) fn slice(&self, offset: usize, length: usize) -> Result { + let end = offset + .checked_add(length) + .filter(|&end| end <= self.num_rows()) + .ok_or_else(|| { + exec_datafusion_err!( + "slice with offset {offset} and length {length} out of bounds for a batch of {} rows", + self.num_rows() + ) + })?; + if length == 0 { + return Ok(RecordBatch::new_empty(self.schema())); + } + + let first = self.segment_of(offset, 0); + let last = self.segment_of(end - 1, first); + let head = &self.segments[first]; + let head_offset = offset - self.offsets[first]; + if first == last { + return Ok(head.slice(head_offset, length)); + } + + // The range spans several segments: only the first and the last one + // need trimming, the ones in between are taken whole. + let head = head.slice(head_offset, head.num_rows() - head_offset); + let tail = self.segments[last].slice(0, end - self.offsets[last]); + let pieces = std::iter::once(&head) + .chain(&self.segments[first + 1..last]) + .chain(std::iter::once(&tail)); + Ok(concat_batches(&self.schema, pieces)?) + } + + /// Index of the segment holding global row `row`, which must be in + /// bounds. `hint` is checked before falling back to a binary search. + fn segment_of(&self, row: usize, hint: usize) -> usize { + debug_assert!(row < self.num_rows()); + if hint + 1 < self.offsets.len() + && (self.offsets[hint]..self.offsets[hint + 1]).contains(&row) + { + return hint; + } + // `offsets[0] == 0 <= row`, so at least one offset precedes the row + self.offsets.partition_point(|&start| start <= row) - 1 + } +} + +impl From for LogicalBatch { + fn from(batch: RecordBatch) -> Self { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Self::new_empty(batch.schema()); + } + Self { + schema: batch.schema(), + segments: vec![batch], + offsets: vec![0, num_rows], + } + } +} + +/// One row of a [`LogicalBatch`], addressed by the [`RecordBatch`] holding +/// it and the row's index within that batch. Obtained from +/// [`LogicalBatch::row`]. +#[derive(Debug, Clone, Copy)] +pub(crate) struct BatchRow<'a> { + batch: &'a RecordBatch, + index: usize, +} + +impl<'a> BatchRow<'a> { + /// Row `index` of `batch`. + /// + /// # Errors + /// + /// Returns an execution error if `index >= batch.num_rows()`. + pub(crate) fn new(batch: &'a RecordBatch, index: usize) -> Result { + if index >= batch.num_rows() { + return exec_err!( + "row index {index} out of bounds for a batch of {} rows", + batch.num_rows() + ); + } + Ok(Self { batch, index }) + } + + /// The array holding the row's value for `column`; the value is at + /// [`Self::index`]. + /// + /// # Errors + /// + /// Returns an execution error if the column index is out of bounds. + pub(crate) fn column(&self, column: usize) -> Result<&'a ArrayRef> { + self.batch.columns().get(column).ok_or_else(|| { + exec_datafusion_err!( + "column index {column} out of bounds for a batch of {} columns", + self.batch.num_columns() + ) + }) + } + + /// The row's index within the arrays returned by [`Self::column`]. + pub(crate) fn index(&self) -> usize { + self.index + } +} + +/// Global row indices of a [`LogicalBatch`], resolved to the segments that +/// hold them. Built by [`LogicalBatch::row_indices`] and consumed by +/// [`LogicalBatch::take_column`]. +#[derive(Debug)] +pub(crate) struct RowIndices(IndicesVariant); + +impl RowIndices { + /// Number of rows selected. + pub(crate) fn len(&self) -> usize { + match &self.0 { + IndicesVariant::Empty => 0, + IndicesVariant::Single(single) => single.indices.len(), + IndicesVariant::Multi(groups) => { + groups.iter().map(|single| single.indices.len()).sum() + } + } + } +} + +/// Row indices grouped by the segments from which they will be gathered. +#[derive(Debug)] +enum IndicesVariant { + /// No rows selected + Empty, + /// Every row lives in the same segment, so a gather is a plain [`take`]. + /// + /// ```text + /// Segment 0: [10, 20, 30] + /// Global rows: [ 0, 1, 2] + /// Selected rows: [1, 2] + /// + /// Single { segment: 0, indices: [1, 2] } + /// Output: [20, 30] + /// ``` + Single(Single), + /// Rows are gathered with [`take`] for each group, then [`concat()`] combines + /// the results in group order. A segment may appear in several groups. + /// + /// ```text + /// Segment 0: [10, 20] (global rows 0, 1) + /// Segment 1: [30, 40] (global rows 2, 3) + /// Selected rows: [1, 2] + /// + /// [ + /// Single { segment: 0, indices: [1] }, // take [20] + /// Single { segment: 1, indices: [0] }, // take [30] + /// ] + /// Output: [20, 30] + /// ``` + Multi(Vec), +} + +/// One consecutive group of output rows gathered from the same segment. +#[derive(Debug)] +struct Single { + /// Index into [`LogicalBatch::segments`]. + segment: usize, + /// Row indices within that segment, in output order. + indices: UInt32Array, +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::ops::Range; + + use arrow::array::{AsArray, Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use datafusion_common::DataFusionError; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ])) + } + + /// One batch holding `rows` with `a = i` and `b = "s{i}"` (`b` is NULL for + /// odd `i`). + fn batch(rows: Range) -> RecordBatch { + let a = Int32Array::from_iter_values(rows.clone()); + let b = + StringArray::from_iter(rows.map(|i| (i % 2 == 0).then(|| format!("s{i}")))); + RecordBatch::try_new(schema(), vec![Arc::new(a), Arc::new(b)]).unwrap() + } + + /// Segments of 3, 2 and 4 rows, plus empty ones that must be ignored, + /// alongside the batch they are logically equal to. + fn logical_batch() -> (LogicalBatch, RecordBatch) { + let segments = vec![ + RecordBatch::new_empty(schema()), + batch(0..3), + batch(3..5), + RecordBatch::new_empty(schema()), + batch(5..9), + RecordBatch::new_empty(schema()), + ]; + let expected = concat_batches(&schema(), &segments).unwrap(); + (LogicalBatch::new(schema(), segments).unwrap(), expected) + } + + #[test] + fn construction() { + let (logical_batch, expected) = logical_batch(); + assert_eq!(logical_batch.num_rows(), 9); + assert_eq!(logical_batch.offsets, vec![0, 3, 5, 9]); + assert_eq!(logical_batch.schema(), schema()); + let segments = logical_batch.into_batches(); + assert_eq!(segments.len(), 3); + assert_eq!(concat_batches(&schema(), &segments).unwrap(), expected); + + let empty = LogicalBatch::new_empty(schema()); + assert_eq!(empty.num_rows(), 0); + assert_eq!(empty.slice(0, 0).unwrap(), RecordBatch::new_empty(schema())); + assert_eq!( + empty + .take_column(0, &empty.row_indices([]).unwrap()) + .unwrap() + .len(), + 0 + ); + + let single = LogicalBatch::from(batch(0..4)); + assert_eq!(single.num_rows(), 4); + assert_eq!(single.slice(0, 4).unwrap(), batch(0..4)); + } + + #[test] + fn row_lookup() { + let (logical_batch, _) = logical_batch(); + for i in 0..9 { + let row = logical_batch.row(i).unwrap(); + let a = row.column(0).unwrap().as_primitive::(); + assert_eq!(a.value(row.index()), i as i32, "row {i}"); + } + // Rows are addressed within their own segment + assert_eq!(logical_batch.row(0).unwrap().index(), 0); + assert_eq!(logical_batch.row(2).unwrap().index(), 2); + assert_eq!(logical_batch.row(3).unwrap().index(), 0); + assert_eq!(logical_batch.row(4).unwrap().index(), 1); + assert_eq!(logical_batch.row(5).unwrap().index(), 0); + assert_eq!(logical_batch.row(8).unwrap().index(), 3); + } + + #[test] + fn row_lookup_out_of_bounds() { + let (logical_batch, _) = logical_batch(); + for index in [9, usize::MAX] { + assert!(matches!( + logical_batch.row(index), + Err(DataFusionError::Execution(_)) + )); + } + assert!(matches!( + LogicalBatch::new_empty(schema()).row(0), + Err(DataFusionError::Execution(_)) + )); + } + + #[test] + fn row_indices_out_of_bounds() { + let (logical_batch, _) = logical_batch(); + for rows in [[0, 9], [3, usize::MAX]] { + assert!(matches!( + logical_batch.row_indices(rows), + Err(DataFusionError::Execution(_)) + )); + } + assert!(matches!( + LogicalBatch::new_empty(schema()).row_indices([0]), + Err(DataFusionError::Execution(_)) + )); + } + + #[test] + fn batch_row_out_of_bounds() { + let batch = batch(0..3); + for index in [3, usize::MAX] { + assert!(matches!( + BatchRow::new(&batch, index), + Err(DataFusionError::Execution(_)) + )); + } + let row = BatchRow::new(&batch, 2).unwrap(); + assert_eq!(row.index(), 2); + assert_eq!(row.column(0).unwrap(), batch.column(0)); + assert!(matches!(row.column(2), Err(DataFusionError::Execution(_)))); + assert!(matches!( + BatchRow::new(&RecordBatch::new_empty(schema()), 0), + Err(DataFusionError::Execution(_)) + )); + } + + /// Gathers `rows` from both columns through the logical batch and + /// through `take` on the equivalent plain batch, and asserts they agree. + fn assert_take(logical_batch: &LogicalBatch, expected: &RecordBatch, rows: &[usize]) { + let indices = logical_batch.row_indices(rows.iter().copied()).unwrap(); + assert_eq!(indices.len(), rows.len()); + let plain_indices = UInt32Array::from_iter_values(rows.iter().map(|i| *i as u32)); + for column in 0..2 { + let actual = logical_batch.take_column(column, &indices).unwrap(); + let expected = take(expected.column(column), &plain_indices, None).unwrap(); + assert_eq!(&actual, &expected, "column {column}, rows {rows:?}"); + } + } + + #[test] + fn take_within_one_segment() { + let (logical_batch, expected) = logical_batch(); + let indices = logical_batch.row_indices([3, 4, 4, 3]).unwrap(); + assert!(matches!( + indices.0, + IndicesVariant::Single(Single { segment: 1, .. }) + )); + assert_take(&logical_batch, &expected, &[3, 4, 4, 3]); + assert_take(&logical_batch, &expected, &[8, 5, 6, 7]); + assert_take(&logical_batch, &expected, &[0, 0, 0]); + } + + #[test] + fn take_across_segments() { + let (logical_batch, expected) = logical_batch(); + let indices = logical_batch.row_indices([2, 3, 2, 3]).unwrap(); + assert!(matches!(indices.0, IndicesVariant::Multi(_))); + assert_take(&logical_batch, &expected, &[2, 3, 2, 3]); + assert_take(&logical_batch, &expected, &[8, 0, 4, 5]); + assert_take(&logical_batch, &expected, &(0..9).collect::>()); + // Repeated visits to a segment must stay in selection order. + assert_take(&logical_batch, &expected, &[3, 4, 4, 7, 5, 3]); + // Selections may skip segments. + assert_take(&logical_batch, &expected, &[0, 8]); + } + + #[test] + fn take_nothing() { + let (logical_batch, _) = logical_batch(); + let indices = logical_batch.row_indices([]).unwrap(); + assert_eq!(indices.len(), 0); + let actual = logical_batch.take_column(1, &indices).unwrap(); + assert_eq!(actual.len(), 0); + assert_eq!(actual.data_type(), &DataType::Utf8); + } + + #[test] + fn take_out_of_bounds() { + let (logical_batch, _) = logical_batch(); + // Check empty, single-segment, and multi-segment selections. + for rows in [vec![], vec![0], vec![0, 3]] { + let indices = logical_batch.row_indices(rows).unwrap(); + assert!(matches!( + logical_batch.take_column(2, &indices), + Err(DataFusionError::Execution(_)) + )); + } + + // Indices resolved against a different batch may refer to missing + // segments or rows. Neither case should panic inside Arrow's take. + let smaller = LogicalBatch::from(batch(0..1)); + for row in [2, 8] { + let indices = logical_batch.row_indices([row]).unwrap(); + assert!(matches!( + smaller.take_column(0, &indices), + Err(DataFusionError::Execution(_)) + )); + } + } + + #[test] + fn slice_within_one_segment() { + let (logical_batch, expected) = logical_batch(); + for (offset, length) in [ + (0, 3), + (1, 2), + (3, 2), + (5, 4), + (6, 1), + (8, 1), + (4, 0), + (9, 0), + ] { + let actual = logical_batch.slice(offset, length).unwrap(); + assert_eq!( + actual, + expected.slice(offset, length), + "{offset}..{}", + offset + length + ); + } + } + + #[test] + fn slice_across_segments() { + let (logical_batch, expected) = logical_batch(); + for (offset, length) in [(0, 9), (2, 2), (2, 4), (1, 7), (4, 5), (3, 6)] { + let actual = logical_batch.slice(offset, length).unwrap(); + assert_eq!( + actual, + expected.slice(offset, length), + "{offset}..{}", + offset + length + ); + } + } + + #[test] + fn slice_out_of_bounds() { + let (logical_batch, _) = logical_batch(); + for (offset, length) in [ + (7, 3), + (10, 0), + (usize::MAX, 1), + (1, usize::MAX), + (usize::MAX, usize::MAX), + ] { + assert!(matches!( + logical_batch.slice(offset, length), + Err(DataFusionError::Execution(_)) + )); + } + } + + #[test] + fn construction_errors() { + let empty_schema = Arc::new(Schema::empty()); + assert!(matches!( + LogicalBatch::new(Arc::clone(&empty_schema), vec![batch(0..1)]), + Err(DataFusionError::Execution(_)) + )); + + // Empty-schema batches can represent these row counts without + // allocating a correspondingly large array. + let batch = |rows| { + RecordBatch::try_new_with_options( + Arc::clone(&empty_schema), + vec![], + &arrow::array::RecordBatchOptions::new().with_row_count(Some(rows)), + ) + .unwrap() + }; + assert!(matches!( + LogicalBatch::new( + Arc::clone(&empty_schema), + vec![batch(usize::MAX), batch(1)] + ), + Err(DataFusionError::Execution(_)) + )); + if let Some(index) = (u32::MAX as usize).checked_add(1) { + let large = LogicalBatch::from(batch(usize::MAX)); + assert!(matches!( + large.row_indices([index]), + Err(DataFusionError::Execution(_)) + )); + } + } + + #[test] + fn empty_schema_counts_rows() { + let schema = Arc::new(Schema::empty()); + let batch = |rows| { + RecordBatch::try_new_with_options( + Arc::clone(&schema), + vec![], + &arrow::array::RecordBatchOptions::new().with_row_count(Some(rows)), + ) + .unwrap() + }; + let logical_batch = + LogicalBatch::new(Arc::clone(&schema), vec![batch(2), batch(3)]).unwrap(); + assert_eq!(logical_batch.num_rows(), 5); + assert_eq!(logical_batch.row(4).unwrap().index(), 2); + assert_eq!(logical_batch.slice(1, 3).unwrap().num_rows(), 3); + assert_eq!(logical_batch.slice(3, 2).unwrap().num_rows(), 2); + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 84de98902d2bf..9268fe3b8b1dd 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -37,6 +37,7 @@ mod asof_join; pub mod chain; mod cross_join; mod hash_join; +mod logical_batch; mod nested_loop_join; mod piecewise_merge_join; #[cfg(feature = "proto")] diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index fd2e0ae21c2ea..e58c150aa6fde 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -30,6 +30,7 @@ use super::utils::{ use crate::common::can_project; use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::joins::SharedBitmapBuilder; +use crate::joins::logical_batch::{BatchRow, LogicalBatch}; use crate::joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, OnceAsync, OnceFut, build_join_schema, check_join_is_valid, estimate_join_statistics, @@ -55,9 +56,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; @@ -128,7 +127,9 @@ use crate::spill::spill_manager::SpillManager; /// /// ## 1. Buffering Left Input /// - The operator eagerly buffers all left-side input batches into memory, -/// util a memory limit is reached. +/// util a memory limit is reached. The batches are kept as they arrive and +/// addressed as one contiguous batch (see `LogicalBatch`), so buffering +/// does not copy them into a merged batch. /// Currently, an out-of-memory error will be thrown if all the left-side input batches /// cannot fit into memory at once. /// In the future, it's possible to make this case finish execution. (see @@ -1084,8 +1085,9 @@ impl EmbeddedProjection for NestedLoopJoinExec { /// Left (build-side) data pub(crate) struct JoinLeftData { - /// Build-side data collected to single batch - batch: RecordBatch, + /// Build-side rows. See [`LogicalBatch`] for details on this layout + /// and why it is used. + batch: LogicalBatch, /// Shared bitmap builder for visited left indices bitmap: SharedBitmapBuilder, /// Counter of running probe-threads, potentially able to update `bitmap` @@ -1099,7 +1101,7 @@ pub(crate) struct JoinLeftData { impl JoinLeftData { pub(crate) fn new( - batch: RecordBatch, + batch: LogicalBatch, bitmap: SharedBitmapBuilder, probe_threads_counter: AtomicUsize, reservation: MemoryReservation, @@ -1112,7 +1114,7 @@ impl JoinLeftData { } } - pub(crate) fn batch(&self) -> &RecordBatch { + pub(crate) fn batch(&self) -> &LogicalBatch { &self.batch } @@ -1187,11 +1189,11 @@ async fn collect_left_input( // polling the child stream above. let build_timer = metrics.build_time.timer(); - let merged_batch = concat_batches(&schema, &batches)?; + let buffered_build_batch = LogicalBatch::new(Arc::clone(&schema), batches)?; // 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 = buffered_build_batch.num_rows(); let buffer_size = n_rows.div_ceil(8); match reservation.try_grow(buffer_size) { Ok(()) => {} @@ -1200,11 +1202,10 @@ async fn collect_left_input( // outside that timer. build_timer.done(); 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], + buffered_build_batch.into_batches(), None, stream, metrics, @@ -1230,7 +1231,7 @@ async fn collect_left_input( }; Ok(LeftLoad::InMemory(Arc::new(JoinLeftData::new( - merged_batch, + buffered_build_batch, Mutex::new(visited_left_side), AtomicUsize::new(probe_threads_count), reservation, @@ -1248,7 +1249,7 @@ fn left_load_from_spill( Some(data) => LeftLoad::Spilled(Arc::new(data)), // No rows means no bitmap either, whatever the join type. None => LeftLoad::InMemory(Arc::new(JoinLeftData::new( - RecordBatch::new_empty(schema), + LogicalBatch::new_empty(schema), Mutex::new(BooleanBufferBuilder::new(0)), AtomicUsize::new(probe_threads_count), reservation, @@ -1694,8 +1695,8 @@ impl FallbackCoordinator { } let _build_timer = build_time.timer(); - let merged_batch = concat_batches(&left_schema, &pending_batches)?; - let n_rows = merged_batch.num_rows(); + let batch = LogicalBatch::new(left_schema, pending_batches)?; + let n_rows = batch.num_rows(); let visited_left_side = if self.with_visited_bitmap { let buffer_size = n_rows.div_ceil(8); reservation.grow(buffer_size); @@ -1722,7 +1723,7 @@ impl FallbackCoordinator { let chunk_reservation = reservation.take(); let data = JoinLeftData::new( - merged_batch, + batch, Mutex::new(visited_left_side), AtomicUsize::new(self.right_partition_count), chunk_reservation, @@ -3022,14 +3023,15 @@ impl NestedLoopJoinStream { // materializes the intermediate batch, and finally applies the join filter // to it. // ----------------------------------------------------------- + let left_batch = left_data.batch(); 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 - 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) - })); + let left_indices = left_batch.row_indices( + (0..l_row_count) + .flat_map(|i| std::iter::repeat_n(l_start_index + i, right_rows)), + )?; let right_indices: UInt32Array = UInt32Array::from_iter_values( (0..l_row_count).flat_map(|_| 0..right_rows as u32), ); @@ -3055,8 +3057,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); - take(col.as_ref(), &left_indices, None)? + left_batch.take_column(column_index.index, &left_indices)? } else { let col = right_batch.column(column_index.index); take(col.as_ref(), &right_indices, None)? @@ -3176,8 +3177,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); - take(col.as_ref(), &left_indices, None)? + left_batch.take_column(column_index.index, &left_indices)? } else { let col = right_batch.column(column_index.index); take(col.as_ref(), &right_indices, None)? @@ -3206,13 +3206,9 @@ impl NestedLoopJoinStream { return Ok(None); } + let left_row = left_data.batch().row(l_index)?; let cur_right_bitmap = if let Some(filter) = &self.join_filter { - apply_filter_to_row_join_batch( - left_data.batch(), - l_index, - right_batch, - filter, - )? + apply_filter_to_row_join_batch(left_row, right_batch, filter)? } else { BooleanArray::from(vec![true; right_row_count]) }; @@ -3240,8 +3236,7 @@ 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_row, right_batch, Some(cur_right_bitmap), &self.column_indices, @@ -3319,8 +3314,8 @@ impl NestedLoopJoinStream { // 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); + let left_batch_sliced = + left_data.batch().slice(start_idx, end_idx - start_idx)?; // Can this be more efficient? let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx); @@ -3458,15 +3453,14 @@ impl NestedLoopJoinStream { // ==== Utilities ==== /// Apply the join filter between: -/// (l_index th row in left buffer) x (right batch) +/// (left_row in left buffer) x (right batch) /// Returns a bitmap, with successfully joined indices set to true fn apply_filter_to_row_join_batch( - left_batch: &RecordBatch, - l_index: usize, + left_row: BatchRow<'_>, right_batch: &RecordBatch, filter: &JoinFilter, ) -> Result { - debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0); + debug_assert!(right_batch.num_rows() != 0); let intermediate_batch = if filter.schema.fields().is_empty() { // If filter is constant (e.g. literal `true`), empty batch can be used @@ -3478,8 +3472,7 @@ fn apply_filter_to_row_join_batch( } else { build_row_join_batch( &filter.schema, - left_batch, - l_index, + left_row, right_batch, None, &filter.column_indices, @@ -3516,21 +3509,20 @@ fn boolean_mask_from_filter(filter_arr: &BooleanArray) -> BooleanArray { /// This function performs the following steps: /// 1. Apply filter to probe-side batch -/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the -/// filtered probe-side batch +/// 2. Broadcast the build row (`build_row`) to the filtered probe-side batch /// 3. Concat them together according to `col_indices`, and return the result /// (None if the result is empty) /// /// Example: -/// build_side_batch: +/// build side batch: /// a /// ---- /// 1 /// 2 /// 3 /// -/// # 0 index element in the build_side_batch (that is `1`) will be used -/// build_side_index: 0 +/// # 0 index row of the build side batch (that is `1`) will be used +/// build_row: row 0 /// /// probe_side_batch: /// b @@ -3563,8 +3555,7 @@ fn boolean_mask_from_filter(filter_arr: &BooleanArray) -> BooleanArray { /// 1 40 fn build_row_join_batch( output_schema: &Schema, - build_side_batch: &RecordBatch, - build_side_index: usize, + build_row: BatchRow<'_>, probe_side_batch: &RecordBatch, probe_side_filter: Option, // See [`NLJStream`] struct's `column_indices` field for more detail @@ -3608,7 +3599,7 @@ fn build_row_join_batch( let array = if column_index.side == build_side { // Broadcast the single build-side row to match the filtered // probe-side batch length - let original_left_array = build_side_batch.column(column_index.index); + let original_left_array = build_row.column(column_index.index)?; // Use `arrow::compute::take` directly for `List(Utf8View)` rather // than going through `ScalarValue::to_array_of_size()`, which @@ -3620,7 +3611,7 @@ fn build_row_join_batch( if field.data_type() == &DataType::Utf8View => { let indices_iter = std::iter::repeat_n( - build_side_index as u64, + build_row.index() as u64, filtered_probe_batch.num_rows(), ); let indices_array = UInt64Array::from_iter_values(indices_iter); @@ -3629,7 +3620,7 @@ fn build_row_join_batch( _ => { let scalar_value = ScalarValue::try_from_array( original_left_array.as_ref(), - build_side_index, + build_row.index(), )?; scalar_value.to_array_of_size(filtered_probe_batch.num_rows())? } @@ -3784,9 +3775,9 @@ fn build_unmatched_batch( .collect::>(), )); let left_null_batch = if nullable_left_schema.fields.is_empty() { - // Left input can be an empty relation, in this case left relation - // won't be used to construct the result batch (i.e. not in `col_indices`) - create_record_batch_with_empty_schema(nullable_left_schema, 0)? + // Keep the placeholder row even when no columns from this + // side are projected, so BatchRow can address row 0. + create_record_batch_with_empty_schema(nullable_left_schema, 1)? } else { RecordBatch::try_new(nullable_left_schema, left_null_columns)? }; @@ -3796,8 +3787,7 @@ fn build_unmatched_batch( build_row_join_batch( output_schema, - &left_null_batch, - 0, + BatchRow::new(&left_null_batch, 0)?, batch, Some(flipped_bitmap), col_indices, @@ -4249,7 +4239,7 @@ pub(crate) mod tests { BooleanBufferBuilder::new(0) }; let chunk = Arc::new(JoinLeftData::new( - left_batch.clone(), + left_batch.clone().into(), Mutex::new(visited), AtomicUsize::new(1), MemoryConsumer::new("NestedLoopJoinFallbackChunk[test]".to_string()) From d523512a78141395a7caed1d7c31bc4c93b38266 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Thu, 17 Sep 2026 16:19:34 +0800 Subject: [PATCH 3/3] review: fix doc --- datafusion/physical-plan/src/joins/logical_batch.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/joins/logical_batch.rs b/datafusion/physical-plan/src/joins/logical_batch.rs index f6bf6762f6e85..133174404a834 100644 --- a/datafusion/physical-plan/src/joins/logical_batch.rs +++ b/datafusion/physical-plan/src/joins/logical_batch.rs @@ -52,10 +52,12 @@ use datafusion_common::{Result, exec_datafusion_err, exec_err}; /// See issue for details: /// - /// -/// # TODO -/// It's named 'logical batch' because it's possible to swap the physical layout -/// and keep the same interface for other usages. For example, segments are aligned -/// at the same size, so it achieves O(1) access speed. +/// # Potential Improvements +/// +/// `LogicalBatch` exposes a logical view of the rows, independent of their +/// physical layout. This allows alternative layouts without changing the +/// interface. For example, fixed-size segments could enable O(1) lookup of +/// the segment containing a given row. #[derive(Debug, Clone)] pub(crate) struct LogicalBatch { schema: SchemaRef,