Fix memory accounting in Datafusion

[cetra3]

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,

#[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:

• Add claim method to recordbatch for memory accounting arrow-rs#9433
• feat(memory-tracking): expose API to NullBuffer, ArrayData, and Array arrow-rs#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

[adriangb]

Benchmark Peak Memory

I already implemented this in the new benchmark runner I'm working on, e.g. adriangb#3 (comment). That said it's using the k8s API - not sure how accurate that's going to be for small spikes. We could switch to use Codspeed, jemalloc high watermark heap dumps (could be useful for debugging the actual cause of allocations) or something else. Partially blocked by #20631.

dhat tests sound nice too!

[adriangb]

I think one thing that might help with codspeed, etc. would be to move the clickbench, etc. benchmarks into criterion or something like that. It might be verbose to do but it seems like something an LLM can auto complete pretty easily.

[mbutrovich]

What granularity do you have in mind for memory accounting? I think for batches flowing through the system and large buffers that are allocated in operators like joins and aggregations that might spill I think the memory pool makes sense (and getting better visibility). But there are tons of heap allocated data structures that I don't think people really think about at all (hash tables, vectors, etc.) that can add up in operators. I'm all for more visibility on those sorts of things, but am terrified of the performance implications of centralizing those allocations. If I want visibility on those I reach for jemalloc.

[cetra3]

I think if there are lots of heap allocated data structures in use, while not part of a RecordBatch, they still account for the overall heap usage of datafusion, and should be accounted for.

Having fully accurate granularity would be an ideal goal, but in lieu of that if performance suffers by doing a lot of accounting, then we should at least strive towards keeping things under a given limit.