Skip to content

Commit

Permalink
Revert the removal of reservation in HashJoin (#13792)
Browse files Browse the repository at this point in the history
* fix: restore memory reservation in JoinLeftData for accurate memory accounting in HashJoin

This commit reintroduces the `_reservation` field in the `JoinLeftData` structure to ensure proper tracking of memory resources during join operations. The absence of this field could lead to inconsistent memory usage reporting and potential out-of-memory issues as upstream operators increase their memory consumption.

* fmt

Signed-off-by: Jay Zhan <[email protected]>

---------

Signed-off-by: Jay Zhan <[email protected]>
  • Loading branch information
jayzhan-synnada authored Dec 16, 2024
1 parent bd2c975 commit 2176330
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions datafusion/physical-plan/src/joins/hash_join.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,11 @@ struct JoinLeftData {
/// Counter of running probe-threads, potentially
/// able to update `visited_indices_bitmap`
probe_threads_counter: AtomicUsize,
/// We need to keep this field to maintain accurate memory accounting, even though we don't directly use it.
/// Without holding onto this reservation, the recorded memory usage would become inconsistent with actual usage.
/// This could hide potential out-of-memory issues, especially when upstream operators increase their memory consumption.
/// The MemoryReservation ensures proper tracking of memory resources throughout the join operation's lifecycle.
_reservation: MemoryReservation,
}

impl JoinLeftData {
Expand All @@ -99,12 +104,14 @@ impl JoinLeftData {
batch: RecordBatch,
visited_indices_bitmap: SharedBitmapBuilder,
probe_threads_counter: AtomicUsize,
reservation: MemoryReservation,
) -> Self {
Self {
hash_map,
batch,
visited_indices_bitmap,
probe_threads_counter,
_reservation: reservation,
}
}

Expand Down Expand Up @@ -897,6 +904,7 @@ async fn collect_left_input(
single_batch,
Mutex::new(visited_indices_bitmap),
AtomicUsize::new(probe_threads_count),
reservation,
);

Ok(data)
Expand Down

0 comments on commit 2176330

Please sign in to comment.