promql/functions/
quantile.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::sync::Arc;
16
17use datafusion::arrow::array::{Float64Array, Float64Builder};
18use datafusion::arrow::datatypes::TimeUnit;
19use datafusion::common::DataFusionError;
20use datafusion::logical_expr::{ScalarUDF, Volatility};
21use datafusion::physical_plan::ColumnarValue;
22use datafusion_common::ScalarValue;
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
31pub struct QuantileOverTime;
32
33impl QuantileOverTime {
34    pub const fn name() -> &'static str {
35        "prom_quantile_over_time"
36    }
37
38    pub fn scalar_udf() -> ScalarUDF {
39        let input_types = vec![
40            // time index column
41            RangeArray::convert_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)),
42            // value column
43            RangeArray::convert_data_type(DataType::Float64),
44            // quantile
45            DataType::Float64,
46        ];
47        create_udf(
48            Self::name(),
49            input_types,
50            DataType::Float64,
51            Volatility::Volatile,
52            Arc::new(Self::quantile_over_time) as _,
53        )
54    }
55
56    fn quantile_over_time(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
57        error::ensure(
58            input.len() == 3,
59            DataFusionError::Plan(
60                "prom_quantile_over_time function should have 3 inputs".to_string(),
61            ),
62        )?;
63
64        let ts_array = extract_array(&input[0])?;
65        let value_array = extract_array(&input[1])?;
66        let quantile_col = &input[2];
67
68        let ts_range: RangeArray = RangeArray::try_new(ts_array.to_data().into())?;
69        let value_range: RangeArray = RangeArray::try_new(value_array.to_data().into())?;
70        error::ensure(
71            ts_range.len() == value_range.len(),
72            DataFusionError::Execution(format!(
73                "{}: input arrays should have the same length, found {} and {}",
74                Self::name(),
75                ts_range.len(),
76                value_range.len()
77            )),
78        )?;
79        error::ensure(
80            ts_range.value_type() == DataType::Timestamp(TimeUnit::Millisecond, None),
81            DataFusionError::Execution(format!(
82                "{}: expect TimestampMillisecond as time index array's type, found {}",
83                Self::name(),
84                ts_range.value_type()
85            )),
86        )?;
87        error::ensure(
88            value_range.value_type() == DataType::Float64,
89            DataFusionError::Execution(format!(
90                "{}: expect Float64 as value array's type, found {}",
91                Self::name(),
92                value_range.value_type()
93            )),
94        )?;
95
96        let value_array = value_range.values();
97        let value_array = value_array.as_any().downcast_ref::<Float64Array>().unwrap();
98        // A NULL field value means the series has no sample at that timestamp, so a window's
99        // samples are not simply its slots.
100        let has_nulls = value_array.null_count() > 0;
101        let mut result_builder = Float64Builder::with_capacity(ts_range.len());
102        let mut scratch = Vec::new();
103        let mut samples = Vec::new();
104
105        match quantile_col {
106            ColumnarValue::Scalar(quantile_scalar) => {
107                let quantile = if let ScalarValue::Float64(Some(q)) = quantile_scalar {
108                    *q
109                } else {
110                    // For `ScalarValue::Float64(None)` or other scalar types, use NAN,
111                    // which conforms to PromQL's behavior.
112                    f64::NAN
113                };
114
115                for index in 0..ts_range.len() {
116                    let (_, ts_len) = ts_range.get_offset_length(index).unwrap();
117                    let (value_offset, value_len) = value_range.get_offset_length(index).unwrap();
118                    error::ensure(
119                        ts_len == value_len,
120                        DataFusionError::Execution(format!(
121                            "{}: time and value arrays in a group should have the same length, found {} and {}",
122                            Self::name(),
123                            ts_len,
124                            value_len
125                        )),
126                    )?;
127
128                    let window = window_samples(
129                        value_array,
130                        has_nulls,
131                        value_offset,
132                        value_len,
133                        &mut samples,
134                    );
135                    match window_quantile(window, quantile, &mut scratch) {
136                        Some(value) => result_builder.append_value(value),
137                        None => result_builder.append_null(),
138                    }
139                }
140            }
141            ColumnarValue::Array(quantile_array) => {
142                let quantile_array = quantile_array
143                    .as_any()
144                    .downcast_ref::<Float64Array>()
145                    .ok_or_else(|| {
146                        DataFusionError::Execution(format!(
147                            "{}: expect Float64 as quantile array's type, found {}",
148                            Self::name(),
149                            quantile_array.data_type()
150                        ))
151                    })?;
152
153                error::ensure(
154                    quantile_array.len() == ts_range.len(),
155                    DataFusionError::Execution(format!(
156                        "{}: quantile array should have the same length as other columns, found {} and {}",
157                        Self::name(),
158                        quantile_array.len(),
159                        ts_range.len()
160                    )),
161                )?;
162                for index in 0..ts_range.len() {
163                    let (_, ts_len) = ts_range.get_offset_length(index).unwrap();
164                    let (value_offset, value_len) = value_range.get_offset_length(index).unwrap();
165                    error::ensure(
166                        ts_len == value_len,
167                        DataFusionError::Execution(format!(
168                            "{}: time and value arrays in a group should have the same length, found {} and {}",
169                            Self::name(),
170                            ts_len,
171                            value_len
172                        )),
173                    )?;
174                    let quantile = if quantile_array.is_null(index) {
175                        f64::NAN
176                    } else {
177                        quantile_array.value(index)
178                    };
179                    let window = window_samples(
180                        value_array,
181                        has_nulls,
182                        value_offset,
183                        value_len,
184                        &mut samples,
185                    );
186                    match window_quantile(window, quantile, &mut scratch) {
187                        Some(value) => result_builder.append_value(value),
188                        None => result_builder.append_null(),
189                    }
190                }
191            }
192        }
193
194        let result = ColumnarValue::Array(Arc::new(result_builder.finish()));
195        Ok(result)
196    }
197}
198
199/// Returns the samples of the window `[offset, offset + len)`, collecting the non-null ones
200/// into `samples` when the backing array has nulls and borrowing the slice otherwise.
201fn window_samples<'a>(
202    values: &'a Float64Array,
203    has_nulls: bool,
204    offset: usize,
205    len: usize,
206    samples: &'a mut Vec<f64>,
207) -> &'a [f64] {
208    let raw_values = values.values();
209    if !has_nulls {
210        return &raw_values[offset..offset + len];
211    }
212    samples.clear();
213    samples.extend(
214        (offset..offset + len)
215            .filter(|index| values.is_valid(*index))
216            .map(|index| raw_values[index]),
217    );
218    samples
219}
220
221/// Quantile of one range window, or `None` when the window holds no sample.
222///
223/// Prometheus returns an empty vector for a range without float samples rather than the NaN
224/// that [`quantile_impl`] yields for an empty slice, so the emptiness check belongs here and
225/// not in the shared kernel.
226fn window_quantile(values: &[f64], quantile: f64, scratch: &mut Vec<f64>) -> Option<f64> {
227    if values.is_empty() {
228        return None;
229    }
230    quantile_with_scratch(values, quantile, scratch)
231}
232
233/// Refer to <https://github.com/prometheus/prometheus/blob/6e2905a4d4ff9b47b1f6d201333f5bd53633f921/promql/quantile.go#L357-L386>
234pub(crate) fn quantile_impl(values: &[f64], quantile: f64) -> Option<f64> {
235    let mut scratch = Vec::new();
236    quantile_with_scratch(values, quantile, &mut scratch)
237}
238
239/// Same as [quantile_impl] but reuses a caller-provided scratch buffer to avoid
240/// per-call allocation.
241fn quantile_with_scratch(values: &[f64], quantile: f64, scratch: &mut Vec<f64>) -> Option<f64> {
242    if quantile.is_nan() || values.is_empty() {
243        return Some(f64::NAN);
244    }
245    if quantile < 0.0 {
246        return Some(f64::NEG_INFINITY);
247    }
248    if quantile > 1.0 {
249        return Some(f64::INFINITY);
250    }
251
252    scratch.clear();
253    scratch.extend_from_slice(values);
254    scratch.sort_unstable_by(f64::total_cmp);
255
256    let length = scratch.len();
257    let rank = quantile * (length - 1) as f64;
258
259    let lower_index = rank.floor() as usize;
260    let upper_index = (length - 1).min(lower_index + 1);
261    let weight = rank - rank.floor();
262
263    let result = scratch[lower_index] * (1.0 - weight) + scratch[upper_index] * weight;
264    Some(result)
265}
266
267#[cfg(test)]
268mod tests {
269    use datafusion::arrow::array::TimestampMillisecondArray;
270    use datafusion::arrow::buffer::NullBuffer;
271
272    use super::*;
273
274    #[test]
275    fn test_quantile_impl_empty() {
276        let values = &[];
277        let q = 0.5;
278        assert!(quantile_impl(values, q).unwrap().is_nan());
279    }
280
281    #[test]
282    fn test_quantile_impl_nan() {
283        let values = &[1.0, 2.0, 3.0];
284        let q = f64::NAN;
285        assert!(quantile_impl(values, q).unwrap().is_nan());
286    }
287
288    #[test]
289    fn test_quantile_impl_negative_quantile() {
290        let values = &[1.0, 2.0, 3.0];
291        let q = -0.5;
292        assert_eq!(quantile_impl(values, q).unwrap(), f64::NEG_INFINITY);
293    }
294
295    #[test]
296    fn test_quantile_impl_greater_than_one_quantile() {
297        let values = &[1.0, 2.0, 3.0];
298        let q = 1.5;
299        assert_eq!(quantile_impl(values, q).unwrap(), f64::INFINITY);
300    }
301
302    #[test]
303    fn test_quantile_impl_single_element() {
304        let values = &[1.0];
305        let q = 0.8;
306        assert_eq!(quantile_impl(values, q).unwrap(), 1.0);
307    }
308
309    #[test]
310    fn test_quantile_impl_even_length() {
311        let values = &[3.0, 1.0, 5.0, 2.0];
312        let q = 0.5;
313        assert_eq!(quantile_impl(values, q).unwrap(), 2.5);
314    }
315
316    #[test]
317    fn test_quantile_impl_odd_length() {
318        let values = &[4.0, 1.0, 3.0, 2.0, 5.0];
319        let q = 0.25;
320        assert_eq!(quantile_impl(values, q).unwrap(), 2.0);
321    }
322
323    #[test]
324    fn quantile_over_time_ranks_samples_only() {
325        let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([
326            0i64, 1000, 2000,
327        ]));
328        // Samples are 1.0 and 4.0; ranking the padding too would pull the median down.
329        let values_array = Arc::new(Float64Array::new(
330            vec![1.0, -100.0, 4.0].into(),
331            Some(NullBuffer::from_iter([true, false, true])),
332        ));
333        // The second window holds no sample, the third holds no slot at all.
334        let ranges = [(0, 3), (1, 1), (3, 0)];
335
336        let input = vec![
337            ColumnarValue::Array(Arc::new(
338                RangeArray::from_ranges(ts_array, ranges)
339                    .unwrap()
340                    .into_dict(),
341            )),
342            ColumnarValue::Array(Arc::new(
343                RangeArray::from_ranges(values_array, ranges)
344                    .unwrap()
345                    .into_dict(),
346            )),
347            ColumnarValue::Scalar(ScalarValue::Float64(Some(0.5))),
348        ];
349        let output = extract_array(&QuantileOverTime::quantile_over_time(&input).unwrap()).unwrap();
350        let output = output.as_any().downcast_ref::<Float64Array>().unwrap();
351
352        assert_eq!(
353            output.iter().collect::<Vec<_>>(),
354            vec![Some(2.5), None, None]
355        );
356    }
357
358    #[test]
359    fn quantile_over_time_keeps_nan_for_an_invalid_quantile() {
360        let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([0i64, 1000]));
361        let values_array = Arc::new(Float64Array::from_iter_values([1.0, 4.0]));
362        let ranges = [(0, 2)];
363
364        let input = vec![
365            ColumnarValue::Array(Arc::new(
366                RangeArray::from_ranges(ts_array, ranges)
367                    .unwrap()
368                    .into_dict(),
369            )),
370            ColumnarValue::Array(Arc::new(
371                RangeArray::from_ranges(values_array, ranges)
372                    .unwrap()
373                    .into_dict(),
374            )),
375            ColumnarValue::Scalar(ScalarValue::Float64(None)),
376        ];
377        let output = extract_array(&QuantileOverTime::quantile_over_time(&input).unwrap()).unwrap();
378        let output = output.as_any().downcast_ref::<Float64Array>().unwrap();
379
380        assert!(output.value(0).is_nan());
381    }
382}
promql/functions/quantile.rs

promql/functions/
quantile.rs
