add used()

Author: LiaCastaneda

datafusion/physical-plan/src/aggregates/no_grouping.rs

@@ -25,7 +25,7 @@ use crate::metrics::{BaselineMetrics, RecordOutput};
 use crate::{RecordBatchStream, SendableRecordBatchStream};
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
-use arrow_buffer::TrackingMemoryPool;
+use arrow_buffer::{MemoryPool, TrackingMemoryPool};
 use datafusion_common::{Result, ScalarValue, internal_datafusion_err, internal_err};
 use datafusion_execution::TaskContext;
 use datafusion_expr::Operator;
@@ -437,6 +437,7 @@ fn aggregate_batch(
     arrow_pool: &TrackingMemoryPool,
 ) -> Result<usize> {
     let mut allocated = 0usize;
+    let pool_size_pre = arrow_pool.used();
 
     // 1.1 iterate accumulators and respective expressions together
     // 1.2 filter the batch if necessary
@@ -473,5 +474,8 @@ fn aggregate_batch(
             res
         })?;
 
+    let pool_size_post = arrow_pool.used();
+    allocated += pool_size_post.saturating_sub(pool_size_pre);
+
     Ok(allocated)
 }

datafusion/physical-plan/src/aggregates/row_hash.rs

@@ -38,7 +38,7 @@ use crate::{RecordBatchStream, SendableRecordBatchStream};
 
 use arrow::array::*;
 use arrow::datatypes::SchemaRef;
-use arrow_buffer::TrackingMemoryPool;
+use arrow_buffer::{MemoryPool, TrackingMemoryPool};
 use datafusion_common::{
     DataFusionError, Result, assert_eq_or_internal_err, assert_or_internal_err,
     internal_err,
@@ -1069,7 +1069,8 @@ impl GroupedHashAggregateStream {
                 .accumulators
                 .iter()
                 .map(|x| x.size(Some(&self.arrow_pool)))
-                .sum::<usize>();
+                .sum::<usize>()
+            + self.arrow_pool.used();
 
         let reservation_result = self.reservation.try_resize(total_size);