AggregateFn MinMax

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/src/aggregate_fn/fns/min_max/bool_accumulator.rs

@@ -0,0 +1,124 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::marker::PhantomData;
+use std::ops::BitAnd;
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use super::Direction;
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::aggregate_fn::accumulator::Accumulator;
+use crate::arrays::BoolArray;
+use crate::canonical::ToCanonical;
+use crate::scalar::Scalar;
+
+/// Accumulator for boolean min/max.
+///
+/// - Min is saturated as soon as `false` is seen (since false < true).
+/// - Max is saturated as soon as `true` is seen.
+pub(super) struct BoolExtremumAccumulator<D> {
+    current: Option<bool>,
+    results: Vec<Option<bool>>,
+    _direction: PhantomData<D>,
+}
+
+impl<D: Direction> BoolExtremumAccumulator<D> {
+    pub(super) fn new() -> Self {
+        Self {
+            current: None,
+            results: Vec::new(),
+            _direction: PhantomData,
+        }
+    }
+
+    #[inline]
+    fn consider(&mut self, candidate: bool) {
+        match self.current {
+            None => self.current = Some(candidate),
+            Some(cur) => {
+                if D::should_replace_bool(cur, candidate) {
+                    self.current = Some(candidate);
+                }
+            }
+        }
+    }
+}
+
+/// Count of true and false values in a boolean array, considering validity.
+struct BoolCounts {
+    true_count: u64,
+    false_count: u64,
+}
+
+fn bool_counts(bool_array: &BoolArray) -> VortexResult<BoolCounts> {
+    let validity = bool_array.validity_mask()?;
+    let bits = bool_array.to_bit_buffer();
+
+    match &validity {
+        Mask::AllTrue(_) => {
+            let true_count = bits.true_count() as u64;
+            let false_count = bool_array.len() as u64 - true_count;
+            Ok(BoolCounts {
+                true_count,
+                false_count,
+            })
+        }
+        Mask::AllFalse(_) => Ok(BoolCounts {
+            true_count: 0,
+            false_count: 0,
+        }),
+        Mask::Values(v) => {
+            let valid_bits = bits.bitand(v.bit_buffer());
+            let true_count = valid_bits.true_count() as u64;
+            let valid_count = v.bit_buffer().true_count() as u64;
+            let false_count = valid_count - true_count;
+            Ok(BoolCounts {
+                true_count,
+                false_count,
+            })
+        }
+    }
+}
+
+impl<D: Direction> Accumulator for BoolExtremumAccumulator<D> {
+    fn accumulate(&mut self, batch: &ArrayRef) -> VortexResult<()> {
+        let bool_array = batch.to_bool();
+        let counts = bool_counts(&bool_array)?;
+
+        if counts.true_count > 0 {
+            self.consider(true);
+        }
+        if counts.false_count > 0 {
+            self.consider(false);
+        }
+
+        Ok(())
+    }
+
+    fn merge(&mut self, state: &Scalar) -> VortexResult<()> {
+        if state.is_null() {
+            return Ok(());
+        }
+        if let Some(v) = state.as_bool().value() {
+            self.consider(v);
+        }
+        Ok(())
+    }
+
+    fn is_saturated(&self) -> bool {
+        self.current.is_some_and(D::is_saturated_bool)
+    }
+
+    fn flush(&mut self) -> VortexResult<()> {
+        self.results.push(self.current);
+        self.current = None;
+        Ok(())
+    }
+
+    fn finish(self: Box<Self>) -> VortexResult<ArrayRef> {
+        Ok(BoolArray::from_iter(self.results).into_array())
+    }
+}

vortex-array/src/aggregate_fn/fns/min_max/mod.rs

@@ -0,0 +1,186 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod bool_accumulator;
+mod primitive_accumulator;
+
+use num_traits::Bounded;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+
+use self::bool_accumulator::BoolExtremumAccumulator;
+use self::primitive_accumulator::PrimitiveExtremumAccumulator;
+use crate::aggregate_fn::AggregateFnId;
+use crate::aggregate_fn::AggregateFnVTable;
+use crate::aggregate_fn::accumulator::Accumulator;
+use crate::dtype::DType;
+use crate::dtype::NativePType;
+use crate::dtype::Nullability;
+use crate::match_each_native_ptype;
+use crate::scalar_fn::EmptyOptions;
+
+/// Compile-time direction for extremum accumulators.
+pub(crate) trait Direction: Send + Sync + 'static {
+    /// Returns `true` if `candidate` should replace `current`.
+    fn should_replace<T: NativePType>(current: T, candidate: T) -> bool;
+
+    /// Returns `true` if `value` is at the type's extreme and no further improvement is possible.
+    fn is_saturated<T: NativePType + Bounded>(value: T) -> bool;
+
+    /// Returns `true` if `candidate` should replace `current` for booleans.
+    fn should_replace_bool(current: bool, candidate: bool) -> bool;
+
+    /// Returns `true` if the boolean value is saturated.
+    fn is_saturated_bool(value: bool) -> bool;
+}
+
+/// Seek the minimum value.
+pub(crate) struct FindMin;
+
+impl Direction for FindMin {
+    #[inline]
+    fn should_replace<T: NativePType>(current: T, candidate: T) -> bool {
+        candidate.is_lt(current)
+    }
+
+    #[inline]
+    fn is_saturated<T: NativePType + Bounded>(value: T) -> bool {
+        value.total_compare(T::min_value()).is_eq()
+    }
+
+    #[inline]
+    fn should_replace_bool(_current: bool, candidate: bool) -> bool {
+        !candidate
+    }
+
+    #[inline]
+    fn is_saturated_bool(value: bool) -> bool {
+        !value
+    }
+}
+
+/// Seek the maximum value.
+pub(crate) struct FindMax;
+
+impl Direction for FindMax {
+    #[inline]
+    fn should_replace<T: NativePType>(current: T, candidate: T) -> bool {
+        candidate.is_gt(current)
+    }
+
+    #[inline]
+    fn is_saturated<T: NativePType + Bounded>(value: T) -> bool {
+        value.total_compare(T::max_value()).is_eq()
+    }
+
+    #[inline]
+    fn should_replace_bool(_current: bool, candidate: bool) -> bool {
+        candidate
+    }
+
+    #[inline]
+    fn is_saturated_bool(value: bool) -> bool {
+        value
+    }
+}
+
+/// Computes the minimum of numeric or boolean values.
+///
+/// Nulls and NaN values are skipped. The output dtype matches the input dtype but is always
+/// nullable.
+///
+/// # Flush semantics
+///
+/// - **Empty group** (no accumulate/merge calls): produces **null**.
+/// - **All-null group**: produces **null**.
+/// - `is_saturated()` returns true once the type's minimum value is seen.
+#[derive(Clone)]
+pub struct Min;
+
+/// Computes the maximum of numeric or boolean values.
+///
+/// Nulls and NaN values are skipped. The output dtype matches the input dtype but is always
+/// nullable.
+///
+/// # Flush semantics
+///
+/// - **Empty group** (no accumulate/merge calls): produces **null**.
+/// - **All-null group**: produces **null**.
+/// - `is_saturated()` returns true once the type's maximum value is seen.
+#[derive(Clone)]
+pub struct Max;
+
+fn return_dtype(input_dtype: &DType) -> VortexResult<DType> {
+    match input_dtype {
+        DType::Bool(_) => Ok(DType::Bool(Nullability::Nullable)),
+        DType::Primitive(p, _) => Ok(DType::Primitive(*p, Nullability::Nullable)),
+        _ => vortex_bail!(
+            "Min/Max requires numeric or boolean input, got {}",
+            input_dtype
+        ),
+    }
+}
+
+fn make_accumulator<D: Direction>(input_dtype: &DType) -> VortexResult<Box<dyn Accumulator>> {
+    match input_dtype {
+        DType::Bool(_) => Ok(Box::new(BoolExtremumAccumulator::<D>::new())),
+        DType::Primitive(p, _) => Ok(match_each_native_ptype!(*p, |T| {
+            Box::new(PrimitiveExtremumAccumulator::<T, D>::new()) as Box<dyn Accumulator>
+        })),
+        _ => vortex_bail!(
+            "Min/Max requires numeric or boolean input, got {}",
+            input_dtype
+        ),
+    }
+}
+
+impl AggregateFnVTable for Min {
+    type Options = EmptyOptions;
+
+    fn id(&self) -> AggregateFnId {
+        AggregateFnId::new_ref("vortex.min")
+    }
+
+    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> VortexResult<DType> {
+        return_dtype(input_dtype)
+    }
+
+    fn state_dtype(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType> {
+        self.return_dtype(options, input_dtype)
+    }
+
+    fn accumulator(
+        &self,
+        _options: &Self::Options,
+        input_dtype: &DType,
+    ) -> VortexResult<Box<dyn Accumulator>> {
+        make_accumulator::<FindMin>(input_dtype)
+    }
+}
+
+impl AggregateFnVTable for Max {
+    type Options = EmptyOptions;
+
+    fn id(&self) -> AggregateFnId {
+        AggregateFnId::new_ref("vortex.max")
+    }
+
+    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> VortexResult<DType> {
+        return_dtype(input_dtype)
+    }
+
+    fn state_dtype(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType> {
+        self.return_dtype(options, input_dtype)
+    }
+
+    fn accumulator(
+        &self,
+        _options: &Self::Options,
+        input_dtype: &DType,
+    ) -> VortexResult<Box<dyn Accumulator>> {
+        make_accumulator::<FindMax>(input_dtype)
+    }
+}
+
+#[cfg(test)]
+mod tests;