promql/functions/
idelta.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
15use std::fmt::Display;
16use std::sync::Arc;
17
18use datafusion::arrow::array::{Float64Array, Float64Builder, TimestampMillisecondArray};
19use datafusion::arrow::datatypes::TimeUnit;
20use datafusion::common::DataFusionError;
21use datafusion::logical_expr::{ScalarUDF, Volatility};
22use datafusion::physical_plan::ColumnarValue;
23use datafusion_expr::create_udf;
24use datatypes::arrow::array::Array;
25use datatypes::arrow::datatypes::DataType;
26
27use crate::error;
28use crate::functions::extract_array;
29use crate::range_array::RangeArray;
30
31/// The `funcIdelta` in Promql,
32/// from <https://github.com/prometheus/prometheus/blob/6bdecf377cea8e856509914f35234e948c4fcb80/promql/functions.go#L235>
33#[derive(Debug)]
34pub struct IDelta<const IS_RATE: bool> {}
35
36impl<const IS_RATE: bool> IDelta<IS_RATE> {
37    pub const fn name() -> &'static str {
38        if IS_RATE { "prom_irate" } else { "prom_idelta" }
39    }
40
41    pub fn scalar_udf() -> ScalarUDF {
42        create_udf(
43            Self::name(),
44            Self::input_type(),
45            Self::return_type(),
46            Volatility::Volatile,
47            Arc::new(Self::calc) as _,
48        )
49    }
50
51    // time index column and value column
52    fn input_type() -> Vec<DataType> {
53        vec![
54            RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)),
55            RangeArray::convert_data_type(DataType::Float64),
56        ]
57    }
58
59    fn return_type() -> DataType {
60        DataType::Float64
61    }
62
63    fn calc(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
64        // construct matrix from input
65        assert_eq!(input.len(), 2);
66        let ts_array = extract_array(&input[0])?;
67        let value_array = extract_array(&input[1])?;
68
69        let ts_range: RangeArray = RangeArray::try_new(ts_array.to_data().into())?;
70        let value_range: RangeArray = RangeArray::try_new(value_array.to_data().into())?;
71        error::ensure(
72            ts_range.len() == value_range.len(),
73            DataFusionError::Execution(format!(
74                "{}: input arrays should have the same length, found {} and {}",
75                Self::name(),
76                ts_range.len(),
77                value_range.len()
78            )),
79        )?;
80        error::ensure(
81            ts_range.value_type() == DataType::Timestamp(TimeUnit::Millisecond, None),
82            DataFusionError::Execution(format!(
83                "{}: expect TimestampMillisecond as time index array's type, found {}",
84                Self::name(),
85                ts_range.value_type()
86            )),
87        )?;
88        error::ensure(
89            value_range.value_type() == DataType::Float64,
90            DataFusionError::Execution(format!(
91                "{}: expect Float64 as value array's type, found {}",
92                Self::name(),
93                value_range.value_type()
94            )),
95        )?;
96
97        let ts_values = ts_range.values();
98        let ts_values = ts_values
99            .as_any()
100            .downcast_ref::<TimestampMillisecondArray>()
101            .unwrap()
102            .values();
103
104        let value_array = value_range.values();
105        let value_array = value_array.as_any().downcast_ref::<Float64Array>().unwrap();
106        // A NULL field value means the series has no sample at that timestamp, so the last two
107        // samples are not necessarily the last two slots.
108        let has_nulls = value_array.null_count() > 0;
109        let value_values = value_array.values();
110
111        let mut result_builder = Float64Builder::with_capacity(ts_range.len());
112
113        for index in 0..ts_range.len() {
114            let (ts_offset, len) = ts_range.get_offset_length(index).unwrap();
115            let (value_offset, value_len) = value_range.get_offset_length(index).unwrap();
116            error::ensure(
117                len == value_len,
118                DataFusionError::Execution(format!(
119                    "{}: input arrays should have the same length, found {} and {}",
120                    Self::name(),
121                    len,
122                    value_len
123                )),
124            )?;
125            let (last_position, prev_position) = if has_nulls {
126                match last_two_samples(value_array, value_offset, len) {
127                    Some(positions) => positions,
128                    None => {
129                        result_builder.append_null();
130                        continue;
131                    }
132                }
133            } else {
134                if len < 2 {
135                    result_builder.append_null();
136                    continue;
137                }
138                (len - 1, len - 2)
139            };
140
141            let last_offset = ts_offset + last_position;
142            let prev_offset = ts_offset + prev_position;
143            let sampled_interval_ms = ts_values[last_offset] - ts_values[prev_offset];
144            // Sub-millisecond samples round to the same ms; skip instead of dividing by zero.
145            if sampled_interval_ms == 0 {
146                result_builder.append_null();
147                continue;
148            }
149            let sampled_interval = sampled_interval_ms as f64 / 1000.0;
150
151            let last_value = value_values[value_offset + last_position];
152            let prev_value = value_values[value_offset + prev_position];
153
154            if !IS_RATE {
155                result_builder.append_value(last_value - prev_value);
156                continue;
157            }
158
159            let result_value = if last_value < prev_value {
160                // counter reset
161                last_value
162            } else {
163                last_value - prev_value
164            };
165
166            result_builder.append_value(result_value / sampled_interval);
167        }
168
169        let result = ColumnarValue::Array(Arc::new(result_builder.finish()));
170        Ok(result)
171    }
172}
173
174/// Locates the last two samples inside `[offset, offset + len)`, returning their positions
175/// relative to `offset`. Returns `None` when the window holds fewer than two samples.
176fn last_two_samples(values: &Float64Array, offset: usize, len: usize) -> Option<(usize, usize)> {
177    let mut last = None;
178    for position in (0..len).rev() {
179        if values.is_null(offset + position) {
180            continue;
181        }
182        match last {
183            None => last = Some(position),
184            Some(last) => return Some((last, position)),
185        }
186    }
187    None
188}
189
190impl<const IS_RATE: bool> Display for IDelta<IS_RATE> {
191    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192        write!(f, "PromQL Idelta Function (is_rate: {IS_RATE})",)
193    }
194}
195
196#[cfg(test)]
197mod test {
198
199    use datafusion::arrow::buffer::NullBuffer;
200
201    use super::*;
202    use crate::functions::test_util::simple_range_udf_runner;
203
204    #[test]
205    fn basic_idelta_and_irate() {
206        let ts_array = Arc::new(TimestampMillisecondArray::from_iter(
207            [1000i64, 3000, 5000, 7000, 9000, 11000, 13000, 15000, 17000]
208                .into_iter()
209                .map(Some),
210        ));
211        let ts_ranges = [(0, 2), (0, 5), (1, 1), (3, 3), (8, 1), (9, 0)];
212
213        let values_array = Arc::new(Float64Array::from_iter([
214            1.0, 2.0, 3.0, 5.0, 0.0, 6.0, 7.0, 8.0, 9.0,
215        ]));
216        let values_ranges = [(0, 2), (0, 5), (1, 1), (3, 3), (8, 1), (9, 0)];
217
218        // test idelta
219        let ts_range_array = RangeArray::from_ranges(ts_array.clone(), ts_ranges).unwrap();
220        let value_range_array =
221            RangeArray::from_ranges(values_array.clone(), values_ranges).unwrap();
222        simple_range_udf_runner(
223            IDelta::<false>::scalar_udf(),
224            ts_range_array,
225            value_range_array,
226            vec![],
227            vec![Some(1.0), Some(-5.0), None, Some(6.0), None, None],
228        );
229
230        // test irate
231        let ts_range_array = RangeArray::from_ranges(ts_array, ts_ranges).unwrap();
232        let value_range_array = RangeArray::from_ranges(values_array, values_ranges).unwrap();
233        simple_range_udf_runner(
234            IDelta::<true>::scalar_udf(),
235            ts_range_array,
236            value_range_array,
237            vec![],
238            // the second point represent counter reset
239            vec![Some(0.5), Some(0.0), None, Some(3.0), None, None],
240        );
241    }
242
243    #[test]
244    fn idelta_uses_last_two_samples_not_last_two_slots() {
245        let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([
246            0i64, 1000, 2000, 3000,
247        ]));
248        // Samples are 1.0@0 and 4.0@1000; the trailing slots only carry padding.
249        let values_array = Arc::new(Float64Array::new(
250            vec![1.0, 4.0, 100.0, 200.0].into(),
251            Some(NullBuffer::from_iter([true, true, false, false])),
252        ));
253        let ranges = [(0, 4), (2, 2), (3, 1)];
254
255        simple_range_udf_runner(
256            IDelta::<false>::scalar_udf(),
257            RangeArray::from_ranges(ts_array.clone(), ranges).unwrap(),
258            RangeArray::from_ranges(values_array.clone(), ranges).unwrap(),
259            vec![],
260            vec![Some(3.0), None, None],
261        );
262
263        simple_range_udf_runner(
264            IDelta::<true>::scalar_udf(),
265            RangeArray::from_ranges(ts_array, ranges).unwrap(),
266            RangeArray::from_ranges(values_array, ranges).unwrap(),
267            vec![],
268            vec![Some(3.0), None, None],
269        );
270    }
271
272    #[test]
273    fn idelta_skips_samples_in_same_millisecond() {
274        // The last two samples of the first range share 999ms; the second range is 1ms apart.
275        let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([
276            999i64, 999, 999, 1000,
277        ]));
278        let values_array = Arc::new(Float64Array::from_iter([10.0, 13.0, 999.0, 1000.0]));
279        let ranges = [(0, 3), (2, 2)];
280
281        simple_range_udf_runner(
282            IDelta::<false>::scalar_udf(),
283            RangeArray::from_ranges(ts_array.clone(), ranges).unwrap(),
284            RangeArray::from_ranges(values_array.clone(), ranges).unwrap(),
285            vec![],
286            vec![None, Some(1.0)],
287        );
288
289        simple_range_udf_runner(
290            IDelta::<true>::scalar_udf(),
291            RangeArray::from_ranges(ts_array, ranges).unwrap(),
292            RangeArray::from_ranges(values_array, ranges).unwrap(),
293            vec![],
294            vec![None, Some(1000.0)],
295        );
296    }
297}
promql/functions/idelta.rs

promql/functions/
idelta.rs
