Try arrow_pool approach

Author: LiaCastaneda

Cargo.lock

@@ -42,7 +42,8 @@ name = "datafusion_expr_common"
 
 [dependencies]
 arrow = { workspace = true }
-datafusion-common = { workspace = true }
 indexmap = { workspace = true }
+datafusion-common = { workspace = true }
 itertools = { workspace = true }
 paste = { workspace = true }
+arrow-buffer = { workspace = true, features = ["pool"] }

datafusion/expr-common/Cargo.toml

@@ -18,6 +18,7 @@
 //! Accumulator module contains the trait definition for aggregation function's accumulators.
 
 use arrow::array::ArrayRef;
+use arrow_buffer::MemoryPool;
 use datafusion_common::{internal_err, Result, ScalarValue};
 use std::fmt::Debug;
 
@@ -71,16 +72,34 @@ pub trait Accumulator: Send + Sync + Debug {
     /// when possible (for example distinct strings)
     fn evaluate(&mut self) -> Result<ScalarValue>;
 
-    /// Returns the allocated size required for this accumulator, in
-    /// bytes, including `Self`.
+    /// Returns the size of non-Arrow allocations in bytes, including `Self`.
     ///
     /// This value is used to calculate the memory used during
     /// execution so DataFusion can stay within its allotted limit.
     ///
+    /// This includes Vec capacity, BufferBuilder capacity, and other
+    /// non-Arrow data structures. Arrow Buffer memory should be tracked
+    /// separately via [`claim_buffers`].
+    ///
     /// "Allocated" means that for internal containers such as `Vec`,
     /// the `capacity` should be used not the `len`.
+    ///
+    /// [`claim_buffers`]: Self::claim_buffers
     fn size(&self) -> usize;
 
+    /// Claim Arrow buffers with the memory pool for accurate tracking.
+    ///
+    /// This method should be called to register Arrow Buffer instances
+    /// with the pool, enabling automatic deduplication of shared buffers.
+    ///
+    /// # Default Implementation
+    ///
+    /// The default implementation does nothing, which is appropriate for
+    /// accumulators that don't store Arrow arrays (e.g., simple numeric accumulators).
+    fn claim_buffers(&self, _pool: &dyn MemoryPool) {
+        // Default: no-op for accumulators without Arrow buffers
+    }
+
     /// Returns the intermediate state of the accumulator, consuming the
     /// intermediate state.
     ///

datafusion/expr-common/src/accumulator.rs

@@ -18,6 +18,7 @@
 //! Vectorized [`GroupsAccumulator`]
 
 use arrow::array::{ArrayRef, BooleanArray};
+use arrow_buffer::MemoryPool as ArrowMemoryPool;
 use datafusion_common::{not_impl_err, Result};
 
 /// Describes how many rows should be emitted during grouping.
@@ -246,8 +247,33 @@ pub trait GroupsAccumulator: Send {
 
     /// Amount of memory used to store the state of this accumulator,
     /// in bytes.
+    /// Returns the size of non-Arrow allocations in bytes.
+    ///
+    /// This includes Vec capacity, BufferBuilder capacity, and other
+    /// non-Arrow data structures. Arrow Buffer memory should be tracked
+    /// separately via [`claim_buffers`].
     ///
     /// This function is called once per batch, so it should be `O(n)` to
     /// compute, not `O(num_groups)`
+    ///
+    /// [`claim_buffers`]: GroupsAccumulator::claim_buffers
     fn size(&self) -> usize;
+
+    /// Claims all internal Arrow buffers with the provided memory pool.
+    ///
+    /// This method should call `.claim()` on all Arrow `Buffer`s held by
+    /// this accumulator to enable accurate memory tracking that automatically
+    /// deduplicates shared buffers.
+    ///
+    /// # Arguments
+    ///
+    /// * `pool` - Arrow memory pool to claim buffers with
+    ///
+    /// # Default Implementation
+    ///
+    /// The default implementation does nothing, which is appropriate for
+    /// accumulators that don't store Arrow arrays (e.g., simple numeric accumulators).
+    fn claim_buffers(&self, _pool: &dyn ArrowMemoryPool) {
+        // Default: no-op for accumulators without Arrow buffers
+    }
 }

datafusion/expr-common/src/groups_accumulator.rs

@@ -43,6 +43,7 @@ name = "datafusion_functions_aggregate_common"
 [dependencies]
 ahash = { workspace = true }
 arrow = { workspace = true }
+arrow-buffer = { workspace = true, features = ["pool"] }
 datafusion-common = { workspace = true }
 datafusion-expr-common = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }

datafusion/functions-aggregate-common/Cargo.toml

@@ -32,6 +32,7 @@ use arrow::{
     compute::take_arrays,
     datatypes::UInt32Type,
 };
+use arrow_buffer::MemoryPool;
 use datafusion_common::{arrow_datafusion_err, DataFusionError, Result, ScalarValue};
 use datafusion_expr_common::accumulator::Accumulator;
 use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
@@ -395,6 +396,13 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
         self.allocation_bytes
     }
 
+    fn claim_buffers(&self, pool: &dyn MemoryPool) {
+        // Delegate to each wrapped Accumulator
+        for state in &self.states {
+            state.accumulator.claim_buffers(pool);
+        }
+    }
+
     fn convert_to_state(
         &self,
         values: &[ArrayRef],

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs

@@ -43,6 +43,7 @@ name = "datafusion_functions_aggregate"
 [dependencies]
 ahash = { workspace = true }
 arrow = { workspace = true }
+arrow-buffer = { workspace = true, features = ["pool"] }
 datafusion-common = { workspace = true }
 datafusion-doc = { workspace = true }
 datafusion-execution = { workspace = true }

datafusion/functions-aggregate/Cargo.toml

@@ -27,6 +27,7 @@ use arrow::array::{
 };
 use arrow::compute::{filter, SortOptions};
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
+use arrow_buffer::MemoryPool;
 
 use datafusion_common::cast::as_list_array;
 use datafusion_common::utils::{
@@ -391,25 +392,17 @@ impl Accumulator for ArrayAggAccumulator {
     fn size(&self) -> usize {
         size_of_val(self)
             + (size_of::<ArrayRef>() * self.values.capacity())
-            + self
-                .values
-                .iter()
-                // Each ArrayRef might be just a reference to a bigger array, and many
-                // ArrayRefs here might be referencing exactly the same array, so if we
-                // were to call `arr.get_array_memory_size()`, we would be double-counting
-                // the same underlying data many times.
-                //
-                // Instead, we do an approximation by estimating how much memory each
-                // ArrayRef would occupy if its underlying data was fully owned by this
-                // accumulator.
-                //
-                // Note that this is just an estimation, but the reality is that this
-                // accumulator might not own any data.
-                .map(|arr| arr.to_data().get_slice_memory_size().unwrap_or_default())
-                .sum::<usize>()
             + self.datatype.size()
             - size_of_val(&self.datatype)
     }
+
+    fn claim_buffers(&self, pool: &dyn MemoryPool) {
+        for arr in &self.values {
+            for buffer in arr.to_data().buffers() {
+                buffer.claim(pool);
+            }
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -1073,10 +1066,25 @@ mod tests {
         acc2.update_batch(&[data(["b", "c", "a"])])?;
         acc1 = merge(acc1, acc2)?;
 
-        assert_eq!(acc1.size(), 266);
+        // size() returns only non-Arrow allocations (Vec capacity, etc.)
+        let non_arrow_size = acc1.size();
+        assert_eq!(non_arrow_size, 236);
+
+        // Verify Arrow buffer tracking - claims full buffer capacity
+        let pool = arrow_buffer::TrackingMemoryPool::default();
+        acc1.claim_buffers(&pool);
+        let arrow_buffer_size = pool.used();
+
+        // With small arrays, buffer over-allocation creates high overhead
+        // This tracks actual physical memory (capacity), not logical slice size
+        assert_eq!(arrow_buffer_size, 2080);
+
+        // Total memory = non-Arrow + Arrow buffers (tracks actual RAM usage)
+        assert_eq!(non_arrow_size + arrow_buffer_size, 2316);
 
         Ok(())
     }
+
     #[test]
     fn does_not_over_account_memory_distinct() -> Result<()> {
         let (mut acc1, mut acc2) = ArrayAggAccumulatorBuilder::string()

datafusion/functions-aggregate/src/array_agg.rs

@@ -48,6 +48,7 @@ name = "datafusion_physical_plan"
 [dependencies]
 ahash = { workspace = true }
 arrow = { workspace = true }
+arrow-buffer = { workspace = true, features = ["pool"] }
 arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }
 async-trait = { workspace = true }

datafusion/physical-plan/Cargo.toml

@@ -24,6 +24,7 @@ use arrow::array::types::{
 };
 use arrow::array::{downcast_primitive, ArrayRef, RecordBatch};
 use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
+use arrow_buffer::MemoryPool;
 use datafusion_common::Result;
 
 use datafusion_expr::EmitTo;
@@ -99,9 +100,22 @@ pub trait GroupValues: Send {
     /// assigned.
     fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()>;
 
-    /// Returns the number of bytes of memory used by this [`GroupValues`]
+    /// Returns the size of non-Arrow allocations in bytes.
+    ///
+    /// This includes Vec capacity, BufferBuilder capacity, and other
+    /// non-Arrow data structures. Arrow Buffer memory should be tracked
+    /// separately via [`claim_buffers`].
+    ///
+    /// [`claim_buffers`]: GroupValues::claim_buffers
     fn size(&self) -> usize;
 
+    /// Claims all internal Arrow buffers with the provided memory pool.
+    ///
+    /// Default implementation does nothing (for builders that don't store arrays with shared buffers).
+    fn claim_buffers(&self, _pool: &dyn MemoryPool) {
+        // Default: no-op
+    }
+
     /// Returns true if this [`GroupValues`] is empty
     fn is_empty(&self) -> bool;