apache / apache/datafusion

Force compact in topK when we hit memory limit

Open
#19,386 7 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Is your feature request related to a problem or challenge?

More context in https://github.com/apache/datafusion/issues/19216#issuecomment-3633472687
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/apache/datafusion/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.batches.capacity() * (size_of::() + size_of::())
// + self.batches_size

let sizeOfSelf = size_of::();
let capacity = self.batches.capacity();
let u32RecordBatch = size_of::() + size_of::();
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 https://github.com/apache/datafusion/issues/9417#issuecomment-2431943283

@alamb also mentioned https://github.com/apache/datafusion/issues/9417#issuecomment-2432677248 , https://github.com/apache/datafusion/issues/16841#issuecomment-3113393092 - 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.
- https://github.com/apache/datafusion/pull/18928 - will this help ?
- https://github.com/apache/datafusion/issues/16841 - I see multiple approaches being discussed for fixing double counting issue.

Contributor guide

Open the contributing guide

Research direction

Start with datafusion/datafusion/datafusion/physical-plan/src/topk/mod.rs and reproduce the memory failure using the provided datafusion-cli ClickBench query. Read the linked discussions on force compaction and memory accounting, then verify that TopK compacts when the memory limit is reached and the query no longer fails from duplicate batch accounting.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data-engineering, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.