Skip to content

Commit 24ef31c

Browse files
fix: account retained merge rows incrementally
Keep each stream charged while its row buffer is taken for reuse, and update only its contribution when the buffer is saved or released. This avoids rescanning every merge input for each batch. Preserve reservation failure cleanup and cover rejected growth while another stream remains retained.
1 parent 90371f9 commit 24ef31c

1 file changed

Lines changed: 46 additions & 7 deletions

File tree

‎datafusion/physical-plan/src/sorts/stream.rs‎

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@ impl FusedStreams {
103103
#[derive(Debug)]
104104
struct ReusableRows {
105105
inner: Vec<Option<Arc<Rows>>>,
106+
/// Each stream's charge stays reserved while its buffer is taken for reuse.
107+
/// The sum of these sizes equals `reservation.size()`.
108+
reserved_sizes: Vec<usize>,
106109
reservation: MemoryReservation,
107110
}
108111

@@ -130,22 +133,25 @@ impl ReusableRows {
130133
/// reservation is free for the same reason. A pool that cannot cover the buffer
131134
/// fails the query here, as it did before the buffer was cached at all.
132135
fn save(&mut self, stream_idx: usize, rows: &Arc<Rows>) -> Result<()> {
133-
self.inner[stream_idx] = Some(Arc::clone(rows));
134-
let retained = self.retained_size();
136+
let size = rows.size();
137+
let retained = self.reservation.size() - self.reserved_sizes[stream_idx] + size;
135138
debug_assert!(retained >= self.reservation.size());
136139
if let Err(e) = self.reservation.try_resize(retained) {
137-
self.inner[stream_idx] = None;
138-
self.reservation.try_resize(self.retained_size())?;
140+
self.release(stream_idx);
139141
return Err(e);
140142
}
143+
self.inner[stream_idx] = Some(Arc::clone(rows));
144+
self.reserved_sizes[stream_idx] = size;
141145
Ok(())
142146
}
143147

144148
// drop whatever a finished stream was holding
145149
fn release(&mut self, stream_idx: usize) {
146150
debug_assert!(self.reservation.size() >= self.retained_size());
147-
if let Some(rows) = self.inner[stream_idx].take() {
148-
self.reservation.shrink(rows.size());
151+
self.inner[stream_idx] = None;
152+
let reserved = mem::take(&mut self.reserved_sizes[stream_idx]);
153+
if reserved != 0 {
154+
self.reservation.shrink(reserved);
149155
}
150156
}
151157

@@ -200,6 +206,7 @@ impl RowCursorStream {
200206
streams: FusedStreams(streams),
201207
rows: ReusableRows {
202208
inner: vec![None; stream_count],
209+
reserved_sizes: vec![0; stream_count],
203210
reservation: rows_reservation,
204211
},
205212
})
@@ -442,7 +449,7 @@ mod tests {
442449
use datafusion_common::DataFusionError;
443450
use datafusion_execution::RecordBatchStream;
444451
use datafusion_execution::memory_pool::{
445-
GreedyMemoryPool, MemoryConsumer, MemoryPool,
452+
GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool,
446453
};
447454
use datafusion_physical_expr::expressions::col;
448455
use futures::Stream;
@@ -907,4 +914,36 @@ mod tests {
907914

908915
Ok(())
909916
}
917+
918+
// The buffer has already been taken out of the cache when its reservation fails
919+
// to grow. Cleanup must still release its previous charge.
920+
#[test]
921+
fn failed_buffer_growth_releases_only_its_stream_reservation() -> Result<()> {
922+
let small = (64, 8);
923+
let large = (2048, 128);
924+
let mut harness =
925+
RowCursorHarness::with_shapes(vec![vec![small, large], vec![small]])?;
926+
927+
// Cache stream 1's buffer and record the reservation that must survive.
928+
assert!(harness.poll_and_drop_cursor(1)?);
929+
let expected_remaining = harness.reserved();
930+
assert!(harness.poll_and_drop_cursor(0)?);
931+
932+
// Leave no space for stream 0's larger batch.
933+
let MemoryLimit::Finite(limit) = harness.pool.memory_limit() else {
934+
unreachable!("the harness uses a bounded memory pool");
935+
};
936+
let competitor = MemoryConsumer::new("competitor").register(&harness.pool);
937+
competitor.try_grow(limit - harness.reserved())?;
938+
939+
assert!(matches!(
940+
harness.poll_cursor(0),
941+
Err(DataFusionError::ResourcesExhausted(_))
942+
));
943+
drop(competitor);
944+
945+
// Only stream 1's buffer and the shared converter remain reserved.
946+
assert_eq!(harness.reserved(), expected_remaining);
947+
Ok(())
948+
}
910949
}

0 commit comments

Comments
 (0)