refactor: move metrics module to datafusion-common crate (#19247)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

## Rationale for this change

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
#18456 requires this
refactor. Now we have dependency physical-plan -> physical-expr ->
common, and metrics module lives in physical-plan. We want metrics
module to directly accessible from physical-expr crate, so they should
be moved to common crate.

For compile time concern, I've checked `cargo build --profile
release-nonlto --timings` benchmark, compiling `common` crate takes 5s
in 3min total, so it has no noticeable impact. They're all keep in this
project wide root `datafusion-common` dependency for simplicity.

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
- Metrics module depends on memory util `get_record_batch_memory_size`,
so it's first moved to `common` crate
- Move metrics module to `common` crate

Above moves are all re-exported in the original place to keep the public
APIs backward-compatible

## Are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

Author: 2010YOUY01

Cargo.lock

@@ -17,9 +17,11 @@
 
 //! This module provides a function to estimate the memory size of a HashTable prior to allocation
 
-use crate::Result;
 use crate::error::_exec_datafusion_err;
-use std::mem::size_of;
+use crate::{HashSet, Result};
+use arrow::array::ArrayData;
+use arrow::record_batch::RecordBatch;
+use std::{mem::size_of, ptr::NonNull};
 
 /// Estimates the memory size required for a hash table prior to allocation.
 ///
@@ -99,6 +101,74 @@ pub fn estimate_memory_size<T>(num_elements: usize, fixed_size: usize) -> Result
         })
 }
 
+/// Calculate total used memory of this batch.
+///
+/// This function is used to estimate the physical memory usage of the `RecordBatch`.
+/// It only counts the memory of large data `Buffer`s, and ignores metadata like
+/// types and pointers.
+/// The implementation will add up all unique `Buffer`'s memory
+/// size, due to:
+/// - The data pointer inside `Buffer` are memory regions returned by global memory
+///   allocator, those regions can't have overlap.
+/// - The actual used range of `ArrayRef`s inside `RecordBatch` can have overlap
+///   or reuse the same `Buffer`. For example: taking a slice from `Array`.
+///
+/// Example:
+/// For a `RecordBatch` with two columns: `col1` and `col2`, two columns are pointing
+/// to a sub-region of the same buffer.
+///
+/// {xxxxxxxxxxxxxxxxxxx} <--- buffer
+///       ^    ^  ^    ^
+///       |    |  |    |
+/// col1->{    }  |    |
+/// col2--------->{    }
+///
+/// In the above case, `get_record_batch_memory_size` will return the size of
+/// the buffer, instead of the sum of `col1` and `col2`'s actual memory size.
+///
+/// Note: Current `RecordBatch`.get_array_memory_size()` will double count the
+/// buffer memory size if multiple arrays within the batch are sharing the same
+/// `Buffer`. This method provides temporary fix until the issue is resolved:
+/// <https://github.com/apache/arrow-rs/issues/6439>
+pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize {
+    // Store pointers to `Buffer`'s start memory address (instead of actual
+    // used data region's pointer represented by current `Array`)
+    let mut counted_buffers: HashSet<NonNull<u8>> = HashSet::new();
+    let mut total_size = 0;
+
+    for array in batch.columns() {
+        let array_data = array.to_data();
+        count_array_data_memory_size(&array_data, &mut counted_buffers, &mut total_size);
+    }
+
+    total_size
+}
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(
+    array_data: &ArrayData,
+    counted_buffers: &mut HashSet<NonNull<u8>>,
+    total_size: &mut usize,
+) {
+    // Count memory usage for `array_data`
+    for buffer in array_data.buffers() {
+        if counted_buffers.insert(buffer.data_ptr()) {
+            *total_size += buffer.capacity();
+        } // Otherwise the buffer's memory is already counted
+    }
+
+    if let Some(null_buffer) = array_data.nulls()
+        && counted_buffers.insert(null_buffer.inner().inner().data_ptr())
+    {
+        *total_size += null_buffer.inner().inner().capacity();
+    }
+
+    // Count all children `ArrayData` recursively
+    for child in array_data.child_data() {
+        count_array_data_memory_size(child, counted_buffers, total_size);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::{collections::HashSet, mem::size_of};
@@ -132,3 +202,129 @@ mod tests {
         assert!(estimated.is_err());
     }
 }
+
+#[cfg(test)]
+mod record_batch_tests {
+    use super::*;
+    use arrow::array::{Float64Array, Int32Array, ListArray};
+    use arrow::datatypes::{DataType, Field, Int32Type, Schema};
+    use std::sync::Arc;
+
+    #[test]
+    fn test_get_record_batch_memory_size() {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("ints", DataType::Int32, true),
+            Field::new("float64", DataType::Float64, false),
+        ]));
+
+        let int_array =
+            Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
+        let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
+
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(int_array), Arc::new(float64_array)],
+        )
+        .unwrap();
+
+        let size = get_record_batch_memory_size(&batch);
+        assert_eq!(size, 60);
+    }
+
+    #[test]
+    fn test_get_record_batch_memory_size_with_null() {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("ints", DataType::Int32, true),
+            Field::new("float64", DataType::Float64, false),
+        ]));
+
+        let int_array = Int32Array::from(vec![None, Some(2), Some(3)]);
+        let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0]);
+
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(int_array), Arc::new(float64_array)],
+        )
+        .unwrap();
+
+        let size = get_record_batch_memory_size(&batch);
+        assert_eq!(size, 100);
+    }
+
+    #[test]
+    fn test_get_record_batch_memory_size_empty() {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "ints",
+            DataType::Int32,
+            false,
+        )]));
+
+        let int_array: Int32Array = Int32Array::from(vec![] as Vec<i32>);
+        let batch = RecordBatch::try_new(schema, vec![Arc::new(int_array)]).unwrap();
+
+        let size = get_record_batch_memory_size(&batch);
+        assert_eq!(size, 0, "Empty batch should have 0 memory size");
+    }
+
+    #[test]
+    fn test_get_record_batch_memory_size_shared_buffer() {
+        let original = Int32Array::from(vec![1, 2, 3, 4, 5]);
+        let slice1 = original.slice(0, 3);
+        let slice2 = original.slice(2, 3);
+
+        let schema_origin = Arc::new(Schema::new(vec![Field::new(
+            "origin_col",
+            DataType::Int32,
+            false,
+        )]));
+        let batch_origin =
+            RecordBatch::try_new(schema_origin, vec![Arc::new(original)]).unwrap();
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("slice1", DataType::Int32, false),
+            Field::new("slice2", DataType::Int32, false),
+        ]));
+
+        let batch_sliced =
+            RecordBatch::try_new(schema, vec![Arc::new(slice1), Arc::new(slice2)])
+                .unwrap();
+
+        let size_origin = get_record_batch_memory_size(&batch_origin);
+        let size_sliced = get_record_batch_memory_size(&batch_sliced);
+
+        assert_eq!(size_origin, size_sliced);
+    }
+
+    #[test]
+    fn test_get_record_batch_memory_size_nested_array() {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new(
+                "nested_int",
+                DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
+                false,
+            ),
+            Field::new(
+                "nested_int2",
+                DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
+                false,
+            ),
+        ]));
+
+        let int_list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+            Some(vec![Some(1), Some(2), Some(3)]),
+        ]);
+
+        let int_list_array2 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+            Some(vec![Some(4), Some(5), Some(6)]),
+        ]);
+
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(int_list_array), Arc::new(int_list_array2)],
+        )
+        .unwrap();
+
+        let size = get_record_batch_memory_size(&batch);
+        assert_eq!(size, 8208);
+    }
+}

datafusion/common/src/utils/memory.rs

@@ -51,6 +51,7 @@ sql = []
 [dependencies]
 arrow = { workspace = true }
 async-trait = { workspace = true }
+chrono = { workspace = true }
 dashmap = { workspace = true }
 datafusion-common = { workspace = true, default-features = false }
 datafusion-expr = { workspace = true, default-features = false }
@@ -64,5 +65,4 @@ tempfile = { workspace = true }
 url = { workspace = true }
 
 [dev-dependencies]
-chrono = { workspace = true }
 insta = { workspace = true }

datafusion/execution/Cargo.toml

@@ -32,6 +32,7 @@ pub mod cache;
 pub mod config;
 pub mod disk_manager;
 pub mod memory_pool;
+pub mod metrics;
 pub mod object_store;
 #[cfg(feature = "parquet_encryption")]
 pub mod parquet_encryption;

datafusion/execution/src/lib.rs

@@ -20,18 +20,16 @@
 use std::task::Poll;
 
 use arrow::record_batch::RecordBatch;
-
-use crate::spill::get_record_batch_memory_size;
+use datafusion_common::{Result, utils::memory::get_record_batch_memory_size};
 
 use super::{Count, ExecutionPlanMetricsSet, MetricBuilder, Time, Timestamp};
-use datafusion_common::Result;
 
 /// Helper for creating and tracking common "baseline" metrics for
 /// each operator
 ///
 /// Example:
 /// ```
-/// use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
+/// use datafusion_execution::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
 /// let metrics = ExecutionPlanMetricsSet::new();
 ///
 /// let partition = 2;
@@ -203,7 +201,7 @@ impl SpillMetrics {
     }
 }
 
-/// Metrics for tracking [`crate::stream::BatchSplitStream`] activity
+/// Metrics for tracking batch splitting activity
 #[derive(Debug, Clone)]
 pub struct SplitMetrics {
     /// Number of times an input [`RecordBatch`] was split

…on/physical-plan/src/metrics/baseline.rs …fusion/execution/src/metrics/baseline.rsdatafusion/physical-plan/src/metrics/baseline.rs renamed to datafusion/execution/src/metrics/baseline.rs datafusion/physical-plan/src/metrics/baseline.rs renamed to datafusion/execution/src/metrics/baseline.rs

@@ -34,7 +34,7 @@ use super::{
 /// case of constant strings
 ///
 /// ```rust
-/// use datafusion_physical_plan::metrics::*;
+/// use datafusion_execution::metrics::*;
 ///
 /// let metrics = ExecutionPlanMetricsSet::new();
 /// let partition = 1;

…ion/physical-plan/src/metrics/builder.rs …afusion/execution/src/metrics/builder.rsdatafusion/physical-plan/src/metrics/builder.rs renamed to datafusion/execution/src/metrics/builder.rs datafusion/physical-plan/src/metrics/builder.rs renamed to datafusion/execution/src/metrics/builder.rs

@@ -44,7 +44,7 @@ use std::{any::Any, fmt::Debug, fmt::Display, sync::Arc};
 /// # use std::any::Any;
 /// # use std::sync::atomic::{AtomicUsize, Ordering};
 ///
-/// # use datafusion_physical_plan::metrics::CustomMetricValue;
+/// # use datafusion_execution::metrics::CustomMetricValue;
 ///
 /// #[derive(Debug, Default)]
 /// struct MyCounter {

…sion/physical-plan/src/metrics/custom.rs …tafusion/execution/src/metrics/custom.rsdatafusion/physical-plan/src/metrics/custom.rs renamed to datafusion/execution/src/metrics/custom.rs datafusion/physical-plan/src/metrics/custom.rs renamed to datafusion/execution/src/metrics/custom.rs

@@ -22,16 +22,16 @@ mod builder;
 mod custom;
 mod value;
 
+use datafusion_common::HashMap;
 use parking_lot::Mutex;
 use std::{
     borrow::Cow,
     fmt::{Debug, Display},
     sync::Arc,
 };
 
-use datafusion_common::HashMap;
-
 // public exports
+
 pub use baseline::{BaselineMetrics, RecordOutput, SpillMetrics, SplitMetrics};
 pub use builder::MetricBuilder;
 pub use custom::CustomMetricValue;
@@ -40,15 +40,14 @@ pub use value::{
     ScopedTimerGuard, Time, Timestamp,
 };
 
-/// Something that tracks a value of interest (metric) of a DataFusion
-/// [`ExecutionPlan`] execution.
+/// Something that tracks a value of interest (metric) during execution.
 ///
 /// Typically [`Metric`]s are not created directly, but instead
 /// are created using [`MetricBuilder`] or methods on
 /// [`ExecutionPlanMetricsSet`].
 ///
 /// ```
-/// use datafusion_physical_plan::metrics::*;
+/// use datafusion_execution::metrics::*;
 ///
 /// let metrics = ExecutionPlanMetricsSet::new();
 /// assert!(metrics.clone_inner().output_rows().is_none());
@@ -66,8 +65,6 @@ pub use value::{
 /// // As well as from the metrics set
 /// assert_eq!(metrics.clone_inner().output_rows(), Some(13));
 /// ```
-///
-/// [`ExecutionPlan`]: super::ExecutionPlan
 
 #[derive(Debug)]
 pub struct Metric {
@@ -204,9 +201,7 @@ impl Metric {
     }
 }
 
-/// A snapshot of the metrics for a particular ([`ExecutionPlan`]).
-///
-/// [`ExecutionPlan`]: super::ExecutionPlan
+/// A snapshot of the metrics for a particular execution plan.
 #[derive(Default, Debug, Clone)]
 pub struct MetricsSet {
     metrics: Vec<Arc<Metric>>,
@@ -401,17 +396,14 @@ impl Display for MetricsSet {
     }
 }
 
-/// A set of [`Metric`]s for an individual "operator" (e.g. `&dyn
-/// ExecutionPlan`).
+/// A set of [`Metric`]s for an individual operator.
 ///
-/// This structure is intended as a convenience for [`ExecutionPlan`]
+/// This structure is intended as a convenience for execution plan
 /// implementations so they can generate different streams for multiple
 /// partitions but easily report them together.
 ///
 /// Each `clone()` of this structure will add metrics to the same
 /// underlying metrics set
-///
-/// [`ExecutionPlan`]: super::ExecutionPlan
 #[derive(Default, Debug, Clone)]
 pub struct ExecutionPlanMetricsSet {
     inner: Arc<Mutex<MetricsSet>>,

…afusion/physical-plan/src/metrics/mod.rs datafusion/execution/src/metrics/mod.rsdatafusion/physical-plan/src/metrics/mod.rs renamed to datafusion/execution/src/metrics/mod.rs datafusion/physical-plan/src/metrics/mod.rs renamed to datafusion/execution/src/metrics/mod.rs

@@ -19,9 +19,8 @@
 
 use super::CustomMetricValue;
 use chrono::{DateTime, Utc};
-use datafusion_common::instant::Instant;
 use datafusion_common::{
-    human_readable_count, human_readable_duration, human_readable_size,
+    human_readable_count, human_readable_duration, human_readable_size, instant::Instant,
 };
 use parking_lot::Mutex;
 use std::{

…usion/physical-plan/src/metrics/value.rs datafusion/execution/src/metrics/value.rsdatafusion/physical-plan/src/metrics/value.rs renamed to datafusion/execution/src/metrics/value.rs datafusion/physical-plan/src/metrics/value.rs renamed to datafusion/execution/src/metrics/value.rs

@@ -52,7 +52,6 @@ arrow = { workspace = true }
 arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }
 async-trait = { workspace = true }
-chrono = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-common-runtime = { workspace = true, default-features = true }
 datafusion-execution = { workspace = true }