geoffreyclaude commented on code in PR #26366:
URL: https://github.com/apache/datafusion/pull/26366#discussion_r4284201216
##########
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:
That makes sense! How does the following look to you?
```rust
fn release(&mut self, stream_idx: usize) {
debug_assert!(self.reservation.size() >= self.retained_size());
self.inner[stream_idx] = None;
let released = mem::take(&mut self.reserved_sizes[stream_idx]);
self.reservation.resize(self.reservation.size() - released);
}
```
This handles the zero-change case too (I used `resize` rather than
`try_resize` since we’re only releasing memory).
--
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]
