common_function/aggrs/
extremum.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 `min` and `max` aggregation operators.
16//!
17//! PromQL ignores NaN samples unless every sample of a group is NaN. DataFusion's
18//! `min`/`max` order NaN differently depending on the code path (its grouped
19//! accumulator lets NaN replace the current value and be replaced by the next one),
20//! so the result depends on the order of the input.
21
22use std::sync::Arc;
23
24use datafusion::arrow::array::{ArrayRef, AsArray};
25use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Float64Type};
26use datafusion::error::Result as DfResult;
27use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
28use datafusion::logical_expr::{
29    Accumulator, AggregateUDF, AggregateUDFImpl, GroupsAccumulator, Signature, Volatility,
30};
31use datafusion_common::ScalarValue;
32use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator;
33
34use crate::function_registry::FunctionRegistry;
35
36/// `max` or `min` with PromQL NaN semantics.
37#[derive(Debug, PartialEq, Eq, Hash)]
38pub struct Extremum {
39    name: &'static str,
40    is_max: bool,
41    signature: Signature,
42}
43
44impl Extremum {
45    pub const MAX_NAME: &'static str = "prom_max";
46    pub const MIN_NAME: &'static str = "prom_min";
47
48    /// Registers both functions; the registry also derives their state functions, which
49    /// distributed plans need to split the aggregation between datanodes and the frontend.
50    pub fn register(registry: &FunctionRegistry) {
51        registry.register_aggr(Self::max_udaf());
52        registry.register_aggr(Self::min_udaf());
53    }
54
55    pub fn max_udaf() -> AggregateUDF {
56        AggregateUDF::from(Self::new(Self::MAX_NAME, true))
57    }
58
59    pub fn min_udaf() -> AggregateUDF {
60        AggregateUDF::from(Self::new(Self::MIN_NAME, false))
61    }
62
63    fn new(name: &'static str, is_max: bool) -> Self {
64        Self {
65            name,
66            is_max,
67            signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
68        }
69    }
70}
71
72/// Folds `value` into `current` like Prometheus' aggregation: a NaN is replaced by any
73/// sample, and a NaN sample never replaces a number.
74fn fold(is_max: bool, current: &mut f64, value: f64) {
75    let better = if is_max {
76        value > *current
77    } else {
78        value < *current
79    };
80    if current.is_nan() || better {
81        *current = value;
82    }
83}
84
85impl AggregateUDFImpl for Extremum {
86    fn name(&self) -> &str {
87        self.name
88    }
89
90    fn signature(&self) -> &Signature {
91        &self.signature
92    }
93
94    fn return_type(&self, _arg_types: &[DataType]) -> DfResult<DataType> {
95        Ok(DataType::Float64)
96    }
97
98    fn accumulator(&self, _acc_args: AccumulatorArgs) -> DfResult<Box<dyn Accumulator>> {
99        Ok(Box::new(ExtremumAccumulator {
100            is_max: self.is_max,
101            value: None,
102        }))
103    }
104
105    fn state_fields(&self, args: StateFieldsArgs) -> DfResult<Vec<FieldRef>> {
106        Ok(vec![Arc::new(Field::new(
107            format!("{}[{}]", args.name, self.name),
108            DataType::Float64,
109            true,
110        ))])
111    }
112
113    fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
114        true
115    }
116
117    fn create_groups_accumulator(
118        &self,
119        _args: AccumulatorArgs,
120    ) -> DfResult<Box<dyn GroupsAccumulator>> {
121        Ok(groups_accumulator(self.is_max))
122    }
123}
124
125fn groups_accumulator(is_max: bool) -> Box<dyn GroupsAccumulator> {
126    Box::new(
127        PrimitiveGroupsAccumulator::<Float64Type, _>::new(&DataType::Float64, move |cur, v| {
128            fold(is_max, cur, v)
129        })
130        .with_starting_value(f64::NAN),
131    )
132}
133
134#[derive(Debug)]
135struct ExtremumAccumulator {
136    is_max: bool,
137    value: Option<f64>,
138}
139
140impl ExtremumAccumulator {
141    fn update(&mut self, values: &ArrayRef) {
142        for v in values.as_primitive::<Float64Type>().iter().flatten() {
143            match &mut self.value {
144                Some(current) => fold(self.is_max, current, v),
145                None => self.value = Some(v),
146            }
147        }
148    }
149}
150
151impl Accumulator for ExtremumAccumulator {
152    fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> {
153        self.update(&values[0]);
154        Ok(())
155    }
156
157    fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> {
158        self.update(&states[0]);
159        Ok(())
160    }
161
162    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
163        Ok(vec![ScalarValue::Float64(self.value)])
164    }
165
166    fn evaluate(&mut self) -> DfResult<ScalarValue> {
167        Ok(ScalarValue::Float64(self.value))
168    }
169
170    fn size(&self) -> usize {
171        size_of_val(self)
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use datafusion::arrow::array::Float64Array;
178    use datafusion::logical_expr::EmitTo;
179
180    use super::*;
181
182    fn groups_result(
183        is_max: bool,
184        values: Vec<Option<f64>>,
185        groups: Vec<usize>,
186    ) -> Vec<Option<f64>> {
187        let num_groups = groups.iter().max().map_or(0, |g| g + 1);
188        let mut acc = groups_accumulator(is_max);
189        let values: ArrayRef = Arc::new(Float64Array::from(values));
190        acc.update_batch(&[values], &groups, None, num_groups)
191            .unwrap();
192        let out = acc.evaluate(EmitTo::All).unwrap();
193        out.as_primitive::<Float64Type>().iter().collect()
194    }
195
196    fn is_nan(v: Option<f64>) -> bool {
197        v.is_some_and(f64::is_nan)
198    }
199
200    #[test]
201    fn nan_is_ignored_unless_the_group_is_all_nan() {
202        // Group 0 has NaN on both sides of a number; group 1 is all NaN; group 2 has only NULL.
203        let values = vec![
204            Some(f64::NAN),
205            Some(3.0),
206            Some(f64::NAN),
207            Some(f64::NAN),
208            None,
209        ];
210        let groups = vec![0, 0, 0, 1, 2];
211        for is_max in [true, false] {
212            let out = groups_result(is_max, values.clone(), groups.clone());
213            assert_eq!(out[0], Some(3.0));
214            assert!(is_nan(out[1]));
215            assert_eq!(out[2], None);
216        }
217    }
218
219    #[test]
220    fn result_does_not_depend_on_input_order() {
221        let forward = vec![Some(1.0), Some(f64::NAN), Some(4.0), Some(2.0)];
222        let mut backward = forward.clone();
223        backward.reverse();
224        assert_eq!(
225            groups_result(true, forward.clone(), vec![0; 4]),
226            vec![Some(4.0)]
227        );
228        assert_eq!(
229            groups_result(true, backward.clone(), vec![0; 4]),
230            vec![Some(4.0)]
231        );
232        assert_eq!(groups_result(false, forward, vec![0; 4]), vec![Some(1.0)]);
233        assert_eq!(groups_result(false, backward, vec![0; 4]), vec![Some(1.0)]);
234    }
235
236    #[test]
237    fn merged_partial_states_keep_nan_semantics() {
238        let mut partial_nan = ExtremumAccumulator {
239            is_max: true,
240            value: None,
241        };
242        partial_nan.update(&(Arc::new(Float64Array::from(vec![f64::NAN])) as ArrayRef));
243        let mut partial_num = ExtremumAccumulator {
244            is_max: true,
245            value: None,
246        };
247        partial_num.update(&(Arc::new(Float64Array::from(vec![2.0, 5.0])) as ArrayRef));
248
249        let mut merged = ExtremumAccumulator {
250            is_max: true,
251            value: None,
252        };
253        for partial in [&mut partial_nan, &mut partial_num] {
254            let state = partial.state().unwrap()[0].to_array().unwrap();
255            merged.merge_batch(&[state]).unwrap();
256        }
257        assert_eq!(merged.evaluate().unwrap(), ScalarValue::Float64(Some(5.0)));
258    }
259}
common_function/aggrs/extremum.rs

common_function/aggrs/
extremum.rs
