EmilyMatt commented on code in PR #26366:
URL: https://github.com/apache/datafusion/pull/26366#discussion_r4283962456
##########
datafusion/physical-plan/src/sorts/stream.rs:
##########
@@ -130,22 +133,25 @@ impl ReusableRows {
/// reservation is free for the same reason. A pool that cannot cover the
buffer
/// fails the query here, as it did before the buffer was cached at all.
fn save(&mut self, stream_idx: usize, rows: &Arc<Rows>) -> Result<()> {
- self.inner[stream_idx] = Some(Arc::clone(rows));
- let retained = self.retained_size();
+ let size = rows.size();
+ let retained = self.reservation.size() -
self.reserved_sizes[stream_idx] + size;
debug_assert!(retained >= self.reservation.size());
if let Err(e) = self.reservation.try_resize(retained) {
- self.inner[stream_idx] = None;
- self.reservation.try_resize(self.retained_size())?;
+ self.release(stream_idx);
return Err(e);
}
+ self.inner[stream_idx] = Some(Arc::clone(rows));
+ self.reserved_sizes[stream_idx] = size;
Ok(())
}
// drop whatever a finished stream was holding
fn release(&mut self, stream_idx: usize) {
debug_assert!(self.reservation.size() >= self.retained_size());
- if let Some(rows) = self.inner[stream_idx].take() {
- self.reservation.shrink(rows.size());
+ self.inner[stream_idx] = None;
+ let reserved = mem::take(&mut self.reserved_sizes[stream_idx]);
+ if reserved != 0 {
+ self.reservation.shrink(reserved);
Review Comment:
Personally I always prefer an accumulative sizing, as in, calculate the
total size that is held right now(after setting this stream to 0), then do a
try_resize, it both avoids bugs and mis-accounting, and also handles the != 0
part on its own
--
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]
