[DISCUSSION] Memory accounting model discussion

[gabotechs]

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

• fix: overcounting of memory in first/last. #15924
• Fix array_agg memory over use #16346
• fix: Incorrect memory accounting in array_agg function #16519
• Address memory over-accounting in array_agg #16816

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

[notfilippo]

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.

[wj-stack]

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

[fmonjalet]

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!

[gabotechs]

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.

[alamb]

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

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

[alamb]

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

[alamb]

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.