/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 | | } |