Repository navigation
Conversation
Add uncount_batch, uncount_array, and uncount_batch_with_array_overhead methods to RecordBatchMemoryCounter, enabling operators that retain batches incrementally and drop them later (sort, window, sort-merge join, TopK) to use accurate buffer-level memory accounting. The implementation replaces the internal BufferIdSet (insert-only set) with BufferIdMap (reference-counted map) that tracks per-buffer counts. A buffer's capacity is added to memory_usage when its count goes from 0 to 1 and subtracted when it returns to 0. The inline fast path for <= 16 distinct buffers is preserved to avoid heap allocation for typical batches. The per-type buffer walk (null buffers, offsets, view buffers, variadic data buffers, dictionaries, nested children, ArrayData fallback) is shared between counting and uncounting via a unified visit_array_buffers method with a BufferOp direction parameter. Existing count_* methods return exactly what they returned before -- no caller changes needed. All existing tests pass unchanged. New tests: - count/uncount round trip for batch, array, and batch_with_array_overhead - two-slice example from the issue (the exact table from apache#26140) - view arrays sharing data buffers - dictionaries sharing values - nested struct types - overflow promotion (>16 distinct buffers) with removals - uncount of never-counted buffer (no-op, no panic) - randomized count/uncount sequence checked against reference model Benchmark: record_batch_memory count_uncount case added; no regression on existing counting benchmarks. Closes apache#26140
Author
|
Closing for now — will reopen after further review. |
Contributor
|
You can turn it to draft instead of closing it |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #26140.
Rationale for this change
RecordBatchMemoryCountercan 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, anduncount_batch_with_array_overheadmethods toRecordBatchMemoryCounter. The internalBufferIdSet(insert-only set) is replaced withBufferIdMap(reference-counted map):memory_usageonly when its count goes from 0 to 1 (identical to today's behavior).The per-type buffer walk (null buffers, offsets, view buffers including variadic data buffers, dictionaries, nested children, the
ArrayDatafallback) is shared between counting and uncounting via a unifiedvisit_array_buffersmethod with aBufferOpdirection 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_uncount_batch_round_triptest_uncount_two_slices_example_from_issuetest_uncount_array_round_triptest_uncount_batch_with_array_overhead_round_triptest_uncount_view_arrays_sharing_data_bufferstest_uncount_dictionaries_sharing_valuestest_uncount_nested_structtest_uncount_with_overflow_promotiontest_uncount_never_counted_is_nooptest_randomized_count_uncount_sequenceBenchmark
New
count_uncountbenchmark group added. No regression on existing counting benchmarks: