EmilyMatt commented on code in PR #26366:
URL: https://github.com/apache/datafusion/pull/26366#discussion_r4285749410
##########
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:
I think while it's tempting to use the mem::take()'s value, it's more
resilient to just set to 0 manually instead and 1. Still use a try_resize(), to
not rely on the assumption that memory has only shrunk(logic could change and
break that assumption) and 2. Call the functiom that calculates the total size
of the stream(the perf impact of this is negligible), as in the future more
stuff could be added to that function.
I.e., never fully "trust" that the current reservation is correct, and
always strive to fix it
--
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]
