andygrove opened a new pull request, #27097: URL: https://github.com/apache/datafusion/pull/27097
## Which issue does this PR close? - Closes #27069. ## 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: ``` assertion failed: reservation.size() + shared.size() >= reserved_bytes ``` 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. `RecordBatchMemoryCounter` identifies a buffer by its start address and sizes it with `Buffer::capacity()`, keeping the first capacity it sees for an address. A buffer from `Buffer::from_custom_allocation`, which is what `arrow::ffi::from_ffi` builds, 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 from `sort_run_sizes` ends up below a later run's own `get_record_batch_memory_size`, and the check in `sort_run` fails. `RecordBatchMemoryTracker` has 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_usage` underflows. ## 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. `insert` returns 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. `increment` returns the bytes added the same way, and `decrement` returns the recorded capacity when the count reaches 0, so the tracker always releases exactly what it added. - The inline addresses stay in their own contiguous array, so the lookup scan is unchanged. The capacities go in a parallel array. Buffers that agree on capacity at a shared address, which covers every non-FFI case, 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_capacities` in `sorts/sort.rs` sorts two runs over one custom allocation that declare 16 and 1024 values, with `sort_in_place_threshold_bytes` set to 0 so each batch stays its own run. Without the fix it panics with the assertion above. - `test_record_batch_memory_counter_counts_largest_capacity_at_address` in `utils/memory.rs` counts 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_address` counts 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. Memory accounting for buffers that share an address and disagree on capacity is now at least as large as before. No public API changes. -- 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]
