Fix memory accounting in Datafusion
- 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?
Right now the memory accounting in datafusion appears to be inaccurate. Often times if you limit memory within the application, this isn't honoured very well, and looking at the process in `top` or similar, you can see it uses way more RSS than you would expect.
### Describe the solution you'd like
There are a few things we probably need to solve this:
## Benchmark Peak Memory
We should introduce benchmarks for peak memory on a variety of different query types, so that we can ensure we don't go comically over our memory allowances.
There are two ways that I can see this being introduced
* Utilise `codspeed` with `--mode memory`: https://github.com/CodSpeedHQ/codspeed?tab=readme-ov-file#memory. This would allow us to run a few `datafusion-cli` like queries against clickhouse benchmarks etc..
* Utilise `dhat-rs` with specific tests for specific queries, ensuring that we don't go above peak memory + some memory slop:
I.e,
```rust
#[global_allocator]
static ALLOC: dhat::Alloc = dhat::Alloc;
#[test]
fn test() {
let _profiler = dhat::Profiler::builder().testing().build();
// run a 1gb limited query here
let stats = dhat::HeapStats::get();
dhat::assert_eq!(stats.max_bytes, 1024 * 1024 * 1024 * 1.1); // 1GB + 10% overhead
}
```
## Introduce the Claim API
Allow the memory pool in datafusion to use the new pool feature in arrow-rs.
Before this is useful, we need to ensure that we can claim full recordbatches. There are two competing PRs that bring us closer to that goal:
- https://github.com/apache/arrow-rs/pull/9433
- https://github.com/apache/arrow-rs/pull/8918
I am biased towards my PR of course, but whatever the `arrow-rs` team thinks is appropriate
With that in place, we could add:
* A new executor that sits around `DataSourceExec` that `claims` recordbatches as they come in.
* Sort/Spill infra to ensure that when we spill to disk or read from disk, we reclaim the input record batches
* Any other `ExecutionPlan` nodes that generate batches, such as row hash streams etc..
### Describe alternatives you've considered
_No response_
### Additional context
_No response_
Contributor guide
Research direction
No specific file or test is named. Start by reading DataFusion's memory-pool implementation and the DataSourceExec execution path, then review the referenced arrow-rs claim PRs and the proposed CodSpeed or dhat-rs benchmark approaches. Done would require agreed peak-memory benchmarks plus coordinated claiming and reclaiming across record batches, spill paths, and other batch-generating execution plans.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- databases
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 28/100