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