diff --git a/encodings/uuid/src/lib.rs b/encodings/uuid/src/lib.rs index 6ad4bc23d89..580e7763edc 100644 --- a/encodings/uuid/src/lib.rs +++ b/encodings/uuid/src/lib.rs @@ -13,8 +13,9 @@ //! is stored as 16 bytes in big-endian (network) byte order, matching [RFC 4122] and Arrow's //! [canonical UUID extension]. //! -//! [`initialize`] registers the extension dtype and the Arrow import/export plugins that map it -//! to Arrow's `arrow.uuid` extension type over `FixedSizeBinary[16]` storage. +//! [`initialize`] registers the extension dtype, the Arrow import/export plugins that map it to +//! Arrow's `arrow.uuid` extension type over `FixedSizeBinary[16]` storage, and `is_sorted` and +//! `min_max` over its storage, since big-endian bytes order like the UUIDs they encode. //! //! [RFC 4122]: https://www.rfc-editor.org/rfc/rfc4122 //! [canonical UUID extension]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#uuid @@ -26,6 +27,7 @@ mod vtable; use std::sync::Arc; pub use metadata::UuidMetadata; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; @@ -37,9 +39,142 @@ pub struct Uuid; /// Register UUID extension support with a session. /// /// This registers the `vortex.uuid` extension dtype along with the Arrow exporter and importer -/// for Arrow's canonical `arrow.uuid` extension type. +/// for Arrow's canonical `arrow.uuid` extension type. UUIDs are stored big-endian, so they compare +/// like their storage bytes, and `is_sorted` and `min_max` are pushed down into the storage. pub fn initialize(session: &VortexSession) { session.dtypes().register(Uuid); + session.aggregate_fns().register_storage_ordered(Uuid); session.arrow().register_exporter(Arc::new(Uuid)); session.arrow().register_importer(Arc::new(Uuid)); } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::ExtensionStoragePushdown; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; + use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions; + use vortex_array::aggregate_fn::fns::min_max::MinMax; + use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; + use vortex_array::aggregate_fn::session::AggregateFnSessionExt; + use vortex_array::array_session; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::*; + + fn uuid_session() -> VortexSession { + let session = array_session(); + initialize(&session); + session + } + + /// A nullable UUID column, with each UUID given by its leading byte. + fn uuid_column(leading_bytes: &[Option]) -> VortexResult { + let bytes = leading_bytes + .iter() + .flat_map(|byte| { + let mut uuid = [0u8; 16]; + uuid[0] = byte.unwrap_or_default(); + uuid + }) + .collect::>(); + let storage = FixedSizeListArray::new( + PrimitiveArray::new(bytes, Validity::NonNullable).into_array(), + 16, + leading_bytes.iter().map(Option::is_some).collect(), + leading_bytes.len(), + ) + .into_array(); + let ext_dtype = + ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), storage.dtype().clone())?; + Ok(ExtensionArray::new(ext_dtype.erased(), storage).into_array()) + } + + #[test] + fn initialize_registers_order_aggregates() { + let session = uuid_session(); + for id in [ + ExtensionStoragePushdown::new(IsSorted, Uuid).id(), + ExtensionStoragePushdown::new(MinMax, Uuid).id(), + ] { + assert!( + session.aggregate_fns().find_plugin(&id).is_some(), + "{id} is not registered" + ); + } + } + + #[rstest] + // Nulls sort first. + #[case::sorted(&[None, Some(1), Some(2), Some(2)], false, true)] + #[case::unsorted(&[Some(2), Some(1), None, Some(3)], false, false)] + #[case::strict_with_duplicates(&[None, Some(1), Some(2), Some(2)], true, false)] + #[case::strict(&[None, Some(1), Some(2), Some(0xff)], true, true)] + fn is_sorted_orders_by_bytes( + #[case] leading_bytes: &[Option], + #[case] strict: bool, + #[case] expected: bool, + ) -> VortexResult<()> { + let session = uuid_session(); + let column = uuid_column(leading_bytes)?; + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(IsSorted, Uuid), + IsSortedOptions { strict }, + column.dtype().clone(), + )?; + acc.accumulate(&column, &mut session.create_execution_ctx())?; + assert_eq!(bool::try_from(&acc.finish()?)?, expected); + Ok(()) + } + + #[test] + fn min_max_orders_by_bytes() -> VortexResult<()> { + let session = uuid_session(); + let column = uuid_column(&[Some(7), None, Some(0xff), Some(3)])?; + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(MinMax, Uuid), + NumericalAggregateOpts::default(), + column.dtype().clone(), + )?; + acc.accumulate(&column, &mut session.create_execution_ctx())?; + + let uuid = |leading_byte: u8| { + let mut bytes = [0u8; 16]; + bytes[0] = leading_byte; + Scalar::fixed_size_list( + DType::Primitive(PType::U8, Nullability::NonNullable), + bytes + .into_iter() + .map(|byte| Scalar::primitive(byte, Nullability::NonNullable)) + .collect(), + Nullability::NonNullable, + ) + }; + assert_eq!( + MinMaxResult::from_scalar(&acc.finish()?)?, + Some(MinMaxResult { + min: uuid(3), + max: uuid(0xff), + }) + ); + Ok(()) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/fixed_size_bytes.rs b/vortex-array/src/aggregate_fn/fns/fixed_size_bytes.rs new file mode 100644 index 00000000000..55c9f612124 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/fixed_size_bytes.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fixed-size lists of non-nullable `u8`, which order lexicographically by byte. +//! +//! This is the storage of byte-string extension types such as UUID, so order-based aggregates can +//! support it without implementing general list comparison. + +use vortex_error::VortexResult; + +use crate::ExecutionCtx; +use crate::arrays::FixedSizeListArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +/// Whether `dtype` is a fixed-size list of non-nullable `u8`. +pub(crate) fn is_fixed_size_bytes(dtype: &DType) -> bool { + matches!( + dtype, + DType::FixedSizeList(element_dtype, ..) + if matches!(element_dtype.as_ref(), DType::Primitive(PType::U8, Nullability::NonNullable)) + ) +} + +/// The values of a fixed-size list of bytes, each one a slice of the decoded elements. +pub(crate) struct FixedSizeBytes { + elements: PrimitiveArray, + size: usize, + len: usize, +} + +impl FixedSizeBytes { + /// Decode the elements of `array`, which must be a fixed-size list of bytes. + pub(crate) fn try_new( + array: &FixedSizeListArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(Self { + elements: array.elements().clone().execute::(ctx)?, + size: array.list_size() as usize, + len: array.len(), + }) + } + + /// Every value in order, including the values of null lists. + pub(crate) fn values(&self) -> impl Iterator { + let bytes = self.elements.as_slice::(); + (0..self.len).map(move |index| &bytes[index * self.size..(index + 1) * self.size]) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs deleted file mode 100644 index dc69b7c8bf9..00000000000 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/extension.rs +++ /dev/null @@ -1,22 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; - -use crate::ExecutionCtx; -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::arrays::extension::ExtensionArrayExt; - -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) - } -} diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.rs new file mode 100644 index 00000000000..3ef6b2546c4 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.rs @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::IsSortedIteratorExt; +use crate::ExecutionCtx; +use crate::aggregate_fn::fns::fixed_size_bytes::FixedSizeBytes; +use crate::arrays::FixedSizeListArray; + +/// Check whether a fixed-size list of bytes is sorted by byte, with nulls first. +pub(super) fn check_fixed_size_bytes_sorted( + array: &FixedSizeListArray, + strict: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = array + .as_ref() + .validity()? + .execute_mask(array.len(), ctx)? + .to_bit_buffer(); + let bytes = FixedSizeBytes::try_new(array, ctx)?; + let iter = bytes + .values() + .zip(mask.iter()) + .map(|(value, valid)| valid.then_some(value)); + Ok(if strict { + iter.is_strict_sorted() + } else { + iter.is_sorted() + }) +} diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index c1d3d7d1a5a..89775ce2dfa 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -3,7 +3,7 @@ mod bool; mod decimal; -mod extension; +mod fixed_size_list; mod primitive; mod varbin; @@ -18,7 +18,7 @@ use vortex_session::registry::CachedId; use self::bool::check_bool_sorted; use self::decimal::check_decimal_sorted; -use self::extension::check_extension_sorted; +use self::fixed_size_list::check_fixed_size_bytes_sorted; use self::primitive::check_primitive_sorted; use self::varbin::check_varbinview_sorted; use crate::ArrayRef; @@ -31,6 +31,7 @@ use crate::aggregate_fn::AggregateArgs; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::fns::fixed_size_bytes::is_fixed_size_bytes; use crate::arrays::Constant; use crate::arrays::Null; use crate::builtins::ArrayBuiltins; @@ -59,7 +60,8 @@ impl Display for IsSortedOptions { /// Compute whether an array is sorted in non-decreasing order. /// /// Returns `true` for empty arrays and arrays of length 1. -/// Returns `false` for struct, list, and fixed-size list arrays. +/// Returns `false` for struct, list, and fixed-size list arrays, except fixed-size lists of +/// non-nullable `u8`, which are ordered by byte. pub fn is_sorted(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { is_sorted_impl(array, false, ctx) } @@ -67,7 +69,8 @@ pub fn is_sorted(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult /// Compute whether an array is strictly sorted in increasing order (no duplicates). /// /// Returns `true` for empty arrays and arrays of length 1. -/// Returns `false` for struct, list, and fixed-size list arrays. +/// Returns `false` for struct, list, and fixed-size list arrays, except fixed-size lists of +/// non-nullable `u8`, which are ordered by byte. pub fn is_strict_sorted(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { is_sorted_impl(array, true, ctx) } @@ -254,6 +257,9 @@ impl AggregateFnVTable for IsSorted { fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { match input_dtype { + DType::FixedSizeList(..) if is_fixed_size_bytes(input_dtype) => { + Some(DType::Bool(Nullability::NonNullable)) + } DType::Null | DType::List(..) | DType::FixedSizeList(..) @@ -272,6 +278,9 @@ impl AggregateFnVTable for IsSorted { fn partial_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { match input_dtype { + DType::FixedSizeList(..) if is_fixed_size_bytes(input_dtype) => { + Some(make_is_sorted_partial_dtype(input_dtype)) + } DType::Null | DType::List(..) | DType::FixedSizeList(..) @@ -514,9 +523,12 @@ impl AggregateFnVTable for IsSorted { check_varbinview_sorted(v, args.options.strict, ctx)? } Canonical::Decimal(d) => check_decimal_sorted(d, args.options.strict, ctx)?, - Canonical::Extension(e) => check_extension_sorted(e, args.options.strict, ctx)?, + Canonical::FixedSizeList(f) => { + check_fixed_size_bytes_sorted(f, args.options.strict, ctx)? + } Canonical::Null(_) => !args.options.strict, - // Struct, List, FixedSizeList should have been filtered out by return_dtype + // Struct, List, other FixedSizeLists and Extension should have been filtered + // out by return_dtype _ => unreachable!(), }; diff --git a/vortex-array/src/aggregate_fn/fns/min_max/fixed_size_list.rs b/vortex-array/src/aggregate_fn/fns/min_max/fixed_size_list.rs new file mode 100644 index 00000000000..f4e3cc6a51d --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/min_max/fixed_size_list.rs @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::MinMaxPartial; +use super::MinMaxResult; +use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateArgs; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::fixed_size_bytes::FixedSizeBytes; +use crate::arrays::FixedSizeListArray; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::dtype::Nullability::NonNullable; +use crate::scalar::Scalar; + +/// Accumulate the byte-wise extrema of a fixed-size list of bytes. +pub(super) fn accumulate_fixed_size_bytes( + args: AggregateArgs<'_, NumericalAggregateOpts>, + partial: &mut MinMaxPartial, + array: &FixedSizeListArray, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let mask = array + .as_ref() + .validity()? + .execute_mask(array.len(), ctx)? + .to_bit_buffer(); + let bytes = FixedSizeBytes::try_new(array, ctx)?; + let extrema = bytes + .values() + .zip(mask.iter()) + .filter_map(|(value, valid)| valid.then_some(value)) + .fold(None, |extrema: Option<(&[u8], &[u8])>, value| { + Some(match extrema { + None => (value, value), + Some((min, max)) => (min.min(value), max.max(value)), + }) + }); + + let (element_dtype, ..) = array.dtype_parts(); + let make_scalar = |value: &[u8]| { + Scalar::fixed_size_list( + element_dtype.clone(), + value + .iter() + .map(|&byte| Scalar::primitive(byte, NonNullable)) + .collect(), + NonNullable, + ) + }; + partial.merge( + args, + extrema.map(|(min, max)| MinMaxResult { + min: make_scalar(min), + max: make_scalar(max), + }), + ); + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index b8b91e9e13e..7e625cfc27d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -4,6 +4,7 @@ mod bool; mod decimal; mod extension; +mod fixed_size_list; mod primitive; mod varbinview; @@ -18,6 +19,7 @@ use vortex_session::registry::CachedId; use self::bool::accumulate_bool; use self::decimal::accumulate_decimal; use self::extension::accumulate_extension; +use self::fixed_size_list::accumulate_fixed_size_bytes; use self::primitive::accumulate_primitive; use self::varbinview::accumulate_varbinview; use crate::ArrayRef; @@ -31,6 +33,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::fixed_size_bytes::is_fixed_size_bytes; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; @@ -299,7 +302,8 @@ fn minmax_supported_dtype(input_dtype: &DType) -> bool { /// /// This is intentionally narrower than [`minmax_supported_dtype`]. List and fixed-size-list /// extrema have a defined output dtype for aggregate expression lowering, but the accumulator does -/// not yet implement lexicographic list comparison. +/// not yet implement lexicographic list comparison. Fixed-size lists of non-nullable `u8` are the +/// exception: they compare by byte. fn minmax_compute_supported_dtype(input_dtype: &DType) -> bool { matches!( input_dtype, @@ -309,7 +313,7 @@ fn minmax_compute_supported_dtype(input_dtype: &DType) -> bool { | DType::Utf8(..) | DType::Binary(..) | DType::Extension(..) - ) + ) || is_fixed_size_bytes(input_dtype) } impl AggregateFnVTable for MinMax { @@ -473,6 +477,9 @@ impl AggregateFnVTable for MinMax { Canonical::VarBinView(v) => accumulate_varbinview(args, partial, v, ctx), Canonical::Decimal(d) => accumulate_decimal(args, partial, d, ctx), Canonical::Extension(e) => accumulate_extension(args, partial, e, ctx), + Canonical::FixedSizeList(f) if is_fixed_size_bytes(f.dtype()) => { + accumulate_fixed_size_bytes(args, partial, f, ctx) + } Canonical::Null(_) => Ok(()), Canonical::Union(_) => { vortex_bail!("TODO(connor)[Union]: implement min_max for Union arrays") diff --git a/vortex-array/src/aggregate_fn/fns/mod.rs b/vortex-array/src/aggregate_fn/fns/mod.rs index c4e8d1b72ca..2a0119820b1 100644 --- a/vortex-array/src/aggregate_fn/fns/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mod.rs @@ -10,6 +10,7 @@ pub mod bounded_max; pub mod bounded_min; pub mod count; pub mod first; +mod fixed_size_bytes; pub mod is_constant; pub mod is_sorted; pub mod last; diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 669fbb2bb82..9726b5e2ff4 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -14,6 +14,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnPluginRef; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::ExtensionStoragePushdown; use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; @@ -49,6 +50,10 @@ use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; use crate::arrays::dict::compute::min_max::DictMinMaxKernel; use crate::dtype::DType; +use crate::dtype::extension::ExtVTable; +use crate::extension::datetime::Date; +use crate::extension::datetime::Time; +use crate::extension::datetime::Timestamp; /// Session state for aggregate functions and encoding-specific aggregate kernels. /// @@ -117,6 +122,12 @@ impl Default for AggregateFnSession { this.register(SumV2); this.register(UncompressedSizeInBytes); + // Temporal extension types compare like their integer storage, so comparison-based + // aggregates are pushed down into it. + this.register_storage_ordered(Date); + this.register_storage_ordered(Time); + this.register_storage_ordered(Timestamp); + // Register the built-in aggregate kernels. this.register_aggregate_kernel(Chunked.id(), None::, &ChunkedArrayAggregate); this.register_aggregate_kernel(Dict.id(), Some(MinMax.id()), &DictMinMaxKernel); @@ -154,6 +165,17 @@ impl AggregateFnSession { self.registry.insert(id, pluginref); } + /// Push the comparison-based aggregates `is_sorted` and `min_max` down into the storage of + /// extension type `ext`, through [`ExtensionStoragePushdown`]. + /// + /// Only call this when comparing two values of `ext` is the same as comparing their storage, + /// as for the temporal types and for UUIDs, which are stored big-endian. See + /// [`ExtensionStoragePushdown`] for how to check. + pub fn register_storage_ordered(&self, ext: E) { + self.register(ExtensionStoragePushdown::new(IsSorted, ext.clone())); + self.register(ExtensionStoragePushdown::new(MinMax, ext)); + } + /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every /// registered aggregate's `zone_stat_default`. /// diff --git a/vortex-array/src/aggregate_fn/storage_pushdown.rs b/vortex-array/src/aggregate_fn/storage_pushdown.rs index b3f369fc70c..701cdd76d17 100644 --- a/vortex-array/src/aggregate_fn/storage_pushdown.rs +++ b/vortex-array/src/aggregate_fn/storage_pushdown.rs @@ -16,7 +16,6 @@ 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; @@ -275,12 +274,14 @@ impl AggregateFnVTable for ExtensionStorageP Columnar::Canonical(Canonical::Extension(array)) => { self.accumulate_storage(args, state, array.storage_array(), ctx) } + // A constant storage array has no encoding to dispatch on and no statistics worth + // caching, so it goes straight to `A` rather than through a nested accumulator. Columnar::Constant(constant) => { let storage = constant.scalar().as_extension().to_storage_scalar(); - self.accumulate_storage( + self.aggregate.accumulate( args, state, - &ConstantArray::new(storage, constant.len()).into_array(), + &Columnar::Constant(ConstantArray::new(storage, constant.len())), ctx, ) } @@ -314,6 +315,7 @@ impl AggregateFnVTable for ExtensionStorageP #[cfg(test)] mod tests { + use rstest::rstest; use vortex_error::VortexResult; use super::ExtensionStoragePushdown; @@ -323,15 +325,20 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::is_sorted::IsSorted; + use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::session::AggregateFnSessionExt; + use crate::arrays::ConstantArray; 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::Time; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; #[test] fn pushes_down_into_storage() -> VortexResult<()> { @@ -403,6 +410,44 @@ mod tests { Ok(()) } + #[test] + fn accumulates_constant_batches() -> VortexResult<()> { + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let constant = |value: i64| { + ConstantArray::new( + Scalar::extension_ref( + ext_dtype.clone(), + Scalar::primitive(value, Nullability::Nullable), + ), + 3, + ) + .into_array() + }; + + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(MinMax, Timestamp), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype.clone()), + )?; + acc.accumulate(&constant(2_000), &mut ctx)?; + acc.accumulate(&constant(1_000), &mut ctx)?; + + let mut expected = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::default(), + ext_dtype.storage_dtype().clone(), + )?; + expected.accumulate( + &PrimitiveArray::from_option_iter([Some(2_000i64), Some(1_000)]).into_array(), + &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()); @@ -412,4 +457,80 @@ mod tests { .is_none() ); } + + #[test] + fn default_session_registers_temporal_order_aggregates() { + let session = crate::array_session(); + let aggregate_fns = session.aggregate_fns(); + for id in [ + ExtensionStoragePushdown::new(IsSorted, Date).id(), + ExtensionStoragePushdown::new(IsSorted, Time).id(), + ExtensionStoragePushdown::new(IsSorted, Timestamp).id(), + ExtensionStoragePushdown::new(MinMax, Date).id(), + ExtensionStoragePushdown::new(MinMax, Time).id(), + ExtensionStoragePushdown::new(MinMax, Timestamp).id(), + ] { + assert!( + aggregate_fns.find_plugin(&id).is_some(), + "{id} is not registered" + ); + } + } + + #[rstest] + // Nulls sort first. + #[case::sorted([None, Some(1_000), Some(2_000), Some(2_000)], false, true)] + #[case::unsorted([Some(2_000), Some(1_000), None, Some(3_000)], false, false)] + #[case::strict_with_duplicates([None, Some(1_000), Some(2_000), Some(2_000)], true, false)] + fn is_sorted_over_timestamps( + #[case] values: [Option; 4], + #[case] strict: bool, + #[case] expected: bool, + ) -> VortexResult<()> { + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let array = ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_option_iter(values).into_array(), + ) + .into_array(); + + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(IsSorted, Timestamp), + IsSortedOptions { strict }, + DType::Extension(ext_dtype), + )?; + acc.accumulate(&array, &mut ctx)?; + + assert_eq!(bool::try_from(&acc.finish()?)?, expected); + Ok(()) + } + + #[test] + fn min_max_over_dates_matches_storage() -> VortexResult<()> { + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); + let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); + let storage = + PrimitiveArray::from_option_iter([Some(20i32), None, Some(-3), Some(7)]).into_array(); + let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array(); + + let mut acc = Accumulator::try_new( + ExtensionStoragePushdown::new(MinMax, Date), + NumericalAggregateOpts::default(), + DType::Extension(ext_dtype), + )?; + acc.accumulate(&array, &mut ctx)?; + + let mut expected = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::default(), + storage.dtype().clone(), + )?; + expected.accumulate(&storage, &mut ctx)?; + + assert_eq!(acc.finish()?, expected.finish()?); + Ok(()) + } }