datafusion_distributed/metrics/
bytes_metric.rs

1use std::{
2    any::Any,
3    borrow::Cow,
4    fmt::{Display, Formatter},
5    sync::{Arc, atomic::AtomicUsize},
6};
7
8use datafusion::physical_plan::Metric;
9use datafusion::{
10    common::human_readable_size,
11    physical_plan::metrics::{CustomMetricValue, MetricBuilder, MetricValue},
12};
13use std::sync::atomic::Ordering::Relaxed;
14
15/// Extension trait for DataFusion's metric system that adds support for byte count metrics
16/// that display using human-readable byte sizes (KB, MB, GB) instead of plain count notation.
17pub trait BytesMetricExt {
18    fn bytes_counter(self, name: impl Into<Cow<'static, str>>) -> BytesCounterMetric;
19}
20
21impl BytesMetricExt for MetricBuilder<'_> {
22    fn bytes_counter(self, name: impl Into<Cow<'static, str>>) -> BytesCounterMetric {
23        let value = BytesCounterMetric::default();
24        self.build(MetricValue::Custom {
25            name: name.into(),
26            value: Arc::new(value.clone()),
27        });
28        value
29    }
30}
31/// A cumulative counter metric for tracking byte counts.
32///
33/// Unlike DataFusion's built-in [`Count`](datafusion::physical_plan::metrics::Count) which formats
34/// large values using plain count notation (e.g., "1.91 B" meaning 1.91 billion), this metric
35/// formats values using [`human_readable_size`] (e.g., "1.91 GB").
36///
37/// This avoids the confusing display where "B" (billions) looks like "bytes".
38///
39/// Aggregation sums values across partitions/tasks.
40#[derive(Debug, Clone)]
41pub struct BytesCounterMetric {
42    bytes: Arc<AtomicUsize>,
43}
44
45impl Default for BytesCounterMetric {
46    fn default() -> Self {
47        Self {
48            bytes: Arc::new(AtomicUsize::new(usize::MIN)),
49        }
50    }
51}
52
53impl BytesCounterMetric {
54    pub fn new_metric(name: impl Into<Cow<'static, str>>, bytes: usize) -> Arc<Metric> {
55        Arc::new(Metric::new(
56            MetricValue::Custom {
57                name: name.into(),
58                value: Arc::new(BytesCounterMetric::from_value(bytes)),
59            },
60            None,
61        ))
62    }
63
64    pub fn from_value(bytes: usize) -> Self {
65        Self {
66            bytes: Arc::new(AtomicUsize::new(bytes)),
67        }
68    }
69
70    pub fn value(&self) -> usize {
71        self.bytes.load(Relaxed)
72    }
73
74    pub fn add_bytes(&self, bytes: usize) -> usize {
75        self.bytes.fetch_add(bytes, Relaxed)
76    }
77}
78
79impl Display for BytesCounterMetric {
80    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
81        write!(f, "{}", human_readable_size(self.value()))
82    }
83}
84
85impl CustomMetricValue for BytesCounterMetric {
86    fn new_empty(&self) -> Arc<dyn CustomMetricValue> {
87        Arc::new(BytesCounterMetric::default())
88    }
89
90    fn aggregate(&self, other: Arc<dyn CustomMetricValue + 'static>) {
91        let Some(other) = other.as_any().downcast_ref::<Self>() else {
92            return;
93        };
94        self.bytes.fetch_add(other.bytes.load(Relaxed), Relaxed);
95    }
96
97    fn as_any(&self) -> &dyn Any {
98        self
99    }
100
101    fn as_usize(&self) -> usize {
102        self.value()
103    }
104
105    fn is_eq(&self, other: &Arc<dyn CustomMetricValue>) -> bool {
106        let Some(other) = other.as_any().downcast_ref::<Self>() else {
107            return false;
108        };
109        other.value() == self.value()
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn default_is_zero_and_add_accumulates() {
119        let m = BytesCounterMetric::default();
120        assert_eq!(m.value(), 0);
121        m.add_bytes(1024);
122        m.add_bytes(2048);
123        assert_eq!(m.value(), 3072);
124    }
125
126    #[test]
127    fn from_value_constructs_correctly() {
128        let m = BytesCounterMetric::from_value(1_000_000);
129        assert_eq!(m.value(), 1_000_000);
130    }
131
132    #[test]
133    fn aggregate_sums_values() {
134        let a = BytesCounterMetric::from_value(500);
135        let b = BytesCounterMetric::from_value(300);
136        a.aggregate(Arc::new(b));
137        assert_eq!(a.value(), 800);
138    }
139
140    #[test]
141    fn display_uses_human_readable_size() {
142        // 0 bytes
143        assert_eq!(format!("{}", BytesCounterMetric::from_value(0)), "0.0 B");
144        // 4 MB (>= 2*MB threshold, so displays in MB)
145        assert_eq!(
146            format!("{}", BytesCounterMetric::from_value(4 * 1024 * 1024)),
147            "4.0 MB"
148        );
149        // 4 GB (>= 2*GB threshold, so displays in GB)
150        assert_eq!(
151            format!("{}", BytesCounterMetric::from_value(4 * 1024 * 1024 * 1024)),
152            "4.0 GB"
153        );
154    }
155}
datafusion_distributed/metrics/bytes_metric.rs

datafusion_distributed/metrics/
bytes_metric.rs
