This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new ea3c0509bc Add `claim` method to recordbatch for memory accounting
(#9433)
ea3c0509bc is described below
commit ea3c0509bcee34e1e85152db56d085c19ae05e9c
Author: Peter L <[email protected]>
AuthorDate: Thu Mar 19 05:44:35 2026 +1030
Add `claim` method to recordbatch for memory accounting (#9433)
# Which issue does this PR close?
None specifically but aligns with some of the changes in
https://github.com/apache/arrow-rs/issues/8137
# Rationale for this change
It should be easy to claim a `RecordBatch` in totality with an arrow
memory pool
# What changes are included in this PR?
Adds a few methods to bubble up the `claim` to `RecordBatch` level if
the `pool` feature is enabled.
# Are these changes tested?
Yes & new tests added
# Are there any user-facing changes?
If `pool` feature is added, a new `claim` method on `RecordBatch` and
associated structs
---
arrow-array/src/record_batch.rs | 14 ++++++++++++++
1 file changed, 14 insertions(+)
diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs
index 780e14fd4f..f400ac4d0d 100644
--- a/arrow-array/src/record_batch.rs
+++ b/arrow-array/src/record_batch.rs
@@ -795,6 +795,20 @@ impl RecordBatch {
RecordBatch::try_new(schema, columns)
}
+ /// Registers all buffers in this record batch with the provided
[`MemoryPool`].
+ ///
+ /// This claims memory for all columns in the batch by calling
[`Array::claim`]
+ /// on each column.
+ ///
+ /// [`MemoryPool`]: arrow_buffer::MemoryPool
+ /// [`Array::claim`]: crate::Array::claim
+ #[cfg(feature = "pool")]
+ pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
+ for column in self.columns() {
+ column.claim(pool);
+ }
+ }
+
/// Returns the total number of bytes of memory occupied physically by
this batch.
///
/// Note that this does not always correspond to the exact memory usage of
a