feat(aggregate): register is_sorted and min_max for temporal and UUID extension types (#10616)

- New `AggregateFnSession::register_storage_ordered(ext)` pushes the
comparison-based aggregates `is_sorted` and `min_max` down into `ext`'s
storage, registering `ExtensionStoragePushdown<IsSorted, E>` and
`ExtensionStoragePushdown<MinMax, E>`.
- The default aggregate session calls it for `Date`, `Time` and
`Timestamp`, whose storage is an integer.
- `vortex_uuid::initialize` calls it for `Uuid`. UUIDs are stored as 16
big-endian bytes, so byte order is UUID order.
- `is_sorted` and `min_max` now order a `FixedSizeList` of non-nullable
`u8` (UUID storage) lexicographically by byte. A shared
`fns/fixed_size_bytes.rs` decodes the elements once and slices each
value. Other fixed-size lists stay unsupported.

Author: joseph-isaacs

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<u8>]) -> VortexResult<ArrayRef> {
+        let bytes = leading_bytes
+            .iter()
+            .flat_map(|byte| {
+                let mut uuid = [0u8; 16];
+                uuid[0] = byte.unwrap_or_default();
+                uuid
+            })
+            .collect::<Buffer<u8>>();
+        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<u8>],
+        #[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(())
+    }
+}

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<Self> {
+        Ok(Self {
+            elements: array.elements().clone().execute::<PrimitiveArray>(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<Item = &[u8]> {
+        let bytes = self.elements.as_slice::<u8>();
+        (0..self.len).map(move |index| &bytes[index * self.size..(index + 1) * self.size])
+    }
+}

vortex-array/src/aggregate_fn/fns/is_sorted/extension.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<bool> {
+    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()
+    })
+}

vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_list.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,15 +60,17 @@ 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<bool> {
     is_sorted_impl(array, false, ctx)
 }
 
 /// 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<bool> {
     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<DType> {
         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<DType> {
         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!(),
                 };