Validate Extension Arrays

encodings/datetime-parts/src/compute/compare.rs

@@ -321,7 +321,8 @@ mod test {
         .unwrap();
 
         // Timestamp with a value larger than i32::MAX.
-        let rhs = dtp_array_from_timestamp(i64::MAX, rhs_validity);
+        // https://github.com/BurntSushi/jiff/blob/e5b7f0d061e4da9598aed73f6171e78baa8b007f/src/shared/tzif.rs#L23
+        let rhs = dtp_array_from_timestamp(253_402_207_200i64, rhs_validity);
 
         let comparison = lhs
             .clone()

vortex-array/public-api.lock

@@ -10238,6 +10238,8 @@ pub fn vortex_array::dtype::extension::ExtVTable::serialize_metadata(&self, meta
 
 pub fn vortex_array::dtype::extension::ExtVTable::unpack_native<'a>(&self, metadata: &'a Self::Metadata, storage_dtype: &'a vortex_array::dtype::DType, storage_value: &'a vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::dtype::extension::ExtVTable::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::dtype::extension::ExtVTable::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::dtype::extension::ExtVTable::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
@@ -10256,6 +10258,8 @@ pub fn vortex_array::extension::datetime::Date::serialize_metadata(&self, metada
 
 pub fn vortex_array::extension::datetime::Date::unpack_native(&self, metadata: &Self::Metadata, _storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::extension::datetime::Date::validate_array(&self, _metadata: &Self::Metadata, _storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::extension::datetime::Date::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::datetime::Date::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
@@ -10274,6 +10278,8 @@ pub fn vortex_array::extension::datetime::Time::serialize_metadata(&self, metada
 
 pub fn vortex_array::extension::datetime::Time::unpack_native(&self, metadata: &Self::Metadata, _storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::extension::datetime::Time::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::extension::datetime::Time::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::datetime::Time::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
@@ -10292,7 +10298,9 @@ pub fn vortex_array::extension::datetime::Timestamp::serialize_metadata(&self, m
 
 pub fn vortex_array::extension::datetime::Timestamp::unpack_native<'a>(&self, metadata: &'a Self::Metadata, _storage_dtype: &'a vortex_array::dtype::DType, storage_value: &'a vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
-pub fn vortex_array::extension::datetime::Timestamp::validate_dtype(&self, _metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
+pub fn vortex_array::extension::datetime::Timestamp::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::extension::datetime::Timestamp::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::datetime::Timestamp::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
 
@@ -10310,6 +10318,8 @@ pub fn vortex_array::extension::uuid::Uuid::serialize_metadata(&self, metadata:
 
 pub fn vortex_array::extension::uuid::Uuid::unpack_native<'a>(&self, metadata: &'a Self::Metadata, _storage_dtype: &'a vortex_array::dtype::DType, storage_value: &'a vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::extension::uuid::Uuid::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::extension::uuid::Uuid::validate_dtype(&self, _metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::uuid::Uuid::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
@@ -14142,6 +14152,8 @@ pub fn vortex_array::extension::datetime::Date::serialize_metadata(&self, metada
 
 pub fn vortex_array::extension::datetime::Date::unpack_native(&self, metadata: &Self::Metadata, _storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::extension::datetime::Date::validate_array(&self, _metadata: &Self::Metadata, _storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::extension::datetime::Date::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::datetime::Date::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
@@ -14192,6 +14204,8 @@ pub fn vortex_array::extension::datetime::Time::serialize_metadata(&self, metada
 
 pub fn vortex_array::extension::datetime::Time::unpack_native(&self, metadata: &Self::Metadata, _storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::extension::datetime::Time::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::extension::datetime::Time::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::datetime::Time::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
@@ -14244,7 +14258,9 @@ pub fn vortex_array::extension::datetime::Timestamp::serialize_metadata(&self, m
 
 pub fn vortex_array::extension::datetime::Timestamp::unpack_native<'a>(&self, metadata: &'a Self::Metadata, _storage_dtype: &'a vortex_array::dtype::DType, storage_value: &'a vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
-pub fn vortex_array::extension::datetime::Timestamp::validate_dtype(&self, _metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
+pub fn vortex_array::extension::datetime::Timestamp::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::extension::datetime::Timestamp::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::datetime::Timestamp::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>
 
@@ -14320,6 +14336,8 @@ pub fn vortex_array::extension::uuid::Uuid::serialize_metadata(&self, metadata:
 
 pub fn vortex_array::extension::uuid::Uuid::unpack_native<'a>(&self, metadata: &'a Self::Metadata, _storage_dtype: &'a vortex_array::dtype::DType, storage_value: &'a vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
 
+pub fn vortex_array::extension::uuid::Uuid::validate_array(&self, metadata: &Self::Metadata, storage_array: &dyn vortex_array::DynArray) -> vortex_error::VortexResult<()>
+
 pub fn vortex_array::extension::uuid::Uuid::validate_dtype(&self, _metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
 
 pub fn vortex_array::extension::uuid::Uuid::validate_scalar_value(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType, storage_value: &vortex_array::scalar::ScalarValue) -> vortex_error::VortexResult<()>

vortex-array/src/arrays/extension/array.rs

@@ -75,13 +75,7 @@ impl ExtensionArray {
     ///
     /// Returns an error if the storage array in not compatible with the extension dtype.
     pub fn try_new(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> VortexResult<Self> {
-        // TODO(connor): Replace these statements once we add `validate_storage_array`.
-        // ext_dtype.validate_storage_array(&storage_array)?;
-        assert_eq!(
-            ext_dtype.storage_dtype(),
-            storage_array.dtype(),
-            "ExtensionArray: storage_dtype must match storage array DType",
-        );
+        ext_dtype.validate_storage_array(&storage_array)?;
 
         // SAFETY: we validate that the inputs are valid above.
         Ok(unsafe { Self::new_unchecked(ext_dtype, storage_array) })
@@ -95,16 +89,10 @@ impl ExtensionArray {
     /// other words, they must know that `ext_dtype.validate_storage_array(&storage_array)` has been
     /// called successfully on this storage array.
     pub unsafe fn new_unchecked(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> Self {
-        // TODO(connor): Replace these statements once we add `validate_storage_array`.
-        // #[cfg(debug_assertions)]
-        // ext_dtype
-        //     .validate_storage_array(&storage_array)
-        //     .vortex_expect("[Debug Assertion]: Invalid storage array for `ExtensionArray`");
-        debug_assert_eq!(
-            ext_dtype.storage_dtype(),
-            storage_array.dtype(),
-            "ExtensionArray: storage_dtype must match storage array DType",
-        );
+        #[cfg(debug_assertions)]
+        ext_dtype
+            .validate_storage_array(&storage_array)
+            .vortex_expect("[Debug Assertion]: Invalid storage array for `ExtensionArray`");
 
         Self {
             dtype: DType::Extension(ext_dtype),

vortex-array/src/arrays/extension/compute/rules.rs

@@ -111,6 +111,14 @@ mod tests {
         ) -> VortexResult<Self::NativeValue<'a>> {
             Ok("")
         }
+
+        fn validate_array(
+            &self,
+            _metadata: &Self::Metadata,
+            _storage_array: &dyn DynArray,
+        ) -> VortexResult<()> {
+            Ok(())
+        }
     }
 
     fn test_ext_dtype() -> ExtDTypeRef {
@@ -213,6 +221,14 @@ mod tests {
             ) -> VortexResult<Self::NativeValue<'a>> {
                 Ok("")
             }
+
+            fn validate_array(
+                &self,
+                _metadata: &Self::Metadata,
+                _storage_array: &dyn DynArray,
+            ) -> VortexResult<()> {
+                Ok(())
+            }
         }
 
         let ext_dtype1 = ExtDType::<TestExt>::try_new(

vortex-array/src/dtype/dtype_impl.rs

@@ -487,7 +487,7 @@ mod tests {
             Timestamp::new_with_tz(TimeUnit::Seconds, Some("UTC".into()), Nullable).erased(),
         );
         let t2 = DType::Extension(
-            Timestamp::new_with_tz(TimeUnit::Seconds, Some("ET".into()), Nullable).erased(),
+            Timestamp::new_with_tz(TimeUnit::Seconds, Some("EST".into()), Nullable).erased(),
         );
         assert!(!t1.eq_ignore_nullability(&t2));
     }

vortex-array/src/dtype/extension/erased.rs

@@ -13,6 +13,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 
+use crate::DynArray;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::extension::ExtDType;
@@ -100,6 +101,11 @@ impl ExtDTypeRef {
     pub(crate) fn validate_storage_value(&self, storage_value: &ScalarValue) -> VortexResult<()> {
         self.0.value_validate(storage_value)
     }
+
+    /// Validates that the given storage scalar value is valid for this dtype.
+    pub(crate) fn validate_storage_array(&self, storage_array: &dyn DynArray) -> VortexResult<()> {
+        self.0.array_validate(storage_array)
+    }
 }
 
 /// Methods for downcasting type-erased extension dtypes.

vortex-array/src/dtype/extension/typed.rs

@@ -15,7 +15,9 @@ use std::sync::Arc;
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure_eq;
 
+use crate::DynArray;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::extension::ExtDTypeRef;
@@ -128,11 +130,13 @@ pub(super) trait DynExtDType: 'static + Send + Sync + super::sealed::Sealed {
     fn metadata_serialize(&self) -> VortexResult<Vec<u8>>;
     /// Returns a new [`ExtDTypeRef`] with the given nullability.
     fn with_nullability(&self, nullability: Nullability) -> ExtDTypeRef;
-    /// Validates that the given storage scalar value is valid for this dtype.
+    /// Validates that the given storage scalar value is valid for this extension type.
     fn value_validate(&self, storage_value: &ScalarValue) -> VortexResult<()>;
     /// Formats an extension scalar value using the current dtype for metadata context.
     fn value_display(&self, f: &mut fmt::Formatter<'_>, storage_value: &ScalarValue)
     -> fmt::Result;
+    /// Validates that the given storage array is valid for this extension type.
+    fn array_validate(&self, storage_array: &dyn DynArray) -> VortexResult<()>;
 }
 
 impl<V: ExtVTable> DynExtDType for ExtDTypeInner<V> {
@@ -208,4 +212,14 @@ impl<V: ExtVTable> DynExtDType for ExtDTypeInner<V> {
             ),
         }
     }
+
+    fn array_validate(&self, storage_array: &dyn DynArray) -> VortexResult<()> {
+        vortex_ensure_eq!(
+            &self.storage_dtype,
+            storage_array.dtype(),
+            "tried to construct an extension array with a storage array that has the wrong dtype"
+        );
+
+        self.vtable.validate_array(&self.metadata, storage_array)
+    }
 }

vortex-array/src/dtype/extension/vtable.rs

@@ -7,6 +7,7 @@ use std::hash::Hash;
 
 use vortex_error::VortexResult;
 
+use crate::DynArray;
 use crate::dtype::DType;
 use crate::dtype::extension::ExtId;
 use crate::scalar::ScalarValue;
@@ -36,9 +37,11 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
     fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata>;
 
     /// Validate that the given storage type is compatible with this extension type.
+    ///
+    /// Implementors should additionally validate that the metadata makes sense for the give dtype.
     fn validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &DType) -> VortexResult<()>;
 
-    // Methods related to the extension scalar values.
+    // Methods related to extension scalar values.
 
     /// Validate the given storage value is compatible with the extension type.
     ///
@@ -72,4 +75,17 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
         storage_dtype: &'a DType,
         storage_value: &'a ScalarValue,
     ) -> VortexResult<Self::NativeValue<'a>>;
+
+    // Methods related to extension arrays.
+
+    /// Validates that the given storage array is compatible with this extension type and type
+    /// medatada.
+    ///
+    /// Note that [`ExtVTable::validate_dtype()`] is always called first on the dtype of the storage
+    /// array to validate the storage [`DType`].
+    fn validate_array(
+        &self,
+        metadata: &Self::Metadata,
+        storage_array: &dyn DynArray,
+    ) -> VortexResult<()>;
 }

vortex-array/src/extension/datetime/date.rs

@@ -7,9 +7,10 @@ use jiff::Span;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
-use vortex_error::vortex_ensure;
+use vortex_error::vortex_ensure_eq;
 use vortex_error::vortex_err;
 
+use crate::DynArray;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
@@ -95,11 +96,10 @@ impl ExtVTable for Date {
         let ptype = date_ptype(metadata)
             .ok_or_else(|| vortex_err!("Date type does not support time unit {}", metadata))?;
 
-        vortex_ensure!(
-            storage_dtype.as_ptype() == ptype,
-            "Date storage dtype for {} must be {}",
-            metadata,
-            ptype
+        vortex_ensure_eq!(
+            storage_dtype.as_ptype(),
+            ptype,
+            "Date storage dtype for {metadata} must be {ptype}",
         );
 
         Ok(())
@@ -119,6 +119,15 @@ impl ExtVTable for Date {
             _ => vortex_bail!("Date type does not support time unit {}", metadata),
         }
     }
+
+    fn validate_array(
+        &self,
+        _metadata: &Self::Metadata,
+        _storage_array: &dyn DynArray,
+    ) -> VortexResult<()> {
+        // All i64 and i32 values are valid dates.
+        Ok(())
+    }
 }
 
 #[cfg(test)]