[DISCUSSION] Memory accounting model discussion

The current model used in DataFusion for measuring memory consumption assumes that the different entities that can consume memory (accumulators, joins, etc...) are the ones owning the data. This does not exactly match with how memory is managed in arrow-rs, where underlying data buffers might be referenced more than 1 time in different parts of DataFusion.

For example: ArrayAggAccumulator accumulates data by storing ArrayRefs, so even if it's accumulating and retaining data, the actual memory was allocated before ArrayAggAccumulator came into play, and the accumulator is only adding a reference to it:

datafusion/datafusion/functions-aggregate/src/array_agg.rs

Line 214 in ac407a1

values: Vec<ArrayRef>,

For that case, one could argue ArrayAggAccumulator is not really consuming any memory, as it's not performing new allocations, and the underlying data was there potentially even before the ArrayAggAccumulator was instantiated.

There has been several attempts in the past towards addressing memory accounting issues in DataFusion code

Some imply copying/compacting just the necessary slice of data from the underlying buffer (ScalarValue::compact) so that it's actually owned by the consumer, but in certain cases that could take a hit to performance.

The main point about this issue is to start a conversation around what could be the ideal approach for memory counting:

Is copying/compacting accumulated data and calling get_array_memory_size() or storing array references and calling get_slice_memory_size() acceptable for measuring memory consumption?
Should the memory counting model in DataFusion be expanded so that it does not take into account just memory consumed, but also memory retained because of references to shared buffers

According to the docs in MemoryPool:

DataFusion’s design ONLY limits operators that require “large” amounts of memory (proportional to number of input rows), such as GroupByHashExec. It does NOT track and limit memory used internally by other operators such as DataSourceExec or the RecordBatches that flow between operators. Furthermore, operators should not reserve memory for the batches they produce. Instead, if a parent operator needs to hold batches from its children in memory for an extended period, it is the parent operator’s responsibility to reserve the necessary memory for those batches.

I would argue that when accounting anything that is not owned (like ArrayRef, ScalarValue::List(...), ...), the memory reserved should be equal to the bytes occupied by the reference itself (size_of::<Arc>() or potentially even 0) and not the size of referenced object.

The caveat I see with this approach is that in all the places where we actually copy data (via something like ScalarValue::compact) we need to be able to correctly account the size of the new copy in addition to the size of the reference.

I fully support this proposal，I encountered a situation where there was only a very small amount of data, but RepartitionExec still needed to alloc 19G of memory

@gabotechs

I would argue that when accounting anything that is not owned (like ArrayRef, ScalarValue::List(...), ...), the memory reserved should be equal to the bytes occupied by the reference itself (size_of::() or potentially even 0) and not the size of referenced object.

The tricky part is that there is no single place that properly owns the data: it's a shared ownership (Arc). Meaning that even if you only use one row, keeping the pointer alive keeps the whole Array alive.
A good example may be if you have 1000 record batches of 1MiB each and retain one value per record batch as an aggregation key:

In theory you are only retaining 1000 rows, so not a lot of data
But in practice you are keeping 1GiB of data alive, which breaks the assumptions of the streaming model

So this data has to be accounted for. In this example, 1000 values indeed retain 1GiB.

But the issue is when there are multiple reference to the same Array retained. If you have 1 record batch of 1MiB and retain 1000 rows in it, without copying the values, the current memory accounting will count the same Array 1000 times, effectively counting 1GiB of used memory instead of 1MiB. This is the main memory over accounting we face when aggregating on arrays (as in: a column in which each row is an Array<String> for example).

I had written a longer analysis there: #16055, although this particular instance was solved thanks to the new compact() API.

The solution of copying can be less efficient in terms of CPU but is the best in terms of retained memory (it keeps the goodies of the streaming model even in nasty corner cases). A good trade-off would probably be to avoid double counting Arrow array buffers, but I am not sure how this could be integrated.

Anyway, strong +1 for this issue as it's been a recurring problem that makes the (otherwise super great) memory tracking fail queries that should easily pass. Thanks @gabotechs!

@notfilippo

One idea @notfilippo mentioned is that arrow-rs could offer some kind of API for tracking allocations. As it's arrow-rs the one who knows when a buffer is allocated, when it's reference count increases, and when it's freed, maybe it's a good candidate for offering an API that allows consumers to track memory usage. Just throwing the idea for if something sticks, but I imagine that going down that path would require a fair amount of non-trivial work.

A more immediate problem that would be nice to get a shorter term solution for is the problem of double-counting memory. There are several places in the codebase that are calling .get_array_memory_size() over arrays baked by shared buffers that can result in over accounting. Some places I have identified:

datafusion/datafusion/common/src/scalar/mod.rs

Lines 3384 to 3386 in 3869857

    
           ScalarValue::List(arr) => arr.get_array_memory_size(), 
        
           ScalarValue::LargeList(arr) => arr.get_array_memory_size(), 
        
           ScalarValue::FixedSizeList(arr) => arr.get_array_memory_size(),

datafusion/datafusion/functions-aggregate/src/array_agg.rs

Line 381 in 47f75ef

.map(|arr| arr.get_array_memory_size())

datafusion/datafusion/physical-plan/src/joins/nested_loop_join.rs

Line 625 in d3cacac

let batch_size = batch.get_array_memory_size();

datafusion/datafusion/physical-plan/src/joins/cross_join.rs

Line 202 in d3cacac

let batch_size = batch.get_array_memory_size();

datafusion/datafusion/physical-plan/src/joins/sort_merge_join.rs

Lines 818 to 821 in d3cacac

    
           let size_estimation = batch.get_array_memory_size() 
        
               + join_arrays 
        
                   .iter() 
        
                   .map(|arr| arr.get_array_memory_size())

datafusion/datafusion/physical-plan/src/joins/symmetric_hash_join.rs

Line 1158 in d3cacac

size += self.input_buffer.get_array_memory_size();

datafusion/datafusion/physical-plan/src/repartition/mod.rs

Line 950 in 4bc66c8

let size = batch.get_array_memory_size();

Probably there are some more.

I wonder if an acceptable short term solution could be to account for memory occupied by an ArrayRef using .get_slice_memory_size() rather than .get_array_memory_size(). It's still not going to be correct, but it might be less wrong than double-counting memory.

Some imply copying/compacting just the necessary slice of data from the underlying buffer (ScalarValue::compact) so that it's actually owned by the consumer, but in certain cases that could take a hit to performance.

The solution of copying can be less efficient in terms of CPU but is the best in terms of retained memory (it keeps the goodies of the streaming model even in nasty corner cases).

I think you have identified the key tradeoff: It is often faster CPU wise to reuse the existing allocation but that risks holding on to more memory than necessary. Copying only the data needed results in less memory use but consumes CPU time.

I think there are several places in DataFusion that deal with a similar issue, that perhaps we can take inspiration from. I think at a high level they all basically employ a heuristic to decide when copying is better than reusing the allocation

TopK: maybe_compact
StringView (and calling "gc" on it), for example in BatchCoalescer

@notfilippo

One idea @notfilippo mentioned is that arrow-rs could offer some kind of API for tracking allocations. As it's arrow-rs the one who knows when a buffer is allocated, when it's reference count increases, and when it's freed, maybe it's a good candidate for offering an API that allows consumers to track memory usage. Just throwing the idea for if something sticks, but I imagine that going down that path would require a fair amount of non-trivial work.

I believe @waynexia added just such an API in

memory tracking with memory pool arrow-rs#7303

It hasn't been released yet but that should happen soon

@Dandandan

I think @Dandandan and @zhuqi-lucas have also been thinking about / trying to improve memory efficiency and might have some perspectives to offer

@waynexia

I believe @waynexia added just such an API in

Datafusion is a very good example and use case for that API. I haven't fully figured out how to evolve that API, like whether we need and how to integrate it with higher-level Array interfaces, and how downstream projects like us here would like to make use of it in practice. Would love to learn more from helping land it here.

One place that IMO would really benefit from using the new arrow APIs is IO pipeline width. We've found that being able to go wide on IO is extremely useful for query latency so that we can saturate the bandwidth of our nodes on bandwidth constrained queries (e.g. point queries where it is necessary to filter GBs of data in memory to produce only 1 row). Currently in DataFusion that is (1) tied to the CPU work width and (2) has no way to account for the selectivity of the query in terms of rows or columns. I started some work on this in #17758 but really I think the ideal scenario is a config along the lines of "buffer up to 1GB of data from DataSourceExec" and then that can be opening 1024 files that produce 1kB of data each or 1 file that produces a 2GB RecordBatch on each poll.

I started some work on this in #17758 but really I think the ideal scenario is a config along the lines of "buffer up to 1GB of data from DataSourceExec" and then that can be opening 1024 files that produce 1kB of data each or 1 file that produces a 2GB RecordBatch on each poll.

I have been thinking a lot about this too in relation to a push parquet decoder

Decouple IO and CPU operations in the Parquet Reader (push decoder) arrow-rs#7983

Currently, the arrow-rs parquet decoder decodes in the unit of a row group (basically it fetches all data pages in a row group needed for reading a column, after accounting for predicates). This does limit the number of IOs, which is important for reading from object storage, but it also means that the minimum memory usage is a function of how large the row group was, which means it can't be limited to some set amount easily

I am hoping that once we have a push decoder working (I hope to get this ready for review later this week) it will finally be feasible to allow decoding a subset of the row groups to trade off # of requests with memory usage

related thread in discord: https://discord.com/channels/885562378132000778/1422843060500299938

Hello,

Currently in DataFusion that is (1) tied to the CPU work width

This is essentially datafusion.execution.target_partitions, right?

has no way to account for the selectivity of the query in terms of rows or columns.

Could you please elaborate on this? Should the outcome (selects 1 row out of 10000 rows spread across 100 files) of the query matter here at all? In any scenario, regardless of the selectivity of the query, we would like to saturate the network bandwidth. Or is the problem particularly visible in selective queries as you described? If this is off topic, I could ask this on your PR.

This is essentially datafusion.execution.target_partitions, right?

Yes exactly.

The issue is that say you set target partitions to 8 because you have 8 CPU cores. For a very selective query that is going to return 1 row but have to prune hundreds of files / row groups (using parquet stats, predicate pushdown, etc.) the bottleneck is going to be IO / latency to object storage. Setting target partitions to say 128 will make the query go a lot faster. But if you set that permanently queries that are CPU bound (e.g. same query but with no filter) will be much slower or OOM your machine.

There is also a tuning guide that tries to help explain the tradeoff between target_partitions and memory usage here: https://datafusion.apache.org/user-guide/configs.html#tuning-guide

feat: Add `output_bytes` to baseline metrics (#18268)

feat: Add `output_bytes` to baseline metrics (apache#18268)

👋 I was looking into this issue and thinking around ways to integrate the new arrow memory tracking API to the operators, instead of doing compact() or get_slice_memory_size() to avoid overaccounting.

Scoping it down to GroupedHashAggregateStream for simplicity (and since it is where people have seen this issue the most), I imagine we could include the Arrow MemoryPool arrow_pool: Arc<TrackingMemoryPool>, in the struct in order to track the actual Buffer size in this operator overall, then we wouldn't need to call size() of each accumulator in update_memory_reservation:
let acc = self.accumulators.iter().map(|x| x.size()).sum::<usize>(); <- where some accumulators might actually share the same underlying Buffer.

I would even argue that accumulators used by GroupedHashAggregateStream don't need a size() function, since at the end what holds the accumulators is GroupedHashAggregateStream and acumulators within that operator won't handle more memory than batches processed while this stream polls.

So I imagine a flow like:

a RecordBatch arrives in group_aggregate_batch()
Before processing: Claim all buffers from the batch into arrow_pool (using bytes.claim(&arrow_pool))
Process the batch normally, creating the ScalarValues and updating accumulators (the ScalarValues reference the already-claimed buffers)
Everytime we call update_memory_reservation(): Use arrow_pool.used() to get the actual buffer memory, I guess here we should also include anyother overhead memory

I think this is more or less what compact() does but without copying data and from the operator level, not from the accumulator level.

Edit: I came to understand that the solution I proposed before would break assumptions of the API 😿, for instance nothing stops API users of having extra memory inside a different Accumulator implementation, that should also be accounted in update_memory_reservation. I will try to rethink this...

For GroupedHashAggregate stream in particular, another potential solution would be to implement a GroupsAccumulator for whatever aggregate you are working on, rather than rely on ScalarValue (which can hold pointers into shared buffers)

They are trickey to do, but are often faster and could be much more precise with reporting their allocation

I think the challenge of any sort of arrow_pool based solution is that the arrow buffers themselves are shared, so it will be very hard to know when to know the memory is "reclaimed" and not to double count it

I think the challenge of any sort of arrow_pool based solution is that the arrow buffers themselves are shared, so it will be very hard to know when to know the memory is "reclaimed" and not to double count it

🤔 wouldn't that already solved by the arrow memory tracking API itself? From the a test in arrow-buffer/src/bytes.rs

// Reserve memory and assign to buffer. Claim twice.
buffer.claim(&pool);
assert_eq!(pool.used(), 1024);
buffer.claim(&pool);
assert_eq!(pool.used(), 1024);

When you call claim() twice on the same ~~buffer~~ Bytes (actually is not on Buffer but on Bytes which essetially holds the data of Buffer so works as well I guess) it replaces the old reservation:

First claim() -> pool counter = 1024
Second claim() -> Creates new reservation: 1024 + 1024 = 2048, inmediately drops old reservation: 2048 - 1024 = 1024
we keep 1024 only, avoiding the overaccounting

I'm probably missing something but could we tie a memory reservation to an array buffer itself so that it gets reserved once when the array is created and freed whenever the array gets dropped? We could even implement as_any() as self.inner.as_any() or something so that downcasts that do read only access to the data continue to work.

I'm probably missing something but could we tie a memory reservation to an array buffer itself so that it gets reserved once when the array is created and freed whenever the array gets dropped?

iiuc the MemoryReservation arrow API does that in Buffer->Bytes here:

So it seems then that if we moved to a world where we use the new Arrow memory tracking, register the memory when we first create the array and then when we slice / get a reference in accumulators, etc. we... do nothing? Aside from maybe GC'ing if we've got the only reference and we have a small slice?

How would we know or ensure all places where data is created call claim()? the biggest Array creator might be the scan node but i imagine there are other places during execution that create arrays too.

I think we'd have to force it, e.g. make ArrayRef a struct and require a memory pool to initialize it?

Makes sense, that seems like the ideal final solution. Although, apart from redefining the MemoryPool model as it is, I imagine this would require multiple PRs across Arrow and DataFusion, maybe its worth breaking into steps or middle ground solutions? For instance, for now we're probably OK not accounting for every single Array used in DataFusion, since we don't do that in the current model anyways - only memory in specific operators is reported to the DF MemoryPool.

I think a middle ground solution could be to use something like an ArrowMemoryPool integration (proposed in #18928) in operators that already track memory (like GroupedHashAggregateStream, HashJoinExec, etc.) to fix in the short/middle term the buffer overaccounting problem we are seeing

feat(memory-tracking): implement arrow_buffer::MemoryPool for MemoryP…

I filed an epic with a bunch of related memory items to try and coordinate some response

[EPIC] A collection of Memory Accounting Limitations and Improvements #22758

Navigation Menu

Repository navigation

[DISCUSSION] Memory accounting model discussion #16841

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	ScalarValue::List(arr) => arr.get_array_memory_size(),
	ScalarValue::LargeList(arr) => arr.get_array_memory_size(),
	ScalarValue::FixedSizeList(arr) => arr.get_array_memory_size(),

	let size_estimation = batch.get_array_memory_size()
	+ join_arrays
	.iter()
	.map(\|arr\| arr.get_array_memory_size())

[DISCUSSION] Memory accounting model discussion #16841

Description

Activity

notfilippo commented on Jul 22, 2025

wj-stack commented on Jul 23, 2025

fmonjalet commented on Jul 23, 2025

gabotechs commented on Jul 24, 2025

alamb commented on Jul 24, 2025

alamb commented on Jul 24, 2025

alamb commented on Jul 24, 2025

waynexia commented on Jul 24, 2025

adriangb commented on Sep 30, 2025

alamb commented on Oct 1, 2025

alamb commented on Oct 1, 2025

debajyoti-truefoundry commented on Oct 1, 2025

adriangb commented on Oct 1, 2025

alamb commented on Oct 6, 2025

LiaCastaneda commented on Nov 21, 2025

alamb commented on Nov 22, 2025

LiaCastaneda commented on Nov 24, 2025

adriangb commented on Nov 24, 2025

LiaCastaneda commented on Nov 24, 2025

adriangb commented on Nov 24, 2025

LiaCastaneda commented on Nov 24, 2025

adriangb commented on Nov 25, 2025

LiaCastaneda commented on Nov 25, 2025

alamb commented on Jun 4, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions