Repository navigation
[EPIC] A collection of Memory Accounting Limitations and Improvements #22758
Description
As DataFusion gets used more in production systems there is more of an interest in improving its memory accounting model. This epic lists related issues to try and get some organized discussion going
- Fix memory accounting in Datafusion #20714 from @cetra3
- [DISCUSSION] Memory accounting model discussion #16841 from @gabotechs
- NestedLoopJoinExec spill path: untracked allocation overshoots memory pool #22723 from @avantgardnerio
- GroupValuesRows::emit untracked decode buffer leaks past MemoryReservation #22739 from @avantgardnerio
- More accurate memory accounting in external sort #14748 from @2010YOUY01
- Downstream consumers of AggregateExec significantly overcount memory usage #22757 from @Samyak2
- Hash aggregation produces batches reporting huge memory size #22526 from @ariel-miculas
- Incorrect memory accounting in
array_aggfunction #16517 from @sfluor - Accurately reserve memory in the build side of hash joins #22861 from @jordepic
Activity
ariel-miculas commentedon Jun 10, 2026
Adding #22861 by https://github.com/jordepic
wirybeaver commentedon Jun 11, 2026
I've been working on memory accounting improvements and filed #22898 proposing two approaches to replace manual
get_array_memory_size()+try_resize()tracking:claim()viaArrowMemoryPool— use Arrow's idempotentbuffer.claim(pool)API at output boundaries. Eliminates both under-accounting (can't miss sites) and over-accounting (sharedArc<Buffer>counted once). Shippable today on arrow-buffer 58.3. Prototype: [RFC]feat: post-construction claim() for hash aggregation memory tracking wirybeaver/datafusion#3. Limitation: claim() is post-construction. A burst allocation can temporarily exceed pool limits before the claim catches it.with_pool()for real-time enforcement — upstream arrow-rs change addingTrackedVec<T>so builder growth goes through the pool, enforcing limits before allocation. Prototype: [RFC]feat(pool): add with_pool() to Arrow builders for real-time memory tracking wirybeaver/arrow-rs#1The key finding is that Arrow builders use
Vec<T>internally (notMutableBuffer), so the existingMutableBufferreservation field is bypassed during the build phase.The alternative of patching each missed
try_growcall site manually (wirybeaver#1) works but is inherently fragile and doesn't solve over-accounting.Would love feedback on the approach — details and tradeoff comparison in #22898.
cetra3 commentedon Jun 11, 2026
I think the only real way you can do this correctly is with a custom counting allocator with an artificial limit. I've been trying to get some traction in rust to get this stabilised but it's an uphill battle. There has been some recent work put into it though.
There is definitely a path using allocator_api2 today but would mean that we would need to swap out all arrow allocations and vecs with custom ones & probably use the fallible versions (i.e,
try_reserveinstead ofreserve, etc.. so we can do something if allocation fails). In other words: not a simple refactorThe
claimapi comes close, but at that point the allocation has already happened so you'll always need a bit of headroom.avantgardnerio commentedon Jun 11, 2026
@cetra3 that is a fantastic write up you did on the state of the allocator API.
I don't understand this part though. I know you can swap them all, but why not use a global allocator?
avantgardnerio commentedon Jun 11, 2026
After a little thought, it makes sense. DataFusion is a library and global allocators are none of it's business. What we and arrow could do is track all our own memory, and I think the non-global allocator would be the Rust way. It's a shame they are only stable in linux.
LiaCastaneda commentedon Jul 7, 2026
I'm not super familiar with the Allocator API, but it looks like they've already experimented with integrating it into arrow-rs -- apache/arrow-rs#6336.
Reading through that PR, I think the main blockers that led to closing it were that Allocator isn't compatible with dyn structs, so Buffer/Bytes can't carry a custom allocator, which they'd need to, since whatever allocated the memory has to free it on drop. Another concern was the fact that the allocator would have to be threaded up through
Buffer,Array, and every kernel, making it a very breaking change.Another approach suggested here that avoids introducing breaking changes in the arrow apis was to build wrapper APIs in DataFusion around the arrow kernels that account for the memory of the arrays they create. It's still less accurate than the Allocator (it claims memory after the allocation happens and can miss a kernel's intermediate buffers), but it's fully scoped to DataFusion and needs no arrow changes. I think it could work as a temporary solution until Allocator is stable and mature, although the arrow API breakage concern still stands whenever we revisit the real fix with Allocator.
cetra3 commentedon Jul 7, 2026
I will speak about a couple of paths we have taken in production to curb memory usage, because I think it's probably important context for this discussion.
The first attempt at limiting memory (besides setting the DF memory limit) was to introduce a record batch memory limiter. This is an
ExecutionPlanwrapper struct that wraps an inner plan node, usesclaimon all record batches in the tree, and applies backpressure/limits when memory reaches a certain threshold. There is an optimizer that does atransform_downwrapping every node in the tree with this limiter.I.e, this is essentially:
And so with this at every node in the tree, because reclaiming old memory will free the reservation, this gives a pretty holistic approach that doesn't depend on certain execution plans to limit memory (There is one outstanding race condition in arrow which makes this inaccurate).
You would think that such a blunt tool would be enough to capture memory usage, or at least a large swathe of it. This is not the case:
GroupsAccumulatorAdapteretcBytesstruct from the bytes crate, which is used forobject_storeamong other things, famously does not include acapacity(), onlylength(), and so you have no real way of knowing how much memory it is actually taking up.Most of the time, the allocations for queries is much higher than the reported watermark.
So, the next thing we've tried, and had better success, albeit not perfect, is to use a custom global allocator, and essentially do the same thing we were doing at the record batch level:
With this approach, this is once again not perfect, because it happens after allocations already, but it does catch a lot more memory pressure events, that the pure record batch one does not catch. In fact, I hardly ever see the recordbatch-level limiter hit its threshold, because most allocations are happening elsewhere.
I feel that, given our production experiences, the only way to really knock this problem on its head is that we need to account for all allocations within datafusion and arrow. It is a bit of a painful, but probably needed, refactor. If we don't do this, we are going to end up in our current whack-a-mole state.
LiaCastaneda commentedon Jul 8, 2026
Thanks for sharing, about what you said here
If a single query is causing the memory pressure, you'd end up cancelling all the other queries running on that node too. Have you found this problematic? In a scenario where all the queries (including the heavy one) retry on a different node, wouldn't you just see the same issue again, cancelling more unrelated queries each time?
I agree, ideally this would live in Arrow, although the community has expressed concern over how breaking it would be to integrate a
claim()API into each Arrow buffer creation call and/or to adopt the still-unstable Allocator API. It looks like the main concern is that we shouldn't leave this in a mid-state, and that we should get it right the first time, given how breaking a change like this would be to public APIs like kernels, constructors, etc.I wonder if in the medium term, it would make sense to come up with a solution in DataFusion that doesn't require changes in arrow-rs, at least until we're able to use the rust Allocator (such as
claim, even if it's not 100% accurate since it reports after allocation) -- which would be the permanent approach.cetra3 commentedon Jul 8, 2026
Yes, and because the memory story in DF is not great, even prior to this we isolate queries to be one per process at a time to prevent any out of memory events having a high blast radius.
Yeah it sucks that there isn't a standard trait we can rely on yet and the risk is that the allocator trait looks very different when standardised so would create churn. I still think pursuing it is worthwhile and is actually a pretty well trodden path. PostgreSQL basically highly suggests if not mandates you use their
pallocfor allocation for everything.avantgardnerio commentedon Jul 8, 2026
FWIW, the solution I mentioned above only kills "the victim". At that point the framework is free to sort currently running queries by memory use descending and decide to kill "the perpetrator" as well.
LiaCastaneda commentedon Jul 10, 2026
@avantgardnerio are you refering to this PR? #22834 -- I took a look at it, the approach is interesting and something perhaps people can implement in their own code as a workaround to the DF memory accounting.
Have you thought about integrating this with the DF
MemoryPoolsomehow? I suppose we will need per query attribution for that. iiuc your PR, right now the guard tracks a single global balance and kill one query by panicing when it cannot alloc (which makes it possible to kill only that query and not all in the node), but it doesn't actually keep track of which query it killed nor which query was responsible for the pressure. To target the heavy query (instead of just whichever one happened to allocate when memory ran out) I guess you'd need per-query accounting. Maybe we would need a query id carried as a task_local var that propagates into every spawned task, and have the allocator read it per alloc or somthing simillar.avantgardnerio commentedon Jul 10, 2026
@LiaCastaneda that's the 2nd PR in my series. In this one, I demonstrated something closer to what you are suggesting (but it was only for testing). Unfortunately we can only kill whatever unlucky query goes over the limit (victim), but my intention was that when a kill was detected, we could use the memory tracker to sort queries by amount of tracked memory descending, and kill the perpatrator as well.
I think that's what's in my OomGuard PR. It's definitely part of what we run at Coralogix. If the tokio task that gets killed is one we spawned, we know the query ID. Unfortunately operators like RepartitionExec spawn their own tasks with their own default panic handlers and those sadly just die.
LiaCastaneda commentedon Jul 10, 2026
Perhaps an abstraction in arrow-rs that uses the unstable API on the inside, if the API changes when it stabilizes, we just update the abstraction's internal implementation and not its public methods, so the API doesn't break. Easier said than done though 😄 this would require a lot more thought to check if it's feasible.
alamb commentedon Jul 10, 2026
I think you can intercept (and install handler/ids) using this https://docs.rs/datafusion/latest/datafusion/common/runtime/struct.JoinSet.html
Minor: Fix docs for JoinSet (#23448)
LiaCastaneda commentedon Jul 13, 2026
I think it uses a
thread_localthat doesn't keep an ID (the first PR does keep acontext_id, but it's thread local too). Given that multiple threads can run the tasks of a given query, I was imagining the ID granularity would need to be per task, all tasks of a given query carrying the same ID? Then you can use a global map like the one in your first PR, but keyed by a query ID instead of a context ID/thread ID. With per query IDs you could then rank queries by memory and cancel the perpetrator.If a task from
RepartitionExecpanics, the error would get propagated to the rest of tasks and the query will fail with a generic error, but it's true it would not be possible to know if it errored because of anallocfailure.ariel-miculas commentedon Jul 13, 2026
Since the revert of custom allocation auditing, are there current ways of writing memory tests in datafusion? I mean tests that track the actual memory allocated, since tracking only the pool's memory doesn't show the discrepancies between the configured memory limits and the actual memory usage.
rluvaton commentedon Jul 13, 2026
I'll give my 2 cents as we saw a couple of OOM and memory problem in our company with datafusion.
we became very sensitive regarding memory so we started creating tests that check each operator in isolation with different data profile (number of columns, size of each row, etc.) and stress those in term of memory.
the assertions on memory is by using custom allocator and checking if the maximum allocated memory across the entire run is over the limit by some config.
we found several issues already in the few operators we started testing, I believe we will see much more.
being conservative about memory and having good performance is a fine balance, for example, how can you know if you can preallocate an entire batch size values for each column in
RepartitionExecimpl if you have a lot of partitions or there is a lot of skew.@LiaCastaneda There are inherent problems with tracking allocation with arrow memory pool feature
because the batches can have shared buffers and with arrow memory pool you avoid double counting memory, no operator is owning that batch if couple are holding it, which means that having one operator spill that batch may not free the memory.
LiaCastaneda commentedon Jul 14, 2026
🤔 Isn't that technically correct though? If another operator still holds a reference to the shared buffer it genuinely needs it in memory, so a spill by one of the holders shouldn't free it -- the arrow memory pool only drops (and deducts from the pool) only when no one references it. I would assume this happens with or without the arrow memory pool. Or is the pool making this worse somehow?
Report peak MemoryPool reservation per query in benchmarks (apache#23985
Report peak MemoryPool reservation per query in benchmarks (apache#23985