Repository navigation
fix: count the largest capacity seen for a shared buffer address - #27097
Conversation
RecordBatchMemoryCounter and RecordBatchMemoryTracker identify a buffer by its start address and size it by Buffer::capacity(). A buffer created by Buffer::from_custom_allocation, which is how arrow::ffi::from_ffi imports buffers, reports the length it was imported with, so zero-copy slices of one batch exported and imported separately share an address but report different capacities. The counter kept the first capacity it saw, so the shared input reservation in ExternalSorter::sort_run_sizes could be smaller than a later run's own size, failing the assertion in sort_run. The tracker added the first capacity and subtracted the last one, which can underflow memory_usage. Record the counted capacity per address, raise it when a larger buffer appears at that address, and release the recorded capacity when the address is no longer retained. Closes apache#27069.
github-actions
Bot
added
common
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
The fix looks right. Each address is now charged the largest capacity seen, so the shared total from sort_run_sizes is always at least any single run's get_record_batch_memory_size. The check in sort_run now holds by construction instead of depending on which slice is counted first.
- The
assert!atsorts/sort.rs:1022still aborts the process in release builds if the accounting ever drifts again. That hits embedders like Comet, where the panic happens inside the JVM. Now that the invariant is proven, consider turning it intoassert_or_internal_err!(or adebug_assert!), here or in a follow-up, so a future mismatch fails the query and leaves the process running.
| /// count per buffer: its capacity is added when it is first counted and | ||
| /// released when its count returns to 0. |
There was a problem hiding this comment.
This is no longer exact. A later count can raise the charge for a buffer.
| /// count per buffer: its capacity is added when it is first counted and | |
| /// released when its count returns to 0. | |
| /// count per buffer: it is charged the largest capacity seen while its count | |
| /// is positive, and released when its count returns to 0. |
| /// Keeps the largest capacity seen per address: FFI-imported slices of one | ||
| /// allocation share an address but report the length they were imported with. | ||
| #[derive(Debug)] | ||
| struct BufferIdSet { |
There was a problem hiding this comment.
With capacities, this is BufferIdMap without counts, and insert repeats the scan, append and promotion logic from increment. Did you measure the counter on BufferIdMap? That would remove this type and one of the two promotion paths. The cost is initializing 64 more bytes per counter, so keeping both is fine if record_batch_memory shows a difference.
If they stay separate, use one scan form in both. insert now uses the indexed loop, increment still uses .position(), and the BufferIdMap doc still refers to the position() scan.
There was a problem hiding this comment.
I measured it: with the counter on BufferIdMap, a batch of four columns took 7–9% longer, about 30 ns against 32 ns in column_count/4, row_count/* and array_layout/primitive, in each of 8 alternating rounds. The other cases were within 3.5%. So I kept both, and both scans now use the indexed loop, as decrement already did: .position() in both took 4–7% longer on those batches. On count_uncount, moving increment to the loop is a wash, 5% faster with 4 columns and 5% slower with 16. The BufferIdMap doc no longer mentions position().
| if let Some((count, counted)) = overflow.get_mut(&buffer_id) { | ||
| *count -= 1; | ||
| if *count == 0 { | ||
| let counted = *counted; | ||
| overflow.remove(&buffer_id); | ||
| return true; | ||
| return counted; | ||
| } | ||
| } | ||
| return false; | ||
| return 0; |
There was a problem hiding this comment.
This path now reads the counted capacity, so it can use one hash lookup instead of get_mut followed by remove (needs use crate::hash_map::Entry):
| if let Some((count, counted)) = overflow.get_mut(&buffer_id) { | |
| *count -= 1; | |
| if *count == 0 { | |
| let counted = *counted; | |
| overflow.remove(&buffer_id); | |
| return true; | |
| return counted; | |
| } | |
| } | |
| return false; | |
| return 0; | |
| let Entry::Occupied(mut entry) = overflow.entry(buffer_id) else { | |
| return 0; | |
| }; | |
| let (count, _) = entry.get_mut(); | |
| *count -= 1; | |
| return if *count == 0 { entry.remove().1 } else { 0 }; |
| /// Keeps the largest capacity seen per address: FFI-imported slices of one | ||
| /// allocation share an address but report the length they were imported with. |
There was a problem hiding this comment.
This holds for any buffer built with Buffer::from_custom_allocation, not only FFI imports. We hit the same pattern from a helper that re-wraps trimmed slices to report their real size. Naming the general case here keeps a later change from narrowing the fix to FFI.
| /// Keeps the largest capacity seen per address: FFI-imported slices of one | |
| /// allocation share an address but report the length they were imported with. | |
| /// Keeps the largest capacity seen per address: buffers built with | |
| /// `Buffer::from_custom_allocation`, such as FFI imports, can share an address | |
| /// while reporting different capacities. |
There was a problem hiding this comment.
Done, thanks. I've also fixed the PR description, which said that buffers outside FFI always agree on capacity.
- Return internal errors from the reservation checks in `ExternalSorter::sort_run` instead of panicking, so that an accounting mismatch fails the query rather than aborting the process. - Scan the inline addresses in `BufferIdMap::increment` with the indexed loop that `BufferIdSet::insert` and `BufferIdMap::decrement` use. The counter stays on `BufferIdSet`: on `BufferIdMap` it took about 8% longer for a batch of four columns. - Release an overflow entry in `BufferIdMap::decrement` with one lookup. - Name `Buffer::from_custom_allocation` as the general case in the `BufferIdSet` docs, and describe the tracker's charge as the largest capacity seen while a buffer is counted.
|
Thanks @comphead and @avantgardnerio for the reviews. I've pushed the changes and replied to each comment inline. The The benchmark numbers. Measuring the counter on |
xudong963
left a comment
xudong963
left a comment
There was a problem hiding this comment.
Thanks, this is a clean fix
|
looks good to me to -- thanks everyone |
Which issue does this PR close?
Rationale for this change
A sort over batches that came through the Arrow C Data Interface as zero-copy slices of one larger batch panics, in release builds too:
This is a regression in 56 (it comes from the shared-buffer accounting added in #25820). Apache DataFusion Comet hits it on TPC-DS q67a, where the sort input arrives from the JVM through FFI.
RecordBatchMemoryCounteridentifies a buffer by its start address and sizes it withBuffer::capacity(), keeping the first capacity it sees for an address. A buffer fromBuffer::from_custom_allocation, which is whatarrow::ffi::from_ffibuilds, reports the length it was imported with rather than the size of the allocation. Slices of one batch that are exported and imported separately therefore share a start address but report different capacities. When a smaller one is counted first, the shared input size fromsort_run_sizesends up below a later run's ownget_record_batch_memory_size, and the check insort_runfails.RecordBatchMemoryTrackerhas the same problem in its uncount path. It adds the capacity of the buffer that brings an address's count from 0 to 1, but subtracts the capacity of whichever buffer brings it back to 0. If that one is larger,memory_usageunderflows, and TopK'schecked_sub(..).unwrap()on its batch store size panics.What changes are included in this PR?
BufferIdSet(used by the counter) stores the capacity counted for each address, both inline and in the overflow map.insertreturns the bytes not counted yet: the full capacity for a new address, the growth when a larger capacity shows up at a known address, and 0 otherwise.BufferIdMap(used by the tracker) stores the counted capacity next to each reference count.incrementreturns the bytes added the same way, anddecrementreturns the recorded capacity when the count reaches 0, so the tracker always releases exactly what it added.ExternalSorter::sort_runreturns an internal error instead of panicking when either of its reservation checks fails, so a future accounting mismatch fails the query rather than aborting the process, which for Comet is the JVM.record_batch_memoryagainstmain, as medians of 8 alternating rounds on an M3 Max: a one-column batch goes from 8.5 ns to 12.8 ns because the counter is larger to construct (7.8 ns to 13.1 ns on its own),shared_slices/16is 8–11% slower,count_uncount/16andcount_uncount/64are up to 6% slower, and the rest are within ±4% or within the variation between rounds.Buffers that agree on capacity at a shared address, which covers every buffer not built with
Buffer::from_custom_allocation, are counted exactly as before.One behavior to note: the tracker keeps an address charged at the largest capacity seen until every buffer at that address is released, even after the largest one is released. The allocation is at least that large and stays alive while any slice of it is held, so the overestimate is safe, and it keeps count and uncount symmetric.
What is the testing strategy for this PR?
sort_runs_sharing_address_with_different_capacitiesinsorts/sort.rssorts two runs over one custom allocation that declare 16 and 1024 values, withsort_in_place_threshold_bytesset to 0 so each batch stays its own run. Without the fix, the check above fails, now as an internal error rather than a panic.test_record_batch_memory_counter_counts_largest_capacity_at_addressinutils/memory.rscounts two arrays over one allocation in both orders, with the inline set and after promotion to the overflow map, and checks the bytes each call returns and the total.test_uncount_releases_largest_capacity_at_addresscounts the smaller array and then the larger one, releases them in that order, and checks that the tracker returns to its starting usage, inline and after promotion.Are there any user-facing changes?
No. Sorting FFI-imported sliced batches no longer panics, and if the sort's reservation check fails again, the query fails with an internal error instead of a panic. Memory accounting for buffers that share an address and disagree on capacity is now at least as large as before. No public API changes.