Repository navigation
Keep the counter add-only, add RecordBatchMemoryTracker, use it in TopK - #1
Merged
txwyy123 merged 3 commits intoOct 18, 2026
Merged
Conversation
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.
This is a proposal for apache#26149, written as commits on top of your branch so you can see exactly what I mean. Merge it into
support-uncount-batchif it makes sense to you, or take what's useful.1. Keep
RecordBatchMemoryCounteradd-only, move release intoRecordBatchMemoryTrackerNone of the current users of
RecordBatchMemoryCounterever release a batch:get_record_batch_memory_size(one counter per batch)ExternalSorter, which drops the counter after a spillGiving them the refcount map is where the
column_count/64regression comes from. In this change:RecordBatchMemoryCounteris unchanged frommain(a set of buffer ids, add-only).RecordBatchMemoryTrackerhascount_*/uncount_*and uses yourBufferIdMap.visit_array_buffers(array, &mut impl BufferVisitor). I kept yourvisit_*names.For primitive columns, the counting path compiles to the same instructions as
main; only addresses and argument registers differ. Benchmark timings shift with code alignment (identical code measured 1.00x to 1.2x depending on the build). With functions and blocks aligned to 64 bytes (-C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6) it runs at 0.96–1.04xmain, against 1.00–1.02x formainversus a copy of itself.2. TopK as the first caller
RecordBatchStore::insertnow callscount_batch,unuse(id)callsuncount_batch, andclear()resets the tracker. @DamiaoCastro's repro from apache#26140 (2,000 one-row slices) is added asmemory_limit::topk_reserves_shared_input_buffers_once:mainit fails under 16 MB, with each slice charged 279.5 KB. With this change it passes.... GROUP BY k) ORDER BY c, k LIMIT 1000000) needs about 1.5 GB onmainand runs at-m 275Mwith this change.LIMIT 100000over 20–50M rows 0.96–1.04x; forLIMIT 10, TopK'selapsed_computeis unchanged.3. Tests trimmed to the acceptance criteria in apache#26140
The tests cover the two-slice example, buffers shared through views, dictionaries and lists, overflow past 16 buffers with removals, uncounting untracked arrays, and array overhead held by two batches. The randomized test now compares against a fresh
RecordBatchMemoryCounterover the arrays still counted, instead of a hand-written model.cargo fmt,cargo clippy -p datafusion-common -p datafusion-physical-plan -p datafusion --all-targets --all-features -D warnings, and theutils::memory,topkandmemory_limittests pass on this branch.