sunchao commented on code in PR #24820:
URL: https://github.com/apache/datafusion/pull/24820#discussion_r4027943725
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -1286,21 +1374,14 @@ async fn spill_left_input(
let mut spill_file =
spill_manager.create_in_progress_file("NestedLoopJoin left spill")?;
- for batch in buffered {
- if batch.num_rows() > 0 {
- spill_file.append_batch(&batch)?;
- }
+ for batch in batches {
+ spill_file.append_batch(&batch)?;
}
Review Comment:
**[P2] Preserve empty-batch filtering during spill**
Removing the `num_rows() > 0` guard allows an entirely empty build side to
create a spill file. `append_batch` initializes the writer even for a zero-row
batch, so `finish()` returns a file; replay then skips every batch and returns
no chunk. This bypasses the assignment to `active.left_schema`, and final
right-side emission panics with `left_schema must be set`.
I reproduced this against base `c5257f054` and head `16593136d` using three
empty `Int32` build batches, right-side row `7`, and a 97-byte memory pool
(each empty batch accounts for 96 bytes). For a `Right` join, base returns
`(NULL, 7)` while head panics. `Full`, `RightSemi`, `RightAnti`, and
`RightMark` also panic on head while base succeeds. The same `Right`-join
failure reproduces with empty slices retaining 4 MiB buffers and a roughly 4
MiB pool, so it also occurs beyond tiny test budgets.
Please retain the previous empty-batch guard here and add coverage for an
empty build side that exhausts its reservation after buffering an empty batch.
##########
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<AtomicUsize>,
+}
+
+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<Option<RecordBatch>> {
+ 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<RwLock<dyn LazyBatchGenerator>> {
+ 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<dyn PhysicalExpr> = 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> =
+ 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;
Review Comment:
**[P2] Initialize backtrace diagnostics before the allocation baseline**
The baseline is taken before the first reservation failure. With DataFusion
backtraces enabled and debug information retained, that failure initializes
process-wide symbolization caches, which this test counts against the join's 12
MiB threshold. The new test therefore fails even when build-side spilling stays
bounded; the [AMD64 CI
job](https://github.com/apache/datafusion/actions/runs/34962482370/job/104365026384)
reports 14.21 MiB against 12 MiB.
In a controlled copy of this test linked to the current head, the measured
peak is **8.14 MiB** with `RUST_BACKTRACE=0` and **13.75 MiB** with
`RUST_BACKTRACE=1`. Calling and dropping `DataFusionError::get_back_trace()`
before recording the baseline initializes about 5.61 MiB of retained diagnostic
caches; the measured peak then returns to **8.14 MiB** with backtraces still
enabled and the same join execution. Retaining line-table debug information
gives the same failure/pass distinction.
Please warm the backtrace machinery before taking the baseline, or isolate
the measurement in a process with backtraces disabled. The existing 12 MiB
threshold can remain unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
