txwyy123 opened a new pull request, #26147: URL: https://github.com/apache/datafusion/pull/26147
## Which issue does this PR close? Closes #26140. ## Rationale for this change `RecordBatchMemoryCounter` can only add buffers — there is no way to stop counting a batch. This means operators that retain batches incrementally and drop them later (sort, window, sort-merge join, TopK) cannot use it and must keep their own estimates instead, leading to the over-counting and not-counted problems described in #26140. ## What changes are included in this PR? Add `uncount_batch`, `uncount_array`, and `uncount_batch_with_array_overhead` methods to `RecordBatchMemoryCounter`. The internal `BufferIdSet` (insert-only set) is replaced with `BufferIdMap` (reference-counted map): - **Counting** a batch increments the count of each of its buffers; a buffer's capacity is added to `memory_usage` only when its count goes from 0 to 1 (identical to today's behavior). - **Uncounting** a batch decrements the counts and returns the bytes released: the capacities of buffers whose count reaches 0. The per-type buffer walk (null buffers, offsets, view buffers including variadic data buffers, dictionaries, nested children, the `ArrayData` fallback) is shared between counting and uncounting via a unified `visit_array_buffers` method with a `BufferOp` direction parameter — no duplication. The inline fast path for ≤16 distinct buffers is preserved (allocation-free for typical batches). ### Existing behavior preserved All existing `count_*` methods return exactly what they returned before. No caller changes needed. All 22 existing tests pass unchanged. ## How was this tested? ### Unit tests (8 new) | Test | Coverage | |------|----------| | `test_uncount_batch_round_trip` | count then uncount returns to 0 | | `test_uncount_two_slices_example_from_issue` | The exact table from #26140 | | `test_uncount_array_round_trip` | Single array count/uncount | | `test_uncount_batch_with_array_overhead_round_trip` | Full round-trip with array overhead | | `test_uncount_view_arrays_sharing_data_buffers` | StringView with data buffers | | `test_uncount_dictionaries_sharing_values` | Two dictionaries sharing values array | | `test_uncount_nested_struct` | Nested struct type | | `test_uncount_with_overflow_promotion` | >16 distinct buffers with removals | | `test_uncount_never_counted_is_noop` | Uncount of never-counted buffer | | `test_randomized_count_uncount_sequence` | Randomized sequence checked against reference model | ### Benchmark New `count_uncount` benchmark group added. No regression on existing counting benchmarks: ``` record_batch_memory_size/column_count/1 time: [21.9 ns] record_batch_memory_size/column_count/4 time: [63.3 ns] record_batch_memory_size/column_count/16 time: [285.6 ns] record_batch_memory_size/column_count/64 time: [2.12 µs] record_batch_memory_size/shared_slices/4 time: [1.84 µs] record_batch_memory_size/shared_slices/16 time: [8.72 µs] record_batch_memory_size/shared_slices/64 time: [30.1 µs] record_batch_memory_size/count_uncount/4 time: [3.74 µs] (new) record_batch_memory_size/count_uncount/16 time: [~15 µs] (new) record_batch_memory_size/count_uncount/64 time: [~52 µs] (new) ``` -- 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]
