diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 2812e43b92f..5a6e70266ca 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -15,6 +15,7 @@ use crate::ExecutionCtx; use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; use crate::columnar::AnyColumnar; use crate::dtype::DType; @@ -196,6 +197,10 @@ pub trait DynAccumulator: 'static + Send { /// Whether the accumulator's result is fully determined. fn is_saturated(&self) -> bool; + /// Whether the current result can satisfy `requested`, taking the accumulated state into + /// account. See [`AggregateFnVTable::partial_can_satisfy`]. + fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction; + /// Reset the accumulator's state to the empty group. fn reset(&mut self); @@ -322,6 +327,15 @@ impl DynAccumulator for Accumulator { }) } + fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction { + let options = self.aggregate_fn.as_::(); + // Without a partial nothing was accumulated yet, so only the options can be consulted. + match &self.partial { + Some(partial) => self.vtable.partial_can_satisfy(options, partial, requested), + None => self.vtable.can_satisfy(options, requested), + } + } + fn reset(&mut self) { self.partial = None; } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 4fedcc48c74..4615931e995 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cmp::Ordering; use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; @@ -33,7 +34,6 @@ use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; -use crate::partial_ord::partial_max; use crate::scalar::Scalar; use crate::scalar::ScalarTruncation; use crate::scalar::upper_bound; @@ -65,27 +65,59 @@ pub struct BoundedMax; enum BoundedMaxState { Empty, - Value(Scalar), + /// The true maximum: it fit within the byte bound, so it was never truncated. + Exact(Scalar), + /// An upper bound of the true maximum, either because the maximum was truncated or because + /// it was combined from a serialized partial, which doesn't record exactness. + Inexact(Scalar), Unknown, } +impl BoundedMaxState { + fn bound(bound: Scalar, exact: bool) -> Self { + if exact { + Self::Exact(bound) + } else { + Self::Inexact(bound) + } + } +} + /// Partial accumulator state for the bounded maximum aggregate. pub struct BoundedMaxPartial { state: BoundedMaxState, } impl BoundedMaxPartial { - fn merge_bound(&mut self, max: Scalar) { - if max.is_null() { + fn merge_bound(&mut self, bound: Scalar, exact: bool) { + if bound.is_null() { return; } - self.state = match std::mem::replace(&mut self.state, BoundedMaxState::Empty) { - BoundedMaxState::Empty => BoundedMaxState::Value(max), - BoundedMaxState::Value(current) => BoundedMaxState::Value( - partial_max(max, current).vortex_expect("incomparable bounded max scalars"), - ), - BoundedMaxState::Unknown => BoundedMaxState::Unknown, + let (current, current_exact) = + match std::mem::replace(&mut self.state, BoundedMaxState::Empty) { + BoundedMaxState::Empty => { + self.state = BoundedMaxState::bound(bound, exact); + return; + } + BoundedMaxState::Unknown => { + self.state = BoundedMaxState::Unknown; + return; + } + BoundedMaxState::Exact(current) => (current, true), + BoundedMaxState::Inexact(current) => (current, false), + }; + + // The larger bound wins. It is the true maximum iff it was not truncated: every other + // value is at most its own bound, which is below the winner. On a tie, either side being + // exact means the true maximum reaches the bound. + self.state = match bound + .partial_cmp(¤t) + .vortex_expect("incomparable bounded max scalars") + { + Ordering::Greater => BoundedMaxState::bound(bound, exact), + Ordering::Less => BoundedMaxState::bound(current, current_exact), + Ordering::Equal => BoundedMaxState::bound(current, exact || current_exact), }; } @@ -96,7 +128,7 @@ impl BoundedMaxPartial { fn final_scalar(&self, args: AggregateArgs<'_, BoundedMaxOptions>) -> VortexResult { let dtype = args.return_dtype.clone(); match &self.state { - BoundedMaxState::Value(max) => max.cast(&dtype), + BoundedMaxState::Exact(bound) | BoundedMaxState::Inexact(bound) => bound.cast(&dtype), BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), } } @@ -183,6 +215,27 @@ impl AggregateFnVTable for BoundedMax { } } + fn partial_can_satisfy( + &self, + options: &Self::Options, + partial: &Self::Partial, + requested: &AggregateFnRef, + ) -> AggregateFnSatisfaction { + let satisfaction = self.can_satisfy(options, requested); + let untruncated = matches!( + partial.state, + BoundedMaxState::Empty | BoundedMaxState::Exact(_) + ); + if satisfaction == AggregateFnSatisfaction::Approximate + && untruncated + && requested.is::() + { + AggregateFnSatisfaction::Exact + } else { + satisfaction + } + } + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { supported_dtype(options, input_dtype).map(make_bounded_max_partial_dtype) } @@ -226,7 +279,8 @@ impl AggregateFnVTable for BoundedMax { } else if bound.is_null() { BoundedMaxState::Empty } else { - BoundedMaxState::Value(bound) + // Serialized partials don't record exactness. + BoundedMaxState::Inexact(bound) } }; Ok(BoundedMaxPartial { state }) @@ -240,7 +294,8 @@ impl AggregateFnVTable for BoundedMax { ) -> VortexResult { match second.state { BoundedMaxState::Empty => {} - BoundedMaxState::Value(max) => first.merge_bound(max), + BoundedMaxState::Exact(bound) => first.merge_bound(bound, true), + BoundedMaxState::Inexact(bound) => first.merge_bound(bound, false), BoundedMaxState::Unknown => first.unknown(), } Ok(first) @@ -255,10 +310,10 @@ impl AggregateFnVTable for BoundedMax { let bound_dtype = args.dtype.as_nullable(); match &partial.state { BoundedMaxState::Empty => Ok(Scalar::null(dtype)), - BoundedMaxState::Value(max) => Ok(Scalar::struct_( + BoundedMaxState::Exact(bound) | BoundedMaxState::Inexact(bound) => Ok(Scalar::struct_( dtype, vec![ - max.cast(&bound_dtype)?, + bound.cast(&bound_dtype)?, Scalar::bool(false, Nullability::NonNullable), ], )), @@ -293,7 +348,7 @@ impl AggregateFnVTable for BoundedMax { return Ok(()); }; match truncate_max(result.max, args.options.max_bytes.get())? { - Some(bound) => partial.merge_bound(bound), + Some((bound, truncated)) => partial.merge_bound(bound, !truncated), None => partial.unknown(), } Ok(()) @@ -322,22 +377,21 @@ fn supported_dtype<'a>(_options: &BoundedMaxOptions, input_dtype: &'a DType) -> .map(|_| input_dtype) } -fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult> { +/// Returns the upper bound and whether it was truncated, or `None` if no bound fits `max_bytes`. +fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult> { let nullability = value.dtype().nullability(); match value.dtype() { - DType::Utf8(_) => { - Ok( - upper_bound(BufferString::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - DType::Binary(_) => { - Ok( - upper_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - _ => Ok(Some(value)), + DType::Utf8(_) => Ok(upper_bound( + BufferString::from_scalar(value)?, + max_bytes, + nullability, + )), + DType::Binary(_) => Ok(upper_bound( + ByteBuffer::from_scalar(value)?, + max_bytes, + nullability, + )), + _ => Ok(Some((value, false))), } } @@ -584,4 +638,67 @@ mod tests { assert_eq!(roundtrip, options); Ok(()) } + + #[test] + fn bounded_max_exactly_satisfies_max_only_when_untruncated() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let requested = Max.bind(NumericalAggregateOpts::skip_nans()); + let new_acc = || { + Accumulator::try_new( + BoundedMax, + BoundedMaxOptions { + max_bytes: max_bytes(4), + }, + VarBinViewArray::from_iter_str(["x"]).dtype().clone(), + ) + }; + + // Every value fits the bound, so the bound is the true extremum. + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["b", "abc"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated value that isn't the extremum doesn't affect exactness. + acc.accumulate( + &VarBinViewArray::from_iter_str(["aaaaa"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated extremum makes the bound approximate. + acc.accumulate( + &VarBinViewArray::from_iter_str(["bbbbb"]).into_array(), + &mut ctx, + )?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + + // Exactness isn't serialized, so a winning partial combined from its scalar form is + // approximate even though it was exact where it was computed. + let mut source = new_acc()?; + source.accumulate( + &VarBinViewArray::from_iter_str(["b"]).into_array(), + &mut ctx, + )?; + assert_eq!( + source.can_satisfy(&requested), + AggregateFnSatisfaction::Exact + ); + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["a"]).into_array(), + &mut ctx, + )?; + acc.combine_partials(&source.partial_scalar()?)?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 8c4e91178b0..05dfabbc4c8 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cmp::Ordering; use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; @@ -27,7 +28,6 @@ use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::columnar_min_max; use crate::dtype::DType; -use crate::partial_ord::partial_min; use crate::scalar::Scalar; use crate::scalar::ScalarTruncation; use crate::scalar::lower_bound; @@ -51,7 +51,21 @@ pub struct BoundedMin; enum BoundedMinState { Empty, - Value(Scalar), + /// The true minimum: it fit within the byte bound, so it was never truncated. + Exact(Scalar), + /// A lower bound of the true minimum, either because the minimum was truncated or because + /// it was combined from a serialized partial, which doesn't record exactness. + Inexact(Scalar), +} + +impl BoundedMinState { + fn bound(bound: Scalar, exact: bool) -> Self { + if exact { + Self::Exact(bound) + } else { + Self::Inexact(bound) + } + } } /// Partial accumulator state for the bounded minimum aggregate. @@ -60,16 +74,31 @@ pub struct BoundedMinPartial { } impl BoundedMinPartial { - fn merge(&mut self, min: Scalar) { - if min.is_null() { + fn merge(&mut self, bound: Scalar, exact: bool) { + if bound.is_null() { return; } - self.state = match std::mem::replace(&mut self.state, BoundedMinState::Empty) { - BoundedMinState::Empty => BoundedMinState::Value(min), - BoundedMinState::Value(current) => BoundedMinState::Value( - partial_min(min, current).vortex_expect("incomparable bounded min scalars"), - ), + let (current, current_exact) = + match std::mem::replace(&mut self.state, BoundedMinState::Empty) { + BoundedMinState::Empty => { + self.state = BoundedMinState::bound(bound, exact); + return; + } + BoundedMinState::Exact(current) => (current, true), + BoundedMinState::Inexact(current) => (current, false), + }; + + // The smaller bound wins. It is the true minimum iff it was not truncated: every other + // value is at least its own bound, which is above the winner. On a tie, either side being + // exact means the true minimum reaches the bound. + self.state = match bound + .partial_cmp(¤t) + .vortex_expect("incomparable bounded min scalars") + { + Ordering::Less => BoundedMinState::bound(bound, exact), + Ordering::Greater => BoundedMinState::bound(current, current_exact), + Ordering::Equal => BoundedMinState::bound(current, exact || current_exact), }; } } @@ -137,6 +166,27 @@ impl AggregateFnVTable for BoundedMin { } } + fn partial_can_satisfy( + &self, + options: &Self::Options, + partial: &Self::Partial, + requested: &AggregateFnRef, + ) -> AggregateFnSatisfaction { + let satisfaction = self.can_satisfy(options, requested); + let untruncated = matches!( + partial.state, + BoundedMinState::Empty | BoundedMinState::Exact(_) + ); + if satisfaction == AggregateFnSatisfaction::Approximate + && untruncated + && requested.is::() + { + AggregateFnSatisfaction::Exact + } else { + satisfaction + } + } + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { self.return_dtype(options, input_dtype) } @@ -155,11 +205,12 @@ impl AggregateFnVTable for BoundedMin { _args: AggregateArgs<'_, Self::Options>, scalar: &Scalar, ) -> VortexResult { - // A null partial means the producing accumulator saw nothing valid. + // A null partial means the producing accumulator saw nothing valid. Serialized partials + // don't record exactness. let state = if scalar.is_null() { BoundedMinState::Empty } else { - BoundedMinState::Value(scalar.clone()) + BoundedMinState::Inexact(scalar.clone()) }; Ok(BoundedMinPartial { state }) } @@ -170,8 +221,10 @@ impl AggregateFnVTable for BoundedMin { mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { - if let BoundedMinState::Value(min) = second.state { - first.merge(min); + match second.state { + BoundedMinState::Empty => {} + BoundedMinState::Exact(bound) => first.merge(bound, true), + BoundedMinState::Inexact(bound) => first.merge(bound, false), } Ok(first) } @@ -184,7 +237,7 @@ impl AggregateFnVTable for BoundedMin { let dtype = args.dtype.as_nullable(); match &partial.state { BoundedMinState::Empty => Ok(Scalar::null(dtype)), - BoundedMinState::Value(min) => min.cast(&dtype), + BoundedMinState::Exact(bound) | BoundedMinState::Inexact(bound) => bound.cast(&dtype), } } @@ -208,8 +261,8 @@ impl AggregateFnVTable for BoundedMin { let Some(result) = columnar_min_max(batch, NumericalAggregateOpts::default(), ctx)? else { return Ok(()); }; - if let Some(bound) = truncate_min(result.min, args.options.max_bytes.get())? { - partial.merge(bound); + if let Some((bound, truncated)) = truncate_min(result.min, args.options.max_bytes.get())? { + partial.merge(bound, !truncated); } Ok(()) } @@ -237,22 +290,21 @@ fn supported_dtype<'a>(_options: &BoundedMinOptions, input_dtype: &'a DType) -> .map(|_| input_dtype) } -fn truncate_min(value: Scalar, max_bytes: usize) -> VortexResult> { +/// Returns the lower bound and whether it was truncated. +fn truncate_min(value: Scalar, max_bytes: usize) -> VortexResult> { let nullability = value.dtype().nullability(); match value.dtype() { - DType::Utf8(_) => { - Ok( - lower_bound(BufferString::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - DType::Binary(_) => { - Ok( - lower_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - _ => Ok(Some(value)), + DType::Utf8(_) => Ok(lower_bound( + BufferString::from_scalar(value)?, + max_bytes, + nullability, + )), + DType::Binary(_) => Ok(lower_bound( + ByteBuffer::from_scalar(value)?, + max_bytes, + nullability, + )), + _ => Ok(Some((value, false))), } } #[cfg(test)] @@ -418,4 +470,67 @@ mod tests { assert_eq!(roundtrip, options); Ok(()) } + + #[test] + fn bounded_min_exactly_satisfies_min_only_when_untruncated() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let requested = Min.bind(NumericalAggregateOpts::skip_nans()); + let new_acc = || { + Accumulator::try_new( + BoundedMin, + BoundedMinOptions { + max_bytes: max_bytes(4), + }, + VarBinViewArray::from_iter_str(["x"]).dtype().clone(), + ) + }; + + // Every value fits the bound, so the bound is the true extremum. + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["c", "dde"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated value that isn't the extremum doesn't affect exactness. + acc.accumulate( + &VarBinViewArray::from_iter_str(["ddddd"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated extremum makes the bound approximate. + acc.accumulate( + &VarBinViewArray::from_iter_str(["aaaaa"]).into_array(), + &mut ctx, + )?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + + // Exactness isn't serialized, so a winning partial combined from its scalar form is + // approximate even though it was exact where it was computed. + let mut source = new_acc()?; + source.accumulate( + &VarBinViewArray::from_iter_str(["a"]).into_array(), + &mut ctx, + )?; + assert_eq!( + source.can_satisfy(&requested), + AggregateFnSatisfaction::Exact + ); + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["c"]).into_array(), + &mut ctx, + )?; + acc.combine_partials(&source.partial_scalar()?)?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index ba236ab1d5f..ade0cdd0889 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -164,6 +164,21 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { } } + /// Return whether the result accumulated in `partial` can satisfy `requested`. + /// + /// Refines [`Self::can_satisfy`], which only sees the options, with what was actually + /// accumulated: e.g. a byte-bounded maximum whose inputs all fit within the bound exactly + /// satisfies the unbounded maximum. Must never report a stronger satisfaction than is true for + /// `partial`. The default implementation ignores `partial`. + fn partial_can_satisfy( + &self, + options: &Self::Options, + _partial: &Self::Partial, + requested: &AggregateFnRef, + ) -> AggregateFnSatisfaction { + self.can_satisfy(options, requested) + } + /// The return [`DType`] of the aggregate. /// /// Returns `None` if the aggregate function cannot be applied to the input dtype.