Skip to content

fix: enforce fair memory limits across sibling reservations - #25172

Draft
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:codex/upstream-consumer-accounting
Draft

fix: enforce fair memory limits across sibling reservations#25172
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:codex/upstream-consumer-accounting

Conversation

@sunchao

@sunchao sunchao commented Sep 10, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Fixes fair-pool admission for consumers with multiple reservations. Related to the merge-workspace accounting in #24740.

Rationale for this change

A spilling operator should have one memory allowance, regardless of how many buffers it uses. DataFusion represents that operator as a consumer, while each reservation accounts for some of its memory. Today, FairSpillPool checks the reservation making an allocation request without counting its siblings. An operator with a 100 MiB allowance can hold 60 MiB in one reservation and 40 MiB in another, then obtain still more through an empty sibling. Splitting or transferring reservations can expose the same gap.

A fair-share check also needs a separate limit on total pool usage. If one consumer fills a 100 MiB pool before another registers, the new consumer's 50 MiB share does not represent free memory. There is no room to admit its growth until existing reservations release space.

Closing these gaps affects aggregate spill replay. To finish an aggregation, DataFusion merges the sorted spill files and feeds the merged rows into an aggregate table. Those two components already share a consumer and run together. Once their reservations correctly count toward one allowance, letting the merge occupy all of it can leave the aggregate table unable to process the rows it receives. The accounting fix therefore needs to give replay room to work.

What changes are included in this PR?

The fair pool now admits growth against the consumer's combined reservations and the pool's remaining capacity. Reservation splitting, transfer, and release preserve that combined accounting, so creating another handle cannot create another allowance.

Aggregate replay uses an advisory consumer allowance to choose how many spill files to merge at once. In all four replay paths, it normally limits merge buffers to half that allowance, leaving room for the aggregate table. This is a planning budget: every allocation must still be accepted by the actual memory pool, whose available capacity can change while the query runs.

The half-share budget must also allow progress with very wide rows. For example, the smallest useful merge may need 60 MiB even though its preferred budget is 50 MiB. If its input rows cannot be split further, the merge may exceed that preference only when the actual pool grants the memory. That retry uses the smallest possible merge to leave room for the aggregate table. Checking an indivisible spill file avoids rewriting it unnecessarily, which matters when the existing files already fill the disk quota. If the actual memory cannot be obtained, execution reports a resource error and releases its resources.

What is the testing strategy for this PR?

PR CI at head 4a5cb8724 is green, including extended tests, forced hash-collision tests, Linux and macOS tests, SQL logic suites, Clippy, and documentation checks. The added regressions cover sibling lifetimes, changing fair shares, all four aggregate replay paths, indivisible rows, full disk quota, and cleanup under permanent memory pressure.

Local validation on upstream base ac7b18de6 passed 11,298 Rust tests and all 512 SQL logic files, plus Clippy and formatting. That older base retained its original dependency lockfile; the patch also applied cleanly to main 1ec9ede5. No targeted performance benchmark was run, and this PR makes no performance claim.

Are there any user-facing changes?

Allocations through try_grow can now fail when a consumer's combined reservations exhaust its allowance or the pool is full. Infallible grow retains its existing contract.

The additive MemoryPool::memory_limit_for API exposes an advisory allowance for planning. It defaults to Unknown; transparent custom wrappers should forward their inner pool's answer so aggregate replay can use it. The spill-file backend and ownership APIs are unchanged.

@github-actions github-actions Bot added core Core DataFusion crate execution Related to the execution crate physical-plan Changes to the physical-plan crate labels Sep 10, 2026
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 91.53713% with 49 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.92%. Comparing base (ac7b18d) to head (4a5cb87).
⚠️ Report is 42 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/aggregates/mod.rs 85.95% 4 Missing and 13 partials ⚠️
...usion/physical-plan/src/sorts/multi_level_merge.rs 91.48% 6 Missing and 10 partials ⚠️
datafusion/execution/src/memory_pool/pool.rs 97.78% 3 Missing and 2 partials ⚠️
...ion/execution/src/memory_pool/merge_memory_pool.rs 0.00% 4 Missing ⚠️
datafusion/execution/src/memory_pool/mod.rs 0.00% 3 Missing ⚠️
...fusion/physical-plan/src/aggregates/hash_stream.rs 87.50% 1 Missing ⚠️
...ysical-plan/src/aggregates/ordered_final_stream.rs 83.33% 1 Missing ⚠️
...sical-plan/src/aggregates/ordered_single_stream.rs 83.33% 1 Missing ⚠️
...sion/physical-plan/src/aggregates/single_stream.rs 87.50% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25172      +/-   ##
==========================================
+ Coverage   81.74%   81.92%   +0.17%     
==========================================
  Files        1128     1132       +4     
  Lines      416782   421114    +4332     
  Branches   416782   421114    +4332     
==========================================
+ Hits       340712   344989    +4277     
+ Misses      55999    55792     -207     
- Partials    20071    20333     +262     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

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

Labels

core Core DataFusion crate execution Related to the execution crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants