jayzhan211 commented on PR #24740:
URL: https://github.com/apache/datafusion/pull/24740#issuecomment-5460002438
### The fix doesn't apply when spilling with a single in-memory batch
`in_mem_sort_stream` is called with `is_output_stream = false` only from
`sort_and_spill_in_mem_batches`, but the three `release_unused()` calls at
`sort.rs:608`, `sort.rs:626` and `sort.rs:635` are unconditional. On the
output
path they're already redundant — `sort()` released at `sort.rs:389` — so they
take effect **only on the spill path**, handing the workspace back to the
execution pool right before it's needed.
Probing the `borrow()` at `sort.rs:775`:
```
spill with 2 in-mem batches: want=950272 got=950272 # works
spill with 1 in-mem batch: want=950272 got=0 # workspace already
released
```
Two consequences on that path: the chunked-output `try_resize` falls back to
the
parent pool, and the `reserve_memory_for_merge()` at the end of the spill
has to
re-acquire the full `sort_spill_reservation_bytes` from the parent — so it
can
still fail with `ResourcesExhausted` under contention, which is what #24739
is
about. `in_mem_batches.len() == 1` at spill time is common: a spill
triggered by
the second batch, or `sort()` flushing a single leftover batch after an
earlier
spill.
Suggested fix — drop the two calls that are dead-or-harmful:
```diff
if self.in_mem_batches.is_empty() {
- self.merge_pool.release_unused();
let empty_stream =
```
```diff
if self.in_mem_batches.len() == 1 {
- self.merge_pool.release_unused();
let batch = self.in_mem_batches.swap_remove(0);
```
The third one (the `sort_in_place_threshold_bytes` concat branch) is
genuinely
different and should stay: `concat_batches` grows `self.reservation`, an
execution-pool consumer that can't borrow merge workspace, so the idle floor
has
to be returned for that growth to have anywhere to come from. Worth a
comment,
since it now reads as inconsistent with the two above:
```diff
if self.reservation.size() < self.sort_in_place_threshold_bytes {
+ // Unlike the paths above, `concat_batches` grows
`self.reservation`,
+ // an execution-pool consumer that cannot borrow merge
workspace.
+ // Return the idle floor so that growth has somewhere to come
from.
self.merge_pool.release_unused();
```
With those changes the probe reads `got=950272` on both spill passes and all
108
`sorts::` tests stay green.
Note that none of the three calls is currently covered — removing all three
leaves the suite at 109/109 green — so this needs a regression test. Roughly
`check_chunked_string_view_workspace` with the pool sized for one batch
instead
of two, so the spill happens while exactly one batch is buffered:
```rust
// ... same Utf8View batches / ordering as
check_chunked_string_view_workspace ...
let input_bytes = get_reserved_bytes_for_record_batch(&batches[0])?;
// Room for exactly one input batch plus the merge workspace.
let capacity = options.sort_spill_reservation_bytes + input_bytes;
// ... build sorter over `pool` ...
sorter.insert_batch(batches[0].clone()).await?;
assert_eq!(pool.reserved(), capacity);
sorter.insert_batch(batches[1].clone()).await?; // spills with one buffered
batch
assert!(sorter.spilled_before());
assert_eq!(sorter.in_mem_batches.len(), 1);
let stream = sorter.sort().await?;
drop(sorter);
let output: Vec<RecordBatch> = stream.try_collect().await?;
assert_eq!(concat_batches(&schema, &output)?.num_rows(), 2 * rows);
assert_released(&pool, &runtime).await;
```
To make it a true regression test it needs an assertion that the borrow
actually
happened rather than just that the sort succeeded — e.g. that
`pool.state.lock().unwrap().denied` doesn't increase across the spill, or
exposing the loan size the way
`test_chunked_sort_returns_live_workspace_loan_on_drop`
does.
--
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]
