Skip to content

feat: support releasing batches in RecordBatchMemoryCounter - #26147

Closed
txwyy123 wants to merge 1 commit into
apache:mainfrom
txwyy123:support-uncount-batch
Closed

txwyy123 wants to merge 1 commit into
apache:mainfrom
txwyy123:support-uncount-batch

Conversation

@txwyy123

@txwyy123 txwyy123 commented Oct 9, 2026

Copy link
Copy Markdown

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)

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
@github-actions github-actions Bot added the common Related to common crate label Oct 9, 2026
@txwyy123

txwyy123 commented Oct 9, 2026

Copy link
Copy Markdown
Author

Closing for now — will reopen after further review.

@txwyy123 txwyy123 closed this Oct 9, 2026
@jayzhan211

Copy link
Copy Markdown
Contributor

You can turn it to draft instead of closing it

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support releasing batches in RecordBatchMemoryCounter

2 participants