Tighten pruning aggregate semantics

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

Author: gatesn

vortex-array/public-api.lock

@@ -468,7 +468,7 @@ pub fn vortex_array::aggregate_fn::fns::bounded_max::BoundedMax::try_accumulate(
 
 pub struct vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions
 
-pub vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions::max_bytes: usize
+pub vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions::max_bytes: core::num::nonzero::NonZeroUsize
 
 impl core::clone::Clone for vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions
 
@@ -546,7 +546,7 @@ pub fn vortex_array::aggregate_fn::fns::bounded_min::BoundedMin::try_accumulate(
 
 pub struct vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions
 
-pub vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions::max_bytes: usize
+pub vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions::max_bytes: core::num::nonzero::NonZeroUsize
 
 impl core::clone::Clone for vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions
 

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

@@ -17,6 +17,8 @@ use crate::scalar::Scalar;
 
 /// Compute whether every value in an array is NaN.
 ///
+/// Like other `all` aggregates, this is vacuously true for empty input.
+///
 /// This is a pruning aggregate, not just a convenience wrapper around
 /// [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount]. Pruning aggregates must prove a
 /// row-wise fact for every value in the scope, so their partials remain valid when a stats column is
@@ -39,7 +41,8 @@ impl AggregateFnVTable for AllNan {
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
-        has_nans(input_dtype).then_some(DType::Bool(Nullability::Nullable))
+        matches!(input_dtype, DType::Primitive(ptype, _) if ptype.is_float())
+            .then_some(DType::Bool(Nullability::Nullable))
     }
 
     fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
@@ -77,7 +80,7 @@ impl AggregateFnVTable for AllNan {
         batch: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<bool> {
-        if !has_nans(batch.dtype()) {
+        if !matches!(batch.dtype(), DType::Primitive(ptype, _) if ptype.is_float()) {
             *state = false;
             return Ok(true);
         }
@@ -96,7 +99,7 @@ impl AggregateFnVTable for AllNan {
             Columnar::Constant(c) => c.clone().into_array(),
             Columnar::Canonical(c) => c.clone().into_array(),
         };
-        if !has_nans(array.dtype()) {
+        if !matches!(array.dtype(), DType::Primitive(ptype, _) if ptype.is_float()) {
             *partial = false;
             return Ok(());
         }
@@ -114,10 +117,6 @@ impl AggregateFnVTable for AllNan {
     }
 }
 
-fn has_nans(dtype: &DType) -> bool {
-    matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float())
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_error::VortexResult;

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

@@ -51,6 +51,7 @@ use crate::validity::Validity;
 /// Returns `true` if and only if:
 /// - Both arrays have the same dtype and length
 /// - At every position, both are null or both are non-null with the same value
+/// - The arrays are empty, vacuously
 ///
 /// This is a fused `bool_all(non_distinct(lhs, rhs))` aggregate that allows early
 /// termination via accumulator saturation as soon as a mismatch is found.
@@ -102,6 +103,8 @@ static NAMES: LazyLock<FieldNames> = LazyLock::new(|| FieldNames::from(["lhs", "
 /// as the first distinct pair is found, the accumulator is saturated and remaining batches
 /// are skipped.
 ///
+/// Like other `all` aggregates, this is vacuously true for empty input.
+///
 /// The input is a `Struct{lhs: T, rhs: T}` and the result is `Bool(NonNullable)`.
 #[derive(Clone, Debug)]
 pub struct AllNonDistinct;

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

@@ -17,6 +17,8 @@ use crate::scalar::Scalar;
 
 /// Compute whether every value in an array is not NaN.
 ///
+/// Like other `all` aggregates, this is vacuously true for empty input.
+///
 /// This is a pruning aggregate, not just a convenience wrapper around
 /// [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount]. Pruning aggregates must prove a
 /// row-wise fact for every value in the scope, so their partials remain valid when a stats column is
@@ -39,7 +41,8 @@ impl AggregateFnVTable for AllNonNan {
     }
 
     fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
-        has_nans(input_dtype).then_some(DType::Bool(Nullability::Nullable))
+        matches!(input_dtype, DType::Primitive(ptype, _) if ptype.is_float())
+            .then_some(DType::Bool(Nullability::Nullable))
     }
 
     fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
@@ -104,10 +107,6 @@ impl AggregateFnVTable for AllNonNan {
     }
 }
 
-fn has_nans(dtype: &DType) -> bool {
-    matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float())
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_error::VortexResult;

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

@@ -15,6 +15,8 @@ use crate::dtype::Nullability;
 use crate::scalar::Scalar;
 
 /// Compute whether every value in an array is non-null.
+///
+/// Like other `all` aggregates, this is vacuously true for empty input.
 #[derive(Clone, Debug)]
 pub struct AllNonNull;
 
@@ -136,4 +138,17 @@ mod tests {
         assert!(!bool::try_from(&acc.finish()?)?);
         Ok(())
     }
+
+    #[test]
+    fn all_non_null_true_for_empty_input() -> VortexResult<()> {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
+        let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?;
+
+        let batch = PrimitiveArray::empty::<i32>(Nullability::Nullable).into_array();
+        acc.accumulate(&batch, &mut ctx)?;
+
+        assert!(bool::try_from(&acc.finish()?)?);
+        Ok(())
+    }
 }

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

@@ -15,6 +15,8 @@ use crate::dtype::Nullability;
 use crate::scalar::Scalar;
 
 /// Compute whether every value in an array is null.
+///
+/// Like other `all` aggregates, this is vacuously true for empty input.
 #[derive(Clone, Debug)]
 pub struct AllNull;
 
@@ -139,4 +141,17 @@ mod tests {
         assert!(!bool::try_from(&acc.finish()?)?);
         Ok(())
     }
+
+    #[test]
+    fn all_null_true_for_empty_input() -> VortexResult<()> {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
+        let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?;
+
+        let batch = PrimitiveArray::empty::<i32>(Nullability::Nullable).into_array();
+        acc.accumulate(&batch, &mut ctx)?;
+
+        assert!(bool::try_from(&acc.finish()?)?);
+        Ok(())
+    }
 }

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

@@ -3,6 +3,7 @@
 
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::num::NonZeroUsize;
 
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
@@ -30,12 +31,12 @@ use crate::scalar::upper_bound;
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct BoundedMaxOptions {
     /// Maximum byte length for UTF8/Binary bounds.
-    pub max_bytes: usize,
+    pub max_bytes: NonZeroUsize,
 }
 
 impl Display for BoundedMaxOptions {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.max_bytes)
+        write!(f, "{}", self.max_bytes.get())
     }
 }
 
@@ -53,7 +54,7 @@ enum BoundedMaxState {
 pub struct BoundedMaxPartial {
     state: BoundedMaxState,
     element_dtype: DType,
-    max_bytes: usize,
+    max_bytes: NonZeroUsize,
 }
 
 impl BoundedMaxPartial {
@@ -86,7 +87,7 @@ impl AggregateFnVTable for BoundedMax {
     }
 
     fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
-        let max_bytes = u64::try_from(options.max_bytes)?;
+        let max_bytes = u64::try_from(options.max_bytes.get())?;
         Ok(Some(max_bytes.to_le_bytes().to_vec()))
     }
 
@@ -105,7 +106,9 @@ impl AggregateFnVTable for BoundedMax {
         bytes.copy_from_slice(metadata);
         let max_bytes = usize::try_from(u64::from_le_bytes(bytes))?;
         vortex_ensure!(max_bytes > 0, "BoundedMax requires max_bytes > 0");
-        Ok(BoundedMaxOptions { max_bytes })
+        Ok(BoundedMaxOptions {
+            max_bytes: NonZeroUsize::new(max_bytes).vortex_expect("checked non-zero max_bytes"),
+        })
     }
 
     fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
@@ -162,7 +165,7 @@ impl AggregateFnVTable for BoundedMax {
         let Some(result) = min_max(&array, ctx)? else {
             return Ok(());
         };
-        match truncate_max(result.max, partial.max_bytes)? {
+        match truncate_max(result.max, partial.max_bytes.get())? {
             Some(bound) => partial.merge(bound),
             None => partial.unknown(),
         }
@@ -178,11 +181,7 @@ impl AggregateFnVTable for BoundedMax {
     }
 }
 
-fn supported_dtype<'a>(options: &BoundedMaxOptions, input_dtype: &'a DType) -> Option<&'a DType> {
-    if options.max_bytes == 0 {
-        return None;
-    }
-
+fn supported_dtype<'a>(_options: &BoundedMaxOptions, input_dtype: &'a DType) -> Option<&'a DType> {
     MinMax
         .return_dtype(&EmptyOptions, input_dtype)
         .map(|_| input_dtype)
@@ -209,7 +208,10 @@ fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult<Option<Scalar>>
 
 #[cfg(test)]
 mod tests {
+    use std::num::NonZeroUsize;
+
     use vortex_buffer::buffer;
+    use vortex_error::VortexExpect;
     use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
@@ -227,13 +229,19 @@ mod tests {
     use crate::scalar::Scalar;
     use crate::validity::Validity;
 
+    fn max_bytes(value: usize) -> NonZeroUsize {
+        NonZeroUsize::new(value).vortex_expect("non-zero max_bytes")
+    }
+
     #[test]
     fn bounded_max_truncates_utf8_to_upper_bound() -> VortexResult<()> {
         let mut ctx = LEGACY_SESSION.create_execution_ctx();
         let array = VarBinViewArray::from_iter_str(["aardvark", "char🪩"]).into_array();
         let mut acc = Accumulator::try_new(
             BoundedMax,
-            BoundedMaxOptions { max_bytes: 5 },
+            BoundedMaxOptions {
+                max_bytes: max_bytes(5),
+            },
             array.dtype().clone(),
         )?;
 
@@ -249,7 +257,9 @@ mod tests {
         let array = VarBinViewArray::from_iter_bin([&[255u8, 255, 255][..]]).into_array();
         let mut acc = Accumulator::try_new(
             BoundedMax,
-            BoundedMaxOptions { max_bytes: 2 },
+            BoundedMaxOptions {
+                max_bytes: max_bytes(2),
+            },
             array.dtype().clone(),
         )?;
 
@@ -259,13 +269,58 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn bounded_max_empty_does_not_poison_later_values() -> VortexResult<()> {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let empty = VarBinViewArray::from_iter_bin(Vec::<&[u8]>::new()).into_array();
+        let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array();
+        let mut acc = Accumulator::try_new(
+            BoundedMax,
+            BoundedMaxOptions {
+                max_bytes: max_bytes(2),
+            },
+            empty.dtype().clone(),
+        )?;
+
+        acc.accumulate(&empty, &mut ctx)?;
+        acc.accumulate(&values, &mut ctx)?;
+
+        assert_eq!(
+            acc.finish()?,
+            Scalar::binary(buffer![1u8], Nullability::Nullable)
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn bounded_max_unknown_poisons_later_values() -> VortexResult<()> {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let unknown = VarBinViewArray::from_iter_bin([&[255u8, 255, 255][..]]).into_array();
+        let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array();
+        let mut acc = Accumulator::try_new(
+            BoundedMax,
+            BoundedMaxOptions {
+                max_bytes: max_bytes(2),
+            },
+            unknown.dtype().clone(),
+        )?;
+
+        acc.accumulate(&unknown, &mut ctx)?;
+        acc.accumulate(&values, &mut ctx)?;
+
+        assert_eq!(acc.finish()?, Scalar::null(unknown.dtype().as_nullable()));
+        Ok(())
+    }
+
     #[test]
     fn bounded_max_keeps_fixed_width_values_exact() -> VortexResult<()> {
         let mut ctx = LEGACY_SESSION.create_execution_ctx();
         let array = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array();
         let mut acc = Accumulator::try_new(
             BoundedMax,
-            BoundedMaxOptions { max_bytes: 9 },
+            BoundedMaxOptions {
+                max_bytes: max_bytes(9),
+            },
             array.dtype().clone(),
         )?;
 
@@ -280,7 +335,9 @@ mod tests {
 
     #[test]
     fn bounded_max_options_round_trip() -> VortexResult<()> {
-        let options = BoundedMaxOptions { max_bytes: 64 };
+        let options = BoundedMaxOptions {
+            max_bytes: max_bytes(64),
+        };
         let metadata = BoundedMax
             .serialize(&options)?
             .expect("serializable options");