datafusion_python/
metrics.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
18use std::collections::HashMap;
19use std::sync::Arc;
20
21use chrono::{Datelike, Timelike};
22use datafusion::physical_plan::metrics::{Metric, MetricValue, MetricsSet, Timestamp};
23use pyo3::prelude::*;
24
25#[pyclass(from_py_object, frozen, name = "MetricsSet", module = "datafusion")]
26#[derive(Debug, Clone)]
27pub struct PyMetricsSet {
28    metrics: MetricsSet,
29}
30
31impl PyMetricsSet {
32    pub fn new(metrics: MetricsSet) -> Self {
33        Self { metrics }
34    }
35}
36
37#[pymethods]
38impl PyMetricsSet {
39    fn metrics(&self) -> Vec<PyMetric> {
40        self.metrics
41            .iter()
42            .map(|m| PyMetric::new(Arc::clone(m)))
43            .collect()
44    }
45
46    fn output_rows(&self) -> Option<usize> {
47        self.metrics.output_rows()
48    }
49
50    fn elapsed_compute(&self) -> Option<usize> {
51        self.metrics.elapsed_compute()
52    }
53
54    fn spill_count(&self) -> Option<usize> {
55        self.metrics.spill_count()
56    }
57
58    fn spilled_bytes(&self) -> Option<usize> {
59        self.metrics.spilled_bytes()
60    }
61
62    fn spilled_rows(&self) -> Option<usize> {
63        self.metrics.spilled_rows()
64    }
65
66    fn sum_by_name(&self, name: &str) -> Option<usize> {
67        self.metrics.sum_by_name(name).map(|v| v.as_usize())
68    }
69
70    fn __repr__(&self) -> String {
71        format!("{}", self.metrics)
72    }
73}
74
75#[pyclass(from_py_object, frozen, name = "Metric", module = "datafusion")]
76#[derive(Debug, Clone)]
77pub struct PyMetric {
78    metric: Arc<Metric>,
79}
80
81impl PyMetric {
82    pub fn new(metric: Arc<Metric>) -> Self {
83        Self { metric }
84    }
85
86    fn timestamp_to_pyobject<'py>(
87        py: Python<'py>,
88        ts: &Timestamp,
89    ) -> PyResult<Option<Bound<'py, PyAny>>> {
90        match ts.value() {
91            Some(dt) => {
92                let datetime_mod = py.import("datetime")?;
93                let datetime_cls = datetime_mod.getattr("datetime")?;
94                let tz_utc = datetime_mod.getattr("timezone")?.getattr("utc")?;
95                let result = datetime_cls.call1((
96                    dt.year(),
97                    dt.month(),
98                    dt.day(),
99                    dt.hour(),
100                    dt.minute(),
101                    dt.second(),
102                    dt.timestamp_subsec_micros(),
103                    tz_utc,
104                ))?;
105                Ok(Some(result))
106            }
107            None => Ok(None),
108        }
109    }
110}
111
112#[pymethods]
113impl PyMetric {
114    #[getter]
115    fn name(&self) -> String {
116        self.metric.value().name().to_string()
117    }
118
119    #[getter]
120    fn value<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
121        match self.metric.value() {
122            MetricValue::OutputRows(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())),
123            MetricValue::OutputBytes(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())),
124            MetricValue::ElapsedCompute(t) => Ok(Some(t.value().into_pyobject(py)?.into_any())),
125            MetricValue::SpillCount(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())),
126            MetricValue::SpilledBytes(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())),
127            MetricValue::SpilledRows(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())),
128            MetricValue::CurrentMemoryUsage(g) => Ok(Some(g.value().into_pyobject(py)?.into_any())),
129            MetricValue::Count { count, .. } => {
130                Ok(Some(count.value().into_pyobject(py)?.into_any()))
131            }
132            MetricValue::Gauge { gauge, .. } => {
133                Ok(Some(gauge.value().into_pyobject(py)?.into_any()))
134            }
135            MetricValue::Time { time, .. } => Ok(Some(time.value().into_pyobject(py)?.into_any())),
136            MetricValue::StartTimestamp(ts) | MetricValue::EndTimestamp(ts) => {
137                Self::timestamp_to_pyobject(py, ts)
138            }
139            _ => Ok(None),
140        }
141    }
142
143    #[getter]
144    fn value_as_datetime<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
145        match self.metric.value() {
146            MetricValue::StartTimestamp(ts) | MetricValue::EndTimestamp(ts) => {
147                Self::timestamp_to_pyobject(py, ts)
148            }
149            _ => Ok(None),
150        }
151    }
152
153    #[getter]
154    fn partition(&self) -> Option<usize> {
155        self.metric.partition()
156    }
157
158    fn labels(&self) -> HashMap<String, String> {
159        self.metric
160            .labels()
161            .iter()
162            .map(|l| (l.name().to_string(), l.value().to_string()))
163            .collect()
164    }
165
166    fn __repr__(&self) -> String {
167        format!("{}", self.metric.value())
168    }
169}
datafusion_python/metrics.rs

datafusion_python/
metrics.rs
