common_function/scalars/
uddsketch_rank.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 the scalar function `uddsketch_rank`.
16
17use std::fmt;
18use std::fmt::Display;
19use std::sync::Arc;
20
21use datafusion_common::DataFusionError;
22use datafusion_common::arrow::array::{Array, AsArray, Float64Builder};
23use datafusion_common::arrow::datatypes::{DataType, Float64Type};
24use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
25
26use crate::function::{Function, extract_args};
27use crate::function_registry::FunctionRegistry;
28
29const NAME: &str = "uddsketch_rank";
30
31/// Implements the scalar function `uddsketch_rank`.
32///
33/// It accepts a value and a serialized UDDSketch state, then estimates the
34/// value's quantile rank by counting half of the matching bucket. Both current
35/// and legacy state encodings are supported. Null arguments, empty or invalid
36/// states, and invalid values produce null results.
37#[derive(Debug)]
38pub(crate) struct UddSketchRankFunction {
39    signature: Signature,
40}
41
42impl UddSketchRankFunction {
43    pub fn register(registry: &FunctionRegistry) {
44        registry.register_scalar(UddSketchRankFunction::default());
45    }
46}
47
48impl Default for UddSketchRankFunction {
49    fn default() -> Self {
50        Self {
51            signature: Signature::exact(
52                vec![DataType::Float64, DataType::Binary],
53                Volatility::Immutable,
54            ),
55        }
56    }
57}
58
59impl Display for UddSketchRankFunction {
60    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
61        write!(f, "{}", NAME.to_ascii_uppercase())
62    }
63}
64
65impl Function for UddSketchRankFunction {
66    fn name(&self) -> &str {
67        NAME
68    }
69
70    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
71        Ok(DataType::Float64)
72    }
73
74    fn signature(&self) -> &Signature {
75        &self.signature
76    }
77
78    fn invoke_with_args(
79        &self,
80        args: ScalarFunctionArgs,
81    ) -> datafusion_common::Result<ColumnarValue> {
82        let [arg0, arg1] = extract_args(self.name(), &args)?;
83
84        let Some(values) = arg0.as_primitive_opt::<Float64Type>() else {
85            return Err(DataFusionError::Execution(format!(
86                "'{}' expects 1st argument to be Float64 datatype, got {}",
87                self.name(),
88                arg0.data_type()
89            )));
90        };
91        let Some(sketches) = arg1.as_binary_opt::<i32>() else {
92            return Err(DataFusionError::Execution(format!(
93                "'{}' expects 2nd argument to be Binary datatype, got {}",
94                self.name(),
95                arg1.data_type()
96            )));
97        };
98        let mut builder = Float64Builder::with_capacity(sketches.len());
99
100        for i in 0..sketches.len() {
101            if values.is_null(i) || sketches.is_null(i) {
102                builder.append_null();
103                continue;
104            }
105
106            match crate::uddsketch_compat::rank(sketches.value(i), values.value(i)) {
107                Ok(Some(rank)) => builder.append_value(rank),
108                Ok(None) => builder.append_null(),
109                Err(error) => {
110                    common_telemetry::trace!("Failed to calculate UDDSketch rank: {}", error);
111                    builder.append_null();
112                }
113            }
114        }
115
116        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use std::sync::Arc;
123
124    use arrow_schema::Field;
125    use datafusion_common::arrow::array::{Array, AsArray, BinaryArray, Float64Array};
126    use datafusion_common::arrow::datatypes::{DataType, Float64Type};
127    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
128    use uddsketch::UddSketch;
129
130    use super::UddSketchRankFunction;
131    use crate::function::Function;
132    use crate::uddsketch_compat;
133
134    fn invoke(
135        function: &UddSketchRankFunction,
136        values: Float64Array,
137        states: BinaryArray,
138    ) -> Float64Array {
139        let number_rows = values.len();
140        let result = function
141            .invoke_with_args(ScalarFunctionArgs {
142                args: vec![
143                    ColumnarValue::Array(Arc::new(values)),
144                    ColumnarValue::Array(Arc::new(states)),
145                ],
146                arg_fields: vec![],
147                number_rows,
148                return_field: Arc::new(Field::new("x", DataType::Float64, true)),
149                config_options: Arc::new(Default::default()),
150            })
151            .unwrap();
152        let ColumnarValue::Array(result) = result else {
153            unreachable!()
154        };
155        result.as_primitive::<Float64Type>().clone()
156    }
157
158    #[test]
159    fn test_uddsketch_rank_function_name_and_return_type() {
160        let function = UddSketchRankFunction::default();
161
162        assert_eq!("uddsketch_rank", function.name());
163        assert_eq!(
164            DataType::Float64,
165            function
166                .return_type(&[DataType::Float64, DataType::Binary])
167                .unwrap()
168        );
169    }
170
171    #[test]
172    fn test_uddsketch_rank_function_canonical_state() {
173        let function = UddSketchRankFunction::default();
174        let mut sketch = UddSketch::new(128, 0.01).unwrap();
175        sketch
176            .add_batch(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0])
177            .unwrap();
178        let values = [5.0, 10.0, 55.0, 100.0, 110.0];
179        let expected = values.map(|value| sketch.rank(value).unwrap().unwrap());
180        let encoded = sketch.encode().unwrap();
181        let states = BinaryArray::from_iter_values((0..values.len()).map(|_| encoded.as_slice()));
182
183        let result = invoke(&function, Float64Array::from(values.to_vec()), states);
184
185        assert_eq!(result.values(), &expected);
186    }
187
188    #[test]
189    fn test_uddsketch_rank_function_reads_legacy_state() {
190        let function = UddSketchRankFunction::default();
191        let values = [f64::NEG_INFINITY, 0.0, 0.99, f64::INFINITY];
192        let states = BinaryArray::from_iter_values(
193            (0..values.len()).map(|_| uddsketch_compat::LEGACY_STATE),
194        );
195
196        let result = invoke(&function, Float64Array::from(values.to_vec()), states);
197
198        assert_eq!(result.null_count(), 0);
199        assert_eq!(result.values(), &[0.0, 0.375, 0.625, 1.0]);
200    }
201
202    #[test]
203    fn test_uddsketch_rank_function_returns_null_for_invalid_rows() {
204        let function = UddSketchRankFunction::default();
205        let empty = UddSketch::new(128, 0.01).unwrap().encode().unwrap();
206        let mut populated = UddSketch::new(128, 0.01).unwrap();
207        populated.add(1.0).unwrap();
208        let populated = populated.encode().unwrap();
209        let malformed = [1, 2, 3];
210        let states = BinaryArray::from_iter([
211            Some(populated.as_slice()),
212            None,
213            Some(empty.as_slice()),
214            Some(malformed.as_slice()),
215            Some(populated.as_slice()),
216        ]);
217        let values =
218            Float64Array::from(vec![None, Some(1.0), Some(1.0), Some(1.0), Some(f64::NAN)]);
219
220        let result = invoke(&function, values, states);
221
222        assert_eq!(result.null_count(), 5);
223    }
224}
common_function/scalars/uddsketch_rank.rs

common_function/scalars/
uddsketch_rank.rs
