Enable automatic memory accounting using the memory pool when creating arrays

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

When tracking memory usage with arrow memory pools (introduced in #7303, tracked by EPIC #8137), users must manually call claim() on buffers after array creation. It could be easy to forget to call claim() on newly created arrays, leading to potentially under-reporting memory usage, also it requires having direct access to ArrayData just to access buffers for claiming, which in some cases is complicated.

It would be much more ergonomic and reliable to automatically claim buffers at array creation time.

Describe the solution you'd like

Add optional MemoryPool parameters to array constructors and builders, so buffers are automatically claimed at creation time. For example we could have: PrimitiveArray:: try_new_with_pool() in array costructors, PrimitiveBuilder::new_with_pool that claim buffers once we call finish() etc.

Additional context

This feature is needed to improve memory accounting and tracking in downstream projects like DataFusion. Related discussions:

[DISCUSSION] Memory accounting model discussion datafusion#16841

I think this is something arrow GO currently support on their builders:
https://pkg.go.dev/github.com/apache/arrow/go/arrow/array#NewTimestampBuilder

Given you have to switch the constructor I'm not sure how this is less likely to be missed? It's also unclear how this would work when arrays are constructed by kernels? I'm very lukewarm on this.

You're right, using this approach we would need to create a new kernel function to use the memory pool which doesn't feel super ergonomic. The ultimately goal is to be able keep track of the size of each Array created. fwiw, I looked into how Arrow Go and Arrow cpp handle this, and it seems to pass a context object everywhere (cpp example), and inside the context is the memory pool/allocator. Maybe we could do something similar in arrow? although I imagine that requires a considerable refactor and still wouldn’t resolve the “don’t rely on passing an extra parameter everywhere” concern

It could be interesting to explore an approach were you can register optionally a global memory pool. If the pool is present then it gets to own the memory of all the newly created arrays without the need of refactoring to add an extra parameter anywhere.

you mean keeping the memory pool in something like a OnceLock? that would make it easier, I was thinking what would be the concurrency overhead but I think the contention bottleneck woud be pretty which the same as passing it through parameters, which is increasing the pool size here

I guess the only downside is that it would force users to have only one memory pool.

you mean keeping the memory pool in something like a OnceLock

Yep!

users to have only one memory pool.

I would argue you could keep the claim methods arounds. The Optional global memory pool (which should solve 99% of use cases) will be used by default but users will be free to not set it and use claim instead.

IMO a solution similar to the one in the Arrow C++ implementation has very good chances to be a proper long term solution to this https://arrow.apache.org/docs/cpp/memory.html#memory-pools.

There, a memory pool is not optional, and rather than passively tracking the buffer allocations with a separate mechanism, is the memory pool itself the one responsible for allocating/deallocating Arrow buffers.

Not sure if the C++ Arrow implementation can be considered the "canonical one", but it definitely seems the most advanced and mature one, so following their lead might be a good path forward

FWIW the current design is a compromise to avoid breaking the API.

Unfortunately arrow-rs has never had such a context object, adding at this stage would be a tremendous undertaking, that would involve churning a huge portion of the public interface of this library.

That's not to say it isn't possible, but it would take multiple highly motivated individuals to see it through. Whilst the approach implemented in #6590 is not without downsides, IMO it's a pragmatic compromise. Ultimately what is ideal matters less than what is practical.

@LiaCastaneda

My impression is that memory tracking is a pillar too fundamental to this project to be built on top of a compromise, although I do acknowledge that changes would be relatively big.

Maybe one step forward would be to quantify and report here how much of change that would actually be, and whether following an approach that diverges as little as possible from https://github.com/apache/arrow makes sense at all. I'd not be surprised if we found out that it's not as bad as it seems.

@LiaCastaneda do you have time to explore what would it take?

Yes, I can gauge how much work this would require on the Arrow side. I’m not super familiar with the repo yet, but this could be a good opportunity to dig into it 👍

TBC whilst I do have concerns about the size of the change to the arrow codebase, my larger concern with this initiative is with making large-scale breaking changes to the public API. This has the potential to be extremely disruptive to the ecosystem and will require a lot more buy-in from the various stakeholders. Unfortunately I don't have the time to help out, but perhaps one of the other maintainers will be able to step-in and offer guidance on how this might be best communicated / consensus built.

My impression is that memory tracking is a pillar too fundamental to this project to be built on top of a compromise, although I do acknowledge that changes would be relatively big.

Memory tracking is inherently always going to be a compromise because Rust doesn't have stable allocator support. The question is therefore whether tracking the allocation once it emerges from the kernel, which is what is currently supported, is materially different from tracking it after the Vec / array is allocated. They're not all that dissimilar IMO.

That's fair, my proposal is not really to move forward with any change, just to evaluate the actual impact on the public API, if there's any chance to make it backwards compatible, and what benefits it would bring compared to the current solutions. The answer might be "it does not make sense" and that's completely fine, but I'd be curious to at least measure it.

TBC whilst I do have concerns about the size of the change to the arrow codebase, my larger concern with this initiative is with making large-scale breaking changes to the public API. This has the potential to be extremely disruptive to the ecosystem and will require a lot more buy-in from the various stakeholders.

I don't think the features contemplated here require a large breaking change

We could, for example, have a default global memory pool that the existing stateless kernels use, as well as new kernels that take an explicit context:

// new kernel with context
fn new_binary(context: &Context, arr: &dyn Array, arr2: &dyn Array)  -> ArrayRef { ... }

// Old kernel / backwards compatible API
fn binary(arr: &dyn Array, arr2: &dyn Array)  -> ArrayRef {
  // calls through to new kernel with context using a global allocator
  new_binary(GLOBAL_CONTEXT.get(), arr, arr2)
}

That's not to say it isn't possible, but it would take multiple highly motivated individuals to see it through.

Yes, this is my biggest concern. I think we would need to muster a significant amount of effort to get this feature in -- in contributors and reviewers. And leaving it half way done would be pretty bad.

That being said, something I have been thinking about for a long time is a kernel API that allows buffer reuse (for example, producing the output in the same allocation as the input, when possible). The current signature of taking a & doesn't really allow for this

@notfilippo

I'm thinking if maybe we could implement this in steps.

As a first step we could do something like what @notfilippo suggested. We can have a global Context with a Memory Pool that if it's set, we can automatically track all created Arrays by claiming inside MutableBuffer. I don't think this change would require much effort, since every Buffer creation goes through MutableBuffer, and we can probably directly test this in DataFusion and see if there are any perf regressions and validate the approach. This could be easily reverted if we see it's not feasible.
For a second step (if we see step 1 proves feasible) we could have explicit APIs with the context, like what @alamb suggested. I imagine we could have something like this to also avoid breaking the existing API:

pub fn add_with_context(
    context: &Context,
    lhs: &dyn Datum,
    rhs: &dyn Datum
) -> Result<ArrayRef, ArrowError> {
    arithmetic_op_with_context(Op::Add, context, lhs, rhs)
}

pub fn add(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
    match global_context() {
        Some(ctx) => add_with_context(ctx, lhs, rhs),
        None => {
            // Fallback: use default context
            let ctx = Context::default();
            add_with_context(&ctx, lhs, rhs)
        }
    }
}

This step will be complex because of the number of APIs involved: ~16 array types with constructors (both new and try_new). For the kernels it's more complicated because there are >100 and each would require its counterpart with a context. In addition, it would require passing the context down from the kernel function to the Array constructor function, changing signatures at every level. For example, for the add kernel it would involve passing the context down 6 levels until we reach the array constructor: add → arithmetic_op → integer_op → try_op → try_binary → PrimitiveArray::new.

Maybe we can do this in batches, starting with the most common kernels. However, as @alamb says, the danger would be leaving it halfway done and having a bunch of kernels claiming created data while others don't.

Step 2 would be the long term solution since it will allow more granular memory accounting, and we can have as many memory pools/contexts as we want. Also, I think adding a Context struct now (even if it only contains MemoryPool initially) sets up the infra for future features. For instance, arrow cpp has vars to tweak execution parameters like chunk sizes, threading etc

@alamb

we can automatically track all created Arrays by claiming inside MutableBuffer

Buffers can also be allocated from external allocations, most commonly Vec.

For a second step (if we see step 1 proves feasible) we could have explicit APIs with the context, like what @alamb suggested.

My confusion remains how this is practically any different from users registering the memory allocation after it emerges from the kernel. It is the difference between:

Memory allocation(s) performed
Kernel constructs array(s)
Caller registers array with pool

And

Memory allocation(s) performed
Kernel registers buffers with pool
Kernel returns arrays

Ultimately until Rust stabilises allocator support, our hands are somewhat tied here. That being said, if the desire is just to track global allocations, you can currently define a global allocator - see https://doc.rust-lang.org/std/alloc/index.html.

Also, I think adding a Context struct now (even if it only contains MemoryPool initially) sets up the infra for future features.

We prefer to have this configured at the kernel level with options, rather than a single global config object - keeps things decoupled. We also don't have threading, chunk sizes, etc... - these are execution engine concerns.

In addition, it would require passing the context down from the kernel function to the Array constructor function, changing signatures at every level. For example, for the add kernel it would involve passing the context down 6 levels until we reach the array constructor: add → arithmetic_op → integer_op → try_op → try_binary → PrimitiveArray::new.

This touches on the key issue for me, whilst there are ways to make the feature technically not breaking, in practice for the feature to work it will function like a breaking change. Every operator in DF, etc... will need to be updated to propogate this context, etc...

My confusion remains how this is practically any different from users registering the memory allocation after it emerges from the kernel.

The main advantage is that callers creating Arrays in multiple locations of their project won't need to remember to claim data manually at each location. This also helps track intermediate allocations. For instance, if a kernel merges two arrays, it will probably need to allocate intermediate buffers (like null buffers or offset buffers). Without automatic claiming, the pool will only report that memory is exhausted after the kernel returns and intermediate allocations are complete, rather than detecting exhaustion while attempting the intermediate allocation itself.

Another option I thought of is exposing the MemoryPool API at the RecordBatch level. However, this would miss any Array creations that happen within the operator itself (for example, Accumulators in DataFusion), giving incomplete accounting.

The main advantage is that callers creating Arrays in multiple locations of their project won't need to remember to claim data manually at each location. This also helps track intermediate allocations. For instance, if a kernel merges two arrays, it will probably need to allocate intermediate buffers (like null buffers or offset buffers). Without automatic claiming, the pool will only report that memory is exhausted after the kernel returns and intermediate allocations are complete, rather than detecting exhaustion while attempting the intermediate allocation itself.

I agree with this assesment

Another potential API we could try is implement wrapper APIs in datafusion -- for example, we could implement the

fn cast(context: & MemoryContet, array: &dyn Array, target_type: &DataType) -> Result<ArrayRef> {
  // call arrow cast, then claim memory with pool
  // return arr
}

Then we could audit the DataFusion code and ensure that there are no explicit calls to the arrow kernels directly in Datafusion code but we routed them through the wrappers

Depending on how that works out, we can then contemplate porting such APIs upstream into Arrow.

There is already some version of these wrappers here:

https://github.com/apache/datafusion/blob/7716cae50403c68c176afbb3987bd38abbbaeac0/datafusion/physical-expr/src/expressions/binary/kernels.rs#L18-L19

(I would also selfishly like the wrappers as an excuse / way to test out kernels that skip new allocations -- that reuse existing buffer allocations when possible)

@LiaCastaneda

@LiaCastaneda and I discussed in person a similar approach to that, and how adding a clippy rule that forbids the usage of the raw kernels and allocators from arrow-rs in favor of DataFusion wrappers could prevent other devs from shooting themselves in the foot.

I think it's worth a try, WDYT @LiaCastaneda?

yeah, I can give it a try.

Are the wrapper APIs implemented in datafusion the way forward for this issue or is there another design considered?

@ariel-miculas

Hi @ariel-miculas, I’m not familiar with the wrapper APIs. The approach I considered was just wrapping the Arrow kernels used by DataFusion to use Arrow’s memory pool (whenever an array was created on each DataFusion ExecutionPlan node) while also accounting for it in the DataFusion memory pool. I have a rough draft of my attempt on this branch.

The main blocker was getting both the Arrow memory pool and DataFusion’s MemoryPool, MemoryReservation, and memory registration to work together in a sound way. I ended up dropping the implementation since it was becoming burdensome and would require deeper thought into DataFusion’s memory accounting design.

I’m not familiar with the wrapper APIs.

I was referring to what Andrew Lamb mentioned a few comments above: #8938 (comment)

I see this PR that implements arrow_buffer::MemoryPool for MemoryPool, not sure how it's supposed to fit into the bigger picture

Sorry for the late response, I think that PR wires the arrow MemoryPool together with the DataFusion MemoryPool, so when you call arrow's claim on any array it gets also counted into the DF memory pool. We should probably find a way to make it automatic since you'd still need to call claim on the arrays.

Navigation Menu

Repository navigation

Enable automatic memory accounting using the memory pool when creating arrays #8938

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Enable automatic memory accounting using the memory pool when creating arrays #8938

Description

Activity

LiaCastaneda commented on Dec 1, 2025

tustvold commented on Dec 1, 2025

LiaCastaneda commented on Dec 1, 2025

notfilippo commented on Dec 2, 2025

LiaCastaneda commented on Dec 3, 2025

notfilippo commented on Dec 3, 2025

gabotechs commented on Dec 31, 2025

tustvold commented on Dec 31, 2025

gabotechs commented on Jan 5, 2026

LiaCastaneda commented on Jan 5, 2026

tustvold commented on Jan 5, 2026

gabotechs commented on Jan 5, 2026

alamb commented on Jan 5, 2026

LiaCastaneda commented on Jan 12, 2026

tustvold commented on Jan 12, 2026

LiaCastaneda commented on Jan 12, 2026

alamb commented on Jan 12, 2026

alamb commented on Jan 12, 2026

gabotechs commented on Jan 12, 2026

LiaCastaneda commented on Jan 12, 2026

ariel-miculas commented on Apr 21, 2026

LiaCastaneda commented on Apr 21, 2026

ariel-miculas commented on Apr 21, 2026

LiaCastaneda commented on May 22, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions