Repository navigation
Force compact in topK when we hit memory limit #19386
Description
Is your feature request related to a problem or challenge?
More context in #19216 (comment)
How to reproduce ?
Take clickbench workload
datafusion-cli -m 8G
DataFusion CLI v51.0.0
> SET datafusion.execution.target_partitions=1;
0 row(s) fetched.
Elapsed 0.001 seconds.
> CREATE EXTERNAL TABLE hits
STORED AS PARQUET
LOCATION '/home/ec2-user/clickdata/partitioned/hits/*.parquet';
0 row(s) fetched.
Elapsed 0.054 seconds.
> SELECT "URL", COUNT(*) AS c FROM hits GROUP BY "URL" ORDER BY c DESC LIMIT 10;
Resources exhausted: Additional allocation failed for TopK[0] with top memory consumers (across reservations) as:
TopK[0]#4(can spill: false) consumed 7.8 GB, peak 7.8 GB,
GroupedHashAggregateStream[0] (count(1))#3(can spill: true) consumed 80.4 KB, peak 4.6 GB,
DataFusion-Cli#2(can spill: false) consumed 0.0 B, peak 0.0 B.
Error: Failed to allocate additional 3.9 GB for TopK[0] with 7.8 GB already allocated for this reservation - 187.0 MB remain available for the total pool
>
With 8 GB and 1 partition, the above query SELECT "URL", COUNT(*) AS c FROM hits GROUP BY "URL" ORDER BY c DESC LIMIT 10; fails as the hash aggregate step for URL field is quite heavy.
For each batch that gets inserted and referred to in TopK , the entire batch size gets added during the size estimation. [ as mentioned in https://github.com/https://github.com//issues/9562 ]
Path : datafusion/datafusion/datafusion/physical-plan/src/topk/mod.rs
/// Insert a record batch entry into this store, tracking its
/// memory use, if it has any uses
pub fn insert(&mut self, entry: RecordBatchEntry) {
// uses of 0 means that none of the rows in the batch were stored in the topk
if entry.uses > 0 {
let size = get_record_batch_memory_size(&entry.batch);
self.batches_size += size;
println!("size during insert : {}", size);
self.batches.insert(entry.id, entry);
}
}
/// returns the size of memory used by this store, including all
/// referenced `RecordBatch`es, in bytes
pub fn size(&self) -> usize {
// size_of::<Self>()
// + self.batches.capacity() * (size_of::<u32>() + size_of::<RecordBatchEntry>())
// + self.batches_size
let sizeOfSelf = size_of::<Self>();
let capacity = self.batches.capacity();
let u32RecordBatch = size_of::<u32>() + size_of::<RecordBatchEntry>();
let batchesSize = self.batches_size;
let size = sizeOfSelf + capacity * u32RecordBatch + batchesSize;
println!("self size : {} , capacity : {} , heap size : {}, batch size : {}",
sizeOfSelf, capacity, u32RecordBatch, batchesSize);
println!("size during get : {}", size);
size
}
Record batch size during insert : 4196909056
self size : 48 , capacity : 3 , heap size : 60, batch size : 4196909056
size during get : 4196909284
Record batch size during insert size during insert : 4196909056
self size : 48 , capacity : 3 , heap size : 60, batch size : 8393818112
size during get : 8393818340
Record batch size during insert size during insert : 4196909056
self size : 48 , capacity : 3 , heap size : 60, batch size : 12590727168
size during get : 12590727396
Describe the solution you'd like
Implement force compaction in topK once we hit memory limit as mentioned in #9417 (comment)
@alamb also mentioned #9417 (comment) , #16841 (comment) - potentially there are other places where we could similar optimization ?
Describe alternatives you've considered
Tried spill but force compact is better option here.
Additional context
The main problem is multiple counting as that's the reason why this particular query fails with memory limit.
- feat(memory-tracking): implement arrow_buffer::MemoryPool for MemoryPool #18928 - will this help ?
- [DISCUSSION] Memory accounting model discussion #16841 - I see multiple approaches being discussed for fixing double counting issue.
Activity
bharath-techie commentedon Dec 18, 2025
take
2010YOUY01 commentedon Dec 19, 2025
This looks like a memory leak in TopK executor. In the worst case and without compaction, in this query the TopK node should only hold 10 batches. However, TopK is using 8 GB of memory, which appears to be far more than what 10 batches should require.
bharath-techie commentedon Dec 19, 2025
I don't think its memory leak. The record batch size in GroupHashExecStream is
4196909056and since we do slice on the record batch and send it to consumers, each sliced batch's size is still calculated as4196909056bytes.#9562 - check this issue for more.
bharath-techie commentedon Dec 19, 2025
I also saw another issue while working on this issue - say you limit the memory to 4 GB and now
GroupedHashAggregateStreamspills - forURLfield or any binaryView/StringView field - for each spill , it writes the record batch + entire original record batch's string view buffers - so the amount it spills to disk exceeds 100 GB and results in resource exhausted exception.There as well , in spill manager - if I
gcthe string view - and then spill only the compacted data , it only spills the current record batch and its view.Easily reproducible :
Take one clickbench partitioned parquet file ~120 mb :
820 MB spill file -> 33 MB spill file
Edit :
#19414 - raised an issue for the same
LiaCastaneda commentedon Dec 29, 2025
I think compacting would be a straightforward solution, however in some situations we’ve observed that compacting has a performance impact (e.g., #16519). We’ve seen a similar issue when DataFusion calculates the size of accumulators, and I found that using the Arrow memory pool API solves this (It correctly counts unique Arrow Buffers). However, I’m not yet sure how to integrate it properly with each operator. I have an experimental draft that addresses the over accounting issue for accumulators using the memory pool (#19501), but it’s probably not the cleanest solution, and its isolated to Aggregations only, so TopK would need another Arrow memory pool.
bharath-techie commentedon Dec 29, 2025
Yes agreed on compact causing some performance impact, we also tried out one more approach of using jemalloc based memory pool like the one used in influxdb for accounting [ instead we directly used it as memory pool for allocations ] - that also seemed to work out well in initial testing and was able to overcome the multiple counting issue. Though we didn't do extensive testing.
I'm happy to try out a working solution on top of #19501 as well for topK once its ready. I was already tracking the associated github issues and it looks promising :)
I still feel that force compact on topK instead of throwing error is still a good fallback solution. As atleast the query will go through after doing significant amount of work in groupby etc.
Plus with your change, we might not end up force compacting in many scenarios as we'll be able to correctly account for memory.
adriangb commentedon Dec 29, 2025
I agree with both of you:
So right now I'm +1 for a change to that takes us from error -> continue slowly.
I also am tracking the Arrow memory pool API changes and look forward to seeing them happen, but both things can happen at the same time and like @bharath-techie says are complementary.
Compact more aggressively in TopK based upon memory usage (apache#20381)
Compact more aggressively in TopK based upon memory usage (apache#20381)