kosiew commented on code in PR #25542:
URL: https://github.com/apache/datafusion/pull/25542#discussion_r4063105737
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -7235,179 +5914,184 @@ pub(crate) mod tests {
assert_eq!(
pool.reserved(),
0,
- "{join_type}: the coordinator still holds the final chunk's memory
\
- while the plan is alive"
+ "{join_type}: memory is still reserved after every partition
finished"
);
Ok(())
}
- #[tokio::test]
- async fn test_nlj_memory_limited_releases_final_chunk_left_join() ->
Result<()> {
- assert_final_chunk_released(JoinType::Left).await
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_nlj_memory_limited_releases_memory_left_join() -> Result<()>
{
+ assert_memory_released_after_completion(JoinType::Left).await
}
- #[tokio::test]
- async fn test_nlj_memory_limited_releases_final_chunk_full_join() ->
Result<()> {
- assert_final_chunk_released(JoinType::Full).await
- }
-
- /// Run a NLJ across 4 right partitions, collecting every output
- /// partition CONCURRENTLY. This is required for the multi-chunk
- /// coordinator path: a chunk is not released until all partitions
- /// finish probing it, and a partition cannot advance to the next chunk
- /// until the current one is released. Collecting partitions
- /// sequentially would therefore deadlock; concurrent collection mirrors
- /// how partitions actually run under the runtime.
- async fn multi_partition_memory_limited_join_collect_concurrent(
- left: Arc<dyn ExecutionPlan>,
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_nlj_memory_limited_releases_memory_full_join() -> Result<()>
{
+ assert_memory_released_after_completion(JoinType::Full).await
+ }
+
+ /// A two-partition LEFT join over a left side of several chunks, under a
+ /// memory limit that makes it spill.
+ fn dropped_partition_test_plan(
right: Arc<dyn ExecutionPlan>,
- join_type: &JoinType,
- join_filter: Option<JoinFilter>,
- context: Arc<TaskContext>,
- ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> {
- let partition_count = 4;
+ ) -> Result<(Arc<NestedLoopJoinExec>, Arc<TaskContext>)> {
+ let task_ctx = task_ctx_with_memory_limit(50, 1)?;
let right = Arc::new(RepartitionExec::try_new(
right,
- Partitioning::RoundRobinBatch(partition_count),
+ Partitioning::RoundRobinBatch(2),
)?) as Arc<dyn ExecutionPlan>;
-
- let nested_loop_join = Arc::new(NestedLoopJoinExec::try_new(
- left,
+ let plan = Arc::new(NestedLoopJoinExec::try_new(
+ build_left_table_multi_chunk(),
right,
- join_filter,
- join_type,
+ Some(prepare_join_filter()),
+ &JoinType::Left,
None,
)?);
- let columns = columns(&nested_loop_join.schema());
+ Ok((plan, task_ctx))
+ }
+
+ /// The partitions move through the left chunks together, so one that is
Review Comment:
Could we also add dropped-partition survivor coverage for `FULL` and `LEFT
MARK`?
The current drop tests exercise `LEFT`, but these two join types have
additional final-output behavior involving the shared left bitmap together with
right-side unmatched or mark handling. A couple of focused regressions here
would give me more confidence that the new departure behavior is correct across
all of the left-emitting paths.
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -1431,649 +1410,352 @@ pub(crate) enum LeftLoad {
/// The left side fit the memory budget and is buffered as one batch.
InMemory(Arc<JoinLeftData>),
/// 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.
+ /// shares this handle, and the chunks that are read back from the file.
Spilled(Arc<LeftSpillData>),
}
+/// The spill file [`spill_left_input`] wrote, before it is wrapped in a
[`LeftSpillData`].
+struct SpilledLeftFile {
+ spill_manager: SpillManager,
+ spill_file: Arc<dyn SpillFile>,
+ /// Total number of rows written to `spill_file`
+ num_rows: usize,
+}
+
/// The spilled left side, shared by every output partition.
+///
+/// This is the memory-limited counterpart of [`JoinLeftData`]: the rows live
in
+/// a spill file instead of memory, but the visited bitmap and the
probe-threads
+/// counter cover the whole left side and are shared by all partitions in the
+/// same way.
+///
+/// The rows come back one chunk at a time (see [`LeftChunkBarrier`]), while
the
+/// bitmap spans all of them: bits are addressed by the row's position in the
+/// file. Keeping match tracking apart from the chunks is what lets a chunk be
+/// dropped as soon as it has been probed, with the final left rows emitted in
+/// one pass at the very end.
pub(crate) struct LeftSpillData {
/// SpillManager used to read the spill file (has the left schema)
spill_manager: SpillManager,
/// The spill file containing all left-side batches
spill_file: Arc<dyn SpillFile>,
/// Left-side schema
schema: SchemaRef,
+ /// Total number of rows in `spill_file`
+ num_rows: usize,
+ /// The pass over `spill_file` that chunks are read from. Each chunk load
+ /// takes it and hands it back for the load of the following chunk.
+ reader: Arc<Mutex<Option<LeftChunkReader>>>,
+ /// Visited bitmap over every row of `spill_file`. Empty when the join type
+ /// does not need it.
+ visited: SharedBitmapBuilder,
+ /// Counter of partitions that have not finished probing every chunk
+ probe_threads_counter: AtomicUsize,
+ /// Memory reservation for `visited`
+ reservation: MemoryReservation,
}
-/// Per-chunk shared state in the memory-limited fallback path.
-///
-/// Each chunk's `JoinLeftData` is loaded once by a "leader" partition and
-/// shared (via `Arc`) with every right-side output partition. The
-/// `probe_threads_counter` inside the `JoinLeftData` is initialized to
-/// `right_partition_count`, so `report_probe_completed` returns `true`
-/// only when the *last* partition has finished probing the chunk. That
-/// last partition is then responsible for emitting unmatched left rows
-/// for the chunk, mirroring the single-pass path's coordination via
-/// `collect_left_input(..., probe_threads_count)`.
-struct CurrentChunk {
- /// 0-based monotonically increasing chunk index.
- chunk_index: usize,
- /// Shared per-chunk left data. Cloned by every partition that probes
- /// this chunk; the last to call `report_probe_completed` emits
- /// unmatched left rows.
- data: Arc<JoinLeftData>,
- /// True if the left stream was exhausted while loading this chunk —
- /// no further chunks will be produced after it.
- is_last: bool,
-}
+impl LeftSpillData {
+ fn new(
+ spilled: SpilledLeftFile,
+ schema: SchemaRef,
+ with_visited_left_side: bool,
+ probe_threads_count: usize,
+ reservation: MemoryReservation,
+ ) -> Self {
+ let SpilledLeftFile {
+ spill_manager,
+ spill_file,
+ num_rows,
+ } = spilled;
+ let visited = if with_visited_left_side {
+ // Use infallible `grow`: one bit per row is all that stays in
+ // memory, and the fallback path has no other recourse.
+ reservation.grow(num_rows.div_ceil(8));
+ let mut buffer = BooleanBufferBuilder::new(num_rows);
+ buffer.append_n(num_rows, false);
+ buffer
+ } else {
+ BooleanBufferBuilder::new(0)
+ };
+ Self {
+ spill_manager,
+ spill_file,
+ schema,
+ num_rows,
+ reader: Arc::new(Mutex::new(None)),
+ visited: Mutex::new(visited),
+ probe_threads_counter: AtomicUsize::new(probe_threads_count),
+ reservation,
+ }
+ }
-/// Inner state of [`FallbackCoordinator`], guarded by a synchronous mutex.
-///
-/// Synchronous because cancellation and chunk release have to complete without
-/// another poll or await -- cancellation runs from `Drop`, which has neither
--
-/// rather than depending on a future a dropped stream would take with it. No
-/// critical section awaits: the one slow operation, reading a chunk, runs
after
-/// the guard is released.
-struct FallbackCoordinatorInner {
- /// Reservation the leader borrows to bound one chunk load.
- ///
- /// On a successful load `load_one_chunk` moves the accounted bytes into
the
- /// chunk's `JoinLeftData` with `take()`, so the accounting follows the
data
- /// rather than staying with this slot. Lazily registered by the first
- /// leader, once a runtime context is available.
- reservation: Option<MemoryReservation>,
- /// The shared left spill stream from which chunks are read. Owned by
- /// the coordinator so only one partition reads it at a time.
- left_stream: Option<SendableRecordBatchStream>,
- /// One batch carried over from the previous chunk's load: when
- /// reservation `try_grow` failed for chunk N, the offending batch is
- /// recorded here and becomes the first batch of chunk N+1.
- carryover: Option<RecordBatch>,
- /// True once the left spill stream has produced `None`.
- left_exhausted: bool,
- /// Index of the next chunk to be loaded.
- next_chunk_index: usize,
- /// The currently-loaded chunk, or `None` if no chunk is currently
- /// loaded (initial state, or the last partition has just released
- /// chunk `next_chunk_index - 1` and the next leader hasn't taken
- /// over yet).
- current: Option<CurrentChunk>,
- /// True while a partition has claimed leader role for the next
- /// chunk and is loading it; prevents two partitions from racing.
- loader_in_flight: bool,
- /// A partition was dropped while still `Pending`, before the shared load
- /// decided whether the left side spills.
- ///
- /// `Pending` is entered by every eligible execution, including those whose
- /// left side ends up fitting in memory -- and those never build a shared
- /// chunk counter, so their right partitions stay independent and a dropped
- /// peer has nothing to coordinate. Cancelling on such a drop would fail a
- /// query that had no fallback at all, so it is only recorded here.
- ///
- /// Paired with `coordination_started`: whichever of the two happens second
- /// performs the cancellation, so the drop is honoured whether it precedes
or
- /// follows the execution becoming coordinated. A bool suffices -- one lost
- /// partition is enough to cancel, and nothing reads a count.
- pending_drop: bool,
- /// Set once any partition has entered the coordinated path.
- ///
- /// Remembered rather than checked in the moment, because a `Pending` drop
can
- /// arrive after a peer is already coordinating; that drop must cancel, and
- /// without this flag there would be nobody left to notice.
- coordination_started: bool,
- /// Set when a stream is dropped before finishing, which cancels the whole
- /// coordinated fallback.
- ///
- /// The partitions of a coordinated fallback are not independent: chunk
- /// advancement requires every one of them to report, so a partition that
- /// disappears mid-probe would otherwise leave the survivors waiting on a
- /// release nobody will ever make. Once set, chunk state is dropped,
waiters
- /// are woken with an error, and a loader that is still reading must
discard
- /// its result instead of publishing it.
- cancelled: bool,
-}
+ /// Open a new pass over the spilled left rows
+ fn open_pass(&self) -> Result<SendableRecordBatchStream> {
+ self.spill_manager
+ .read_spill_as_stream(Arc::clone(&self.spill_file), None)
+ }
-/// Plan-level shared coordinator for the memory-limited fallback path.
-///
-/// All right-side output partitions share one of these. It serializes
-/// access to the left spill stream (so each chunk is read exactly once),
-/// publishes the loaded chunk as an `Arc<JoinLeftData>` for every
-/// partition to clone, and uses a `Notify` so partitions waiting for the
-/// next chunk can sleep without busy-looping.
-pub(crate) struct FallbackCoordinator {
- /// Number of right-side partitions; equals the
- /// `probe_threads_counter` initial value for each chunk.
- right_partition_count: usize,
- /// Whether `JoinLeftData` should carry a left visited bitmap (for
- /// join types that emit unmatched left rows in the final output).
- with_visited_bitmap: bool,
- inner: Mutex<FallbackCoordinatorInner>,
- /// Notified when a new chunk becomes available, when the left stream
- /// is exhausted, or when a chunk is released.
- notify: tokio::sync::Notify,
- /// Broadcast signalled when the fallback is cancelled.
- ///
- /// This carries no state of its own -- `cancelled` is what persists. A
- /// delivered broadcast is itself sufficient to establish cancellation;
- /// observers read the flag before awaiting, so neither alone is relied on.
- /// Kept separate from `notify`
- /// because cancellation has to reach tasks that are not waiting on chunk
- /// progress at all: a stream parked on its right input, or a loader parked
- /// on a spill read. Waiters enable their `Notified` before reading
- /// `cancelled`, so a cancellation landing between those two steps is
- /// delivered rather than lost.
- cancel_notify: tokio::sync::Notify,
- /// Test seam that reproduces one cancellation interleaving
deterministically.
+ /// Load the next chunk, accounting for it in `reservation`.
///
- /// Set to `1` to make the leader cancel after claiming the load but before
- /// it registers its cancellation watcher. A cancellation lost in that
window
- /// strands the loader itself -- other observers may already be returning
- /// errors -- which is what the paired test checks. Consumed when it fires,
- /// so one store arms it once.
- #[cfg(test)]
- cancel_at_leader_claim: AtomicUsize,
-}
-
-impl FallbackCoordinator {
- fn new(right_partition_count: usize, with_visited_bitmap: bool) -> Self {
- Self {
- right_partition_count,
- with_visited_bitmap,
- inner: Mutex::new(FallbackCoordinatorInner {
- reservation: None,
- left_stream: None,
- carryover: None,
- left_exhausted: false,
- next_chunk_index: 0,
- current: None,
- loader_in_flight: false,
- pending_drop: false,
- coordination_started: false,
- cancelled: false,
- }),
- notify: tokio::sync::Notify::new(),
- cancel_notify: tokio::sync::Notify::new(),
- #[cfg(test)]
- cancel_at_leader_claim: AtomicUsize::new(0),
- }
+ /// The load is a shared future, the way the whole left side is a shared
+ /// [`OnceFut`]: every partition waiting for the chunk holds a clone, and
+ /// whichever of them is polled drives it. So it does not matter which
+ /// partition started the load, or whether that one is still around.
+ fn load_chunk(
+ &self,
+ reservation: MemoryReservation,
+ build_time: &Time,
+ ) -> LeftChunkFut {
+ load_left_chunk(
+ Arc::clone(&self.reader),
+ self.spill_manager.clone(),
+ Arc::clone(&self.spill_file),
+ Arc::clone(&self.schema),
+ reservation,
+ build_time.clone(),
+ )
+ .map(|chunk| chunk.map(Arc::new).map_err(Arc::new))
+ .boxed()
+ .shared()
}
- /// After the last partition finishes processing chunk
- /// `released_chunk_index`, drop the slot so the next leader can
- /// load chunk `released_chunk_index + 1`.
- fn release_chunk(self: &Arc<Self>, released_chunk_index: usize) {
+ /// Record the matches of a finished chunk, whose first row is the
+ /// `row_offset`-th row of the spill file.
+ fn merge_visited(&self, row_offset: usize, chunk_visited:
&BooleanBufferBuilder) {
+ let mut visited = self.visited.lock();
+ for idx in BitIndexIterator::new(chunk_visited.as_slice(), 0,
chunk_visited.len())
{
- let mut inner = self.inner.lock();
- if let Some(cur) = &inner.current
- && cur.chunk_index == released_chunk_index
- {
- inner.current = None;
- inner.next_chunk_index = released_chunk_index + 1;
- }
+ visited.set_bit(row_offset + idx, true);
}
- // Always notify: waiters may be blocked because they couldn't
- // become leader while a previous chunk was current.
- self.notify.notify_waiters();
}
- /// True once a partition was dropped unfinished, cancelling the fallback.
- ///
- /// Production code observes cancellation through `cancellation_watcher`,
so
- /// that a waker is registered; this plain read is for assertions only.
- #[cfg(test)]
- fn is_cancelled(&self) -> bool {
- self.inner.lock().cancelled
+ /// Decrements counter of running threads, and returns `true`
+ /// if caller is the last running thread
+ fn report_probe_completed(&self) -> bool {
+ self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
}
- /// A future that resolves when the fallback is cancelled.
- ///
- /// Registration happens on the future's **first poll**, not at
construction:
- /// that poll enables the `Notified` and then reads `cancelled`, so
whichever
- /// happens first is observed. Callers must therefore poll it, not merely
hold
- /// it.
+ /// Take the visited bitmap for the final emission. Only the last running
+ /// thread may call this, after which the bitmap is complete.
///
- /// Callers that park on something other than chunk progress -- a stream
- /// waiting on its right input, for instance -- need one of these polled
- /// alongside their own work, or a peer's cancellation never reaches their
- /// waker.
- fn cancellation_watcher(self: &Arc<Self>) -> BoxFuture<'static, ()> {
- let coordinator = Arc::clone(self);
- async move {
- let notified = coordinator.cancel_notify.notified();
- let mut notified = std::pin::pin!(notified);
- notified.as_mut().enable();
- if coordinator.inner.lock().cancelled {
- return;
- }
- notified.await;
- }
- .boxed()
+ /// The bitmap's memory is no longer accounted for here afterwards, so that
+ /// a plan that outlives its execution does not keep it reserved. The
+ /// caller accounts for the returned buffer instead.
+ fn take_visited(&self) -> BooleanBuffer {
Review Comment:
I think there is still a memory-lifetime issue here when one of the streams
is dropped early.
`take_visited()` is currently the only path that releases the global bitmap
reservation before `LeftSpillData` itself is dropped. However, an unfinished
stream deliberately does not decrement `probe_threads_counter`; its `Drop`
implementation only calls `LeftChunkBarrier::depart()`.
That means the remaining streams can finish probing, but none of them can
become the unmatched-left emitter. As a result, `take_visited()` is never
called and the bitmap reservation remains attached to `LeftSpillData` until the
plan is dropped.
The new dropped-partition test appears to encode this behavior by explicitly
dropping `plan` before asserting that the pool reservation reaches zero. For a
retained or cached physical plan, that means memory from a partially consumed
query can remain reserved after execution and potentially cause later queries
to hit `ResourcesExhausted`.
Could we track incomplete execution separately from normal probe completion,
and release or discard the global bitmap once the surviving streams have
finished and we know there will be no unmatched-left emitter? It would also be
good to add a regression test that asserts the memory reservation reaches zero
while the plan itself is still alive.
--
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]
