promql/
functions.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
15mod aggr_over_time;
16mod changes;
17mod deriv;
18mod double_exponential_smoothing;
19mod edge_count;
20mod extrapolate_rate;
21mod idelta;
22mod ieee_float;
23mod native_histogram;
24mod predict_linear;
25mod quantile;
26mod quantile_aggr;
27mod resets;
28mod round;
29#[cfg(test)]
30mod test_util;
31mod vector_matching;
32
33pub use aggr_over_time::{
34    AbsentOverTime, AvgOverTime, CountOverTime, LastOverTime, MaxOverTime, MinOverTime,
35    PresentOverTime, StddevOverTime, StdvarOverTime, SumOverTime,
36};
37pub use changes::Changes;
38use datafusion::arrow::array::{
39    ArrayRef, DictionaryArray, Float64Array, TimestampMillisecondArray,
40};
41use datafusion::error::DataFusionError;
42use datafusion::physical_plan::ColumnarValue;
43use datatypes::arrow::array::Array;
44use datatypes::arrow::datatypes::{DataType, Int64Type};
45pub use deriv::Deriv;
46pub use double_exponential_smoothing::DoubleExponentialSmoothing;
47pub use extrapolate_rate::{Delta, Increase, Rate};
48pub use idelta::IDelta;
49pub use ieee_float::{IeeeComparison, IeeeSqrt};
50pub use native_histogram::{
51    MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg,
52    NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges,
53    NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta,
54    NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction,
55    NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
56    NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
57    NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile,
58    NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev,
59    NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime,
60    NativeHistogramToString, PromqlFloatToString,
61};
62pub use predict_linear::PredictLinear;
63pub use quantile::QuantileOverTime;
64pub use quantile_aggr::{QUANTILE_NAME, quantile_udaf};
65pub use resets::Resets;
66pub use round::Round;
67pub use vector_matching::{MatchGroupViolation, UniqueMatchGroup};
68
69use crate::range_array::RangeArray;
70
71/// Extracts an array from a `ColumnarValue`.
72///
73/// If the `ColumnarValue` is a scalar, it converts it to an array of size 1.
74pub(crate) fn extract_array(columnar_value: &ColumnarValue) -> Result<ArrayRef, DataFusionError> {
75    match columnar_value {
76        ColumnarValue::Array(array) => Ok(array.clone()),
77        ColumnarValue::Scalar(scalar) => Ok(scalar.to_array_of_size(1)?),
78    }
79}
80
81/// Extracts and validates a range dictionary with the expected value type.
82pub(crate) fn extract_range_dict(
83    columnar_value: &ColumnarValue,
84    func_name: &str,
85    arg_name: &str,
86    expected_value_type: &DataType,
87) -> Result<DictionaryArray<Int64Type>, DataFusionError> {
88    let array = extract_array(columnar_value)?;
89    let dict = array
90        .as_any()
91        .downcast_ref::<DictionaryArray<Int64Type>>()
92        .ok_or_else(|| {
93            DataFusionError::Execution(format!(
94                "{func_name}: expect {arg_name} as DictionaryArray<Int64>, found {}",
95                array.data_type()
96            ))
97        })?
98        .clone();
99
100    if &dict.value_type() != expected_value_type {
101        return Err(DataFusionError::Execution(format!(
102            "{func_name}: expect {arg_name} values of type {expected_value_type}, found {}",
103            dict.value_type()
104        )));
105    }
106
107    RangeArray::try_new(dict.clone()).map_err(DataFusionError::from)?;
108    Ok(dict)
109}
110
111/// Extracts a validated [RangeArray] from a [ColumnarValue].
112pub(crate) fn extract_range_array(
113    columnar_value: &ColumnarValue,
114) -> Result<RangeArray, DataFusionError> {
115    let array = extract_array(columnar_value)?;
116    let dict = array
117        .as_any()
118        .downcast_ref::<DictionaryArray<Int64Type>>()
119        .ok_or_else(|| {
120            DataFusionError::Execution(format!(
121                "expected DictionaryArray<Int64>, found {}",
122                array.data_type()
123            ))
124        })?
125        .clone();
126    RangeArray::try_new(dict).map_err(DataFusionError::from)
127}
128
129/// compensation(Kahan) summation algorithm - a technique for reducing the numerical error
130/// in floating-point arithmetic. The algorithm also includes the modification ("Neumaier improvement")
131/// that reduces the numerical error further in cases
132/// where the numbers being summed have a large difference in magnitude
133/// Prometheus's implementation:
134/// <https://github.com/prometheus/prometheus/blob/f55ab2217984770aa1eecd0f2d5f54580029b1c0/promql/functions.go#L782>
135pub(crate) fn compensated_sum_inc(inc: f64, sum: f64, mut compensation: f64) -> (f64, f64) {
136    let new_sum = sum + inc;
137    if sum.abs() >= inc.abs() {
138        compensation += (sum - new_sum) + inc;
139    } else {
140        compensation += (inc - new_sum) + sum;
141    }
142    (new_sum, compensation)
143}
144
145/// linear_regression performs a least-square linear regression analysis on the
146/// times and values. It return the slope and intercept based on times and values.
147/// Prometheus's implementation: <https://github.com/prometheus/prometheus/blob/90b2f7a540b8a70d8d81372e6692dcbb67ccbaaa/promql/functions.go#L793-L837>
148pub(crate) fn linear_regression(
149    times: &TimestampMillisecondArray,
150    values: &Float64Array,
151    intercept_time: i64,
152) -> (Option<f64>, Option<f64>) {
153    linear_regression_slice(times.values(), values, 0, values.len(), intercept_time)
154}
155
156pub(crate) fn linear_regression_slice(
157    times: &[i64],
158    values: &Float64Array,
159    offset: usize,
160    len: usize,
161    intercept_time: i64,
162) -> (Option<f64>, Option<f64>) {
163    linear_regression_slices(times, offset, values, offset, len, intercept_time)
164}
165
166pub(crate) fn linear_regression_slices(
167    times: &[i64],
168    time_offset: usize,
169    values: &Float64Array,
170    value_offset: usize,
171    len: usize,
172    intercept_time: i64,
173) -> (Option<f64>, Option<f64>) {
174    let raw_values = values.values();
175    let has_nulls = values.null_count() > 0;
176    let mut count: f64 = 0.0;
177    let mut sum_x: f64 = 0.0;
178    let mut sum_y: f64 = 0.0;
179    let mut sum_xy: f64 = 0.0;
180    let mut sum_x2: f64 = 0.0;
181    let mut comp_x: f64 = 0.0;
182    let mut comp_y: f64 = 0.0;
183    let mut comp_xy: f64 = 0.0;
184    let mut comp_x2: f64 = 0.0;
185
186    let mut const_y = true;
187    let mut init_y = None;
188
189    for i in 0..len {
190        let time_idx = time_offset + i;
191        let value_idx = value_offset + i;
192        if has_nulls && values.is_null(value_idx) {
193            continue;
194        }
195        let value = raw_values[value_idx];
196        let time = times[time_idx] as f64;
197        let initial = init_y.get_or_insert(value);
198        if const_y && count > 0.0 && value != *initial {
199            const_y = false;
200        }
201        count += 1.0;
202        let x = (time - intercept_time as f64) / 1e3f64;
203        (sum_x, comp_x) = compensated_sum_inc(x, sum_x, comp_x);
204        (sum_y, comp_y) = compensated_sum_inc(value, sum_y, comp_y);
205        (sum_xy, comp_xy) = compensated_sum_inc(x * value, sum_xy, comp_xy);
206        (sum_x2, comp_x2) = compensated_sum_inc(x * x, sum_x2, comp_x2);
207    }
208
209    if count < 2.0 {
210        return (None, None);
211    }
212
213    if const_y {
214        let init_y = init_y.unwrap();
215        if !init_y.is_finite() {
216            return (None, None);
217        }
218        return (Some(0.0), Some(init_y));
219    }
220
221    sum_x += comp_x;
222    sum_y += comp_y;
223    sum_xy += comp_xy;
224    sum_x2 += comp_x2;
225
226    let cov_xy = sum_xy - sum_x * sum_y / count;
227    let var_x = sum_x2 - sum_x * sum_x / count;
228
229    let slope = cov_xy / var_x;
230    let intercept = sum_y / count - slope * sum_x / count;
231
232    (Some(slope), Some(intercept))
233}
234
235#[cfg(test)]
236mod test {
237    use std::sync::Arc;
238
239    use datafusion::physical_plan::ColumnarValue;
240    use datatypes::arrow::array::Int64Array;
241    use datatypes::arrow::datatypes::Int64Type;
242
243    use super::*;
244    use crate::range_array::RangeArray;
245
246    #[test]
247    fn calculate_linear_regression_none() {
248        let ts_array = TimestampMillisecondArray::from_iter(
249            [
250                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
251            ]
252            .into_iter()
253            .map(Some),
254        );
255        let values_array = Float64Array::from_iter([
256            1.0 / 0.0,
257            1.0 / 0.0,
258            1.0 / 0.0,
259            1.0 / 0.0,
260            1.0 / 0.0,
261            1.0 / 0.0,
262            1.0 / 0.0,
263            1.0 / 0.0,
264            1.0 / 0.0,
265            1.0 / 0.0,
266        ]);
267        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
268        assert_eq!(slope, None);
269        assert_eq!(intercept, None);
270    }
271
272    #[test]
273    fn calculate_linear_regression_value_is_const() {
274        let ts_array = TimestampMillisecondArray::from_iter(
275            [
276                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
277            ]
278            .into_iter()
279            .map(Some),
280        );
281        let values_array =
282            Float64Array::from_iter([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]);
283        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
284        assert_eq!(slope, Some(0.0));
285        assert_eq!(intercept, Some(10.0));
286    }
287
288    #[test]
289    fn calculate_linear_regression() {
290        let ts_array = TimestampMillisecondArray::from_iter(
291            [
292                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
293            ]
294            .into_iter()
295            .map(Some),
296        );
297        let values_array = Float64Array::from_iter([
298            0.0, 10.0, 20.0, 30.0, 40.0, 0.0, 10.0, 20.0, 30.0, 40.0, 50.0,
299        ]);
300        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
301        assert_eq!(slope, Some(10.606060606060607));
302        assert_eq!(intercept, Some(6.818181818181815));
303
304        let (slope, intercept) = linear_regression(&ts_array, &values_array, 3000);
305        assert_eq!(slope, Some(10.606060606060607));
306        assert_eq!(intercept, Some(38.63636363636364));
307    }
308
309    #[test]
310    fn calculate_linear_regression_value_have_none() {
311        let ts_array = TimestampMillisecondArray::from_iter(
312            [
313                0i64, 300, 600, 900, 1200, 1350, 1500, 1800, 2100, 2400, 2550, 2700, 3000,
314            ]
315            .into_iter()
316            .map(Some),
317        );
318        let values_array: Float64Array = [
319            Some(0.0),
320            Some(10.0),
321            Some(20.0),
322            Some(30.0),
323            Some(40.0),
324            None,
325            Some(0.0),
326            Some(10.0),
327            Some(20.0),
328            Some(30.0),
329            None,
330            Some(40.0),
331            Some(50.0),
332        ]
333        .into_iter()
334        .collect();
335        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
336        assert_eq!(slope, Some(10.606060606060607));
337        assert_eq!(intercept, Some(6.818181818181815));
338    }
339
340    #[test]
341    fn calculate_linear_regression_value_all_none() {
342        let ts_array = TimestampMillisecondArray::from_iter([0i64, 300, 600].into_iter().map(Some));
343        let values_array: Float64Array = [None, None, None].into_iter().collect();
344        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
345        assert_eq!(slope, None);
346        assert_eq!(intercept, None);
347    }
348
349    // From prometheus `promql/functions_test.go` case `TestKahanSum`
350    #[test]
351    fn test_kahan_sum() {
352        let inputs = vec![1.0, 10.0f64.powf(100.0), 1.0, -10.0f64.powf(100.0)];
353
354        let mut sum = 0.0;
355        let mut c = 0f64;
356
357        for v in inputs {
358            (sum, c) = compensated_sum_inc(v, sum, c);
359        }
360        assert_eq!(sum + c, 2.0)
361    }
362
363    #[test]
364    fn extract_range_array_rejects_external_dictionary_with_null_keys() {
365        let keys = Int64Array::from_iter([Some(0), None]);
366        let values = Arc::new(Float64Array::from_iter([1.0, 2.0]));
367        let dict = DictionaryArray::<Int64Type>::try_new(keys, values).unwrap();
368
369        let err = extract_range_array(&ColumnarValue::Array(Arc::new(dict))).unwrap_err();
370        assert!(err.to_string().contains("Empty range is not expected"));
371    }
372
373    #[test]
374    fn extract_range_array_accepts_internal_packed_ranges() {
375        let values = Arc::new(Float64Array::from_iter([1.0, 2.0, 3.0]));
376        let range_array = RangeArray::from_ranges(values, [(0, 2), (1, 2)]).unwrap();
377
378        let extracted =
379            extract_range_array(&ColumnarValue::Array(Arc::new(range_array.into_dict()))).unwrap();
380
381        assert_eq!(extracted.get_offset_length(0), Some((0, 2)));
382        assert_eq!(extracted.get_offset_length(1), Some((1, 2)));
383    }
384}
promql/functions.rs

promql/
functions.rs
