[arrow-datafusion] branch main updated: Minor: deprecate `batch_byte_size` (#7245)

[yjshen]

This is an automated email from the ASF dual-hosted git repository.

yjshen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 37b6516efd Minor: deprecate `batch_byte_size` (#7245)
37b6516efd is described below

commit 37b6516efded94af1019518e7c35520712ea42ca
Author: Andrew Lamb <and...@nerdnetworks.org>
AuthorDate: Wed Aug 9 11:16:22 2023 -0500

    Minor: deprecate `batch_byte_size` (#7245)
---
 datafusion/core/src/physical_plan/common.rs     | 15 ++++++++-------
 datafusion/core/src/physical_plan/sorts/sort.rs |  4 ++--
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/datafusion/core/src/physical_plan/common.rs 
b/datafusion/core/src/physical_plan/common.rs
index 46dbc9ef62..9e94fe4e25 100644
--- a/datafusion/core/src/physical_plan/common.rs
+++ b/datafusion/core/src/physical_plan/common.rs
@@ -131,7 +131,11 @@ pub fn compute_record_batch_statistics(
 ) -> Statistics {
     let nb_rows = batches.iter().flatten().map(RecordBatch::num_rows).sum();
 
-    let total_byte_size = batches.iter().flatten().map(batch_byte_size).sum();
+    let total_byte_size = batches
+        .iter()
+        .flatten()
+        .map(|b| b.get_array_memory_size())
+        .sum();
 
     let projection = match projection {
         Some(p) => p,
@@ -340,7 +344,7 @@ impl IPCWriter {
         self.writer.write(batch)?;
         self.num_batches += 1;
         self.num_rows += batch.num_rows() as u64;
-        let num_bytes: usize = batch_byte_size(batch);
+        let num_bytes: usize = batch.get_array_memory_size();
         self.num_bytes += num_bytes as u64;
         Ok(())
     }
@@ -357,12 +361,9 @@ impl IPCWriter {
 }
 
 /// Returns the total number of bytes of memory occupied physically by this 
batch.
+#[deprecated(since = "28.0.0", note = "RecordBatch::get_array_memory_size")]
 pub fn batch_byte_size(batch: &RecordBatch) -> usize {
-    batch
-        .columns()
-        .iter()
-        .map(|array| array.get_array_memory_size())
-        .sum()
+    batch.get_array_memory_size()
 }
 
 #[cfg(test)]
diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs 
b/datafusion/core/src/physical_plan/sorts/sort.rs
index 411a425b51..1813fd6a15 100644
--- a/datafusion/core/src/physical_plan/sorts/sort.rs
+++ b/datafusion/core/src/physical_plan/sorts/sort.rs
@@ -19,7 +19,7 @@
 //! It will do in-memory sorting if it has enough memory budget
 //! but spills to disk if needed.
 
-use crate::physical_plan::common::{batch_byte_size, spawn_buffered, IPCWriter};
+use crate::physical_plan::common::{spawn_buffered, IPCWriter};
 use crate::physical_plan::expressions::PhysicalSortExpr;
 use crate::physical_plan::metrics::{
     BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
@@ -279,7 +279,7 @@ impl ExternalSorter {
         }
         self.reserve_memory_for_merge()?;
 
-        let size = batch_byte_size(&input);
+        let size = input.get_array_memory_size();
         if self.reservation.try_grow(size).is_err() {
             let before = self.reservation.size();
             self.in_mem_sort().await?;