Report peak MemoryPool reservation per query in benchmarks - #23985
Merged
gabotechs merged 7 commits intoJul 30, 2026
Conversation
DataFusion's `MemoryPool` deliberately accounts for only the "large" allocations that scale with input size; intermediate batches flowing between operators are assumed to be small and left untracked. The `MemoryPool` docs therefore advise reserving "some overhead (e.g. 10%)" on top of the configured limit. Nothing measures what that overhead actually is. `MemoryPool::reserved` is a live value that has usually fallen back to zero by the time a query finishes, so the peak is never observed, and benchmarks report peak RSS with nothing to compare it against. Add `PeakRecordingPool`, a delegating `MemoryPool` wrapper that records the high-water mark of `reserved()`. `CommonOpt::runtime_env_builder` installs it around the pool it already builds, so every benchmark that runs with a memory limit reports the peak without further changes. - `BenchQuery` gains `pool_peak_bytes`, reset per case by `start_new_case`, so each query gets its own reading rather than inheriting the high-water mark of the queries before it. - `print_memory_stats` prints the run-wide peak after the existing mimalloc line, leaving `mem_profile`'s output parser untouched. This is measurement only. Nothing enforces a relationship between the pool's accounting and actual allocation, and no threshold or CI check is added. The wrapper is installed only when a memory limit is configured, since without one there is no pool to wrap. `pool_peak_bytes` is omitted from the results JSON in that case, leaving the output byte-identical to before for runs without a limit. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`ArrowMemoryPool` implements Arrow's `MemoryPool` by growing a DataFusion `MemoryReservation` against the pool it wraps, so a buffer claimed through it reaches `PeakRecordingPool::grow` and is recorded like any other reservation. Nothing in DataFusion claims buffers today, but that makes the peak follow the accounting as it changes rather than fixing it to the current set of manually tracked consumers. That property was assumed rather than tested. Add a test that builds an `ArrowMemoryPool` over the recording pool and asserts an Arrow-side reservation both raises the peak and releases on drop. `arrow-buffer/pool` and `datafusion-execution/arrow_buffer_pool` are enabled as dev-dependencies only, so the benchmark binaries are built with exactly the features they were before. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
adriangb
force-pushed
the
claude/peak-pool-memory-accounting-3b4c26
branch
from
July 29, 2026 21:30
920656e to
5bb2c36
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23985 +/- ##
==========================================
+ Coverage 80.71% 80.85% +0.14%
==========================================
Files 1089 1098 +9
Lines 368760 374236 +5476
Branches 368760 374236 +5476
==========================================
+ Hits 297647 302607 +4960
- Misses 53372 53586 +214
- Partials 17741 18043 +302 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
adriangb
marked this pull request as ready for review
July 29, 2026 21:57
Contributor
Author
|
@avantgardnerio @LiaCastaneda any interest in this change? |
gabotechs
reviewed
Jul 30, 2026
The high-water marks lived in process-wide statics, which meant the recorded peak could not be attributed to a particular pool, the `Debug` impl reported a number that was not the instance's, and the unit tests needed a mutex to keep from clobbering each other. Move both marks onto `PeakRecordingPool` and reach them through `RuntimeEnv::memory_pool`, which `dyn MemoryPool::downcast_ref` already supports. `BenchmarkRun::set_memory_pool` takes the pool the queries run against, so benchmarks that build a runtime per query now report each query against the pool it actually ran on. No change to what is measured or emitted.
Recording the peak by calling `inner.reserved()` after every grow put the wrapped pool's bookkeeping on the allocation path. `--mem-pool-type` defaults to `fair`, and `FairSpillPool::reserved()` takes the same state lock `try_grow` just released, so every accounted allocation acquired it twice — in a harness whose job is measuring time, on exactly the spilling benchmarks where the number matters. Keep a running total in the wrapper instead. It stays exact because the trait grants exactly what is asked for: `grow` is infallible and `try_grow` either grants `additional` or leaves the reservation alone. Resetting now sets the mark to what is currently reserved rather than to zero, so a query that starts with data already held reports that as its floor instead of under-reporting until the next grow.
`pool_peak_bytes` was only assigned in `write_iter`, which a query that failed never reaches. Its absence was documented as meaning "no memory limit was configured", so an out-of-memory query — the case where the peak is most worth seeing — serialized identically to a run with no pool at all. Assign it in `mark_failed` as well.
`external_aggr` builds its pool directly instead of going through `CommonOpt::runtime_env_builder`, so it was the one benchmark that never got a recorder installed — despite being the benchmark that exists to run queries under a memory limit. Wrap the pool it builds and hand it to the run, so each (query, limit) case reports its peak like the rest.
The README's "Collecting data" section is the contract benchmarks follow to produce results JSON, and it did not mention the new field or the `set_memory_pool` call that populates it.
avantgardnerio
approved these changes
Jul 30, 2026
avantgardnerio
left a comment
avantgardnerio
left a comment
Contributor
There was a problem hiding this comment.
Great idea! Thanks @adriangb
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jul 30, 2026
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
5 participants
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?
Part of #22758. Doesn't close an issue.
Rationale for this change
DataFusion's
MemoryPooldeliberately accounts for only the "large" allocations that scale with input size; intermediate batches flowing between operators are assumed to be small and left untracked. TheMemoryPooldocs therefore advise reserving "some overhead (e.g. 10%)" on top of the configured limit.Nothing measures what that overhead actually is.
MemoryPool::reserved()is a live value that has usually fallen back to zero by the time a query finishes, so the peak is never observed, and benchmarks report peak RSS with nothing to compare it against.This adds the missing number so the two can be compared. It is measurement only — nothing enforces a relationship between the pool's accounting and actual allocation, and no threshold or CI check is added.
This is complementary to the accounting work discussed in #22898. That issue proposes changing what gets accounted (Arrow
claim()/ builderwith_pool()); becauseArrowMemoryPoolgrows a DataFusionMemoryReservationagainst the pool it wraps,reserved()is where both models converge. This PR just makes the peak of that value observable per query, so the effect of any such change — or of a regression — is visible as a number rather than inferred.Concretely, that means the peak follows the accounting instead of being pinned to the current set of manually tracked consumers: nothing claims buffers today, but if #22898 lands in either form, those bytes show up here without further changes. There is a test covering that path.
What changes are included in this PR?
PeakRecordingPool, a delegatingMemoryPoolwrapper recording the high-water mark ofreserved(). Every method delegates; wrapping does not change how memory is granted, limited, or reported.CommonOpt::runtime_env_builderinstalls it around the pool it already builds, so every benchmark run with a memory limit reports the peak without further changes.BenchQuerygainspool_peak_bytes, reset per case bystart_new_case, so each query gets its own reading rather than inheriting the high-water mark of the queries before it.print_memory_statsprints the run-wide peak after the existing mimalloc line, leavingmem_profile's output parser untouched.Everything is confined to the
benchmarkscrate. No trait changes.arrow-buffer/poolanddatafusion-execution/arrow_buffer_poolare enabled as dev-dependencies only, so the benchmark binaries build with exactly the features they had before — confirmed absent fromcargo tree --no-dev-dependencies.Are these changes tested?
Yes — 5 unit tests plus a doctest, and verified end-to-end with
dfbench nlj --query 1 -i 2 --memory-limit 512M, which emits"pool_peak_bytes": 262528per query.One of those tests builds an
ArrowMemoryPoolover the recording pool and asserts an Arrow-side reservation both raises the peak and releases on drop, pinning the interaction described above.The pre-existing
test_runtime_env_builder_reads_env_varstill assertsMemoryLimit::Finite(2G)through the wrapper, which is direct evidence the delegation is transparent.The wrapper is installed only when a memory limit is configured, since without one there is no pool to wrap.
pool_peak_bytesis omitted from the results JSON in that case, leaving output byte-identical to before for runs without a limit — confirmedcompare.pyparses both old and new files.Are there any user-facing changes?
No. Confined to the benchmarks crate; adds one optional field to the benchmark results JSON.