From 0683b359edec3ed6401258ae652f5d7fa25a2afe Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 10:11:14 +0000 Subject: [PATCH 1/9] feat(aggregate): add ExtensionDelegate to lift aggregates over extension storage Add a StorageLifter trait and an ExtensionDelegate adapter that implements AggregateFnVTable for an extension dtype by running the lifter's storage aggregate over the extension's storage array and lifting the final result back. Route the existing hand-written storage delegations through it, with one lifter per aggregate: min_max, is_constant, is_sorted, uncompressed_size_in_bytes and the zoned bloom filter. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- vortex-array/src/aggregate_fn/extension.rs | 314 ++++++++++++++++++ .../aggregate_fn/fns/is_constant/extension.rs | 37 ++- .../aggregate_fn/fns/is_sorted/extension.rs | 51 ++- .../src/aggregate_fn/fns/min_max/extension.rs | 159 ++++++++- .../uncompressed_size_in_bytes/extension.rs | 38 ++- vortex-array/src/aggregate_fn/mod.rs | 3 + .../bloom_filter/canonical/extension.rs | 31 +- .../aggregates/bloom_filter/canonical/mod.rs | 4 +- .../zoned/aggregates/bloom_filter/mod.rs | 4 +- 9 files changed, 609 insertions(+), 32 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/extension.rs diff --git a/vortex-array/src/aggregate_fn/extension.rs b/vortex-array/src/aggregate_fn/extension.rs new file mode 100644 index 00000000000..562c230aad2 --- /dev/null +++ b/vortex-array/src/aggregate_fn/extension.rs @@ -0,0 +1,314 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Generic adapter for aggregates over an extension dtype that are computed by running an +//! aggregate over the extension's storage dtype. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::Accumulator; +use crate::aggregate_fn::AggregateArgs; +use crate::aggregate_fn::AggregateDTypes; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynAccumulator; +use crate::arrays::ConstantArray; +use crate::arrays::Extension; +use crate::arrays::ExtensionArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::dtype::DType; +use crate::dtype::extension::ExtDTypeRef; +use crate::scalar::Scalar; + +type StorageOptions = <::Storage as AggregateFnVTable>::Options; + +/// Declare an aggregate over an extension dtype in terms of an aggregate over its storage dtype. +/// +/// Wrap an implementation in [`ExtensionDelegate`] to obtain an [`AggregateFnVTable`]. Every batch +/// is unwrapped to its storage array and accumulated by [`Self::Storage`], so storage-encoding +/// kernels registered for the storage aggregate still fire. Partial states are the storage +/// aggregate's partials; only the final result is lifted back with [`Self::lift_scalar`] and +/// [`Self::lift_array`]. +/// +/// The defaults return the storage result unchanged, which suits aggregates whose result does not +/// depend on the input dtype, e.g. `is_constant` or `is_sorted`. +pub trait StorageLifter: 'static + Send + Sync + Clone { + /// The aggregate computed over the storage dtype. + type Storage: AggregateFnVTable; + + /// Stable identifier of the lifted aggregate. + /// + /// It must differ from the storage aggregate's id: the lifted partials and results have + /// different dtypes, so kernels and cached statistics of one cannot stand in for the other. + fn id(&self) -> AggregateFnId; + + /// Construct the storage aggregate vtable. + fn storage(&self) -> Self::Storage; + + /// Compute the final storage result over a whole storage array. + /// + /// Used by [`ExtensionDelegate::aggregate`]. Override it to reuse a storage entry point that + /// consults cached statistics, e.g. `is_constant`. + fn aggregate_storage( + &self, + options: &StorageOptions, + storage: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let mut acc = Accumulator::try_new(self.storage(), options.clone(), storage.dtype().clone())?; + acc.accumulate(storage, ctx)?; + acc.finish() + } + + /// The return dtype of the lifted aggregate, given the storage aggregate's return dtype. + fn lift_return_dtype(&self, ext_dtype: &ExtDTypeRef, storage_dtype: DType) -> Option { + let _ = ext_dtype; + Some(storage_dtype) + } + + /// Lift a storage result scalar into a result of the lifted return dtype. + fn lift_scalar(&self, ext_dtype: &ExtDTypeRef, storage_result: Scalar) -> VortexResult { + let _ = ext_dtype; + Ok(storage_result) + } + + /// Lift an array of storage results into an array of the lifted return dtype. + fn lift_array( + &self, + ext_dtype: &ExtDTypeRef, + storage_results: ArrayRef, + ) -> VortexResult { + let _ = ext_dtype; + Ok(storage_results) + } +} + +/// Adapter that exposes any [`StorageLifter`] as an [`AggregateFnVTable`] over extension dtypes. +#[derive(Clone, Debug)] +pub struct ExtensionDelegate(pub L); + +impl ExtensionDelegate { + /// Construct a new extension delegate vtable. + pub fn new(lifter: L) -> Self { + Self(lifter) + } + + /// Aggregate an extension `array` into the lifted final result. + /// + /// This is the entry point for aggregates that handle extension inputs inside their own + /// `accumulate` by delegating to the storage array. + pub fn aggregate( + &self, + options: &StorageOptions, + array: &ExtensionArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let storage_result = self + .0 + .aggregate_storage(options, array.storage_array(), ctx)?; + self.0.lift_scalar(array.ext_dtype(), storage_result) + } + + /// Accumulate an extension `array` straight into a partial state of the storage aggregate. + /// + /// This suits aggregates whose partial state does not depend on the input dtype, e.g. a Bloom + /// filter, which can then fold extension batches into the partial they already hold. + pub fn accumulate_storage( + &self, + options: &StorageOptions, + partial: &mut ::Partial, + array: &ExtensionArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let storage = self.0.storage(); + let dtypes = + AggregateDTypes::try_new(&storage, options, array.storage_array().dtype().clone())?; + let args = dtypes.args(options); + if storage.try_accumulate(args, partial, array.storage_array(), ctx)? { + return Ok(()); + } + let columnar = array.storage_array().clone().execute::(ctx)?; + storage.accumulate(args, partial, &columnar, ctx) + } + + /// Derive the storage aggregate's dtypes from the lifted aggregate's. + /// + /// The partial dtypes are shared, so only the storage return dtype is left to resolve. + fn storage_dtypes<'a>( + &self, + args: AggregateArgs<'a, StorageOptions>, + ) -> VortexResult<(&'a ExtDTypeRef, AggregateDTypes)> { + let ext_dtype = ext_dtype(args.dtype)?; + let storage = self.0.storage(); + let storage_dtype = ext_dtype.storage_dtype().clone(); + let return_dtype = storage + .return_dtype(args.options, &storage_dtype) + .ok_or_else(|| { + vortex_err!( + "Aggregate function {} cannot be applied to dtype {}", + storage.id(), + storage_dtype + ) + })?; + Ok(( + ext_dtype, + AggregateDTypes::new(storage_dtype, return_dtype, args.partial_dtype.clone()), + )) + } +} + +fn ext_dtype(dtype: &DType) -> VortexResult<&ExtDTypeRef> { + dtype + .as_extension_opt() + .ok_or_else(|| vortex_err!("Extension delegate expects an extension dtype, got {dtype}")) +} + +impl AggregateFnVTable for ExtensionDelegate { + type Options = StorageOptions; + // The storage aggregate is held as a fully-fledged `Accumulator` so that storage batches + // consult the kernel registry and the cached statistics of the storage aggregate. + type Partial = Accumulator; + + fn id(&self) -> AggregateFnId { + self.0.id() + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + self.0.storage().serialize(options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + self.0.storage().deserialize(metadata, session) + } + + fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + let ext_dtype = input_dtype.as_extension_opt()?; + let storage_return = self + .0 + .storage() + .return_dtype(options, ext_dtype.storage_dtype())?; + self.0.lift_return_dtype(ext_dtype, storage_return) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + let ext_dtype = input_dtype.as_extension_opt()?; + self.0 + .storage() + .partial_dtype(options, ext_dtype.storage_dtype()) + } + + fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult { + let (_, dtypes) = self.storage_dtypes(args)?; + Ok(Accumulator::from_dtypes( + self.0.storage(), + args.options.clone(), + dtypes, + )) + } + + fn partial_from_scalar( + &self, + args: AggregateArgs<'_, Self::Options>, + scalar: &Scalar, + ) -> VortexResult { + let mut partial = self.empty_partial(args)?; + partial.combine_partials(scalar)?; + Ok(partial) + } + + fn merge_partials( + &self, + _args: AggregateArgs<'_, Self::Options>, + mut first: Self::Partial, + mut second: Self::Partial, + ) -> VortexResult { + first.merge_from(&mut second)?; + Ok(first) + } + + fn to_scalar( + &self, + _args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> VortexResult { + partial.partial_scalar() + } + + fn is_saturated( + &self, + _args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> bool { + partial.is_saturated() + } + + /// Unwrap extension arrays to their storage before decoding, so the storage accumulator sees + /// the storage encoding. + fn try_accumulate( + &self, + _args: AggregateArgs<'_, Self::Options>, + state: &mut Self::Partial, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(array) = batch.as_opt::() else { + return Ok(false); + }; + state.accumulate(array.storage_array(), ctx)?; + Ok(true) + } + + fn accumulate( + &self, + _args: AggregateArgs<'_, Self::Options>, + state: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match batch { + Columnar::Canonical(Canonical::Extension(array)) => { + state.accumulate(array.storage_array(), ctx) + } + Columnar::Constant(constant) => { + let storage = constant.scalar().as_extension().to_storage_scalar(); + state.accumulate(&ConstantArray::new(storage, constant.len()).into_array(), ctx) + } + Columnar::Canonical(_) => { + vortex_bail!( + "Extension delegate expects an extension array, got {}", + batch.dtype() + ) + } + } + } + + fn finalize( + &self, + args: AggregateArgs<'_, Self::Options>, + states: ArrayRef, + ) -> VortexResult { + let (ext_dtype, dtypes) = self.storage_dtypes(args)?; + let storage_results = self + .0 + .storage() + .finalize(dtypes.args(args.options), states)?; + self.0.lift_array(ext_dtype, storage_results) + } + + fn finalize_scalar( + &self, + args: AggregateArgs<'_, Self::Options>, + partial: &Self::Partial, + ) -> VortexResult { + self.0 + .lift_scalar(ext_dtype(args.dtype)?, partial.final_scalar()?) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs b/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs index fa9393ede64..f834717cf45 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs @@ -2,16 +2,49 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use super::IsConstant; use super::is_constant; +use crate::ArrayRef; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::ExtensionDelegate; +use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; +use crate::scalar::Scalar; + +/// Lifts [`IsConstant`] over an extension's storage; constancy does not depend on the extension. +#[derive(Clone, Debug)] +pub(super) struct IsConstantLifter; + +impl StorageLifter for IsConstantLifter { + type Storage = IsConstant; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.is_constant.extension"); + *ID + } + + fn storage(&self) -> Self::Storage { + IsConstant + } + + fn aggregate_storage( + &self, + _options: &EmptyOptions, + storage: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(is_constant(storage, ctx)?.into()) + } +} /// Check if an extension array is constant by delegating to its storage array. pub(super) fn check_extension_constant( e: &ExtensionArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - is_constant(e.storage_array(), ctx) + bool::try_from(&ExtensionDelegate::new(IsConstantLifter).aggregate(&EmptyOptions, e, ctx)?) } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs index dc69b7c8bf9..28015d7702e 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs @@ -2,21 +2,56 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use super::IsSorted; +use super::IsSortedOptions; +use super::is_sorted; +use super::is_strict_sorted; +use crate::ArrayRef; use crate::ExecutionCtx; -use crate::aggregate_fn::fns::is_sorted::is_sorted; -use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::ExtensionDelegate; +use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; +use crate::scalar::Scalar; + +/// Lifts [`IsSorted`] over an extension's storage, which defines the extension's order. +#[derive(Clone, Debug)] +pub(super) struct IsSortedLifter; + +impl StorageLifter for IsSortedLifter { + type Storage = IsSorted; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.is_sorted.extension"); + *ID + } + + fn storage(&self) -> Self::Storage { + IsSorted + } + + fn aggregate_storage( + &self, + options: &IsSortedOptions, + storage: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let sorted = if options.strict { + is_strict_sorted(storage, ctx)? + } else { + is_sorted(storage, ctx)? + }; + Ok(sorted.into()) + } +} pub(super) fn check_extension_sorted( array: &ExtensionArray, strict: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { - if strict { - is_strict_sorted(array.storage_array(), ctx) - } else { - is_sorted(array.storage_array(), ctx) - } + let options = IsSortedOptions { strict }; + bool::try_from(&ExtensionDelegate::new(IsSortedLifter).aggregate(&options, array, ctx)?) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs index 97f6b29984b..85a1cc61b1d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs @@ -2,34 +2,171 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use super::MinMax; use super::MinMaxPartial; use super::MinMaxResult; +use super::NAMES; +use super::make_minmax_dtype; use super::min_max; +use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateArgs; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::ExtensionDelegate; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::struct_::StructArrayExt; +use crate::dtype::DType; use crate::dtype::Nullability; +use crate::dtype::extension::ExtDTypeRef; +use crate::legacy_session; use crate::scalar::Scalar; +/// Lifts [`MinMax`] over an extension's storage into extension-typed extrema. +#[derive(Clone, Debug)] +pub(super) struct MinMaxLifter; + +impl StorageLifter for MinMaxLifter { + type Storage = MinMax; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.min_max.extension"); + *ID + } + + fn storage(&self) -> Self::Storage { + MinMax + } + + fn aggregate_storage( + &self, + options: &NumericalAggregateOpts, + storage: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let dtype = make_minmax_dtype(storage.dtype()); + Ok(match min_max(storage, ctx, *options)? { + Some(MinMaxResult { min, max }) => Scalar::struct_(dtype, vec![min, max]), + None => Scalar::null(dtype), + }) + } + + fn lift_return_dtype(&self, ext_dtype: &ExtDTypeRef, _storage_dtype: DType) -> Option { + Some(make_minmax_dtype(&DType::Extension(ext_dtype.clone()))) + } + + fn lift_scalar(&self, ext_dtype: &ExtDTypeRef, storage_result: Scalar) -> VortexResult { + let return_dtype = make_minmax_dtype(&DType::Extension(ext_dtype.clone())); + let Some(MinMaxResult { min, max }) = MinMaxResult::from_scalar(&storage_result)? else { + return Ok(Scalar::null(return_dtype)); + }; + let non_nullable_ext_dtype = ext_dtype.with_nullability(Nullability::NonNullable); + Ok(Scalar::struct_( + return_dtype, + vec![ + Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), + Scalar::extension_ref(non_nullable_ext_dtype, max), + ], + )) + } + + fn lift_array( + &self, + ext_dtype: &ExtDTypeRef, + storage_results: ArrayRef, + ) -> VortexResult { + let storage_results = match storage_results.as_opt::() { + Some(results) => results.into_owned(), + None => storage_results + .execute::(&mut legacy_session().create_execution_ctx())?, + }; + let non_nullable_ext_dtype = ext_dtype.with_nullability(Nullability::NonNullable); + let fields = storage_results.iter_unmasked_fields().map(|field| { + ExtensionArray::new(non_nullable_ext_dtype.clone(), field.clone()).into_array() + }); + Ok(StructArray::try_new( + NAMES.clone(), + fields, + storage_results.len(), + storage_results.struct_validity(), + )? + .into_array()) + } +} + pub(super) fn accumulate_extension( args: AggregateArgs<'_, NumericalAggregateOpts>, partial: &mut MinMaxPartial, array: &ExtensionArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let non_nullable_ext_dtype = array.ext_dtype().with_nullability(Nullability::NonNullable); - let local = min_max( - array.storage_array(), + let result = ExtensionDelegate::new(MinMaxLifter).aggregate( + &NumericalAggregateOpts::default(), + array, ctx, - NumericalAggregateOpts::default(), - )? - .map(|MinMaxResult { min, max }| MinMaxResult { - min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), - max: Scalar::extension_ref(non_nullable_ext_dtype, max), - }); - partial.merge(args, local); + )?; + partial.merge(args, MinMaxResult::from_scalar(&result)?); Ok(()) } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::MinMaxLifter; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::ExtensionDelegate; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::min_max::MinMaxResult; + use crate::arrays::ExtensionArray; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; + + #[test] + fn delegate_lifts_storage_extrema() -> VortexResult<()> { + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let array = ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_option_iter([Some(3_000i64), None, Some(1_000)]).into_array(), + ) + .into_array(); + + let mut ctx = crate::legacy_session().create_execution_ctx(); + let mut acc = Accumulator::try_new( + ExtensionDelegate::new(MinMaxLifter), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype.clone()), + )?; + acc.accumulate(&array, &mut ctx)?; + + let non_nullable = ext_dtype.with_nullability(Nullability::NonNullable); + let extremum = |value: i64| { + Scalar::extension_ref( + non_nullable.clone(), + Scalar::primitive(value, Nullability::NonNullable), + ) + }; + assert_eq!( + MinMaxResult::from_scalar(&acc.finish()?)?, + Some(MinMaxResult { + min: extremum(1_000), + max: extremum(3_000), + }) + ); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs index d2a7a6b5f3a..156885310d6 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs @@ -2,17 +2,51 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use super::UncompressedSizeInBytes; use super::UncompressedSizeOpts; use super::uncompressed_size_in_bytes_u64; +use crate::ArrayRef; use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::ExtensionDelegate; +use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; +use crate::scalar::Scalar; + +/// Lifts [`UncompressedSizeInBytes`] over an extension's storage, which holds all its buffers. +#[derive(Clone, Debug)] +pub(super) struct UncompressedSizeInBytesLifter; + +impl StorageLifter for UncompressedSizeInBytesLifter { + type Storage = UncompressedSizeInBytes; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.uncompressed_size_in_bytes.extension"); + *ID + } + + fn storage(&self) -> Self::Storage { + UncompressedSizeInBytes + } + + fn aggregate_storage( + &self, + options: &UncompressedSizeOpts, + storage: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(uncompressed_size_in_bytes_u64(storage, *options, ctx)?.into()) + } +} pub(super) fn extension_uncompressed_size_in_bytes( array: &ExtensionArray, opts: UncompressedSizeOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { - uncompressed_size_in_bytes_u64(array.storage_array(), opts, ctx) + u64::try_from( + &ExtensionDelegate::new(UncompressedSizeInBytesLifter).aggregate(&opts, array, ctx)?, + ) } diff --git a/vortex-array/src/aggregate_fn/mod.rs b/vortex-array/src/aggregate_fn/mod.rs index 820001d3a58..64a739216d8 100644 --- a/vortex-array/src/aggregate_fn/mod.rs +++ b/vortex-array/src/aggregate_fn/mod.rs @@ -32,6 +32,9 @@ pub use typed::*; mod erased; pub use erased::*; +mod extension; +pub use extension::*; + mod options; pub use options::*; diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs index dc41c5e2e91..ef3e7ca1ce1 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs @@ -1,23 +1,42 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::Canonical; use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::ExtensionDelegate; +use vortex_array::aggregate_fn::StorageLifter; use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use super::BloomOptions; use super::BloomPartial; +use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; + +/// Lifts [`BloomFilter`] over an extension's storage, so extension values hash as their storage. +#[derive(Clone, Debug)] +struct BloomFilterLifter; + +impl StorageLifter for BloomFilterLifter { + type Storage = BloomFilter; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.bloom_filter.sbbf.extension"); + *ID + } + + fn storage(&self) -> Self::Storage { + BloomFilter + } +} pub(super) fn accumulate_extension( + options: &BloomOptions, array: &ExtensionArray, partial: &mut BloomPartial, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let storage = array.storage_array().clone(); - let canonical = storage.execute::(ctx)?; - - super::accumulate_canonical(&canonical, partial, ctx) + ExtensionDelegate::new(BloomFilterLifter).accumulate_storage(options, partial, array, ctx) } #[cfg(test)] diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs index cdc900bd339..253b376ae5b 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs @@ -16,9 +16,11 @@ use extension::accumulate_extension; use primitive::accumulate_primitive; use varbin::accumulate_varbin; +use super::BloomOptions; use super::BloomPartial; pub(super) fn accumulate_canonical( + options: &BloomOptions, canonical: &Canonical, partial: &mut BloomPartial, ctx: &mut ExecutionCtx, @@ -27,7 +29,7 @@ pub(super) fn accumulate_canonical( Canonical::Bool(array) => accumulate_bool(array, partial, ctx)?, Canonical::Primitive(array) => accumulate_primitive(array, partial, ctx)?, Canonical::VarBinView(array) => accumulate_varbin(array, partial, ctx)?, - Canonical::Extension(array) => accumulate_extension(array, partial, ctx)?, + Canonical::Extension(array) => accumulate_extension(options, array, partial, ctx)?, // Nulls are skipped and are not included in any Bloom filter. Canonical::Null(_) => {} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 3986f092ac6..82d6026c239 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -368,7 +368,7 @@ impl AggregateFnVTable for BloomFilter { fn accumulate( &self, - _args: AggregateArgs<'_, Self::Options>, + args: AggregateArgs<'_, Self::Options>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -376,7 +376,7 @@ impl AggregateFnVTable for BloomFilter { match batch { Columnar::Constant(constant) => constant::accumulate_constant(constant, partial)?, Columnar::Canonical(canonical) => { - canonical::accumulate_canonical(canonical, partial, ctx)? + canonical::accumulate_canonical(args.options, canonical, partial, ctx)? } } Ok(()) From 241723d7819fe675ff65f612a5df7478e656250a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 15:36:22 +0000 Subject: [PATCH 2/9] refactor(aggregate): replace storage lifters with ExtensionAggregate Replace the StorageLifter trait and ExtensionDelegate with a single ExtensionAggregate adapter, which defines aggregate A over extension type E by running A over E's storage. It is an aggregate of its own, with id `[]`, so supporting a new pair is one session registration: session.aggregate_fns().register(ExtensionAggregate::new(Sum, Timestamp)); Restore the original storage delegation in min_max, is_constant, is_sorted, uncompressed_size_in_bytes and the zoned bloom filter. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- vortex-array/src/aggregate_fn/extension.rs | 304 ++++++++---------- .../aggregate_fn/fns/is_constant/extension.rs | 37 +-- .../aggregate_fn/fns/is_sorted/extension.rs | 51 +-- .../src/aggregate_fn/fns/min_max/extension.rs | 159 +-------- .../uncompressed_size_in_bytes/extension.rs | 38 +-- .../bloom_filter/canonical/extension.rs | 31 +- .../aggregates/bloom_filter/canonical/mod.rs | 4 +- .../zoned/aggregates/bloom_filter/mod.rs | 4 +- 8 files changed, 171 insertions(+), 457 deletions(-) diff --git a/vortex-array/src/aggregate_fn/extension.rs b/vortex-array/src/aggregate_fn/extension.rs index 562c230aad2..29d6e0973b3 100644 --- a/vortex-array/src/aggregate_fn/extension.rs +++ b/vortex-array/src/aggregate_fn/extension.rs @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Generic adapter for aggregates over an extension dtype that are computed by running an -//! aggregate over the extension's storage dtype. +//! Generic adapter that defines an aggregate over one extension type by running an existing +//! aggregate over the extension's storage. use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -22,195 +22,105 @@ use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::arrays::ConstantArray; use crate::arrays::Extension; -use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::dtype::DType; -use crate::dtype::extension::ExtDTypeRef; +use crate::dtype::extension::ExtVTable; use crate::scalar::Scalar; -type StorageOptions = <::Storage as AggregateFnVTable>::Options; - -/// Declare an aggregate over an extension dtype in terms of an aggregate over its storage dtype. +/// The aggregate `A` applied to extension type `E` through its storage. /// -/// Wrap an implementation in [`ExtensionDelegate`] to obtain an [`AggregateFnVTable`]. Every batch -/// is unwrapped to its storage array and accumulated by [`Self::Storage`], so storage-encoding -/// kernels registered for the storage aggregate still fire. Partial states are the storage -/// aggregate's partials; only the final result is lifted back with [`Self::lift_scalar`] and -/// [`Self::lift_array`]. +/// Inputs must have an extension dtype of type `E`. Every batch is unwrapped to its storage array +/// and accumulated by `A`, so storage-encoding kernels registered for `A` still fire. Partials and +/// results are those of `A` over the storage dtype. /// -/// The defaults return the storage result unchanged, which suits aggregates whose result does not -/// depend on the input dtype, e.g. `is_constant` or `is_sorted`. -pub trait StorageLifter: 'static + Send + Sync + Clone { - /// The aggregate computed over the storage dtype. - type Storage: AggregateFnVTable; - - /// Stable identifier of the lifted aggregate. - /// - /// It must differ from the storage aggregate's id: the lifted partials and results have - /// different dtypes, so kernels and cached statistics of one cannot stand in for the other. - fn id(&self) -> AggregateFnId; - - /// Construct the storage aggregate vtable. - fn storage(&self) -> Self::Storage; - - /// Compute the final storage result over a whole storage array. - /// - /// Used by [`ExtensionDelegate::aggregate`]. Override it to reuse a storage entry point that - /// consults cached statistics, e.g. `is_constant`. - fn aggregate_storage( - &self, - options: &StorageOptions, - storage: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut acc = Accumulator::try_new(self.storage(), options.clone(), storage.dtype().clone())?; - acc.accumulate(storage, ctx)?; - acc.finish() - } - - /// The return dtype of the lifted aggregate, given the storage aggregate's return dtype. - fn lift_return_dtype(&self, ext_dtype: &ExtDTypeRef, storage_dtype: DType) -> Option { - let _ = ext_dtype; - Some(storage_dtype) - } - - /// Lift a storage result scalar into a result of the lifted return dtype. - fn lift_scalar(&self, ext_dtype: &ExtDTypeRef, storage_result: Scalar) -> VortexResult { - let _ = ext_dtype; - Ok(storage_result) - } - - /// Lift an array of storage results into an array of the lifted return dtype. - fn lift_array( - &self, - ext_dtype: &ExtDTypeRef, - storage_results: ArrayRef, - ) -> VortexResult { - let _ = ext_dtype; - Ok(storage_results) - } -} - -/// Adapter that exposes any [`StorageLifter`] as an [`AggregateFnVTable`] over extension dtypes. +/// The adapter is an aggregate of its own, with id `[]`, so supporting a new pair is one +/// registration: +/// +/// ``` +/// # use vortex_array::aggregate_fn::ExtensionAggregate; +/// # use vortex_array::aggregate_fn::fns::sum::Sum; +/// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +/// # use vortex_array::extension::datetime::Timestamp; +/// # let session = vortex_array::array_session(); +/// session.aggregate_fns().register(ExtensionAggregate::new(Sum, Timestamp)); +/// ``` +/// +/// Only register a pair when `A` over the storage is meaningful for `E`, e.g. when `E` orders and +/// compares like its storage. #[derive(Clone, Debug)] -pub struct ExtensionDelegate(pub L); - -impl ExtensionDelegate { - /// Construct a new extension delegate vtable. - pub fn new(lifter: L) -> Self { - Self(lifter) - } +pub struct ExtensionAggregate { + aggregate: A, + ext: E, + id: AggregateFnId, +} - /// Aggregate an extension `array` into the lifted final result. - /// - /// This is the entry point for aggregates that handle extension inputs inside their own - /// `accumulate` by delegating to the storage array. - pub fn aggregate( - &self, - options: &StorageOptions, - array: &ExtensionArray, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let storage_result = self - .0 - .aggregate_storage(options, array.storage_array(), ctx)?; - self.0.lift_scalar(array.ext_dtype(), storage_result) +impl ExtensionAggregate { + /// Define `aggregate` over extension type `ext`. + pub fn new(aggregate: A, ext: E) -> Self { + let id = AggregateFnId::new(&format!("{}[{}]", aggregate.id(), ext.id())); + Self { aggregate, ext, id } } - /// Accumulate an extension `array` straight into a partial state of the storage aggregate. - /// - /// This suits aggregates whose partial state does not depend on the input dtype, e.g. a Bloom - /// filter, which can then fold extension batches into the partial they already hold. - pub fn accumulate_storage( - &self, - options: &StorageOptions, - partial: &mut ::Partial, - array: &ExtensionArray, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let storage = self.0.storage(); - let dtypes = - AggregateDTypes::try_new(&storage, options, array.storage_array().dtype().clone())?; - let args = dtypes.args(options); - if storage.try_accumulate(args, partial, array.storage_array(), ctx)? { - return Ok(()); - } - let columnar = array.storage_array().clone().execute::(ctx)?; - storage.accumulate(args, partial, &columnar, ctx) + /// The storage dtype of `dtype`, if it is an extension dtype of type `E`. + fn storage_dtype<'a>(&self, dtype: &'a DType) -> Option<&'a DType> { + dtype + .as_extension_opt() + .filter(|ext_dtype| ext_dtype.is::()) + .map(|ext_dtype| ext_dtype.storage_dtype()) } - /// Derive the storage aggregate's dtypes from the lifted aggregate's. - /// - /// The partial dtypes are shared, so only the storage return dtype is left to resolve. - fn storage_dtypes<'a>( - &self, - args: AggregateArgs<'a, StorageOptions>, - ) -> VortexResult<(&'a ExtDTypeRef, AggregateDTypes)> { - let ext_dtype = ext_dtype(args.dtype)?; - let storage = self.0.storage(); - let storage_dtype = ext_dtype.storage_dtype().clone(); - let return_dtype = storage - .return_dtype(args.options, &storage_dtype) - .ok_or_else(|| { - vortex_err!( - "Aggregate function {} cannot be applied to dtype {}", - storage.id(), - storage_dtype - ) - })?; - Ok(( - ext_dtype, - AggregateDTypes::new(storage_dtype, return_dtype, args.partial_dtype.clone()), + /// Derive the storage aggregate's dtypes from this aggregate's, which share the partial and + /// return dtypes. + fn storage_dtypes(&self, args: AggregateArgs<'_, A::Options>) -> VortexResult { + let storage_dtype = self.storage_dtype(args.dtype).ok_or_else(|| { + vortex_err!( + "Aggregate function {} expects extension type {}, got {}", + self.id, + self.ext.id(), + args.dtype + ) + })?; + Ok(AggregateDTypes::new( + storage_dtype.clone(), + args.return_dtype.clone(), + args.partial_dtype.clone(), )) } } -fn ext_dtype(dtype: &DType) -> VortexResult<&ExtDTypeRef> { - dtype - .as_extension_opt() - .ok_or_else(|| vortex_err!("Extension delegate expects an extension dtype, got {dtype}")) -} - -impl AggregateFnVTable for ExtensionDelegate { - type Options = StorageOptions; +impl AggregateFnVTable for ExtensionAggregate { + type Options = A::Options; // The storage aggregate is held as a fully-fledged `Accumulator` so that storage batches - // consult the kernel registry and the cached statistics of the storage aggregate. - type Partial = Accumulator; + // consult the kernel registry and the cached statistics of `A`. + type Partial = Accumulator; fn id(&self) -> AggregateFnId { - self.0.id() + self.id } fn serialize(&self, options: &Self::Options) -> VortexResult>> { - self.0.storage().serialize(options) + self.aggregate.serialize(options) } fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { - self.0.storage().deserialize(metadata, session) + self.aggregate.deserialize(metadata, session) } fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - let ext_dtype = input_dtype.as_extension_opt()?; - let storage_return = self - .0 - .storage() - .return_dtype(options, ext_dtype.storage_dtype())?; - self.0.lift_return_dtype(ext_dtype, storage_return) + self.aggregate + .return_dtype(options, self.storage_dtype(input_dtype)?) } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - let ext_dtype = input_dtype.as_extension_opt()?; - self.0 - .storage() - .partial_dtype(options, ext_dtype.storage_dtype()) + self.aggregate + .partial_dtype(options, self.storage_dtype(input_dtype)?) } fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult { - let (_, dtypes) = self.storage_dtypes(args)?; Ok(Accumulator::from_dtypes( - self.0.storage(), + self.aggregate.clone(), args.options.clone(), - dtypes, + self.storage_dtypes(args)?, )) } @@ -283,7 +193,8 @@ impl AggregateFnVTable for ExtensionDelegate { } Columnar::Canonical(_) => { vortex_bail!( - "Extension delegate expects an extension array, got {}", + "Aggregate function {} expects an extension array, got {}", + self.id, batch.dtype() ) } @@ -295,20 +206,83 @@ impl AggregateFnVTable for ExtensionDelegate { args: AggregateArgs<'_, Self::Options>, states: ArrayRef, ) -> VortexResult { - let (ext_dtype, dtypes) = self.storage_dtypes(args)?; - let storage_results = self - .0 - .storage() - .finalize(dtypes.args(args.options), states)?; - self.0.lift_array(ext_dtype, storage_results) + let dtypes = self.storage_dtypes(args)?; + self.aggregate.finalize(dtypes.args(args.options), states) } fn finalize_scalar( &self, - args: AggregateArgs<'_, Self::Options>, + _args: AggregateArgs<'_, Self::Options>, partial: &Self::Partial, ) -> VortexResult { - self.0 - .lift_scalar(ext_dtype(args.dtype)?, partial.final_scalar()?) + partial.final_scalar() + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::ExtensionAggregate; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::Accumulator; + use crate::aggregate_fn::AggregateFnVTable; + use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::session::AggregateFnSessionExt; + use crate::arrays::ExtensionArray; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::extension::datetime::Date; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + + #[test] + fn sums_extension_through_storage() -> VortexResult<()> { + let session = crate::array_session(); + let aggregate = ExtensionAggregate::new(Sum, Timestamp); + session.aggregate_fns().register(aggregate.clone()); + assert!( + session + .aggregate_fns() + .find_plugin(&aggregate.id()) + .is_some() + ); + + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let storage = + PrimitiveArray::from_option_iter([Some(3_000i64), None, Some(1_000)]).into_array(); + let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array(); + + let mut ctx = session.create_execution_ctx(); + let mut acc = Accumulator::try_new( + aggregate, + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype), + )?; + acc.accumulate(&array, &mut ctx)?; + + let mut expected = Accumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + storage.dtype().clone(), + )?; + expected.accumulate(&storage, &mut ctx)?; + + assert_eq!(acc.finish()?, expected.finish()?); + Ok(()) + } + + #[test] + fn rejects_other_extension_types() { + let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()); + assert!( + ExtensionAggregate::new(Sum, Timestamp) + .return_dtype(&NumericalAggregateOpts::default(), &date) + .is_none() + ); } } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs b/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs index f834717cf45..fa9393ede64 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/extension.rs @@ -2,49 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_session::registry::CachedId; -use super::IsConstant; use super::is_constant; -use crate::ArrayRef; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::ExtensionDelegate; -use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::scalar::Scalar; - -/// Lifts [`IsConstant`] over an extension's storage; constancy does not depend on the extension. -#[derive(Clone, Debug)] -pub(super) struct IsConstantLifter; - -impl StorageLifter for IsConstantLifter { - type Storage = IsConstant; - - fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.is_constant.extension"); - *ID - } - - fn storage(&self) -> Self::Storage { - IsConstant - } - - fn aggregate_storage( - &self, - _options: &EmptyOptions, - storage: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(is_constant(storage, ctx)?.into()) - } -} +use crate::arrays::extension::ExtensionArrayExt; /// Check if an extension array is constant by delegating to its storage array. pub(super) fn check_extension_constant( e: &ExtensionArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - bool::try_from(&ExtensionDelegate::new(IsConstantLifter).aggregate(&EmptyOptions, e, ctx)?) + is_constant(e.storage_array(), ctx) } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs index 28015d7702e..dc69b7c8bf9 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs @@ -2,56 +2,21 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_session::registry::CachedId; -use super::IsSorted; -use super::IsSortedOptions; -use super::is_sorted; -use super::is_strict_sorted; -use crate::ArrayRef; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::ExtensionDelegate; -use crate::aggregate_fn::StorageLifter; +use crate::aggregate_fn::fns::is_sorted::is_sorted; +use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; use crate::arrays::ExtensionArray; -use crate::scalar::Scalar; - -/// Lifts [`IsSorted`] over an extension's storage, which defines the extension's order. -#[derive(Clone, Debug)] -pub(super) struct IsSortedLifter; - -impl StorageLifter for IsSortedLifter { - type Storage = IsSorted; - - fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.is_sorted.extension"); - *ID - } - - fn storage(&self) -> Self::Storage { - IsSorted - } - - fn aggregate_storage( - &self, - options: &IsSortedOptions, - storage: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let sorted = if options.strict { - is_strict_sorted(storage, ctx)? - } else { - is_sorted(storage, ctx)? - }; - Ok(sorted.into()) - } -} +use crate::arrays::extension::ExtensionArrayExt; pub(super) fn check_extension_sorted( array: &ExtensionArray, strict: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { - let options = IsSortedOptions { strict }; - bool::try_from(&ExtensionDelegate::new(IsSortedLifter).aggregate(&options, array, ctx)?) + if strict { + is_strict_sorted(array.storage_array(), ctx) + } else { + is_sorted(array.storage_array(), ctx) + } } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs index 85a1cc61b1d..97f6b29984b 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs @@ -2,171 +2,34 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_session::registry::CachedId; -use super::MinMax; use super::MinMaxPartial; use super::MinMaxResult; -use super::NAMES; -use super::make_minmax_dtype; use super::min_max; -use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateArgs; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::ExtensionDelegate; use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::arrays::Struct; -use crate::arrays::StructArray; -use crate::arrays::struct_::StructArrayExt; -use crate::dtype::DType; +use crate::arrays::extension::ExtensionArrayExt; use crate::dtype::Nullability; -use crate::dtype::extension::ExtDTypeRef; -use crate::legacy_session; use crate::scalar::Scalar; -/// Lifts [`MinMax`] over an extension's storage into extension-typed extrema. -#[derive(Clone, Debug)] -pub(super) struct MinMaxLifter; - -impl StorageLifter for MinMaxLifter { - type Storage = MinMax; - - fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.min_max.extension"); - *ID - } - - fn storage(&self) -> Self::Storage { - MinMax - } - - fn aggregate_storage( - &self, - options: &NumericalAggregateOpts, - storage: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let dtype = make_minmax_dtype(storage.dtype()); - Ok(match min_max(storage, ctx, *options)? { - Some(MinMaxResult { min, max }) => Scalar::struct_(dtype, vec![min, max]), - None => Scalar::null(dtype), - }) - } - - fn lift_return_dtype(&self, ext_dtype: &ExtDTypeRef, _storage_dtype: DType) -> Option { - Some(make_minmax_dtype(&DType::Extension(ext_dtype.clone()))) - } - - fn lift_scalar(&self, ext_dtype: &ExtDTypeRef, storage_result: Scalar) -> VortexResult { - let return_dtype = make_minmax_dtype(&DType::Extension(ext_dtype.clone())); - let Some(MinMaxResult { min, max }) = MinMaxResult::from_scalar(&storage_result)? else { - return Ok(Scalar::null(return_dtype)); - }; - let non_nullable_ext_dtype = ext_dtype.with_nullability(Nullability::NonNullable); - Ok(Scalar::struct_( - return_dtype, - vec![ - Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), - Scalar::extension_ref(non_nullable_ext_dtype, max), - ], - )) - } - - fn lift_array( - &self, - ext_dtype: &ExtDTypeRef, - storage_results: ArrayRef, - ) -> VortexResult { - let storage_results = match storage_results.as_opt::() { - Some(results) => results.into_owned(), - None => storage_results - .execute::(&mut legacy_session().create_execution_ctx())?, - }; - let non_nullable_ext_dtype = ext_dtype.with_nullability(Nullability::NonNullable); - let fields = storage_results.iter_unmasked_fields().map(|field| { - ExtensionArray::new(non_nullable_ext_dtype.clone(), field.clone()).into_array() - }); - Ok(StructArray::try_new( - NAMES.clone(), - fields, - storage_results.len(), - storage_results.struct_validity(), - )? - .into_array()) - } -} - pub(super) fn accumulate_extension( args: AggregateArgs<'_, NumericalAggregateOpts>, partial: &mut MinMaxPartial, array: &ExtensionArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let result = ExtensionDelegate::new(MinMaxLifter).aggregate( - &NumericalAggregateOpts::default(), - array, + let non_nullable_ext_dtype = array.ext_dtype().with_nullability(Nullability::NonNullable); + let local = min_max( + array.storage_array(), ctx, - )?; - partial.merge(args, MinMaxResult::from_scalar(&result)?); + NumericalAggregateOpts::default(), + )? + .map(|MinMaxResult { min, max }| MinMaxResult { + min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), + max: Scalar::extension_ref(non_nullable_ext_dtype, max), + }); + partial.merge(args, local); Ok(()) } - -#[cfg(test)] -mod tests { - use vortex_error::VortexResult; - - use super::MinMaxLifter; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::Accumulator; - use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::ExtensionDelegate; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::min_max::MinMaxResult; - use crate::arrays::ExtensionArray; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::scalar::Scalar; - - #[test] - fn delegate_lifts_storage_extrema() -> VortexResult<()> { - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); - let array = ExtensionArray::new( - ext_dtype.clone(), - PrimitiveArray::from_option_iter([Some(3_000i64), None, Some(1_000)]).into_array(), - ) - .into_array(); - - let mut ctx = crate::legacy_session().create_execution_ctx(); - let mut acc = Accumulator::try_new( - ExtensionDelegate::new(MinMaxLifter), - NumericalAggregateOpts::default(), - DType::Extension(ext_dtype.clone()), - )?; - acc.accumulate(&array, &mut ctx)?; - - let non_nullable = ext_dtype.with_nullability(Nullability::NonNullable); - let extremum = |value: i64| { - Scalar::extension_ref( - non_nullable.clone(), - Scalar::primitive(value, Nullability::NonNullable), - ) - }; - assert_eq!( - MinMaxResult::from_scalar(&acc.finish()?)?, - Some(MinMaxResult { - min: extremum(1_000), - max: extremum(3_000), - }) - ); - Ok(()) - } -} diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs index 156885310d6..d2a7a6b5f3a 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs @@ -2,51 +2,17 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_session::registry::CachedId; -use super::UncompressedSizeInBytes; use super::UncompressedSizeOpts; use super::uncompressed_size_in_bytes_u64; -use crate::ArrayRef; use crate::ExecutionCtx; -use crate::aggregate_fn::AggregateFnId; -use crate::aggregate_fn::ExtensionDelegate; -use crate::aggregate_fn::StorageLifter; use crate::arrays::ExtensionArray; -use crate::scalar::Scalar; - -/// Lifts [`UncompressedSizeInBytes`] over an extension's storage, which holds all its buffers. -#[derive(Clone, Debug)] -pub(super) struct UncompressedSizeInBytesLifter; - -impl StorageLifter for UncompressedSizeInBytesLifter { - type Storage = UncompressedSizeInBytes; - - fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.uncompressed_size_in_bytes.extension"); - *ID - } - - fn storage(&self) -> Self::Storage { - UncompressedSizeInBytes - } - - fn aggregate_storage( - &self, - options: &UncompressedSizeOpts, - storage: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(uncompressed_size_in_bytes_u64(storage, *options, ctx)?.into()) - } -} +use crate::arrays::extension::ExtensionArrayExt; pub(super) fn extension_uncompressed_size_in_bytes( array: &ExtensionArray, opts: UncompressedSizeOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { - u64::try_from( - &ExtensionDelegate::new(UncompressedSizeInBytesLifter).aggregate(&opts, array, ctx)?, - ) + uncompressed_size_in_bytes_u64(array.storage_array(), opts, ctx) } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs index ef3e7ca1ce1..dc41c5e2e91 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs @@ -1,42 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::Canonical; use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::AggregateFnId; -use vortex_array::aggregate_fn::ExtensionDelegate; -use vortex_array::aggregate_fn::StorageLifter; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_error::VortexResult; -use vortex_session::registry::CachedId; -use super::BloomOptions; use super::BloomPartial; -use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; - -/// Lifts [`BloomFilter`] over an extension's storage, so extension values hash as their storage. -#[derive(Clone, Debug)] -struct BloomFilterLifter; - -impl StorageLifter for BloomFilterLifter { - type Storage = BloomFilter; - - fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.bloom_filter.sbbf.extension"); - *ID - } - - fn storage(&self) -> Self::Storage { - BloomFilter - } -} pub(super) fn accumulate_extension( - options: &BloomOptions, array: &ExtensionArray, partial: &mut BloomPartial, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - ExtensionDelegate::new(BloomFilterLifter).accumulate_storage(options, partial, array, ctx) + let storage = array.storage_array().clone(); + let canonical = storage.execute::(ctx)?; + + super::accumulate_canonical(&canonical, partial, ctx) } #[cfg(test)] diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs index 253b376ae5b..cdc900bd339 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs @@ -16,11 +16,9 @@ use extension::accumulate_extension; use primitive::accumulate_primitive; use varbin::accumulate_varbin; -use super::BloomOptions; use super::BloomPartial; pub(super) fn accumulate_canonical( - options: &BloomOptions, canonical: &Canonical, partial: &mut BloomPartial, ctx: &mut ExecutionCtx, @@ -29,7 +27,7 @@ pub(super) fn accumulate_canonical( Canonical::Bool(array) => accumulate_bool(array, partial, ctx)?, Canonical::Primitive(array) => accumulate_primitive(array, partial, ctx)?, Canonical::VarBinView(array) => accumulate_varbin(array, partial, ctx)?, - Canonical::Extension(array) => accumulate_extension(options, array, partial, ctx)?, + Canonical::Extension(array) => accumulate_extension(array, partial, ctx)?, // Nulls are skipped and are not included in any Bloom filter. Canonical::Null(_) => {} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 82d6026c239..3986f092ac6 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -368,7 +368,7 @@ impl AggregateFnVTable for BloomFilter { fn accumulate( &self, - args: AggregateArgs<'_, Self::Options>, + _args: AggregateArgs<'_, Self::Options>, partial: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, @@ -376,7 +376,7 @@ impl AggregateFnVTable for BloomFilter { match batch { Columnar::Constant(constant) => constant::accumulate_constant(constant, partial)?, Columnar::Canonical(canonical) => { - canonical::accumulate_canonical(args.options, canonical, partial, ctx)? + canonical::accumulate_canonical(canonical, partial, ctx)? } } Ok(()) From 3db29a5ffeb57e84efa1e2f827d1c2827dce7163 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 19 Oct 2026 14:17:14 +0100 Subject: [PATCH 3/9] fix(aggregate): resolve extension adapter lint failures Signed-off-by: Joe Isaacs --- vortex-array/src/aggregate_fn/extension.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/aggregate_fn/extension.rs b/vortex-array/src/aggregate_fn/extension.rs index 29d6e0973b3..0b3a76bf3c0 100644 --- a/vortex-array/src/aggregate_fn/extension.rs +++ b/vortex-array/src/aggregate_fn/extension.rs @@ -57,6 +57,7 @@ pub struct ExtensionAggregate { impl ExtensionAggregate { /// Define `aggregate` over extension type `ext`. pub fn new(aggregate: A, ext: E) -> Self { + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let id = AggregateFnId::new(&format!("{}[{}]", aggregate.id(), ext.id())); Self { aggregate, ext, id } } @@ -189,7 +190,10 @@ impl AggregateFnVTable for ExtensionAggregat } Columnar::Constant(constant) => { let storage = constant.scalar().as_extension().to_storage_scalar(); - state.accumulate(&ConstantArray::new(storage, constant.len()).into_array(), ctx) + state.accumulate( + &ConstantArray::new(storage, constant.len()).into_array(), + ctx, + ) } Columnar::Canonical(_) => { vortex_bail!( From 3ce423110a49c51eaa53c42fc44482f1aed80527 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 19 Oct 2026 17:19:12 +0100 Subject: [PATCH 4/9] perf(aggregate): hold the storage partial directly in ExtensionAggregate (#10612) Co-authored-by: Claude Signed-off-by: Claude --- vortex-array/src/aggregate_fn/accumulator.rs | 7 + vortex-array/src/aggregate_fn/extension.rs | 137 ++++++++++++++----- 2 files changed, 106 insertions(+), 38 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index d2273a24cca..bd47f6d7b5c 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -70,6 +70,13 @@ impl Accumulator { self.vtable.empty_partial(self.dtypes.args(&self.options)) } + /// Take the accumulated partial state, leaving the accumulator empty. + /// + /// Returns `None` if nothing was accumulated. + pub(crate) fn take_partial(&mut self) -> Option { + self.partial.take() + } + fn ensure_partial(&mut self) -> VortexResult<()> { if self.partial.is_none() { self.partial = Some(self.empty_partial()?); diff --git a/vortex-array/src/aggregate_fn/extension.rs b/vortex-array/src/aggregate_fn/extension.rs index 0b3a76bf3c0..1f0b311fbee 100644 --- a/vortex-array/src/aggregate_fn/extension.rs +++ b/vortex-array/src/aggregate_fn/extension.rs @@ -4,6 +4,8 @@ //! Generic adapter that defines an aggregate over one extension type by running an existing //! aggregate over the extension's storage. +use std::mem; + use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -70,10 +72,15 @@ impl ExtensionAggregate { .map(|ext_dtype| ext_dtype.storage_dtype()) } - /// Derive the storage aggregate's dtypes from this aggregate's, which share the partial and - /// return dtypes. - fn storage_dtypes(&self, args: AggregateArgs<'_, A::Options>) -> VortexResult { - let storage_dtype = self.storage_dtype(args.dtype).ok_or_else(|| { + /// The arguments of the storage aggregate `A`. + /// + /// `A` shares this aggregate's options, return and partial dtypes, and its input dtype is + /// borrowed from the extension dtype, so this allocates nothing. + fn storage_args<'a>( + &self, + args: AggregateArgs<'a, A::Options>, + ) -> VortexResult> { + let dtype = self.storage_dtype(args.dtype).ok_or_else(|| { vortex_err!( "Aggregate function {} expects extension type {}, got {}", self.id, @@ -81,19 +88,40 @@ impl ExtensionAggregate { args.dtype ) })?; - Ok(AggregateDTypes::new( - storage_dtype.clone(), - args.return_dtype.clone(), - args.partial_dtype.clone(), - )) + Ok(AggregateArgs { dtype, ..args }) + } + + /// Accumulate a storage batch into `state` through an [`Accumulator`] of `A`, so that the + /// kernels and cached statistics of `A` on the storage encoding still apply. + fn accumulate_storage( + &self, + args: AggregateArgs<'_, A::Options>, + state: &mut A::Partial, + storage: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let mut accumulator = Accumulator::from_dtypes( + self.aggregate.clone(), + args.options.clone(), + AggregateDTypes::new( + args.dtype.clone(), + args.return_dtype.clone(), + args.partial_dtype.clone(), + ), + ); + accumulator.accumulate(storage, ctx)?; + let Some(batch) = accumulator.take_partial() else { + return Ok(()); + }; + let running = mem::replace(state, self.aggregate.empty_partial(args)?); + *state = self.aggregate.merge_partials(args, running, batch)?; + Ok(()) } } impl AggregateFnVTable for ExtensionAggregate { type Options = A::Options; - // The storage aggregate is held as a fully-fledged `Accumulator` so that storage batches - // consult the kernel registry and the cached statistics of `A`. - type Partial = Accumulator; + type Partial = A::Partial; fn id(&self) -> AggregateFnId { self.id @@ -118,11 +146,7 @@ impl AggregateFnVTable for ExtensionAggregat } fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult { - Ok(Accumulator::from_dtypes( - self.aggregate.clone(), - args.options.clone(), - self.storage_dtypes(args)?, - )) + self.aggregate.empty_partial(self.storage_args(args)?) } fn partial_from_scalar( @@ -130,42 +154,42 @@ impl AggregateFnVTable for ExtensionAggregat args: AggregateArgs<'_, Self::Options>, scalar: &Scalar, ) -> VortexResult { - let mut partial = self.empty_partial(args)?; - partial.combine_partials(scalar)?; - Ok(partial) + self.aggregate + .partial_from_scalar(self.storage_args(args)?, scalar) } fn merge_partials( &self, - _args: AggregateArgs<'_, Self::Options>, - mut first: Self::Partial, - mut second: Self::Partial, + args: AggregateArgs<'_, Self::Options>, + first: Self::Partial, + second: Self::Partial, ) -> VortexResult { - first.merge_from(&mut second)?; - Ok(first) + self.aggregate + .merge_partials(self.storage_args(args)?, first, second) } fn to_scalar( &self, - _args: AggregateArgs<'_, Self::Options>, + args: AggregateArgs<'_, Self::Options>, partial: &Self::Partial, ) -> VortexResult { - partial.partial_scalar() + self.aggregate.to_scalar(self.storage_args(args)?, partial) } fn is_saturated( &self, - _args: AggregateArgs<'_, Self::Options>, + args: AggregateArgs<'_, Self::Options>, partial: &Self::Partial, ) -> bool { - partial.is_saturated() + self.storage_args(args) + .is_ok_and(|args| self.aggregate.is_saturated(args, partial)) } /// Unwrap extension arrays to their storage before decoding, so the storage accumulator sees /// the storage encoding. fn try_accumulate( &self, - _args: AggregateArgs<'_, Self::Options>, + args: AggregateArgs<'_, Self::Options>, state: &mut Self::Partial, batch: &ArrayRef, ctx: &mut ExecutionCtx, @@ -173,24 +197,27 @@ impl AggregateFnVTable for ExtensionAggregat let Some(array) = batch.as_opt::() else { return Ok(false); }; - state.accumulate(array.storage_array(), ctx)?; + self.accumulate_storage(self.storage_args(args)?, state, array.storage_array(), ctx)?; Ok(true) } fn accumulate( &self, - _args: AggregateArgs<'_, Self::Options>, + args: AggregateArgs<'_, Self::Options>, state: &mut Self::Partial, batch: &Columnar, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + let args = self.storage_args(args)?; match batch { Columnar::Canonical(Canonical::Extension(array)) => { - state.accumulate(array.storage_array(), ctx) + self.accumulate_storage(args, state, array.storage_array(), ctx) } Columnar::Constant(constant) => { let storage = constant.scalar().as_extension().to_storage_scalar(); - state.accumulate( + self.accumulate_storage( + args, + state, &ConstantArray::new(storage, constant.len()).into_array(), ctx, ) @@ -210,16 +237,16 @@ impl AggregateFnVTable for ExtensionAggregat args: AggregateArgs<'_, Self::Options>, states: ArrayRef, ) -> VortexResult { - let dtypes = self.storage_dtypes(args)?; - self.aggregate.finalize(dtypes.args(args.options), states) + self.aggregate.finalize(self.storage_args(args)?, states) } fn finalize_scalar( &self, - _args: AggregateArgs<'_, Self::Options>, + args: AggregateArgs<'_, Self::Options>, partial: &Self::Partial, ) -> VortexResult { - partial.final_scalar() + self.aggregate + .finalize_scalar(self.storage_args(args)?, partial) } } @@ -280,6 +307,40 @@ mod tests { Ok(()) } + #[test] + fn combines_flushed_partials() -> VortexResult<()> { + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let batch = |values: [Option; 2]| { + ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_option_iter(values).into_array(), + ) + .into_array() + }; + let accumulator = || { + Accumulator::try_new( + ExtensionAggregate::new(Sum, Timestamp), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype.clone()), + ) + }; + + let mut first = accumulator()?; + first.accumulate(&batch([Some(3_000), None]), &mut ctx)?; + let mut second = accumulator()?; + second.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?; + first.combine_partials(&second.flush()?)?; + + let mut expected = accumulator()?; + expected.accumulate(&batch([Some(3_000), None]), &mut ctx)?; + expected.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?; + + assert_eq!(first.finish()?, expected.finish()?); + Ok(()) + } + #[test] fn rejects_other_extension_types() { let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()); From 9cdad1c19afd32ebf35313296a63924164a45d83 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 16:27:52 +0000 Subject: [PATCH 5/9] docs(aggregate): explain when ExtensionAggregate is correct for each aggregate Document, per aggregate, the property the storage must share with the extension type before a pair is registered, with examples that fit and that do not, and add a doctest that registers and runs ExtensionAggregate. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- vortex-array/src/aggregate_fn/extension.rs | 76 +++++++++++++++++++--- 1 file changed, 67 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/aggregate_fn/extension.rs b/vortex-array/src/aggregate_fn/extension.rs index 1f0b311fbee..e2d163ab54a 100644 --- a/vortex-array/src/aggregate_fn/extension.rs +++ b/vortex-array/src/aggregate_fn/extension.rs @@ -33,22 +33,80 @@ use crate::scalar::Scalar; /// /// Inputs must have an extension dtype of type `E`. Every batch is unwrapped to its storage array /// and accumulated by `A`, so storage-encoding kernels registered for `A` still fire. Partials and -/// results are those of `A` over the storage dtype. +/// results are those of `A` over the storage dtype: they are not converted back to `E`, so e.g. +/// the extrema of a timestamp column come back as its raw `i64` storage values. /// -/// The adapter is an aggregate of its own, with id `[]`, so supporting a new pair is one -/// registration: +/// The adapter is an aggregate of its own, with id `[]` (e.g. +/// `vortex.is_sorted[vortex.timestamp]`). Its options are `A`'s options, and it serializes +/// through `A`, so a registered pair round-trips by id. +/// +/// # When to use it +/// +/// Running `A` over the storage is only correct when the storage has the semantics `A` relies on +/// for `E`. The adapter cannot check this, so only register a pair after checking the property +/// for that aggregate: +/// +/// | Aggregates | `E` must... | Fits | Does not fit | +/// | --- | --- | --- | --- | +/// | `min`, `max`, `min_max`, `bounded_min`, `bounded_max`, `is_sorted` | order exactly like its storage | `Date`, `Time`, `Timestamp` over integers | a case-insensitive string type, a little-endian integer stored as bytes | +/// | `is_constant`, `all_non_distinct` | have one storage value per value, so equal values have equal storage | `Date`, `Time`, `Timestamp`, `Uuid` | JSON text (`{"a":1}` and `{ "a": 1 }` are equal), a case-insensitive string type | +/// | `sum`, `mean` | have values for which adding the storage is meaningful | a duration stored as an integer count | `Timestamp` (the sum of two instants is not an instant), `Uuid` | +/// | `nan_count`, `all_nan`, `all_non_nan` | treat a NaN storage value as NaN | a float-backed measurement type | a type that stores a sentinel in its float storage | +/// | `count`, `null_count`, `all_null`, `all_non_null`, `uncompressed_size_in_bytes` | (always holds: these depend only on validity or storage bytes) | any `E` | | +/// +/// `A` must also accept `E`'s storage dtype. If it does not, the pair still registers, but binding +/// it to a column of `E` fails: `IsSorted` rejects the fixed-size-list storage of `Uuid`, for +/// example. +/// +/// Partials depend on `E`'s metadata as well as `E`. For example, timestamps in milliseconds and +/// in seconds share the `Timestamp` type but not the scale of their storage. Only merge partials +/// accumulated over the same extension dtype. +/// +/// # Example +/// +/// `Timestamp` orders like its `i64` storage, so `is_sorted` and `min_max` can run on the +/// storage. Register the pair once on the session, then bind it like any other aggregate: /// /// ``` +/// # use vortex_array::IntoArray; +/// # use vortex_array::VortexSessionExecute; +/// # use vortex_array::aggregate_fn::Accumulator; +/// # use vortex_array::aggregate_fn::DynAccumulator; /// # use vortex_array::aggregate_fn::ExtensionAggregate; -/// # use vortex_array::aggregate_fn::fns::sum::Sum; +/// # use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; +/// # use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; /// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +/// # use vortex_array::arrays::ExtensionArray; +/// # use vortex_array::arrays::PrimitiveArray; +/// # use vortex_array::dtype::DType; +/// # use vortex_array::dtype::Nullability; +/// # use vortex_array::extension::datetime::TimeUnit; /// # use vortex_array::extension::datetime::Timestamp; -/// # let session = vortex_array::array_session(); -/// session.aggregate_fns().register(ExtensionAggregate::new(Sum, Timestamp)); -/// ``` +/// # fn main() -> vortex_error::VortexResult<()> { +/// let session = vortex_array::array_session(); +/// +/// // Register `vortex.is_sorted[vortex.timestamp]`. +/// let is_sorted = ExtensionAggregate::new(IsSorted, Timestamp); +/// session.aggregate_fns().register(is_sorted.clone()); /// -/// Only register a pair when `A` over the storage is meaningful for `E`, e.g. when `E` orders and -/// compares like its storage. +/// // Aggregate a timestamp column through its storage. +/// let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); +/// let column = ExtensionArray::new( +/// ext_dtype.clone(), +/// PrimitiveArray::from_iter([1_000i64, 2_000, 3_000]).into_array(), +/// ) +/// .into_array(); +/// +/// let mut acc = Accumulator::try_new( +/// is_sorted, +/// IsSortedOptions { strict: true }, +/// DType::Extension(ext_dtype), +/// )?; +/// acc.accumulate(&column, &mut session.create_execution_ctx())?; +/// assert!(bool::try_from(&acc.finish()?)?); +/// # Ok(()) +/// # } +/// ``` #[derive(Clone, Debug)] pub struct ExtensionAggregate { aggregate: A, From 9a0a02fdd4c91549f020e9d5fa8f689627218137 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 16:36:32 +0000 Subject: [PATCH 6/9] refactor(aggregate): rename ExtensionAggregate to StoragePushdown The adapter is an opt-in pushdown of one aggregate into one extension type's storage, not an aggregate over any extension type, so name it for that and move it to aggregate_fn/storage_pushdown.rs. Expand its docs: registering a pair is an unchecked claim about the extension type; a pushdown is correct only when every operation the aggregate uses (comparison, equality, addition, NaN detection) derives from the same operation on the storage; Timestamp supports it for comparison but not addition. Switch the tests and examples from Sum to MinMax, which is correct for Timestamp. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- vortex-array/src/aggregate_fn/mod.rs | 4 +- .../{extension.rs => storage_pushdown.rs} | 116 +++++++++++------- 2 files changed, 77 insertions(+), 43 deletions(-) rename vortex-array/src/aggregate_fn/{extension.rs => storage_pushdown.rs} (69%) diff --git a/vortex-array/src/aggregate_fn/mod.rs b/vortex-array/src/aggregate_fn/mod.rs index 64a739216d8..a6610b010da 100644 --- a/vortex-array/src/aggregate_fn/mod.rs +++ b/vortex-array/src/aggregate_fn/mod.rs @@ -32,8 +32,8 @@ pub use typed::*; mod erased; pub use erased::*; -mod extension; -pub use extension::*; +mod storage_pushdown; +pub use storage_pushdown::*; mod options; pub use options::*; diff --git a/vortex-array/src/aggregate_fn/extension.rs b/vortex-array/src/aggregate_fn/storage_pushdown.rs similarity index 69% rename from vortex-array/src/aggregate_fn/extension.rs rename to vortex-array/src/aggregate_fn/storage_pushdown.rs index e2d163ab54a..3f6a43756a6 100644 --- a/vortex-array/src/aggregate_fn/extension.rs +++ b/vortex-array/src/aggregate_fn/storage_pushdown.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Generic adapter that defines an aggregate over one extension type by running an existing -//! aggregate over the extension's storage. +//! Opt-in pushdown of an aggregate into the storage of one extension type. +//! +//! See [`StoragePushdown`] for when such a pushdown is correct. use std::mem; @@ -29,42 +30,72 @@ use crate::dtype::DType; use crate::dtype::extension::ExtVTable; use crate::scalar::Scalar; -/// The aggregate `A` applied to extension type `E` through its storage. +/// An opt-in pushdown of aggregate `A` into the storage of extension type `E`. /// -/// Inputs must have an extension dtype of type `E`. Every batch is unwrapped to its storage array -/// and accumulated by `A`, so storage-encoding kernels registered for `A` still fire. Partials and -/// results are those of `A` over the storage dtype: they are not converted back to `E`, so e.g. -/// the extrema of a timestamp column come back as its raw `i64` storage values. +/// An extension type gives its storage a meaning: a [`Timestamp`] is an instant, not an `i64`. +/// Aggregates do not see through extension types on their own, so most of them reject an +/// extension dtype. `StoragePushdown` lets you declare, for one aggregate and one extension type +/// at a time, that running the aggregate on the storage gives the right answer for the extension +/// type. /// -/// The adapter is an aggregate of its own, with id `[]` (e.g. -/// `vortex.is_sorted[vortex.timestamp]`). Its options are `A`'s options, and it serializes -/// through `A`, so a registered pair round-trips by id. +/// **Registering a pair is a claim about `E`, and nothing checks it.** The adapter runs `A` on +/// the raw storage values whether or not that is meaningful for `E`. A wrong registration does +/// not fail. It silently returns answers that are wrong for `E`, and query planning and pruning +/// may rely on those answers. /// -/// # When to use it +/// # When a pushdown is correct /// -/// Running `A` over the storage is only correct when the storage has the semantics `A` relies on -/// for `E`. The adapter cannot check this, so only register a pair after checking the property -/// for that aggregate: +/// Every aggregate is built from a few operations on its input values: comparison, equality, +/// addition, or NaN detection. A pushdown of `A` into `E`'s storage is correct only when each +/// operation `A` uses can be derived from the same operation on `E`'s storage. That is a +/// property of the specific type `E` and the specific aggregate `A`, not of extension types in +/// general. One extension type can support pushdown for some aggregates and not others. /// -/// | Aggregates | `E` must... | Fits | Does not fit | -/// | --- | --- | --- | --- | -/// | `min`, `max`, `min_max`, `bounded_min`, `bounded_max`, `is_sorted` | order exactly like its storage | `Date`, `Time`, `Timestamp` over integers | a case-insensitive string type, a little-endian integer stored as bytes | -/// | `is_constant`, `all_non_distinct` | have one storage value per value, so equal values have equal storage | `Date`, `Time`, `Timestamp`, `Uuid` | JSON text (`{"a":1}` and `{ "a": 1 }` are equal), a case-insensitive string type | -/// | `sum`, `mean` | have values for which adding the storage is meaningful | a duration stored as an integer count | `Timestamp` (the sum of two instants is not an instant), `Uuid` | -/// | `nan_count`, `all_nan`, `all_non_nan` | treat a NaN storage value as NaN | a float-backed measurement type | a type that stores a sentinel in its float storage | -/// | `count`, `null_count`, `all_null`, `all_non_null`, `uncompressed_size_in_bytes` | (always holds: these depend only on validity or storage bytes) | any `E` | | +/// [`Timestamp`] shows both cases: /// -/// `A` must also accept `E`'s storage dtype. If it does not, the pair still registers, but binding -/// it to a column of `E` fails: `IsSorted` rejects the fixed-size-list storage of `Uuid`, for -/// example. +/// - Comparing two timestamps of the same unit and time zone is the same as comparing their `i64` +/// storage. So `min`, `max` and `is_sorted` can be pushed down into its storage. +/// - Adding two timestamps is not meaningful: the sum of two instants is not an instant. The `i64` +/// addition on the storage is still defined, though, so a `sum` pushdown would run and return a +/// meaningless number. `sum` must not be registered for `Timestamp`. /// -/// Partials depend on `E`'s metadata as well as `E`. For example, timestamps in milliseconds and -/// in seconds share the `Timestamp` type but not the scale of their storage. Only merge partials -/// accumulated over the same extension dtype. +/// Check each aggregate against this table before registering it for a type: +/// +/// | Aggregates | Operation used | Holds for `E` when... | Fits | Does not fit | +/// | --- | --- | --- | --- | --- | +/// | `min`, `max`, `min_max`, `bounded_min`, `bounded_max`, `is_sorted` | comparison | comparing two values of `E` is comparing their storage | `Date`, `Time`, `Timestamp` | a case-insensitive string type, a little-endian integer stored as bytes | +/// | `is_constant`, `all_non_distinct` | equality | each value of `E` has exactly one storage value, so equal values have equal storage | `Date`, `Time`, `Timestamp`, `Uuid` | JSON text (`{"a":1}` and `{ "a": 1 }` are equal), a case-insensitive string type | +/// | `sum`, `mean` | addition | adding the storage of two values gives the storage of their sum | a duration stored as an integer count | `Timestamp`, `Date`, `Uuid` | +/// | `nan_count`, `all_nan`, `all_non_nan` | NaN detection | a NaN storage value is a NaN value of `E` | a float-backed measurement type | a type that stores a sentinel in its float storage | +/// | `count`, `null_count`, `all_null`, `all_non_null`, `uncompressed_size_in_bytes` | none: validity or storage bytes only | always | any `E` | | +/// +/// If an aggregate is not in the table, work out which operations it uses and check each one. +/// +/// # Behavior +/// +/// - Inputs must have an extension dtype of type `E`. For any other input, `return_dtype` and +/// `partial_dtype` return `None`, so binding fails. +/// - `A` must also accept `E`'s storage dtype. If it does not, the pair still registers but +/// binding it to a column of `E` fails. For example, `IsSorted` rejects the fixed-size-list +/// storage of `Uuid`. +/// - Each batch is unwrapped to its storage array and accumulated by `A`, so kernels and cached +/// statistics of `A` on the storage encoding still apply. +/// - Partials and results are those of `A` over the storage dtype, and are not converted back to +/// `E`. For example, the extrema of a timestamp column come back as raw `i64` storage values, +/// in the column's unit. +/// - Partials depend on `E`'s metadata as well as on `E`. Timestamps in milliseconds and in +/// seconds share the `Timestamp` type but not the scale of their storage. Only merge partials +/// accumulated over the same extension dtype. +/// - The pushdown is an aggregate of its own, with id `[]` (e.g. +/// `vortex.is_sorted[vortex.timestamp]`). It never replaces `A`, and a request for `A` on a +/// column of `E` does not use it. Its options are `A`'s options, and it serializes through +/// `A`, so a registered pair round-trips by id. +/// +/// [`Timestamp`]: crate::extension::datetime::Timestamp /// /// # Example /// -/// `Timestamp` orders like its `i64` storage, so `is_sorted` and `min_max` can run on the +/// Timestamps compare like their `i64` storage, so `is_sorted` can be pushed down into the /// storage. Register the pair once on the session, then bind it like any other aggregate: /// /// ``` @@ -72,7 +103,7 @@ use crate::scalar::Scalar; /// # use vortex_array::VortexSessionExecute; /// # use vortex_array::aggregate_fn::Accumulator; /// # use vortex_array::aggregate_fn::DynAccumulator; -/// # use vortex_array::aggregate_fn::ExtensionAggregate; +/// # use vortex_array::aggregate_fn::StoragePushdown; /// # use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; /// # use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; /// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt; @@ -86,7 +117,7 @@ use crate::scalar::Scalar; /// let session = vortex_array::array_session(); /// /// // Register `vortex.is_sorted[vortex.timestamp]`. -/// let is_sorted = ExtensionAggregate::new(IsSorted, Timestamp); +/// let is_sorted = StoragePushdown::new(IsSorted, Timestamp); /// session.aggregate_fns().register(is_sorted.clone()); /// /// // Aggregate a timestamp column through its storage. @@ -108,14 +139,17 @@ use crate::scalar::Scalar; /// # } /// ``` #[derive(Clone, Debug)] -pub struct ExtensionAggregate { +pub struct StoragePushdown { aggregate: A, ext: E, id: AggregateFnId, } -impl ExtensionAggregate { - /// Define `aggregate` over extension type `ext`. +impl StoragePushdown { + /// Push `aggregate` down into the storage of extension type `ext`. + /// + /// Only do this when `aggregate` over the storage is correct for `ext`; see + /// [`StoragePushdown`] for how to check. pub fn new(aggregate: A, ext: E) -> Self { #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let id = AggregateFnId::new(&format!("{}[{}]", aggregate.id(), ext.id())); @@ -177,7 +211,7 @@ impl ExtensionAggregate { } } -impl AggregateFnVTable for ExtensionAggregate { +impl AggregateFnVTable for StoragePushdown { type Options = A::Options; type Partial = A::Partial; @@ -312,14 +346,14 @@ impl AggregateFnVTable for ExtensionAggregat mod tests { use vortex_error::VortexResult; - use super::ExtensionAggregate; + use super::StoragePushdown; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; - use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::arrays::ExtensionArray; use crate::arrays::PrimitiveArray; @@ -330,9 +364,9 @@ mod tests { use crate::extension::datetime::Timestamp; #[test] - fn sums_extension_through_storage() -> VortexResult<()> { + fn pushes_down_into_storage() -> VortexResult<()> { let session = crate::array_session(); - let aggregate = ExtensionAggregate::new(Sum, Timestamp); + let aggregate = StoragePushdown::new(MinMax, Timestamp); session.aggregate_fns().register(aggregate.clone()); assert!( session @@ -355,7 +389,7 @@ mod tests { acc.accumulate(&array, &mut ctx)?; let mut expected = Accumulator::try_new( - Sum, + MinMax, NumericalAggregateOpts::default(), storage.dtype().clone(), )?; @@ -379,7 +413,7 @@ mod tests { }; let accumulator = || { Accumulator::try_new( - ExtensionAggregate::new(Sum, Timestamp), + StoragePushdown::new(MinMax, Timestamp), NumericalAggregateOpts::default(), DType::Extension(ext_dtype.clone()), ) @@ -403,7 +437,7 @@ mod tests { fn rejects_other_extension_types() { let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()); assert!( - ExtensionAggregate::new(Sum, Timestamp) + StoragePushdown::new(MinMax, Timestamp) .return_dtype(&NumericalAggregateOpts::default(), &date) .is_none() ); From f82b6d013ab2f6dfd30d36903a5e818c7e5fa083 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 17:01:18 +0000 Subject: [PATCH 7/9] docs(aggregate): drop the StoragePushdown row for aggregates that need none count, null_count, all_null, all_non_null and uncompressed_size_in_bytes already accept any extension dtype, so say they need no pushdown instead of listing them as always-correct pushdowns. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- vortex-array/src/aggregate_fn/storage_pushdown.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/aggregate_fn/storage_pushdown.rs b/vortex-array/src/aggregate_fn/storage_pushdown.rs index 3f6a43756a6..dc501d96535 100644 --- a/vortex-array/src/aggregate_fn/storage_pushdown.rs +++ b/vortex-array/src/aggregate_fn/storage_pushdown.rs @@ -67,9 +67,10 @@ use crate::scalar::Scalar; /// | `is_constant`, `all_non_distinct` | equality | each value of `E` has exactly one storage value, so equal values have equal storage | `Date`, `Time`, `Timestamp`, `Uuid` | JSON text (`{"a":1}` and `{ "a": 1 }` are equal), a case-insensitive string type | /// | `sum`, `mean` | addition | adding the storage of two values gives the storage of their sum | a duration stored as an integer count | `Timestamp`, `Date`, `Uuid` | /// | `nan_count`, `all_nan`, `all_non_nan` | NaN detection | a NaN storage value is a NaN value of `E` | a float-backed measurement type | a type that stores a sentinel in its float storage | -/// | `count`, `null_count`, `all_null`, `all_non_null`, `uncompressed_size_in_bytes` | none: validity or storage bytes only | always | any `E` | | /// /// If an aggregate is not in the table, work out which operations it uses and check each one. +/// Aggregates that already accept any extension dtype, such as `count`, `null_count`, +/// `all_null`, `all_non_null` and `uncompressed_size_in_bytes`, need no pushdown. /// /// # Behavior /// From 511a84ae9f7b27be6dc007d2808c45099fa9b3fc Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 17:07:09 +0000 Subject: [PATCH 8/9] refactor(aggregate): rename StoragePushdown to ExtensionStoragePushdown, tighten docs Name the extension in the type, and cut its docs down to the two conditions a registration must meet: the extension type's domain supports the aggregation, and the storage computes it soundly, with the operation each aggregate needs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- .../src/aggregate_fn/storage_pushdown.rs | 104 ++++++------------ 1 file changed, 33 insertions(+), 71 deletions(-) diff --git a/vortex-array/src/aggregate_fn/storage_pushdown.rs b/vortex-array/src/aggregate_fn/storage_pushdown.rs index dc501d96535..1af2ed9e4ff 100644 --- a/vortex-array/src/aggregate_fn/storage_pushdown.rs +++ b/vortex-array/src/aggregate_fn/storage_pushdown.rs @@ -3,7 +3,7 @@ //! Opt-in pushdown of an aggregate into the storage of one extension type. //! -//! See [`StoragePushdown`] for when such a pushdown is correct. +//! See [`ExtensionStoragePushdown`] for when such a pushdown is correct. use std::mem; @@ -32,79 +32,41 @@ use crate::scalar::Scalar; /// An opt-in pushdown of aggregate `A` into the storage of extension type `E`. /// -/// An extension type gives its storage a meaning: a [`Timestamp`] is an instant, not an `i64`. -/// Aggregates do not see through extension types on their own, so most of them reject an -/// extension dtype. `StoragePushdown` lets you declare, for one aggregate and one extension type -/// at a time, that running the aggregate on the storage gives the right answer for the extension -/// type. +/// Registering a pair declares that `A` over `E`'s storage is correct for `E`. Nothing checks +/// this, and a wrong registration silently returns wrong answers. Only register a pair when both +/// of these hold: /// -/// **Registering a pair is a claim about `E`, and nothing checks it.** The adapter runs `A` on -/// the raw storage values whether or not that is meaningful for `E`. A wrong registration does -/// not fail. It silently returns answers that are wrong for `E`, and query planning and pruning -/// may rely on those answers. +/// - **`E`'s domain supports the aggregation.** The result must mean something for values of `E`. +/// The minimum of some timestamps is an instant, but their sum is not. +/// - **The storage computes it soundly.** Every operation `A` uses must give the same answer on +/// the storage as on `E`: +/// - comparison (`min`, `max`, `min_max`, `bounded_*`, `is_sorted`): `E` orders like its +/// storage. +/// - equality (`is_constant`, `all_non_distinct`): each value of `E` has exactly one storage +/// value. +/// - addition (`sum`, `mean`): adding storage values gives the storage of their sum. +/// - NaN detection (`nan_count`, `all_nan`, `all_non_nan`): a NaN in the storage is a NaN of `E`. /// -/// # When a pushdown is correct +/// For example, `Date`, `Time` and `Timestamp` pass for comparison but not for `sum`. Aggregates +/// that already accept any extension dtype, such as `count` and `null_count`, need no pushdown. /// -/// Every aggregate is built from a few operations on its input values: comparison, equality, -/// addition, or NaN detection. A pushdown of `A` into `E`'s storage is correct only when each -/// operation `A` uses can be derived from the same operation on `E`'s storage. That is a -/// property of the specific type `E` and the specific aggregate `A`, not of extension types in -/// general. One extension type can support pushdown for some aggregates and not others. -/// -/// [`Timestamp`] shows both cases: -/// -/// - Comparing two timestamps of the same unit and time zone is the same as comparing their `i64` -/// storage. So `min`, `max` and `is_sorted` can be pushed down into its storage. -/// - Adding two timestamps is not meaningful: the sum of two instants is not an instant. The `i64` -/// addition on the storage is still defined, though, so a `sum` pushdown would run and return a -/// meaningless number. `sum` must not be registered for `Timestamp`. -/// -/// Check each aggregate against this table before registering it for a type: -/// -/// | Aggregates | Operation used | Holds for `E` when... | Fits | Does not fit | -/// | --- | --- | --- | --- | --- | -/// | `min`, `max`, `min_max`, `bounded_min`, `bounded_max`, `is_sorted` | comparison | comparing two values of `E` is comparing their storage | `Date`, `Time`, `Timestamp` | a case-insensitive string type, a little-endian integer stored as bytes | -/// | `is_constant`, `all_non_distinct` | equality | each value of `E` has exactly one storage value, so equal values have equal storage | `Date`, `Time`, `Timestamp`, `Uuid` | JSON text (`{"a":1}` and `{ "a": 1 }` are equal), a case-insensitive string type | -/// | `sum`, `mean` | addition | adding the storage of two values gives the storage of their sum | a duration stored as an integer count | `Timestamp`, `Date`, `Uuid` | -/// | `nan_count`, `all_nan`, `all_non_nan` | NaN detection | a NaN storage value is a NaN value of `E` | a float-backed measurement type | a type that stores a sentinel in its float storage | -/// -/// If an aggregate is not in the table, work out which operations it uses and check each one. -/// Aggregates that already accept any extension dtype, such as `count`, `null_count`, -/// `all_null`, `all_non_null` and `uncompressed_size_in_bytes`, need no pushdown. -/// -/// # Behavior -/// -/// - Inputs must have an extension dtype of type `E`. For any other input, `return_dtype` and -/// `partial_dtype` return `None`, so binding fails. -/// - `A` must also accept `E`'s storage dtype. If it does not, the pair still registers but -/// binding it to a column of `E` fails. For example, `IsSorted` rejects the fixed-size-list -/// storage of `Uuid`. -/// - Each batch is unwrapped to its storage array and accumulated by `A`, so kernels and cached -/// statistics of `A` on the storage encoding still apply. -/// - Partials and results are those of `A` over the storage dtype, and are not converted back to -/// `E`. For example, the extrema of a timestamp column come back as raw `i64` storage values, -/// in the column's unit. -/// - Partials depend on `E`'s metadata as well as on `E`. Timestamps in milliseconds and in -/// seconds share the `Timestamp` type but not the scale of their storage. Only merge partials -/// accumulated over the same extension dtype. -/// - The pushdown is an aggregate of its own, with id `[]` (e.g. -/// `vortex.is_sorted[vortex.timestamp]`). It never replaces `A`, and a request for `A` on a -/// column of `E` does not use it. Its options are `A`'s options, and it serializes through -/// `A`, so a registered pair round-trips by id. -/// -/// [`Timestamp`]: crate::extension::datetime::Timestamp +/// Inputs must be extension dtypes of type `E` whose storage `A` accepts; anything else fails to +/// bind. Results are `A`'s over the storage and are not converted back to `E`, so the extrema of a +/// timestamp column are raw `i64` values in the column's unit. Partials are only comparable for +/// the same extension dtype, including metadata such as a timestamp's unit. The pushdown is an +/// aggregate of its own, with id `[]`, and never replaces `A`. /// /// # Example /// -/// Timestamps compare like their `i64` storage, so `is_sorted` can be pushed down into the -/// storage. Register the pair once on the session, then bind it like any other aggregate: +/// Timestamps compare like their `i64` storage, so `is_sorted` can be pushed down. Register the +/// pair once on the session, then bind it like any other aggregate: /// /// ``` /// # use vortex_array::IntoArray; /// # use vortex_array::VortexSessionExecute; /// # use vortex_array::aggregate_fn::Accumulator; /// # use vortex_array::aggregate_fn::DynAccumulator; -/// # use vortex_array::aggregate_fn::StoragePushdown; +/// # use vortex_array::aggregate_fn::ExtensionStoragePushdown; /// # use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; /// # use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; /// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt; @@ -118,7 +80,7 @@ use crate::scalar::Scalar; /// let session = vortex_array::array_session(); /// /// // Register `vortex.is_sorted[vortex.timestamp]`. -/// let is_sorted = StoragePushdown::new(IsSorted, Timestamp); +/// let is_sorted = ExtensionStoragePushdown::new(IsSorted, Timestamp); /// session.aggregate_fns().register(is_sorted.clone()); /// /// // Aggregate a timestamp column through its storage. @@ -140,17 +102,17 @@ use crate::scalar::Scalar; /// # } /// ``` #[derive(Clone, Debug)] -pub struct StoragePushdown { +pub struct ExtensionStoragePushdown { aggregate: A, ext: E, id: AggregateFnId, } -impl StoragePushdown { +impl ExtensionStoragePushdown { /// Push `aggregate` down into the storage of extension type `ext`. /// /// Only do this when `aggregate` over the storage is correct for `ext`; see - /// [`StoragePushdown`] for how to check. + /// [`ExtensionStoragePushdown`] for how to check. pub fn new(aggregate: A, ext: E) -> Self { #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let id = AggregateFnId::new(&format!("{}[{}]", aggregate.id(), ext.id())); @@ -212,7 +174,7 @@ impl StoragePushdown { } } -impl AggregateFnVTable for StoragePushdown { +impl AggregateFnVTable for ExtensionStoragePushdown { type Options = A::Options; type Partial = A::Partial; @@ -347,7 +309,7 @@ impl AggregateFnVTable for StoragePushdown VortexResult<()> { let session = crate::array_session(); - let aggregate = StoragePushdown::new(MinMax, Timestamp); + let aggregate = ExtensionStoragePushdown::new(MinMax, Timestamp); session.aggregate_fns().register(aggregate.clone()); assert!( session @@ -414,7 +376,7 @@ mod tests { }; let accumulator = || { Accumulator::try_new( - StoragePushdown::new(MinMax, Timestamp), + ExtensionStoragePushdown::new(MinMax, Timestamp), NumericalAggregateOpts::default(), DType::Extension(ext_dtype.clone()), ) @@ -438,7 +400,7 @@ mod tests { fn rejects_other_extension_types() { let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()); assert!( - StoragePushdown::new(MinMax, Timestamp) + ExtensionStoragePushdown::new(MinMax, Timestamp) .return_dtype(&NumericalAggregateOpts::default(), &date) .is_none() ); From 597e2ae803d91379c8821e91bfa7acb5a4213161 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 20:47:53 +0000 Subject: [PATCH 9/9] docs(aggregate): adopt the reviewed ExtensionStoragePushdown doc comment Co-authored-by: connortsui20 <87130162+connortsui20@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hp7ZGEG8mrip6qkAtK4iWD Signed-off-by: Claude --- .../src/aggregate_fn/storage_pushdown.rs | 49 +++++++++++-------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/vortex-array/src/aggregate_fn/storage_pushdown.rs b/vortex-array/src/aggregate_fn/storage_pushdown.rs index 1af2ed9e4ff..b3f369fc70c 100644 --- a/vortex-array/src/aggregate_fn/storage_pushdown.rs +++ b/vortex-array/src/aggregate_fn/storage_pushdown.rs @@ -30,31 +30,38 @@ use crate::dtype::DType; use crate::dtype::extension::ExtVTable; use crate::scalar::Scalar; -/// An opt-in pushdown of aggregate `A` into the storage of extension type `E`. +/// Runs aggregate `A` on the storage values of extension type `E`. /// -/// Registering a pair declares that `A` over `E`'s storage is correct for `E`. Nothing checks -/// this, and a wrong registration silently returns wrong answers. Only register a pair when both -/// of these hold: +/// The input must have extension type `E`, and `A` must support its storage dtype. Binding fails +/// for any other input. The wrapper has a separate aggregate ID, `[]`, formed from the IDs +/// of `A` and `E`. Registering it does not replace `A`. /// -/// - **`E`'s domain supports the aggregation.** The result must mean something for values of `E`. -/// The minimum of some timestamps is an instant, but their sum is not. -/// - **The storage computes it soundly.** Every operation `A` uses must give the same answer on -/// the storage as on `E`: -/// - comparison (`min`, `max`, `min_max`, `bounded_*`, `is_sorted`): `E` orders like its -/// storage. -/// - equality (`is_constant`, `all_non_distinct`): each value of `E` has exactly one storage -/// value. -/// - addition (`sum`, `mean`): adding storage values gives the storage of their sum. -/// - NaN detection (`nan_count`, `all_nan`, `all_non_nan`): a NaN in the storage is a NaN of `E`. +/// Aggregates that already support extension dtypes directly do not need this wrapper. /// -/// For example, `Date`, `Time` and `Timestamp` pass for comparison but not for `sum`. Aggregates -/// that already accept any extension dtype, such as `count` and `null_count`, need no pushdown. +/// # Correctness /// -/// Inputs must be extension dtypes of type `E` whose storage `A` accepts; anything else fails to -/// bind. Results are `A`'s over the storage and are not converted back to `E`, so the extrema of a -/// timestamp column are raw `i64` values in the column's unit. Partials are only comparable for -/// the same extension dtype, including metadata such as a timestamp's unit. The pushdown is an -/// aggregate of its own, with id `[]`, and never replaces `A`. +/// A registration **must** satisfy both conditions: +/// +/// - The aggregation is defined for logical values of `E`. +/// - For every valid input, aggregating the storage values produces the storage representation +/// of the logical aggregate result. This requires storage to preserve whichever properties `A` +/// relies on, such as ordering, equality, arithmetic, or NaN classification. +/// +/// For example, minimum is defined for timestamps, and their integer storage preserves their +/// ordering, so both conditions hold. Summing timestamps fails the first condition even though +/// their storage supports addition. +/// +/// These conditions are not checked. An invalid registration can silently return incorrect results. +/// +/// # Results and partials +/// +/// Results and partial states are those of `A` over the storage dtype. Results are not converted +/// back to `E`. For example, the minimum and maximum of a timestamp column are raw `i64` values +/// in the column's time unit. +/// +/// Partial states can only be combined when the aggregate, options, and full input dtype match. +/// This includes extension metadata, so timestamp partials with different time units must not be +/// combined even if their storage dtypes match. /// /// # Example ///