Skip to content

Keep the counter add-only, add RecordBatchMemoryTracker, use it in TopK - #1

Merged
txwyy123 merged 3 commits into
txwyy123:support-uncount-batchfrom
jayzhan211:uncount-split-topk
Oct 18, 2026
Merged

txwyy123 merged 3 commits into
txwyy123:support-uncount-batchfrom
jayzhan211:uncount-split-topk

Conversation

@jayzhan211

@jayzhan211 jayzhan211 commented Oct 18, 2026 •

Copy link
Copy Markdown

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-batch if it makes sense to you, or take what's useful.

1. Keep RecordBatchMemoryCounter add-only, move release into RecordBatchMemoryTracker

None of the current users of RecordBatchMemoryCounter ever release a batch:

  • get_record_batch_memory_size (one counter per batch)
  • the hash join build side, symmetric hash join and AsOf join
  • ExternalSorter, which drops the counter after a spill

Giving them the refcount map is where the column_count/64 regression comes from. In this change:

  • RecordBatchMemoryCounter is unchanged from main (a set of buffer ids, add-only).
  • The new RecordBatchMemoryTracker has count_* / uncount_* and uses your BufferIdMap.
  • Both share one buffer walk: visit_array_buffers(array, &mut impl BufferVisitor). I kept your visit_* 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.04x main, against 1.00–1.02x for main versus a copy of itself.

2. TopK as the first caller

RecordBatchStore::insert now calls count_batch, unuse(id) calls uncount_batch, and clear() resets the tracker. @DamiaoCastro's repro from apache#26140 (2,000 one-row slices) is added as memory_limit::topk_reserves_shared_input_buffers_once:

  • On main it fails under 16 MB, with each slice charged 279.5 KB. With this change it passes.
  • TopK over aggregate output (... GROUP BY k) ORDER BY c, k LIMIT 1000000) needs about 1.5 GB on main and runs at -m 275M with this change.
  • End to end it is neutral: LIMIT 100000 over 20–50M rows 0.96–1.04x; for LIMIT 10, TopK's elapsed_compute is 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 RecordBatchMemoryCounter over 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 the utils::memory, topk and memory_limit tests pass on this branch.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants