promql/functions/
predict_linear.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 [`predict_linear`](https://prometheus.io/docs/prometheus/latest/querying/functions/#predict_linear) in PromQL. Refer to the [original
16//! implementation](https://github.com/prometheus/prometheus/blob/90b2f7a540b8a70d8d81372e6692dcbb67ccbaaa/promql/functions.go#L859-L872).
17
18use std::sync::Arc;
19
20use datafusion::arrow::array::{Float64Array, Float64Builder, TimestampMillisecondArray};
21use datafusion::arrow::datatypes::TimeUnit;
22use datafusion::common::DataFusionError;
23use datafusion::logical_expr::{ScalarUDF, Volatility};
24use datafusion::physical_plan::ColumnarValue;
25use datafusion_common::ScalarValue;
26use datafusion_expr::create_udf;
27use datatypes::arrow::array::Array;
28use datatypes::arrow::datatypes::DataType;
29
30use crate::error;
31use crate::functions::{extract_range_array, linear_regression_slices};
32use crate::range_array::RangeArray;
33
34pub struct PredictLinear;
35
36impl PredictLinear {
37    pub const fn name() -> &'static str {
38        "prom_predict_linear"
39    }
40
41    pub fn scalar_udf() -> ScalarUDF {
42        let input_types = vec![
43            // time index column
44            RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)),
45            // value column
46            RangeArray::convert_data_type(DataType::Float64),
47            // t
48            DataType::Int64,
49            // evaluation timestamp
50            DataType::Timestamp(TimeUnit::Millisecond, None),
51        ];
52        create_udf(
53            Self::name(),
54            input_types,
55            DataType::Float64,
56            Volatility::Volatile,
57            Arc::new(Self::predict_linear) as _,
58        )
59    }
60
61    fn predict_linear(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
62        error::ensure(
63            input.len() == 4,
64            DataFusionError::Plan("prom_predict_linear function should have 4 inputs".to_string()),
65        )?;
66
67        let t_col = &input[2];
68        let eval_ts_col = &input[3];
69
70        let ts_range = extract_range_array(&input[0])?;
71        let value_range = extract_range_array(&input[1])?;
72        error::ensure(
73            ts_range.len() == value_range.len(),
74            DataFusionError::Execution(format!(
75                "{}: input arrays should have the same length, found {} and {}",
76                Self::name(),
77                ts_range.len(),
78                value_range.len()
79            )),
80        )?;
81        error::ensure(
82            ts_range.value_type() == DataType::Timestamp(TimeUnit::Millisecond, None),
83            DataFusionError::Execution(format!(
84                "{}: expect TimestampMillisecond as time index array's type, found {}",
85                Self::name(),
86                ts_range.value_type()
87            )),
88        )?;
89        error::ensure(
90            value_range.value_type() == DataType::Float64,
91            DataFusionError::Execution(format!(
92                "{}: expect Float64 as value array's type, found {}",
93                Self::name(),
94                value_range.value_type()
95            )),
96        )?;
97
98        let t_iter: Box<dyn Iterator<Item = Option<i64>>> = match t_col {
99            ColumnarValue::Scalar(t_scalar) => {
100                let t = if let ScalarValue::Int64(Some(t_val)) = t_scalar {
101                    *t_val
102                } else {
103                    // For `ScalarValue::Int64(None)` or other scalar types, returns NULL array,
104                    // which conforms to PromQL's behavior.
105                    let null_array = Float64Array::new_null(ts_range.len());
106                    return Ok(ColumnarValue::Array(Arc::new(null_array)));
107                };
108                Box::new((0..ts_range.len()).map(move |_| Some(t)))
109            }
110            ColumnarValue::Array(t_array) => {
111                let t_array = t_array
112                    .as_any()
113                    .downcast_ref::<datafusion::arrow::array::Int64Array>()
114                    .ok_or_else(|| {
115                        DataFusionError::Execution(format!(
116                            "{}: expect Int64 as t array's type, found {}",
117                            Self::name(),
118                            t_array.data_type()
119                        ))
120                    })?;
121                error::ensure(
122                    t_array.len() == ts_range.len(),
123                    DataFusionError::Execution(format!(
124                        "{}: t array should have the same length as other columns, found {} and {}",
125                        Self::name(),
126                        t_array.len(),
127                        ts_range.len()
128                    )),
129                )?;
130
131                Box::new(t_array.iter())
132            }
133        };
134        // The evaluation instant the window above was folded for. It is the instant the regression
135        // is centered on, which is *not* the last sample's timestamp: a window may end before the
136        // step being evaluated, and an `@`-anchored window is replayed at every step while keeping
137        // its own end. See `PromPlanner::create_range_eval_ts_expr`.
138        let eval_ts_iter: Box<dyn Iterator<Item = Option<i64>>> = match eval_ts_col {
139            ColumnarValue::Scalar(eval_ts_scalar) => {
140                let eval_ts = match eval_ts_scalar {
141                    ScalarValue::TimestampMillisecond(Some(eval_ts), _) => *eval_ts,
142                    // For a NULL or otherwise unusable evaluation timestamp, returns NULL array,
143                    // which conforms to PromQL's behavior.
144                    _ => {
145                        let null_array = Float64Array::new_null(ts_range.len());
146                        return Ok(ColumnarValue::Array(Arc::new(null_array)));
147                    }
148                };
149                Box::new((0..ts_range.len()).map(move |_| Some(eval_ts)))
150            }
151            ColumnarValue::Array(eval_ts_array) => {
152                error::ensure(
153                    eval_ts_array.len() == ts_range.len(),
154                    DataFusionError::Execution(format!(
155                        "{}: evaluation timestamp array should have the same length as other columns, found {} and {}",
156                        Self::name(),
157                        eval_ts_array.len(),
158                        ts_range.len()
159                    )),
160                )?;
161                match eval_ts_array.data_type() {
162                    DataType::Timestamp(TimeUnit::Millisecond, _) => Box::new(
163                        eval_ts_array
164                            .as_any()
165                            .downcast_ref::<TimestampMillisecondArray>()
166                            .expect("checked by data type")
167                            .iter(),
168                    ),
169                    other => {
170                        return Err(DataFusionError::Execution(format!(
171                            "{}: expect TimestampMillisecond as evaluation timestamp array's type, found {}",
172                            Self::name(),
173                            other
174                        )));
175                    }
176                }
177            }
178        };
179        let all_timestamps = ts_range
180            .values()
181            .as_any()
182            .downcast_ref::<TimestampMillisecondArray>()
183            .unwrap()
184            .values();
185        let all_values = value_range
186            .values()
187            .as_any()
188            .downcast_ref::<Float64Array>()
189            .unwrap();
190        let mut result_builder = Float64Builder::with_capacity(ts_range.len());
191        for (index, (t, eval_ts)) in t_iter.zip(eval_ts_iter).enumerate() {
192            // A step without an evaluation instant has nothing to extrapolate to.
193            let Some(eval_ts) = eval_ts else {
194                result_builder.append_null();
195                continue;
196            };
197            match predict_linear_impl(
198                &ts_range,
199                &value_range,
200                all_timestamps,
201                all_values,
202                index,
203                t.unwrap(),
204                eval_ts,
205            )? {
206                Some(value) => result_builder.append_value(value),
207                None => result_builder.append_null(),
208            }
209        }
210
211        let result = ColumnarValue::Array(Arc::new(result_builder.finish()));
212        Ok(result)
213    }
214}
215
216fn predict_linear_impl(
217    ts_range: &RangeArray,
218    value_range: &RangeArray,
219    all_timestamps: &[i64],
220    all_values: &Float64Array,
221    index: usize,
222    t: i64,
223    eval_ts: i64,
224) -> Result<Option<f64>, DataFusionError> {
225    let (ts_offset, ts_len) = ts_range.get_offset_length(index).unwrap();
226    let (value_offset, value_len) = value_range.get_offset_length(index).unwrap();
227    error::ensure(
228        ts_len == value_len,
229        DataFusionError::Execution(format!(
230            "{}: time and value arrays in a group should have the same length, found {} and {}",
231            PredictLinear::name(),
232            ts_len,
233            value_len
234        )),
235    )?;
236    if ts_len < 2 {
237        return Ok(None);
238    }
239
240    // Like Prometheus, the regression is centered on the evaluation timestamp: the returned
241    // intercept is the value the window's trend predicts at the step being evaluated, and `t` is
242    // the horizon (in seconds) that is added on top of it. Centering on the last sample instead
243    // would drop the distance between the window's end and the step, which is non-zero both
244    // without `@` (the window can end before the step) and with it (the same anchored window is
245    // replayed at every step).
246    let (slope, intercept) = linear_regression_slices(
247        all_timestamps,
248        ts_offset,
249        all_values,
250        value_offset,
251        value_len,
252        eval_ts,
253    );
254
255    if slope.is_none() || intercept.is_none() {
256        return Ok(None);
257    }
258
259    Ok(Some(slope.unwrap() * t as f64 + intercept.unwrap()))
260}
261
262#[cfg(test)]
263mod test {
264    use std::vec;
265
266    use datafusion::arrow::array::{DictionaryArray, Int64Array};
267    use datatypes::arrow::datatypes::Int64Type;
268
269    use super::*;
270    use crate::functions::test_util::simple_range_udf_runner;
271
272    // build timestamp range and value range arrays for test
273    fn build_test_range_arrays() -> (RangeArray, RangeArray) {
274        let ts_array = Arc::new(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 ranges = [(0, 11)];
282
283        let values_array = Arc::new(Float64Array::from_iter([
284            0.0, 10.0, 20.0, 30.0, 40.0, 0.0, 10.0, 20.0, 30.0, 40.0, 50.0,
285        ]));
286
287        let ts_range_array = RangeArray::from_ranges(ts_array, ranges).unwrap();
288        let value_range_array = RangeArray::from_ranges(values_array, ranges).unwrap();
289
290        (ts_range_array, value_range_array)
291    }
292
293    /// The evaluation timestamp at which the window built by [`build_test_range_arrays`] ends.
294    fn window_end_eval_ts() -> ScalarValue {
295        ScalarValue::TimestampMillisecond(Some(3000), None)
296    }
297
298    #[test]
299    fn calculate_predict_linear_none() {
300        let ts_array = Arc::new(TimestampMillisecondArray::from_iter(
301            [0i64].into_iter().map(Some),
302        ));
303        let ranges = [(0, 0), (0, 1)];
304        let values_array = Arc::new(Float64Array::from_iter([0.0]));
305        let ts_array = RangeArray::from_ranges(ts_array, ranges).unwrap();
306        let value_array = RangeArray::from_ranges(values_array, ranges).unwrap();
307        simple_range_udf_runner(
308            PredictLinear::scalar_udf(),
309            ts_array,
310            value_array,
311            vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()],
312            vec![None, None],
313        );
314    }
315
316    #[test]
317    fn calculate_predict_linear_test1() {
318        let (ts_array, value_array) = build_test_range_arrays();
319        simple_range_udf_runner(
320            PredictLinear::scalar_udf(),
321            ts_array,
322            value_array,
323            vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()],
324            // value at t = 0
325            vec![Some(38.63636363636364)],
326        );
327    }
328
329    #[test]
330    fn calculate_predict_linear_test2() {
331        let (ts_array, value_array) = build_test_range_arrays();
332        simple_range_udf_runner(
333            PredictLinear::scalar_udf(),
334            ts_array,
335            value_array,
336            vec![ScalarValue::Int64(Some(3000)), window_end_eval_ts()],
337            // value at t = 3000
338            vec![Some(31856.818181818187)],
339        );
340    }
341
342    #[test]
343    fn calculate_predict_linear_test3() {
344        let (ts_array, value_array) = build_test_range_arrays();
345        simple_range_udf_runner(
346            PredictLinear::scalar_udf(),
347            ts_array,
348            value_array,
349            vec![ScalarValue::Int64(Some(4200)), window_end_eval_ts()],
350            // value at t = 4200
351            vec![Some(44584.09090909091)],
352        );
353    }
354
355    #[test]
356    fn calculate_predict_linear_test4() {
357        let (ts_array, value_array) = build_test_range_arrays();
358        simple_range_udf_runner(
359            PredictLinear::scalar_udf(),
360            ts_array,
361            value_array,
362            vec![ScalarValue::Int64(Some(6600)), window_end_eval_ts()],
363            // value at t = 6600
364            vec![Some(70038.63636363638)],
365        );
366    }
367
368    #[test]
369    fn calculate_predict_linear_test5() {
370        let (ts_array, value_array) = build_test_range_arrays();
371        simple_range_udf_runner(
372            PredictLinear::scalar_udf(),
373            ts_array,
374            value_array,
375            vec![ScalarValue::Int64(Some(7800)), window_end_eval_ts()],
376            // value at t = 7800
377            vec![Some(82765.9090909091)],
378        );
379    }
380
381    #[test]
382    fn calculate_predict_linear_with_misaligned_offsets() {
383        let ts_values = Arc::new(TimestampMillisecondArray::from_iter(
384            [0i64, 1000, 2000, 3000].into_iter().map(Some),
385        ));
386        let value_values = Arc::new(Float64Array::from_iter([10.0, 20.0, 30.0]));
387        let ts_array = RangeArray::from_ranges(ts_values, [(1, 3)]).unwrap();
388        let value_array = RangeArray::from_ranges(value_values, [(0, 3)]).unwrap();
389
390        simple_range_udf_runner(
391            PredictLinear::scalar_udf(),
392            ts_array,
393            value_array,
394            vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()],
395            vec![Some(30.0)],
396        );
397    }
398
399    #[test]
400    fn predict_linear_rejects_external_dictionary_with_null_keys() {
401        let ts_values = Arc::new(TimestampMillisecondArray::from_iter(
402            [0i64, 1000].into_iter().map(Some),
403        ));
404        let ts_keys = Int64Array::from_iter([Some(0), None]);
405        let ts_dict = DictionaryArray::<Int64Type>::try_new(ts_keys, ts_values).unwrap();
406
407        let value_values = Arc::new(Float64Array::from_iter([1.0, 2.0]));
408        let value_keys = Int64Array::from_iter([Some(0), Some(1)]);
409        let value_dict = DictionaryArray::<Int64Type>::try_new(value_keys, value_values).unwrap();
410
411        let err = PredictLinear::predict_linear(&[
412            ColumnarValue::Array(Arc::new(ts_dict)),
413            ColumnarValue::Array(Arc::new(value_dict)),
414            ColumnarValue::Scalar(ScalarValue::Int64(Some(0))),
415            ColumnarValue::Scalar(window_end_eval_ts()),
416        ])
417        .unwrap_err();
418
419        assert!(err.to_string().contains("Empty range is not expected"));
420    }
421
422    /// The regression is centered on the evaluation timestamp, so the returned intercept follows
423    /// the step being evaluated even though the window stays the same.
424    #[test]
425    fn calculate_predict_linear_centers_on_eval_ts() {
426        // One sample per second with values 0, 1, 2: the trend is exactly one unit per second.
427        for (eval_ts, t, expected) in [
428            // The window ends at 2s: evaluating there returns its last value.
429            (2000, 0, 2.0),
430            // One second past the window's end the trend is at 3.
431            (3000, 0, 3.0),
432            // One second past the window's end, predicting another two seconds ahead.
433            (3000, 2, 5.0),
434        ] {
435            let ts_values = Arc::new(TimestampMillisecondArray::from_iter(
436                [0i64, 1000, 2000].into_iter().map(Some),
437            ));
438            let value_values = Arc::new(Float64Array::from_iter([0.0, 1.0, 2.0]));
439            let ts_array = RangeArray::from_ranges(ts_values, [(0, 3)]).unwrap();
440            let value_array = RangeArray::from_ranges(value_values, [(0, 3)]).unwrap();
441
442            simple_range_udf_runner(
443                PredictLinear::scalar_udf(),
444                ts_array,
445                value_array,
446                vec![
447                    ScalarValue::Int64(Some(t)),
448                    ScalarValue::TimestampMillisecond(Some(eval_ts), None),
449                ],
450                vec![Some(expected)],
451            );
452        }
453    }
454
455    /// The evaluation timestamp arrives as a column, so the UDF accepts it as an array too.
456    #[test]
457    fn calculate_predict_linear_accepts_eval_ts_array() {
458        let (ts_array, value_array) = build_test_range_arrays();
459        let eval_ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([3000]));
460
461        let result = PredictLinear::predict_linear(&[
462            ColumnarValue::Array(Arc::new(ts_array.into_dict())),
463            ColumnarValue::Array(Arc::new(value_array.into_dict())),
464            ColumnarValue::Scalar(ScalarValue::Int64(Some(0))),
465            ColumnarValue::Array(eval_ts_array),
466        ])
467        .unwrap();
468
469        let result = result
470            .into_array(1)
471            .unwrap()
472            .as_any()
473            .downcast_ref::<Float64Array>()
474            .unwrap()
475            .clone();
476        assert_eq!(result.len(), 1);
477        // The same value `calculate_predict_linear_test1` gets from a scalar evaluation timestamp.
478        assert!((result.value(0) - 38.63636363636364).abs() < 0.0001);
479    }
480}
promql/functions/predict_linear.rs

promql/functions/
predict_linear.rs
