vortex_array/aggregate_fn/fns/max/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6use vortex_session::VortexSession;
7use vortex_session::registry::CachedId;
8
9use crate::ArrayRef;
10use crate::Columnar;
11use crate::ExecutionCtx;
12use crate::IntoArray;
13use crate::aggregate_fn::AggregateArgs;
14use crate::aggregate_fn::AggregateFnId;
15use crate::aggregate_fn::AggregateFnRef;
16use crate::aggregate_fn::AggregateFnSatisfaction;
17use crate::aggregate_fn::AggregateFnVTable;
18use crate::aggregate_fn::NumericalAggregateOpts;
19use crate::aggregate_fn::fns::bounded_max::BoundedMax;
20use crate::aggregate_fn::fns::min_max::MinMax;
21use crate::aggregate_fn::fns::min_max::min_max;
22use crate::aggregate_fn::fns::min_max::nan_scalar;
23use crate::aggregate_fn::fns::min_max::scalar_is_nan;
24use crate::dtype::DType;
25use crate::expr::stats::Precision;
26use crate::expr::stats::Stat;
27use crate::expr::stats::StatsProvider;
28use crate::expr::stats::StatsProviderExt;
29use crate::partial_ord::partial_max;
30use crate::scalar::Scalar;
31
32/// Compute the maximum non-null value of an array.
33///
34/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the
35/// default) NaN values are ignored, otherwise any NaN value poisons the maximum to NaN.
36#[derive(Clone, Debug)]
37pub struct Max;
38
39/// Partial accumulator state for the maximum aggregate.
40pub struct MaxPartial {
41    max: Option<Scalar>,
42}
43
44impl MaxPartial {
45    fn merge(&mut self, args: AggregateArgs<'_, NumericalAggregateOpts>, max: Scalar) {
46        if max.is_null() {
47            return;
48        }
49
50        // NaN scalars are incomparable under `partial_max`; they poison the maximum when NaNs
51        // participate, and are dropped when they are skipped.
52        if scalar_is_nan(&max) || self.is_poisoned() {
53            if !args.options.skip_nans {
54                self.poison(args);
55            }
56            return;
57        }
58
59        self.max = Some(match self.max.take() {
60            Some(current) => partial_max(max, current).vortex_expect("incomparable max scalars"),
61            None => max,
62        });
63    }
64
65    fn poison(&mut self, args: AggregateArgs<'_, NumericalAggregateOpts>) {
66        self.max = Some(nan_scalar(args.dtype));
67    }
68
69    fn is_poisoned(&self) -> bool {
70        self.max.as_ref().is_some_and(scalar_is_nan)
71    }
72}
73
74impl AggregateFnVTable for Max {
75    type Options = NumericalAggregateOpts;
76    type Partial = MaxPartial;
77
78    fn id(&self) -> AggregateFnId {
79        static ID: CachedId = CachedId::new("vortex.max");
80        *ID
81    }
82
83    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
84        Ok(Some(options.serialize()))
85    }
86
87    fn deserialize(
88        &self,
89        metadata: &[u8],
90        _session: &VortexSession,
91    ) -> VortexResult<Self::Options> {
92        NumericalAggregateOpts::deserialize(metadata)
93    }
94
95    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
96        MinMax
97            .return_dtype(options, input_dtype)
98            .map(|_| input_dtype.as_nullable())
99    }
100
101    fn can_satisfy(
102        &self,
103        options: &Self::Options,
104        requested: &AggregateFnRef,
105    ) -> AggregateFnSatisfaction {
106        if requested
107            .as_opt::<Self>()
108            .is_some_and(|other| other == options)
109        {
110            AggregateFnSatisfaction::Exact
111        } else if requested.is::<BoundedMax>() && options.skip_nans {
112            // A NaN-including maximum may be NaN, which is not a usable upper bound.
113            AggregateFnSatisfaction::Approximate
114        } else {
115            AggregateFnSatisfaction::No
116        }
117    }
118
119    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
120        self.return_dtype(options, input_dtype)
121    }
122
123    fn empty_partial(
124        &self,
125        _args: AggregateArgs<'_, Self::Options>,
126    ) -> VortexResult<Self::Partial> {
127        Ok(MaxPartial { max: None })
128    }
129
130    fn partial_from_scalar(
131        &self,
132        args: AggregateArgs<'_, Self::Options>,
133        scalar: &Scalar,
134    ) -> VortexResult<Self::Partial> {
135        let mut partial = MaxPartial { max: None };
136        // `merge` normalizes the parsed scalar: nulls stay empty and NaNs poison or drop.
137        partial.merge(args, scalar.clone());
138        Ok(partial)
139    }
140
141    fn merge_partials(
142        &self,
143        args: AggregateArgs<'_, Self::Options>,
144        mut first: Self::Partial,
145        second: Self::Partial,
146    ) -> VortexResult<Self::Partial> {
147        if let Some(max) = second.max {
148            first.merge(args, max);
149        }
150        Ok(first)
151    }
152
153    fn to_scalar(
154        &self,
155        args: AggregateArgs<'_, Self::Options>,
156        partial: &Self::Partial,
157    ) -> VortexResult<Scalar> {
158        let dtype = args.dtype.as_nullable();
159        match &partial.max {
160            Some(max) => max.cast(&dtype),
161            None => Ok(Scalar::null(dtype)),
162        }
163    }
164
165    fn is_saturated(
166        &self,
167        _args: AggregateArgs<'_, Self::Options>,
168        partial: &Self::Partial,
169    ) -> bool {
170        // A poisoned NaN-including maximum is fully determined.
171        partial.is_poisoned()
172    }
173
174    fn try_accumulate(
175        &self,
176        args: AggregateArgs<'_, Self::Options>,
177        partial: &mut Self::Partial,
178        batch: &ArrayRef,
179        _ctx: &mut ExecutionCtx,
180    ) -> VortexResult<bool> {
181        // NaN-aware shortcircuits only apply to the NaN-including float maximum; everything else
182        // takes the default dispatch path.
183        if args.options.skip_nans || !args.dtype.is_float() {
184            return Ok(false);
185        }
186        match batch.statistics().get_as::<u64>(Stat::NaNCount) {
187            Precision::Exact(0) => {
188                // NaN-free batch: the cached NaN-skipping maximum (if any) is valid. `to_scalar`
189                // re-casts to the result dtype, so the cached scalar can merge as-is.
190                if let Some(max) = batch.statistics().get(Stat::Max).as_exact() {
191                    partial.merge(args, max);
192                    return Ok(true);
193                }
194                Ok(false)
195            }
196            Precision::Exact(_) => {
197                partial.poison(args);
198                Ok(true)
199            }
200            _ => Ok(false),
201        }
202    }
203
204    fn accumulate(
205        &self,
206        args: AggregateArgs<'_, Self::Options>,
207        partial: &mut Self::Partial,
208        batch: &Columnar,
209        ctx: &mut ExecutionCtx,
210    ) -> VortexResult<()> {
211        // Delegate to the existing min_max implementation for now. A dedicated max aggregate
212        // would avoid computing min when only max is needed.
213        let array = match batch {
214            Columnar::Canonical(canonical) => canonical.clone().into_array(),
215            Columnar::Constant(constant) => constant.clone().into_array(),
216        };
217        if let Some(result) = min_max(&array, ctx, *args.options)? {
218            partial.merge(args, result.max);
219        }
220        Ok(())
221    }
222
223    fn finalize(
224        &self,
225        _args: AggregateArgs<'_, Self::Options>,
226        partials: ArrayRef,
227    ) -> VortexResult<ArrayRef> {
228        Ok(partials)
229    }
230
231    fn finalize_scalar(
232        &self,
233        args: AggregateArgs<'_, Self::Options>,
234        partial: &Self::Partial,
235    ) -> VortexResult<Scalar> {
236        self.to_scalar(args, partial)
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use vortex_buffer::buffer;
243    use vortex_error::VortexResult;
244
245    use crate::IntoArray as _;
246    use crate::VortexSessionExecute;
247    use crate::aggregate_fn::Accumulator;
248    use crate::aggregate_fn::DynAccumulator;
249    use crate::aggregate_fn::NumericalAggregateOpts;
250    use crate::aggregate_fn::fns::max::Max;
251    use crate::array_session;
252    use crate::arrays::PrimitiveArray;
253    use crate::dtype::DType;
254    use crate::dtype::Nullability;
255    use crate::dtype::PType;
256    use crate::expr::stats::Precision;
257    use crate::expr::stats::Stat;
258    use crate::scalar::Scalar;
259    use crate::scalar::ScalarValue;
260    use crate::validity::Validity;
261
262    #[test]
263    fn max_aggregate_fn() -> VortexResult<()> {
264        let mut ctx = array_session().create_execution_ctx();
265        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
266        let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::default(), dtype)?;
267
268        let batch1 = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array();
269        acc.accumulate(&batch1, &mut ctx)?;
270
271        let batch2 = PrimitiveArray::new(buffer![3i32, 25], Validity::NonNullable).into_array();
272        acc.accumulate(&batch2, &mut ctx)?;
273
274        assert_eq!(
275            acc.finish()?,
276            Scalar::primitive(25i32, Nullability::Nullable)
277        );
278        Ok(())
279    }
280
281    #[test]
282    fn max_empty_group_returns_null() -> VortexResult<()> {
283        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
284        let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::default(), dtype)?;
285
286        assert_eq!(
287            acc.finish()?,
288            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))
289        );
290        Ok(())
291    }
292
293    #[test]
294    fn max_with_nan_not_skipping() -> VortexResult<()> {
295        let mut ctx = array_session().create_execution_ctx();
296        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
297        let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::include_nans(), dtype)?;
298
299        let batch = PrimitiveArray::new(buffer![1.0f64, f64::NAN, -5.0], Validity::NonNullable)
300            .into_array();
301        acc.accumulate(&batch, &mut ctx)?;
302        assert!(acc.is_saturated());
303
304        let result = acc.finish()?;
305        assert!(
306            result
307                .as_primitive()
308                .typed_value::<f64>()
309                .is_some_and(f64::is_nan)
310        );
311        Ok(())
312    }
313
314    #[test]
315    fn max_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> {
316        let mut ctx = array_session().create_execution_ctx();
317        // The array has no NaNs; a planted exact NaNCount stat proves the poisoning came from
318        // the stat rather than a scan.
319        let batch = PrimitiveArray::new(buffer![1.0f64, 2.0], Validity::NonNullable).into_array();
320        batch
321            .statistics()
322            .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64)));
323        let mut acc = Accumulator::try_new(
324            Max,
325            NumericalAggregateOpts::include_nans(),
326            batch.dtype().clone(),
327        )?;
328        acc.accumulate(&batch, &mut ctx)?;
329        let result = acc.finish()?;
330        assert!(
331            result
332                .as_primitive()
333                .typed_value::<f64>()
334                .is_some_and(f64::is_nan)
335        );
336        Ok(())
337    }
338
339    #[test]
340    fn max_nan_including_nullable_cached_stat() -> VortexResult<()> {
341        // A nullable float array's cached Max stat is reconstructed as a nullable scalar. The
342        // NaN-including shortcircuit merges it as-is; `to_scalar` re-casts to the result dtype.
343        let mut ctx = array_session().create_execution_ctx();
344        let array =
345            PrimitiveArray::from_option_iter([Some(1.0f64), Some(2.0), Some(3.0)]).into_array();
346        array
347            .statistics()
348            .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64)));
349        array
350            .statistics()
351            .set(Stat::Max, Precision::Exact(ScalarValue::from(3.0f64)));
352        let mut acc = Accumulator::try_new(
353            Max,
354            NumericalAggregateOpts::include_nans(),
355            array.dtype().clone(),
356        )?;
357        acc.accumulate(&array, &mut ctx)?;
358        assert_eq!(
359            acc.finish()?,
360            Scalar::primitive(3.0f64, Nullability::Nullable)
361        );
362        Ok(())
363    }
364
365    #[test]
366    fn max_casts_nonnullable_legacy_stat_to_nullable_partial() -> VortexResult<()> {
367        let mut ctx = array_session().create_execution_ctx();
368        let batch = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array();
369        batch
370            .statistics()
371            .set(Stat::Max, Precision::Exact(ScalarValue::from(25i32)));
372        let mut acc = Accumulator::try_new(
373            Max,
374            NumericalAggregateOpts::default(),
375            batch.dtype().clone(),
376        )?;
377
378        acc.accumulate(&batch, &mut ctx)?;
379
380        assert_eq!(
381            acc.finish()?,
382            Scalar::primitive(25i32, Nullability::Nullable)
383        );
384        Ok(())
385    }
386}
vortex_array/aggregate_fn/fns/max/mod.rs

vortex_array/aggregate_fn/fns/max/
mod.rs
