geoffreyclaude commented on code in PR #26366:
URL: https://github.com/apache/datafusion/pull/26366#discussion_r4292225619
##########
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 tried using the existing `retained_size()` helper in draft #26389: clear
the finished stream and its recorded size, then call
`try_resize(self.retained_size())`:
```rust
// drop whatever a finished stream was holding
fn release(&mut self, stream_idx: usize) -> Result<()> {
debug_assert!(self.reservation.size() >= self.retained_size());
self.inner[stream_idx] = None;
self.reserved_sizes[stream_idx] = 0;
self.reservation.try_resize(self.retained_size())
}
fn retained_size(&self) -> usize {
self.inner.iter().flatten().map(|rows| rows.size()).sum()
}
```
The [benchmark bot
results](https://github.com/apache/datafusion/pull/26389#issuecomment-6202596131)
show the small-batch sort going from **2.6s to 11.5s**, about **4.4× slower**.
Each completed input triggers a scan of all input slots, so this brings back
quadratic accounting work during stream completion. The cost is substantial for
the workload this PR targets.
Would you be ok with keeping `resize()` with the incremental total, but
checking it against the sum of the recorded sizes in a debug assertion? See
latest commit.
--
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]
