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
33use std::sync::Arc;
34
35pub use aggr_over_time::{
36    AbsentOverTime, AvgOverTime, CountOverTime, LastOverTime, MaxOverTime, MinOverTime,
37    PresentOverTime, StddevOverTime, StdvarOverTime, SumOverTime,
38};
39pub use changes::Changes;
40use datafusion::arrow::array::{
41    ArrayRef, DictionaryArray, Float64Array, TimestampMillisecondArray,
42};
43use datafusion::error::DataFusionError;
44use datafusion::execution::{FunctionRegistry, SessionState};
45use datafusion::physical_plan::ColumnarValue;
46use datatypes::arrow::array::Array;
47use datatypes::arrow::datatypes::{DataType, Int64Type};
48pub use deriv::Deriv;
49pub use double_exponential_smoothing::DoubleExponentialSmoothing;
50pub use extrapolate_rate::{Delta, Increase, Rate};
51pub use idelta::IDelta;
52pub use ieee_float::{IeeeComparison, IeeeSqrt};
53pub use native_histogram::{
54    MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg,
55    NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges,
56    NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta,
57    NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction,
58    NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease,
59    NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg,
60    NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile,
61    NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev,
62    NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime,
63    NativeHistogramToString, PromqlFloatToString,
64};
65pub use predict_linear::PredictLinear;
66pub use quantile::QuantileOverTime;
67pub use quantile_aggr::{QUANTILE_NAME, quantile_udaf};
68pub use resets::Resets;
69pub use round::Round;
70pub use vector_matching::{MatchGroupViolation, UniqueMatchGroup};
71
72use crate::range_array::RangeArray;
73
74/// Extracts an array from a `ColumnarValue`.
75///
76/// If the `ColumnarValue` is a scalar, it converts it to an array of size 1.
77pub(crate) fn extract_array(columnar_value: &ColumnarValue) -> Result<ArrayRef, DataFusionError> {
78    match columnar_value {
79        ColumnarValue::Array(array) => Ok(array.clone()),
80        ColumnarValue::Scalar(scalar) => Ok(scalar.to_array_of_size(1)?),
81    }
82}
83
84/// Extracts and validates a range dictionary with the expected value type.
85pub(crate) fn extract_range_dict(
86    columnar_value: &ColumnarValue,
87    func_name: &str,
88    arg_name: &str,
89    expected_value_type: &DataType,
90) -> Result<DictionaryArray<Int64Type>, DataFusionError> {
91    let array = extract_array(columnar_value)?;
92    let dict = array
93        .as_any()
94        .downcast_ref::<DictionaryArray<Int64Type>>()
95        .ok_or_else(|| {
96            DataFusionError::Execution(format!(
97                "{func_name}: expect {arg_name} as DictionaryArray<Int64>, found {}",
98                array.data_type()
99            ))
100        })?
101        .clone();
102
103    if &dict.value_type() != expected_value_type {
104        return Err(DataFusionError::Execution(format!(
105            "{func_name}: expect {arg_name} values of type {expected_value_type}, found {}",
106            dict.value_type()
107        )));
108    }
109
110    RangeArray::try_new(dict.clone()).map_err(DataFusionError::from)?;
111    Ok(dict)
112}
113
114/// Extracts a validated [RangeArray] from a [ColumnarValue].
115pub(crate) fn extract_range_array(
116    columnar_value: &ColumnarValue,
117) -> Result<RangeArray, DataFusionError> {
118    let array = extract_array(columnar_value)?;
119    let dict = array
120        .as_any()
121        .downcast_ref::<DictionaryArray<Int64Type>>()
122        .ok_or_else(|| {
123            DataFusionError::Execution(format!(
124                "expected DictionaryArray<Int64>, found {}",
125                array.data_type()
126            ))
127        })?
128        .clone();
129    RangeArray::try_new(dict).map_err(DataFusionError::from)
130}
131
132/// compensation(Kahan) summation algorithm - a technique for reducing the numerical error
133/// in floating-point arithmetic. The algorithm also includes the modification ("Neumaier improvement")
134/// that reduces the numerical error further in cases
135/// where the numbers being summed have a large difference in magnitude
136/// Prometheus's implementation:
137/// <https://github.com/prometheus/prometheus/blob/f55ab2217984770aa1eecd0f2d5f54580029b1c0/promql/functions.go#L782>
138pub(crate) fn compensated_sum_inc(inc: f64, sum: f64, mut compensation: f64) -> (f64, f64) {
139    let new_sum = sum + inc;
140    if sum.abs() >= inc.abs() {
141        compensation += (sum - new_sum) + inc;
142    } else {
143        compensation += (inc - new_sum) + sum;
144    }
145    (new_sum, compensation)
146}
147
148/// linear_regression performs a least-square linear regression analysis on the
149/// times and values. It return the slope and intercept based on times and values.
150/// Prometheus's implementation: <https://github.com/prometheus/prometheus/blob/90b2f7a540b8a70d8d81372e6692dcbb67ccbaaa/promql/functions.go#L793-L837>
151pub(crate) fn linear_regression(
152    times: &TimestampMillisecondArray,
153    values: &Float64Array,
154    intercept_time: i64,
155) -> (Option<f64>, Option<f64>) {
156    linear_regression_slice(times.values(), values, 0, values.len(), intercept_time)
157}
158
159pub(crate) fn linear_regression_slice(
160    times: &[i64],
161    values: &Float64Array,
162    offset: usize,
163    len: usize,
164    intercept_time: i64,
165) -> (Option<f64>, Option<f64>) {
166    linear_regression_slices(times, offset, values, offset, len, intercept_time)
167}
168
169pub(crate) fn linear_regression_slices(
170    times: &[i64],
171    time_offset: usize,
172    values: &Float64Array,
173    value_offset: usize,
174    len: usize,
175    intercept_time: i64,
176) -> (Option<f64>, Option<f64>) {
177    let raw_values = values.values();
178    let has_nulls = values.null_count() > 0;
179    let mut count: f64 = 0.0;
180    let mut sum_x: f64 = 0.0;
181    let mut sum_y: f64 = 0.0;
182    let mut sum_xy: f64 = 0.0;
183    let mut sum_x2: f64 = 0.0;
184    let mut comp_x: f64 = 0.0;
185    let mut comp_y: f64 = 0.0;
186    let mut comp_xy: f64 = 0.0;
187    let mut comp_x2: f64 = 0.0;
188
189    let mut const_y = true;
190    let mut init_y = None;
191
192    for i in 0..len {
193        let time_idx = time_offset + i;
194        let value_idx = value_offset + i;
195        if has_nulls && values.is_null(value_idx) {
196            continue;
197        }
198        let value = raw_values[value_idx];
199        let time = times[time_idx] as f64;
200        let initial = init_y.get_or_insert(value);
201        if const_y && count > 0.0 && value != *initial {
202            const_y = false;
203        }
204        count += 1.0;
205        let x = (time - intercept_time as f64) / 1e3f64;
206        (sum_x, comp_x) = compensated_sum_inc(x, sum_x, comp_x);
207        (sum_y, comp_y) = compensated_sum_inc(value, sum_y, comp_y);
208        (sum_xy, comp_xy) = compensated_sum_inc(x * value, sum_xy, comp_xy);
209        (sum_x2, comp_x2) = compensated_sum_inc(x * x, sum_x2, comp_x2);
210    }
211
212    if count < 2.0 {
213        return (None, None);
214    }
215
216    if const_y {
217        let init_y = init_y.unwrap();
218        if !init_y.is_finite() {
219            return (None, None);
220        }
221        return (Some(0.0), Some(init_y));
222    }
223
224    sum_x += comp_x;
225    sum_y += comp_y;
226    sum_xy += comp_xy;
227    sum_x2 += comp_x2;
228
229    let cov_xy = sum_xy - sum_x * sum_y / count;
230    let var_x = sum_x2 - sum_x * sum_x / count;
231
232    let slope = cov_xy / var_x;
233    let intercept = sum_y / count - slope * sum_x / count;
234
235    (Some(slope), Some(intercept))
236}
237
238/// Registers every PromQL scalar function (UDF) and aggregate function (UDAF)
239/// implemented in this module into `session_state`.
240pub fn register_promql_functions(session_state: &mut SessionState) {
241    let _ = session_state.register_udaf(quantile_udaf());
242
243    for udf in IeeeComparison::all() {
244        let _ = session_state.register_udf(Arc::new(udf));
245    }
246    let _ = session_state.register_udf(Arc::new(IeeeSqrt::scalar_udf()));
247
248    let _ = session_state.register_udf(Arc::new(IDelta::<false>::scalar_udf()));
249    let _ = session_state.register_udf(Arc::new(IDelta::<true>::scalar_udf()));
250    let _ = session_state.register_udf(Arc::new(Rate::scalar_udf()));
251    let _ = session_state.register_udf(Arc::new(Increase::scalar_udf()));
252    let _ = session_state.register_udf(Arc::new(Delta::scalar_udf()));
253    let _ = session_state.register_udf(Arc::new(Resets::scalar_udf()));
254    let _ = session_state.register_udf(Arc::new(Changes::scalar_udf()));
255    let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf()));
256    let _ = session_state.register_udf(Arc::new(Round::scalar_udf()));
257    let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf()));
258    let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf()));
259    let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf()));
260    let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf()));
261    let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf()));
262    let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf()));
263    let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf()));
264    let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf()));
265    let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf()));
266    let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf()));
267    let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf()));
268    let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf()));
269    // Legacy Prometheus name of the double exponential smoothing function.
270    let double_exponential_smoothing_udf =
271        DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]);
272    let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf));
273
274    for udf in [
275        NativeHistogramAbsentOverTime::scalar_udf(),
276        NativeHistogramAdd::scalar_udf(),
277        NativeHistogramAvg::scalar_udf(),
278        NativeHistogramAvgOverTime::scalar_udf(),
279        NativeHistogramChanges::scalar_udf(),
280        NativeHistogramCount::scalar_udf(),
281        NativeHistogramCountOverTime::scalar_udf(),
282        NativeHistogramDelta::scalar_udf(),
283        NativeHistogramDivScalar::scalar_udf(),
284        NativeHistogramDrop::bool_false_udf(String::new(), None),
285        NativeHistogramDrop::bool_true_udf(String::new(), None),
286        NativeHistogramDrop::float_null_udf(String::new(), None),
287        NativeHistogramEq::scalar_udf(),
288        NativeHistogramFraction::scalar_udf(),
289        NativeHistogramIDelta::scalar_udf(),
290        NativeHistogramIRate::scalar_udf(),
291        NativeHistogramIncrease::scalar_udf(),
292        NativeHistogramLastOverTime::scalar_udf(),
293        MixedRange::float_udf(None),
294        MixedRange::histogram_udf(None),
295        NativeHistogramMulScalar::scalar_udf(),
296        NativeHistogramNeg::scalar_udf(),
297        NativeHistogramNotEq::scalar_udf(),
298        NativeHistogramPresentOverTime::scalar_udf(),
299        NativeHistogramQuantile::scalar_udf(),
300        NativeHistogramRate::scalar_udf(),
301        NativeHistogramResets::scalar_udf(),
302        NativeHistogramScalarMul::scalar_udf(),
303        NativeHistogramStddev::scalar_udf(),
304        NativeHistogramStdvar::scalar_udf(),
305        NativeHistogramSub::scalar_udf(),
306        NativeHistogramSum::scalar_udf(),
307        NativeHistogramSumOverTime::scalar_udf(),
308        NativeHistogramToString::scalar_udf(),
309        PromqlFloatToString::scalar_udf(),
310    ] {
311        let _ = session_state.register_udf(Arc::new(udf));
312    }
313    for udaf in [
314        NativeHistogramAggAvg::aggregate_udf(),
315        NativeHistogramAggSum::aggregate_udf(),
316    ] {
317        let _ = session_state.register_udaf(Arc::new(udaf));
318    }
319}
320
321#[cfg(test)]
322mod test {
323    use std::collections::BTreeSet;
324    use std::sync::Arc;
325
326    use datafusion::execution::SessionStateBuilder;
327    use datafusion::physical_plan::ColumnarValue;
328    use datatypes::arrow::array::Int64Array;
329    use datatypes::arrow::datatypes::Int64Type;
330
331    use super::*;
332    use crate::range_array::RangeArray;
333
334    #[test]
335    fn calculate_linear_regression_none() {
336        let ts_array = TimestampMillisecondArray::from_iter(
337            [
338                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
339            ]
340            .into_iter()
341            .map(Some),
342        );
343        let values_array = Float64Array::from_iter([
344            1.0 / 0.0,
345            1.0 / 0.0,
346            1.0 / 0.0,
347            1.0 / 0.0,
348            1.0 / 0.0,
349            1.0 / 0.0,
350            1.0 / 0.0,
351            1.0 / 0.0,
352            1.0 / 0.0,
353            1.0 / 0.0,
354        ]);
355        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
356        assert_eq!(slope, None);
357        assert_eq!(intercept, None);
358    }
359
360    #[test]
361    fn calculate_linear_regression_value_is_const() {
362        let ts_array = TimestampMillisecondArray::from_iter(
363            [
364                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
365            ]
366            .into_iter()
367            .map(Some),
368        );
369        let values_array =
370            Float64Array::from_iter([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]);
371        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
372        assert_eq!(slope, Some(0.0));
373        assert_eq!(intercept, Some(10.0));
374    }
375
376    #[test]
377    fn calculate_linear_regression() {
378        let ts_array = TimestampMillisecondArray::from_iter(
379            [
380                0i64, 300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700, 3000,
381            ]
382            .into_iter()
383            .map(Some),
384        );
385        let values_array = Float64Array::from_iter([
386            0.0, 10.0, 20.0, 30.0, 40.0, 0.0, 10.0, 20.0, 30.0, 40.0, 50.0,
387        ]);
388        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
389        assert_eq!(slope, Some(10.606060606060607));
390        assert_eq!(intercept, Some(6.818181818181815));
391
392        let (slope, intercept) = linear_regression(&ts_array, &values_array, 3000);
393        assert_eq!(slope, Some(10.606060606060607));
394        assert_eq!(intercept, Some(38.63636363636364));
395    }
396
397    #[test]
398    fn calculate_linear_regression_value_have_none() {
399        let ts_array = TimestampMillisecondArray::from_iter(
400            [
401                0i64, 300, 600, 900, 1200, 1350, 1500, 1800, 2100, 2400, 2550, 2700, 3000,
402            ]
403            .into_iter()
404            .map(Some),
405        );
406        let values_array: Float64Array = [
407            Some(0.0),
408            Some(10.0),
409            Some(20.0),
410            Some(30.0),
411            Some(40.0),
412            None,
413            Some(0.0),
414            Some(10.0),
415            Some(20.0),
416            Some(30.0),
417            None,
418            Some(40.0),
419            Some(50.0),
420        ]
421        .into_iter()
422        .collect();
423        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
424        assert_eq!(slope, Some(10.606060606060607));
425        assert_eq!(intercept, Some(6.818181818181815));
426    }
427
428    #[test]
429    fn calculate_linear_regression_value_all_none() {
430        let ts_array = TimestampMillisecondArray::from_iter([0i64, 300, 600].into_iter().map(Some));
431        let values_array: Float64Array = [None, None, None].into_iter().collect();
432        let (slope, intercept) = linear_regression(&ts_array, &values_array, ts_array.value(0));
433        assert_eq!(slope, None);
434        assert_eq!(intercept, None);
435    }
436
437    // From prometheus `promql/functions_test.go` case `TestKahanSum`
438    #[test]
439    fn test_kahan_sum() {
440        let inputs = vec![1.0, 10.0f64.powf(100.0), 1.0, -10.0f64.powf(100.0)];
441
442        let mut sum = 0.0;
443        let mut c = 0f64;
444
445        for v in inputs {
446            (sum, c) = compensated_sum_inc(v, sum, c);
447        }
448        assert_eq!(sum + c, 2.0)
449    }
450
451    #[test]
452    fn extract_range_array_rejects_external_dictionary_with_null_keys() {
453        let keys = Int64Array::from_iter([Some(0), None]);
454        let values = Arc::new(Float64Array::from_iter([1.0, 2.0]));
455        let dict = DictionaryArray::<Int64Type>::try_new(keys, values).unwrap();
456
457        let err = extract_range_array(&ColumnarValue::Array(Arc::new(dict))).unwrap_err();
458        assert!(err.to_string().contains("Empty range is not expected"));
459    }
460
461    #[test]
462    fn extract_range_array_accepts_internal_packed_ranges() {
463        let values = Arc::new(Float64Array::from_iter([1.0, 2.0, 3.0]));
464        let range_array = RangeArray::from_ranges(values, [(0, 2), (1, 2)]).unwrap();
465
466        let extracted =
467            extract_range_array(&ColumnarValue::Array(Arc::new(range_array.into_dict()))).unwrap();
468
469        assert_eq!(extracted.get_offset_length(0), Some((0, 2)));
470        assert_eq!(extracted.get_offset_length(1), Some((1, 2)));
471    }
472
473    #[test]
474    fn register_promql_functions_registers_all_promql_functions() {
475        let mut session_state = SessionStateBuilder::new().with_default_features().build();
476        register_promql_functions(&mut session_state);
477
478        // Every PromQL scalar function is registered under its `prom_` name;
479        // DataFusion built-ins never use the `prom_` prefix.
480        let registered: BTreeSet<&str> = session_state
481            .scalar_functions()
482            .keys()
483            .filter(|name| name.starts_with("prom_"))
484            .map(String::as_str)
485            .collect();
486        let expected: BTreeSet<&str> = [
487            "prom_absent_over_time",
488            "prom_avg_over_time",
489            "prom_changes",
490            "prom_count_over_time",
491            "prom_delta",
492            "prom_deriv",
493            "prom_double_exponential_smoothing",
494            "prom_eq",
495            "prom_float_to_string",
496            "prom_ge",
497            "prom_gt",
498            "prom_holt_winters",
499            "prom_idelta",
500            "prom_increase",
501            "prom_irate",
502            "prom_last_over_time",
503            "prom_le",
504            "prom_lt",
505            "prom_max_over_time",
506            "prom_min_over_time",
507            "prom_mixed_range_float",
508            "prom_mixed_range_histogram",
509            "prom_native_histogram_absent_over_time",
510            "prom_native_histogram_add",
511            "prom_native_histogram_avg",
512            "prom_native_histogram_avg_over_time",
513            "prom_native_histogram_changes",
514            "prom_native_histogram_count",
515            "prom_native_histogram_count_over_time",
516            "prom_native_histogram_delta",
517            "prom_native_histogram_div_scalar",
518            "prom_native_histogram_drop_bool",
519            "prom_native_histogram_drop_float",
520            "prom_native_histogram_eq",
521            "prom_native_histogram_fraction",
522            "prom_native_histogram_idelta",
523            "prom_native_histogram_increase",
524            "prom_native_histogram_irate",
525            "prom_native_histogram_keep_bool",
526            "prom_native_histogram_last_over_time",
527            "prom_native_histogram_mul_scalar",
528            "prom_native_histogram_neg",
529            "prom_native_histogram_not_eq",
530            "prom_native_histogram_present_over_time",
531            "prom_native_histogram_quantile",
532            "prom_native_histogram_rate",
533            "prom_native_histogram_resets",
534            "prom_native_histogram_scalar_mul",
535            "prom_native_histogram_stddev",
536            "prom_native_histogram_stdvar",
537            "prom_native_histogram_sub",
538            "prom_native_histogram_sum",
539            "prom_native_histogram_sum_over_time",
540            "prom_native_histogram_to_string",
541            "prom_ne",
542            "prom_predict_linear",
543            "prom_present_over_time",
544            "prom_quantile_over_time",
545            "prom_rate",
546            "prom_resets",
547            "prom_round",
548            "prom_sqrt",
549            "prom_stddev_over_time",
550            "prom_stdvar_over_time",
551            "prom_sum_over_time",
552        ]
553        .into_iter()
554        .collect();
555        assert_eq!(registered, expected);
556
557        for name in [
558            QUANTILE_NAME,
559            NativeHistogramAggAvg::name(),
560            NativeHistogramAggSum::name(),
561        ] {
562            assert!(
563                session_state.aggregate_functions().contains_key(name),
564                "{name} should be registered"
565            );
566        }
567    }
568}
promql/functions.rs

promql/
functions.rs
