Enable automatic memory accounting using the memory pool when creating arrays

[LiaCastaneda]

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

[LiaCastaneda]

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

[tustvold]

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.

[LiaCastaneda]

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

[notfilippo]

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.

[LiaCastaneda]

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.

[notfilippo]

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.

[gabotechs]

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

[tustvold]

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.

[gabotechs]

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?

[LiaCastaneda]

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 👍

[tustvold]

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.

[gabotechs]

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.