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

vortex_array/aggregate_fn/fns/max/
mod.rs
