Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/metrics/builder.rs
Line
Count
Source (jump to first uncovered line)
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
20
use std::{borrow::Cow, sync::Arc};
21
22
use super::{
23
    Count, ExecutionPlanMetricsSet, Gauge, Label, Metric, MetricValue, Time, Timestamp,
24
};
25
26
/// Structure for constructing metrics, counters, timers, etc.
27
///
28
/// Note the use of `Cow<..>` is to avoid allocations in the common
29
/// case of constant strings
30
///
31
/// ```rust
32
///  use datafusion_physical_plan::metrics::*;
33
///
34
///  let metrics = ExecutionPlanMetricsSet::new();
35
///  let partition = 1;
36
///
37
///  // Create the standard output_rows metric
38
///  let output_rows = MetricBuilder::new(&metrics).output_rows(partition);
39
///
40
///  // Create a operator specific counter with some labels
41
///  let num_bytes = MetricBuilder::new(&metrics)
42
///    .with_new_label("filename", "my_awesome_file.parquet")
43
///    .counter("num_bytes", partition);
44
///
45
/// ```
46
pub struct MetricBuilder<'a> {
47
    /// Location that the metric created by this builder will be added do
48
    metrics: &'a ExecutionPlanMetricsSet,
49
50
    /// optional partition number
51
    partition: Option<usize>,
52
53
    /// arbitrary name=value pairs identifiying this metric
54
    labels: Vec<Label>,
55
}
56
57
impl<'a> MetricBuilder<'a> {
58
    /// Create a new `MetricBuilder` that will register the result of `build()` with the `metrics`
59
35.8k
    pub fn new(metrics: &'a ExecutionPlanMetricsSet) -> Self {
60
35.8k
        Self {
61
35.8k
            metrics,
62
35.8k
            partition: None,
63
35.8k
            labels: vec![],
64
35.8k
        }
65
35.8k
    }
66
67
    /// Add a label to the metric being constructed
68
5.59k
    pub fn with_label(mut self, label: Label) -> Self {
69
5.59k
        self.labels.push(label);
70
5.59k
        self
71
5.59k
    }
72
73
    /// Add a label to the metric being constructed
74
2
    pub fn with_new_label(
75
2
        self,
76
2
        name: impl Into<Cow<'static, str>>,
77
2
        value: impl Into<Cow<'static, str>>,
78
2
    ) -> Self {
79
2
        self.with_label(Label::new(name.into(), value.into()))
80
2
    }
81
82
    /// Set the partition of the metric being constructed
83
35.8k
    pub fn with_partition(mut self, partition: usize) -> Self {
84
35.8k
        self.partition = Some(partition);
85
35.8k
        self
86
35.8k
    }
87
88
    /// Consume self and create a metric of the specified value
89
    /// registered with the MetricsSet
90
35.8k
    pub fn build(self, value: MetricValue) {
91
35.8k
        let Self {
92
35.8k
            labels,
93
35.8k
            partition,
94
35.8k
            metrics,
95
35.8k
        } = self;
96
35.8k
        let metric = Arc::new(Metric::new_with_labels(value, partition, labels));
97
35.8k
        metrics.register(metric);
98
35.8k
    }
99
100
    /// Consume self and create a new counter for recording output rows
101
3.44k
    pub fn output_rows(self, partition: usize) -> Count {
102
3.44k
        let count = Count::new();
103
3.44k
        self.with_partition(partition)
104
3.44k
            .build(MetricValue::OutputRows(count.clone()));
105
3.44k
        count
106
3.44k
    }
107
108
    /// Consume self and create a new counter for recording the number of spills
109
    /// triggered by an operator
110
98
    pub fn spill_count(self, partition: usize) -> Count {
111
98
        let count = Count::new();
112
98
        self.with_partition(partition)
113
98
            .build(MetricValue::SpillCount(count.clone()));
114
98
        count
115
98
    }
116
117
    /// Consume self and create a new counter for recording the total spilled bytes
118
    /// triggered by an operator
119
98
    pub fn spilled_bytes(self, partition: usize) -> Count {
120
98
        let count = Count::new();
121
98
        self.with_partition(partition)
122
98
            .build(MetricValue::SpilledBytes(count.clone()));
123
98
        count
124
98
    }
125
126
    /// Consume self and create a new counter for recording the total spilled rows
127
    /// triggered by an operator
128
98
    pub fn spilled_rows(self, partition: usize) -> Count {
129
98
        let count = Count::new();
130
98
        self.with_partition(partition)
131
98
            .build(MetricValue::SpilledRows(count.clone()));
132
98
        count
133
98
    }
134
135
    /// Consume self and create a new gauge for reporting current memory usage
136
0
    pub fn mem_used(self, partition: usize) -> Gauge {
137
0
        let gauge = Gauge::new();
138
0
        self.with_partition(partition)
139
0
            .build(MetricValue::CurrentMemoryUsage(gauge.clone()));
140
0
        gauge
141
0
    }
142
143
    /// Consumes self and creates a new [`Count`] for recording some
144
    /// arbitrary metric of an operator.
145
16.1k
    pub fn counter(
146
16.1k
        self,
147
16.1k
        counter_name: impl Into<Cow<'static, str>>,
148
16.1k
        partition: usize,
149
16.1k
    ) -> Count {
150
16.1k
        self.with_partition(partition).global_counter(counter_name)
151
16.1k
    }
152
153
    /// Consumes self and creates a new [`Gauge`] for reporting some
154
    /// arbitrary metric of an operator.
155
3.24k
    pub fn gauge(
156
3.24k
        self,
157
3.24k
        gauge_name: impl Into<Cow<'static, str>>,
158
3.24k
        partition: usize,
159
3.24k
    ) -> Gauge {
160
3.24k
        self.with_partition(partition).global_gauge(gauge_name)
161
3.24k
    }
162
163
    /// Consumes self and creates a new [`Count`] for recording a
164
    /// metric of an overall operator (not per partition)
165
16.1k
    pub fn global_counter(self, counter_name: impl Into<Cow<'static, str>>) -> Count {
166
16.1k
        let count = Count::new();
167
16.1k
        self.build(MetricValue::Count {
168
16.1k
            name: counter_name.into(),
169
16.1k
            count: count.clone(),
170
16.1k
        });
171
16.1k
        count
172
16.1k
    }
173
174
    /// Consumes self and creates a new [`Gauge`] for reporting a
175
    /// metric of an overall operator (not per partition)
176
3.24k
    pub fn global_gauge(self, gauge_name: impl Into<Cow<'static, str>>) -> Gauge {
177
3.24k
        let gauge = Gauge::new();
178
3.24k
        self.build(MetricValue::Gauge {
179
3.24k
            name: gauge_name.into(),
180
3.24k
            gauge: gauge.clone(),
181
3.24k
        });
182
3.24k
        gauge
183
3.24k
    }
184
185
    /// Consume self and create a new Timer for recording the elapsed
186
    /// CPU time spent by an operator
187
198
    pub fn elapsed_compute(self, partition: usize) -> Time {
188
198
        let time = Time::new();
189
198
        self.with_partition(partition)
190
198
            .build(MetricValue::ElapsedCompute(time.clone()));
191
198
        time
192
198
    }
193
194
    /// Consumes self and creates a new Timer for recording some
195
    /// subset of an operators execution time.
196
12.1k
    pub fn subset_time(
197
12.1k
        self,
198
12.1k
        subset_name: impl Into<Cow<'static, str>>,
199
12.1k
        partition: usize,
200
12.1k
    ) -> Time {
201
12.1k
        let time = Time::new();
202
12.1k
        self.with_partition(partition).build(MetricValue::Time {
203
12.1k
            name: subset_name.into(),
204
12.1k
            time: time.clone(),
205
12.1k
        });
206
12.1k
        time
207
12.1k
    }
208
209
    /// Consumes self and creates a new Timestamp for recording the
210
    /// starting time of execution for a partition
211
195
    pub fn start_timestamp(self, partition: usize) -> Timestamp {
212
195
        let timestamp = Timestamp::new();
213
195
        self.with_partition(partition)
214
195
            .build(MetricValue::StartTimestamp(timestamp.clone()));
215
195
        timestamp
216
195
    }
217
218
    /// Consumes self and creates a new Timestamp for recording the
219
    /// ending time of execution for a partition
220
195
    pub fn end_timestamp(self, partition: usize) -> Timestamp {
221
195
        let timestamp = Timestamp::new();
222
195
        self.with_partition(partition)
223
195
            .build(MetricValue::EndTimestamp(timestamp.clone()));
224
195
        timestamp
225
195
    }
226
}