Force compact in topK when we hit memory limit

[bharath-techie]

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.

[bharath-techie]

take

[2010YOUY01]

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]

I don't think its memory leak. The record batch size in GroupHashExecStream is 4196909056 and since we do slice on the record batch and send it to consumers, each sliced batch's size is still calculated as 4196909056 bytes.

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

#9562 - check this issue for more.

As we see in https://github.com/apache/datafusion/issues/9417, if there are upstream operators like TopK that hold references to any of these sliced RecordBatchs, those slices are treated as though they were an additional allocation that needs to be tracked ([source](https://github.com/apache/arrow-datafusion/blob/e642cc2a94f38518d765d25c8113523aedc29198/datafusion/physical-plan/src/topk/mod.rs#L576))

[bharath-techie]

I also saw another issue while working on this issue - say you limit the memory to 4 GB and now GroupedHashAggregateStream spills - for URL field 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 gc the 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 :

RUST_LOG=datafusion_physical_plan=debug ./datafusion-cli -m 40m --disk-spill-path /home/ec2-user/spilldir --disk-limit 75g

SET datafusion.execution.listing_table_ignore_subdirectory = false;
SET datafusion.execution.target_partitions=1;
SET datafusion.execution.parquet.binary_as_string=true;
CREATE EXTERNAL TABLE hits 
STORED AS PARQUET 
LOCATION '/home/ec2-user/hits_0.parquet';
SELECT "URL", COUNT(*) AS c FROM hits GROUP BY "URL" ORDER BY c DESC LIMIT 10;


Before fix : 
[2025-12-19T07:44:47Z DEBUG datafusion_physical_plan::spill::in_progress_spill_file] [SPILL_FILE] 
Finished spill file: path="/home/ec2-user/spill/datafusion-xKj4Qt/.tmpoxkWIz", 
size=820.54 MB, total_spilled_bytes=820.54 MB, total_spill_files=1


After fix : 
[2025-12-19T07:46:54Z DEBUG datafusion_physical_plan::spill::in_progress_spill_file] [SPILL_FILE] Finished 
spill file: path="home/ec2-user/spill/datafusion-3z9mL6/.tmpF7hNi9",
 size=33.43 MB, total_spilled_bytes=33.43 MB, total_spill_files=1

820 MB spill file -> 33 MB spill file

Edit :
#19414 - raised an issue for the same

[LiaCastaneda]

I don't think its memory leak. The record batch size in GroupHashExecStream is 4196909056 and since we do slice on the record batch and send it to consumers, each sliced batch's size is still calculated as 4196909056 bytes.

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]

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]

I agree with both of you:

1. Arrow memory pool API seems like a better way to address the root cause.
2. Doing some compute and continuing even if slowly is always better than an error (i.e. if we can literally replace and error in code with try to compact and then error if that doesn't help that's a win).
3. Long term we can be more thoughtful about when / what heuristics to use to auto compact.

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.