apache / apache/datafusion

Fix memory accounting in Datafusion

Open
#20,714 5 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?

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.