datafusion_physical_expr_common/metrics/
builder.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Builder for creating arbitrary metrics
19
20use std::{borrow::Cow, sync::Arc};
21
22use crate::metrics::{
23    MetricCategory, MetricType,
24    value::{PruningMetrics, RatioMergeStrategy, RatioMetrics},
25};
26
27use super::{
28    Count, ExecutionPlanMetricsSet, Gauge, Label, LabelValue, Metric, MetricValue, Time,
29    Timestamp,
30};
31
32/// Structure for constructing metrics, counters, timers, etc.
33///
34/// Note the use of `Cow<..>` is to avoid allocations in the common
35/// case of constant strings. Dynamically created label strings are shared when
36/// [`Label`] values are cloned.
37///
38/// ```rust
39/// use datafusion_physical_expr_common::metrics::*;
40///
41/// let metrics = ExecutionPlanMetricsSet::new();
42/// let partition = 1;
43///
44/// // Create the standard output_rows metric
45/// let output_rows = MetricBuilder::new(&metrics).output_rows(partition);
46///
47/// // Create a operator specific counter with some labels
48/// let num_bytes = MetricBuilder::new(&metrics)
49///     .with_new_label("filename", "my_awesome_file.parquet")
50///     .counter("num_bytes", partition);
51/// ```
52#[derive(Clone)]
53pub struct MetricBuilder<'a> {
54    /// Location that the metric created by this builder will be added do
55    metrics: &'a ExecutionPlanMetricsSet,
56
57    /// optional partition number
58    partition: Option<usize>,
59
60    /// arbitrary name=value pairs identifying this metric
61    labels: Vec<Label>,
62
63    /// The type controlling the verbosity/category for this builder
64    /// See comments in [`MetricType`] for details
65    metric_type: MetricType,
66
67    /// Semantic category (rows / bytes / timing).
68    /// `None` means "always include" (the default for custom metrics).
69    metric_category: Option<MetricCategory>,
70}
71
72impl<'a> MetricBuilder<'a> {
73    /// Create a new `MetricBuilder` that will register the result of `build()` with the `metrics`
74    ///
75    /// `self.metric_type` controls when such metric is displayed. See comments in
76    /// [`MetricType`] for details.
77    pub fn new(metrics: &'a ExecutionPlanMetricsSet) -> Self {
78        Self {
79            metrics,
80            partition: None,
81            labels: vec![],
82            metric_type: MetricType::Dev,
83            metric_category: None,
84        }
85    }
86
87    /// Add a label to the metric being constructed
88    pub fn with_label(mut self, label: Label) -> Self {
89        self.labels.push(label);
90        self
91    }
92
93    /// Set the metric type to the metric being constructed
94    pub fn with_type(mut self, metric_type: MetricType) -> Self {
95        self.metric_type = metric_type;
96        self
97    }
98
99    /// Set the semantic category for the metric being constructed.
100    ///
101    /// See [`MetricCategory`] for details on the determinism properties
102    /// of each category.
103    pub fn with_category(mut self, category: MetricCategory) -> Self {
104        self.metric_category = Some(category);
105        self
106    }
107
108    /// Add a label to the metric being constructed
109    pub fn with_new_label(
110        self,
111        name: impl Into<Cow<'static, str>>,
112        value: impl Into<Cow<'static, str>>,
113    ) -> Self {
114        self.with_label(Label::new(
115            LabelValue::from(name.into()),
116            LabelValue::from(value.into()),
117        ))
118    }
119
120    /// Set the partition of the metric being constructed
121    pub fn with_partition(mut self, partition: usize) -> Self {
122        self.partition = Some(partition);
123        self
124    }
125
126    /// Consume self and create a metric of the specified value
127    /// registered with the MetricsSet
128    pub fn build(self, value: MetricValue) {
129        let Self {
130            labels,
131            partition,
132            metrics,
133            metric_type,
134            metric_category,
135        } = self;
136        let mut metric =
137            Metric::new_with_labels(value, partition, labels).with_type(metric_type);
138        if let Some(category) = metric_category {
139            metric = metric.with_category(category);
140        }
141        metrics.register(Arc::new(metric));
142    }
143
144    /// Consume self and create a new counter for recording output rows
145    pub fn output_rows(self, partition: usize) -> Count {
146        let count = Count::new();
147        self.with_category(MetricCategory::Rows)
148            .with_partition(partition)
149            .build(MetricValue::OutputRows(count.clone()));
150        count
151    }
152
153    /// Consume self and create a new counter for recording the number of spills
154    /// triggered by an operator
155    pub fn spill_count(self, partition: usize) -> Count {
156        let count = Count::new();
157        self.with_category(MetricCategory::Rows)
158            .with_partition(partition)
159            .build(MetricValue::SpillCount(count.clone()));
160        count
161    }
162
163    /// Consume self and create a new counter for recording the total spilled bytes
164    /// triggered by an operator
165    pub fn spilled_bytes(self, partition: usize) -> Count {
166        let count = Count::new();
167        self.with_category(MetricCategory::Bytes)
168            .with_partition(partition)
169            .build(MetricValue::SpilledBytes(count.clone()));
170        count
171    }
172
173    /// Consume self and create a new counter for recording the total spilled rows
174    /// triggered by an operator
175    pub fn spilled_rows(self, partition: usize) -> Count {
176        let count = Count::new();
177        self.with_category(MetricCategory::Rows)
178            .with_partition(partition)
179            .build(MetricValue::SpilledRows(count.clone()));
180        count
181    }
182
183    /// Consume self and create a new counter for recording total output bytes
184    pub fn output_bytes(self, partition: usize) -> Count {
185        let count = Count::new();
186        self.with_category(MetricCategory::Bytes)
187            .with_partition(partition)
188            .build(MetricValue::OutputBytes(count.clone()));
189        count
190    }
191
192    /// Consume self and create a new counter for recording total output batches
193    pub fn output_batches(self, partition: usize) -> Count {
194        let count = Count::new();
195        self.with_category(MetricCategory::Rows)
196            .with_partition(partition)
197            .build(MetricValue::OutputBatches(count.clone()));
198        count
199    }
200
201    /// Consume self and create a new gauge for reporting current memory usage
202    pub fn mem_used(self, partition: usize) -> Gauge {
203        let gauge = Gauge::new();
204        self.with_category(MetricCategory::Bytes)
205            .with_partition(partition)
206            .build(MetricValue::CurrentMemoryUsage(gauge.clone()));
207        gauge
208    }
209
210    /// Consumes self and creates a new [`Count`] for recording some
211    /// arbitrary metric of an operator.
212    pub fn counter(
213        self,
214        counter_name: impl Into<Cow<'static, str>>,
215        partition: usize,
216    ) -> Count {
217        self.with_partition(partition).global_counter(counter_name)
218    }
219
220    /// Consumes self and creates a new [`Gauge`] for reporting some
221    /// arbitrary metric of an operator.
222    pub fn gauge(
223        self,
224        gauge_name: impl Into<Cow<'static, str>>,
225        partition: usize,
226    ) -> Gauge {
227        self.with_partition(partition).global_gauge(gauge_name)
228    }
229
230    /// Consumes self and creates a new [`Count`] for recording a
231    /// metric of an overall operator (not per partition)
232    pub fn global_counter(self, counter_name: impl Into<Cow<'static, str>>) -> Count {
233        let count = Count::new();
234        self.build(MetricValue::Count {
235            name: counter_name.into(),
236            count: count.clone(),
237        });
238        count
239    }
240
241    /// Consumes self and creates a new [`Gauge`] for reporting a
242    /// metric of an overall operator (not per partition)
243    pub fn global_gauge(self, gauge_name: impl Into<Cow<'static, str>>) -> Gauge {
244        let gauge = Gauge::new();
245        self.build(MetricValue::Gauge {
246            name: gauge_name.into(),
247            gauge: gauge.clone(),
248        });
249        gauge
250    }
251
252    /// Consumes self and creates a new [`Gauge`] for recording peak memory
253    /// usage in bytes.
254    pub fn peak_memory_usage(
255        self,
256        gauge_name: impl Into<Cow<'static, str>>,
257        partition: usize,
258    ) -> Gauge {
259        let gauge = Gauge::new();
260        self.with_category(MetricCategory::Bytes)
261            .with_partition(partition)
262            .build(MetricValue::PeakMemoryUsage {
263                name: gauge_name.into(),
264                gauge: gauge.clone(),
265            });
266        gauge
267    }
268
269    /// Consume self and create a new Timer for recording the elapsed
270    /// CPU time spent by an operator
271    pub fn elapsed_compute(self, partition: usize) -> Time {
272        let time = Time::new();
273        self.with_category(MetricCategory::Timing)
274            .with_partition(partition)
275            .build(MetricValue::ElapsedCompute(time.clone()));
276        time
277    }
278
279    /// Consumes self and creates a new Timer for recording some
280    /// subset of an operators execution time.
281    pub fn subset_time(
282        self,
283        subset_name: impl Into<Cow<'static, str>>,
284        partition: usize,
285    ) -> Time {
286        let time = Time::new();
287        self.with_category(MetricCategory::Timing)
288            .with_partition(partition)
289            .build(MetricValue::Time {
290                name: subset_name.into(),
291                time: time.clone(),
292            });
293        time
294    }
295
296    /// Consumes self and creates a new Timestamp for recording the
297    /// starting time of execution for a partition
298    pub fn start_timestamp(self, partition: usize) -> Timestamp {
299        let timestamp = Timestamp::new();
300        self.with_category(MetricCategory::Timing)
301            .with_partition(partition)
302            .build(MetricValue::StartTimestamp(timestamp.clone()));
303        timestamp
304    }
305
306    /// Consumes self and creates a new Timestamp for recording the
307    /// ending time of execution for a partition
308    pub fn end_timestamp(self, partition: usize) -> Timestamp {
309        let timestamp = Timestamp::new();
310        self.with_category(MetricCategory::Timing)
311            .with_partition(partition)
312            .build(MetricValue::EndTimestamp(timestamp.clone()));
313        timestamp
314    }
315
316    /// Consumes self and creates a new `PruningMetrics`
317    pub fn pruning_metrics(
318        self,
319        name: impl Into<Cow<'static, str>>,
320        partition: usize,
321    ) -> PruningMetrics {
322        let pruning_metrics = PruningMetrics::new();
323        self.with_category(MetricCategory::Rows)
324            .with_partition(partition)
325            .build(MetricValue::PruningMetrics {
326                name: name.into(),
327                // inner values will be `Arc::clone()`
328                pruning_metrics: pruning_metrics.clone(),
329            });
330        pruning_metrics
331    }
332
333    /// Consumes self and creates a new [`RatioMetrics`]
334    pub fn ratio_metrics(
335        self,
336        name: impl Into<Cow<'static, str>>,
337        partition: usize,
338    ) -> RatioMetrics {
339        self.ratio_metrics_with_strategy(name, partition, RatioMergeStrategy::default())
340    }
341
342    /// Consumes self and creates a new [`RatioMetrics`] with a specific merge strategy
343    pub fn ratio_metrics_with_strategy(
344        self,
345        name: impl Into<Cow<'static, str>>,
346        partition: usize,
347        merge_strategy: RatioMergeStrategy,
348    ) -> RatioMetrics {
349        let ratio_metrics = RatioMetrics::new().with_merge_strategy(merge_strategy);
350        self.with_category(MetricCategory::Rows)
351            .with_partition(partition)
352            .build(MetricValue::Ratio {
353                name: name.into(),
354                ratio_metrics: ratio_metrics.clone(),
355            });
356        ratio_metrics
357    }
358}
datafusion_physical_expr_common/metrics/builder.rs

datafusion_physical_expr_common/metrics/
builder.rs
