promql/functions/
deriv.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Implementation of [`deriv`](https://prometheus.io/docs/prometheus/latest/querying/functions/#deriv) in PromQL. Refer to the [original
16//! implementation](https://github.com/prometheus/prometheus/blob/90b2f7a540b8a70d8d81372e6692dcbb67ccbaaa/promql/functions.go#L839-L856).
17
18use std::sync::Arc;
19
20use common_macro::range_fn;
21use datafusion::arrow::array::{Float64Array, TimestampMillisecondArray};
22use datafusion::common::DataFusionError;
23use datafusion::logical_expr::{ScalarUDF, Volatility};
24use datafusion::physical_plan::ColumnarValue;
25use datatypes::arrow::array::Array;
26use datatypes::arrow::datatypes::DataType;
27
28use crate::functions::{extract_array, linear_regression};
29use crate::range_array::RangeArray;
30
31#[range_fn(name = Deriv, ret = Float64Array, display_name = prom_deriv)]
32pub fn deriv(times: &TimestampMillisecondArray, values: &Float64Array) -> Option<f64> {
33    if values.len() < 2 {
34        None
35    } else {
36        let intercept_time = times.value(0);
37        let (slope, _) = linear_regression(times, values, intercept_time);
38        slope
39    }
40}
41
42#[cfg(test)]
43mod test {
44    use std::sync::Arc;
45
46    use super::*;
47    use crate::functions::test_util::simple_range_udf_runner;
48
49    // build timestamp range and value range arrays for test
50    fn build_test_range_arrays() -> (RangeArray, RangeArray) {
51        let ts_array = Arc::new(TimestampMillisecondArray::from_iter(
52            [
53                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
54            ]
55            .into_iter()
56            .map(Some),
57        ));
58        let ranges = [(0, 11), (0, 1)];
59
60        let values_array = Arc::new(Float64Array::from_iter([
61            0.0, 10.0, 20.0, 30.0, 40.0, 0.0, 10.0, 20.0, 30.0, 40.0, 50.0,
62        ]));
63
64        let ts_range_array = RangeArray::from_ranges(ts_array, ranges).unwrap();
65        let value_range_array = RangeArray::from_ranges(values_array, ranges).unwrap();
66
67        (ts_range_array, value_range_array)
68    }
69
70    #[test]
71    fn calculate_deriv() {
72        let (ts_array, value_array) = build_test_range_arrays();
73        simple_range_udf_runner(
74            Deriv::scalar_udf(),
75            ts_array,
76            value_array,
77            vec![],
78            vec![Some(10.606060606060607), None],
79        );
80    }
81
82    // From prometheus `promql/functions_test.go` case `TestDeriv`
83    #[test]
84    fn complicate_deriv() {
85        let start = 1493712816939;
86        let interval = 30 * 1000;
87        let mut ts_data = vec![];
88        for i in 0..15 {
89            let jitter = 12 * i % 2;
90            ts_data.push(Some(start + interval * i + jitter));
91        }
92        let val_data = vec![Some(1.0); 15];
93        let ts_array = Arc::new(TimestampMillisecondArray::from_iter(ts_data));
94        let val_array = Arc::new(Float64Array::from_iter(val_data));
95        let range = [(0, 15)];
96        let ts_range_array = RangeArray::from_ranges(ts_array, range).unwrap();
97        let value_range_array = RangeArray::from_ranges(val_array, range).unwrap();
98
99        simple_range_udf_runner(
100            Deriv::scalar_udf(),
101            ts_range_array,
102            value_range_array,
103            vec![],
104            vec![Some(0.0)],
105        );
106    }
107}