Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/metrics/baseline.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
//! Metrics common for almost all operators
19
20
use std::task::Poll;
21
22
use arrow::record_batch::RecordBatch;
23
24
use super::{Count, ExecutionPlanMetricsSet, MetricBuilder, Time, Timestamp};
25
use datafusion_common::Result;
26
27
/// Helper for creating and tracking common "baseline" metrics for
28
/// each operator
29
///
30
/// Example:
31
/// ```
32
/// use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
33
/// let metrics = ExecutionPlanMetricsSet::new();
34
///
35
/// let partition = 2;
36
/// let baseline_metrics = BaselineMetrics::new(&metrics, partition);
37
///
38
/// // during execution, in CPU intensive operation:
39
/// let timer = baseline_metrics.elapsed_compute().timer();
40
/// // .. do CPU intensive work
41
/// timer.done();
42
///
43
/// // when operator is finished:
44
/// baseline_metrics.done();
45
/// ```
46
#[derive(Debug, Clone)]
47
pub struct BaselineMetrics {
48
    /// end_time is set when `ExecutionMetrics::done()` is called
49
    end_time: Timestamp,
50
51
    /// amount of time the operator was actively trying to use the CPU
52
    elapsed_compute: Time,
53
54
    /// output rows: the total output rows
55
    output_rows: Count,
56
}
57
58
impl BaselineMetrics {
59
    /// Create a new BaselineMetric structure, and set `start_time` to now
60
192
    pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
61
192
        let start_time = MetricBuilder::new(metrics).start_timestamp(partition);
62
192
        start_time.record();
63
192
64
192
        Self {
65
192
            end_time: MetricBuilder::new(metrics).end_timestamp(partition),
66
192
            elapsed_compute: MetricBuilder::new(metrics).elapsed_compute(partition),
67
192
            output_rows: MetricBuilder::new(metrics).output_rows(partition),
68
192
        }
69
192
    }
70
71
    /// Returns a [`BaselineMetrics`] that updates the same `elapsed_compute` ignoring
72
    /// all other metrics
73
    ///
74
    /// This is useful when an operator offloads some of its intermediate work to separate tasks
75
    /// that as a result won't be recorded by [`Self::record_poll`]
76
7
    pub fn intermediate(&self) -> BaselineMetrics {
77
7
        Self {
78
7
            end_time: Default::default(),
79
7
            elapsed_compute: self.elapsed_compute.clone(),
80
7
            output_rows: Default::default(),
81
7
        }
82
7
    }
83
84
    /// return the metric for cpu time spend in this operator
85
808
    pub fn elapsed_compute(&self) -> &Time {
86
808
        &self.elapsed_compute
87
808
    }
88
89
    /// return the metric for the total number of output rows produced
90
5
    pub fn output_rows(&self) -> &Count {
91
5
        &self.output_rows
92
5
    }
93
94
    /// Records the fact that this operator's execution is complete
95
    /// (recording the `end_time` metric).
96
    ///
97
    /// Note care should be taken to call `done()` manually if
98
    /// `BaselineMetrics` is not `drop`ped immediately upon operator
99
    /// completion, as async streams may not be dropped immediately
100
    /// depending on the consumer.
101
91
    pub fn done(&self) {
102
91
        self.end_time.record()
103
91
    }
104
105
    /// Record that some number of rows have been produced as output
106
    ///
107
    /// See the [`RecordOutput`] for conveniently recording record
108
    /// batch output for other thing
109
1.16k
    pub fn record_output(&self, num_rows: usize) {
110
1.16k
        self.output_rows.add(num_rows);
111
1.16k
    }
112
113
    /// If not previously recorded `done()`, record
114
221
    pub fn try_done(&self) {
115
221
        if self.end_time.value().is_none() {
116
110
            self.end_time.record()
117
111
        }
118
221
    }
119
120
    /// Process a poll result of a stream producing output for an
121
    /// operator, recording the output rows and stream done time and
122
    /// returning the same poll result
123
2.67k
    pub fn record_poll(
124
2.67k
        &self,
125
2.67k
        poll: Poll<Option<Result<RecordBatch>>>,
126
2.67k
    ) -> Poll<Option<Result<RecordBatch>>> {
127
2.67k
        if let Poll::Ready(
maybe_batch1.10k
) = &poll {
128
1.01k
            match maybe_batch {
129
1.01k
                Some(Ok(batch)) => {
130
1.01k
                    batch.record_output(self);
131
1.01k
                }
132
0
                Some(Err(_)) => self.done(),
133
91
                None => self.done(),
134
            }
135
1.57k
        }
136
2.67k
        poll
137
2.67k
    }
138
}
139
140
impl Drop for BaselineMetrics {
141
221
    fn drop(&mut self) {
142
221
        self.try_done()
143
221
    }
144
}
145
146
/// Trait for things that produce output rows as a result of execution.
147
pub trait RecordOutput {
148
    /// Record that some number of output rows have been produced
149
    ///
150
    /// Meant to be composable so that instead of returning `batch`
151
    /// the operator can return `batch.record_output(baseline_metrics)`
152
    fn record_output(self, bm: &BaselineMetrics) -> Self;
153
}
154
155
impl RecordOutput for usize {
156
0
    fn record_output(self, bm: &BaselineMetrics) -> Self {
157
0
        bm.record_output(self);
158
0
        self
159
0
    }
160
}
161
162
impl RecordOutput for RecordBatch {
163
122
    fn record_output(self, bm: &BaselineMetrics) -> Self {
164
122
        bm.record_output(self.num_rows());
165
122
        self
166
122
    }
167
}
168
169
impl RecordOutput for &RecordBatch {
170
1.01k
    fn record_output(self, bm: &BaselineMetrics) -> Self {
171
1.01k
        bm.record_output(self.num_rows());
172
1.01k
        self
173
1.01k
    }
174
}
175
176
impl RecordOutput for Option<&RecordBatch> {
177
0
    fn record_output(self, bm: &BaselineMetrics) -> Self {
178
0
        if let Some(record_batch) = &self {
179
0
            record_batch.record_output(bm);
180
0
        }
181
0
        self
182
0
    }
183
}
184
185
impl RecordOutput for Option<RecordBatch> {
186
0
    fn record_output(self, bm: &BaselineMetrics) -> Self {
187
0
        if let Some(record_batch) = &self {
188
0
            record_batch.record_output(bm);
189
0
        }
190
0
        self
191
0
    }
192
}
193
194
impl RecordOutput for Result<RecordBatch> {
195
0
    fn record_output(self, bm: &BaselineMetrics) -> Self {
196
0
        if let Ok(record_batch) = &self {
197
0
            record_batch.record_output(bm);
198
0
        }
199
0
        self
200
0
    }
201
}