promql/functions/
ieee_float.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//! PromQL comparison operators and `sqrt()`, which follow IEEE 754.
16//!
17//! Arrow compares floats by total order, where NaN equals NaN and is greater than `+Inf`, and
18//! DataFusion's `sqrt` rejects negative input. In PromQL every comparison involving NaN is false
19//! except `!=`, and the square root of a negative number is NaN. NULL input gives NULL output.
20//!
21//! The functions report the column name of the expression they replace (`a < b`, `sqrt(a)`), so
22//! query output columns are unchanged.
23
24use std::sync::Arc;
25
26use datafusion::arrow::array::{AsArray, BooleanArray, Float64Array};
27use datafusion::arrow::datatypes::{DataType, Float64Type};
28use datafusion::error::Result as DfResult;
29use datafusion::logical_expr::{
30    ColumnarValue, Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
31};
32use datafusion_expr::Expr;
33
34/// A PromQL comparison operator on two float operands.
35#[derive(Debug, PartialEq, Eq, Hash)]
36pub struct IeeeComparison {
37    op: Operator,
38    signature: Signature,
39}
40
41impl IeeeComparison {
42    /// The comparison UDF for `op`, which must be a comparison operator.
43    pub fn scalar_udf(op: Operator) -> ScalarUDF {
44        ScalarUDF::new_from_impl(Self {
45            op,
46            signature: Signature::uniform(2, vec![DataType::Float64], Volatility::Immutable),
47        })
48    }
49
50    /// All comparison UDFs, for registering them where plans are decoded.
51    pub fn all() -> Vec<ScalarUDF> {
52        [
53            Operator::Eq,
54            Operator::NotEq,
55            Operator::Lt,
56            Operator::LtEq,
57            Operator::Gt,
58            Operator::GtEq,
59        ]
60        .into_iter()
61        .map(Self::scalar_udf)
62        .collect()
63    }
64
65    fn compare(&self, lhs: f64, rhs: f64) -> bool {
66        match self.op {
67            Operator::Eq => lhs == rhs,
68            Operator::NotEq => lhs != rhs,
69            Operator::Lt => lhs < rhs,
70            Operator::LtEq => lhs <= rhs,
71            Operator::Gt => lhs > rhs,
72            Operator::GtEq => lhs >= rhs,
73            _ => unreachable!("not a comparison operator: {}", self.op),
74        }
75    }
76}
77
78impl ScalarUDFImpl for IeeeComparison {
79    fn name(&self) -> &str {
80        match self.op {
81            Operator::Eq => "prom_eq",
82            Operator::NotEq => "prom_ne",
83            Operator::Lt => "prom_lt",
84            Operator::LtEq => "prom_le",
85            Operator::Gt => "prom_gt",
86            Operator::GtEq => "prom_ge",
87            _ => unreachable!("not a comparison operator: {}", self.op),
88        }
89    }
90
91    fn signature(&self) -> &Signature {
92        &self.signature
93    }
94
95    fn return_type(&self, _arg_types: &[DataType]) -> DfResult<DataType> {
96        Ok(DataType::Boolean)
97    }
98
99    fn schema_name(&self, args: &[Expr]) -> DfResult<String> {
100        Ok(format!(
101            "{} {} {}",
102            args[0].schema_name(),
103            self.op,
104            args[1].schema_name()
105        ))
106    }
107
108    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult<ColumnarValue> {
109        let arrays = ColumnarValue::values_to_arrays(&args.args)?;
110        let lhs = arrays[0].as_primitive::<Float64Type>();
111        let rhs = arrays[1].as_primitive::<Float64Type>();
112        let result: BooleanArray = lhs
113            .iter()
114            .zip(rhs.iter())
115            .map(|(lhs, rhs)| Some(self.compare(lhs?, rhs?)))
116            .collect();
117        Ok(ColumnarValue::Array(Arc::new(result)))
118    }
119}
120
121/// PromQL `sqrt()`.
122#[derive(Debug, PartialEq, Eq, Hash)]
123pub struct IeeeSqrt {
124    signature: Signature,
125}
126
127impl IeeeSqrt {
128    pub fn scalar_udf() -> ScalarUDF {
129        ScalarUDF::new_from_impl(Self {
130            signature: Signature::uniform(1, vec![DataType::Float64], Volatility::Immutable),
131        })
132    }
133}
134
135impl ScalarUDFImpl for IeeeSqrt {
136    fn name(&self) -> &str {
137        "prom_sqrt"
138    }
139
140    fn signature(&self) -> &Signature {
141        &self.signature
142    }
143
144    fn return_type(&self, _arg_types: &[DataType]) -> DfResult<DataType> {
145        Ok(DataType::Float64)
146    }
147
148    fn schema_name(&self, args: &[Expr]) -> DfResult<String> {
149        Ok(format!("sqrt({})", args[0].schema_name()))
150    }
151
152    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult<ColumnarValue> {
153        let arrays = ColumnarValue::values_to_arrays(&args.args)?;
154        let input = arrays[0].as_primitive::<Float64Type>();
155        let result: Float64Array = input.unary(f64::sqrt);
156        Ok(ColumnarValue::Array(Arc::new(result)))
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use datafusion::arrow::array::{Array, ArrayRef};
163    use datafusion::arrow::datatypes::Field;
164    use datafusion_common::config::ConfigOptions;
165
166    use super::*;
167
168    fn invoke(udf: &ScalarUDF, args: Vec<ArrayRef>, return_type: DataType) -> ArrayRef {
169        let number_rows = args[0].len();
170        let arg_fields = args
171            .iter()
172            .map(|array| Arc::new(Field::new("a", array.data_type().clone(), true)))
173            .collect();
174        let result = udf
175            .invoke_with_args(ScalarFunctionArgs {
176                args: args.into_iter().map(ColumnarValue::Array).collect(),
177                arg_fields,
178                number_rows,
179                return_field: Arc::new(Field::new("r", return_type, true)),
180                config_options: Arc::new(ConfigOptions::default()),
181            })
182            .unwrap();
183        match result {
184            ColumnarValue::Array(array) => array,
185            ColumnarValue::Scalar(_) => unreachable!(),
186        }
187    }
188
189    #[test]
190    fn comparisons_follow_ieee_and_propagate_null() {
191        let lhs: ArrayRef = Arc::new(Float64Array::from(vec![
192            Some(1.0),
193            Some(f64::NAN),
194            Some(f64::NAN),
195            None,
196            Some(f64::INFINITY),
197        ]));
198        let rhs: ArrayRef = Arc::new(Float64Array::from(vec![
199            Some(1.0),
200            Some(f64::NAN),
201            Some(1.0),
202            Some(f64::NAN),
203            Some(f64::NAN),
204        ]));
205        let expected = [
206            (
207                Operator::Eq,
208                [Some(true), Some(false), Some(false), None, Some(false)],
209            ),
210            (
211                Operator::NotEq,
212                [Some(false), Some(true), Some(true), None, Some(true)],
213            ),
214            (
215                Operator::Lt,
216                [Some(false), Some(false), Some(false), None, Some(false)],
217            ),
218            (
219                Operator::LtEq,
220                [Some(true), Some(false), Some(false), None, Some(false)],
221            ),
222            (
223                Operator::Gt,
224                [Some(false), Some(false), Some(false), None, Some(false)],
225            ),
226            (
227                Operator::GtEq,
228                [Some(true), Some(false), Some(false), None, Some(false)],
229            ),
230        ];
231        for (op, expected) in expected {
232            let result = invoke(
233                &IeeeComparison::scalar_udf(op),
234                vec![lhs.clone(), rhs.clone()],
235                DataType::Boolean,
236            );
237            let result: Vec<_> = result.as_boolean().iter().collect();
238            assert_eq!(result, expected, "{op}");
239        }
240    }
241
242    #[test]
243    fn sqrt_of_negative_is_nan() {
244        let input: ArrayRef = Arc::new(Float64Array::from(vec![Some(4.0), Some(-1.0), None]));
245        let result = invoke(&IeeeSqrt::scalar_udf(), vec![input], DataType::Float64);
246        let result = result.as_primitive::<Float64Type>();
247        assert_eq!(result.value(0), 2.0);
248        assert!(result.value(1).is_nan());
249        assert!(result.is_null(2));
250    }
251
252    #[test]
253    fn column_names_match_the_replaced_expressions() {
254        let a = datafusion_expr::col("a");
255        let b = datafusion_expr::lit(1.5_f64);
256        let comparison =
257            IeeeComparison::scalar_udf(Operator::LtEq).call(vec![a.clone(), b.clone()]);
258        assert_eq!(
259            comparison.schema_name().to_string(),
260            a.clone().lt_eq(b).schema_name().to_string()
261        );
262        let sqrt = IeeeSqrt::scalar_udf().call(vec![a.clone()]);
263        assert_eq!(sqrt.schema_name().to_string(), "sqrt(a)");
264    }
265}
promql/functions/ieee_float.rs

promql/functions/
ieee_float.rs
