promql/functions/
round.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::error::DataFusionError;
18use datafusion_common::ScalarValue;
19use datafusion_expr::{ColumnarValue, ScalarUDF, Volatility, create_udf};
20use datatypes::arrow::array::{AsArray, Float64Array, PrimitiveArray};
21use datatypes::arrow::datatypes::{DataType, Float64Type};
22use datatypes::arrow::error::ArrowError;
23
24use crate::error;
25use crate::functions::extract_array;
26
27pub struct Round;
28
29impl Round {
30    pub const fn name() -> &'static str {
31        "prom_round"
32    }
33
34    fn input_type() -> Vec<DataType> {
35        vec![DataType::Float64, DataType::Float64]
36    }
37
38    pub fn return_type() -> DataType {
39        DataType::Float64
40    }
41
42    pub fn scalar_udf() -> ScalarUDF {
43        create_udf(
44            Self::name(),
45            Self::input_type(),
46            Self::return_type(),
47            Volatility::Volatile,
48            Arc::new(Self::round) as _,
49        )
50    }
51
52    fn round(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
53        error::ensure(
54            input.len() == 2,
55            DataFusionError::Plan("prom_round function should have 2 inputs".to_string()),
56        )?;
57
58        let value_array = extract_array(&input[0])?;
59        let nearest_col = &input[1];
60
61        match nearest_col {
62            ColumnarValue::Scalar(nearest_scalar) => {
63                let nearest = if let ScalarValue::Float64(Some(val)) = nearest_scalar {
64                    *val
65                } else {
66                    let null_array = Float64Array::new_null(value_array.len());
67                    return Ok(ColumnarValue::Array(Arc::new(null_array)));
68                };
69                let op = |a: f64| round_to_nearest(a, nearest);
70                let result: PrimitiveArray<Float64Type> =
71                    value_array.as_primitive::<Float64Type>().unary(op);
72                Ok(ColumnarValue::Array(Arc::new(result) as _))
73            }
74            ColumnarValue::Array(nearest_array) => {
75                let value_array = value_array.as_primitive::<Float64Type>();
76                let nearest_array = nearest_array.as_primitive::<Float64Type>();
77                error::ensure(
78                    value_array.len() == nearest_array.len(),
79                    DataFusionError::Execution(format!(
80                        "input arrays should have the same length, found {} and {}",
81                        value_array.len(),
82                        nearest_array.len()
83                    )),
84                )?;
85
86                let result: PrimitiveArray<Float64Type> =
87                    datatypes::arrow::compute::binary(value_array, nearest_array, round_to_nearest)
88                        .map_err(|err: ArrowError| {
89                            DataFusionError::ArrowError(Box::new(err), None)
90                        })?;
91
92                Ok(ColumnarValue::Array(Arc::new(result) as _))
93            }
94        }
95    }
96}
97
98/// Rounds like Prometheus: ties round up (towards +Inf), and multiplying by the inverse
99/// of `nearest` follows its float behavior, including NaN for `nearest == 0`.
100fn round_to_nearest(value: f64, nearest: f64) -> f64 {
101    let inverse = 1.0 / nearest;
102    (value * inverse + 0.5).floor() / inverse
103}
104
105#[cfg(test)]
106mod tests {
107    use datafusion_common::config::ConfigOptions;
108    use datafusion_expr::ScalarFunctionArgs;
109    use datatypes::arrow::array::Float64Array;
110    use datatypes::arrow::datatypes::Field;
111
112    use super::*;
113
114    fn test_round_f64(value: Vec<f64>, nearest: f64, expected: Vec<f64>) {
115        let round_udf = Round::scalar_udf();
116        let input = vec![
117            ColumnarValue::Array(Arc::new(Float64Array::from(value))),
118            ColumnarValue::Scalar(ScalarValue::Float64(Some(nearest))),
119        ];
120        let arg_fields = vec![
121            Arc::new(Field::new("a", input[0].data_type(), false)),
122            Arc::new(Field::new("b", input[1].data_type(), false)),
123        ];
124        let return_field = Arc::new(Field::new("c", DataType::Float64, false));
125        let args = ScalarFunctionArgs {
126            args: input,
127            arg_fields,
128            number_rows: 1,
129            return_field,
130            config_options: Arc::new(ConfigOptions::default()),
131        };
132        let result = round_udf.invoke_with_args(args).unwrap();
133        let result_array = extract_array(&result).unwrap();
134        assert_eq!(result_array.len(), 1);
135        assert_eq!(
136            result_array.as_primitive::<Float64Type>().values(),
137            &expected
138        );
139    }
140
141    #[test]
142    fn test_round() {
143        // Expected values are Prometheus 3.15 `round()` results.
144        test_round_f64(vec![123.456], 0.001, vec![123.456]);
145        test_round_f64(vec![123.456], 0.01, vec![123.46]);
146        test_round_f64(vec![123.456], 0.1, vec![123.5]);
147        test_round_f64(vec![123.456], 1.0, vec![123.0]);
148        test_round_f64(vec![123.456], 10.0, vec![120.0]);
149        test_round_f64(vec![123.456], 100.0, vec![100.0]);
150        test_round_f64(vec![123.456], 105.0, vec![104.99999999999999]);
151        test_round_f64(vec![123.456], 1000.0, vec![0.0]);
152    }
153
154    #[test]
155    fn test_round_ties_towards_positive_infinity() {
156        test_round_f64(vec![-0.5], 1.0, vec![0.0]);
157        test_round_f64(vec![-2.5], 1.0, vec![-2.0]);
158        test_round_f64(vec![2.5], 1.0, vec![3.0]);
159    }
160
161    #[test]
162    fn test_round_to_zero_is_nan() {
163        let round_udf = Round::scalar_udf();
164        let input = vec![
165            ColumnarValue::Array(Arc::new(Float64Array::from(vec![123.456]))),
166            ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))),
167        ];
168        let arg_fields = vec![
169            Arc::new(Field::new("a", input[0].data_type(), false)),
170            Arc::new(Field::new("b", input[1].data_type(), false)),
171        ];
172        let args = ScalarFunctionArgs {
173            args: input,
174            arg_fields,
175            number_rows: 1,
176            return_field: Arc::new(Field::new("c", DataType::Float64, false)),
177            config_options: Arc::new(ConfigOptions::default()),
178        };
179        let result = extract_array(&round_udf.invoke_with_args(args).unwrap()).unwrap();
180        assert!(result.as_primitive::<Float64Type>().value(0).is_nan());
181    }
182}
promql/functions/round.rs

promql/functions/
round.rs
