feat: added validator functions for extension arrays

Pratham1812 · Pratham1812 · commit 5b99fa6fd0c2 · 2025-12-27T22:44:46.000+05:30
Signed-off-by: Pratham Agarwal &lt;agarwalpratham1812@gmail.com&gt;
diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs
@@ -662,6 +662,9 @@ pub fn assert_array_eq(
     rhs: &ArrayRef,
     step: usize,
 ) -> crate::error::VortexFuzzResult<()> {
+    use vortex_array::ToCanonical;
+    use vortex_array::arrays::validator_for_ext_type;
+
     use crate::error::Backtrace;
     use crate::error::VortexFuzzError;
 
@@ -699,7 +702,22 @@ pub fn assert_array_eq(
                 Backtrace::capture(),
             ));
         }
+
+        // Also validate the expected array's domain constraints for extension types
+        if matches!(lhs.dtype(), DType::Extension(..)) {
+            let validator = validator_for_ext_type(lhs.to_extension().ext_dtype());
+            if !validator(&l) {
+                return Err(VortexFuzzError::DomainValidationFailed(
+                    l,
+                    idx,
+                    lhs.clone(),
+                    step,
+                    Backtrace::capture(),
+                ));
+            }
+        }
     }
+
     Ok(())
 }
 
diff --git a/fuzz/src/error.rs b/fuzz/src/error.rs
@@ -59,6 +59,8 @@ pub enum VortexFuzzError {
     LengthMismatch(usize, usize, ArrayRef, ArrayRef, usize, Backtrace),
 
     VortexError(VortexError, Backtrace),
+
+    DomainValidationFailed(Scalar, usize, ArrayRef, usize, Backtrace),
 }
 
 impl Debug for VortexFuzzError {
@@ -121,6 +123,13 @@ impl Display for VortexFuzzError {
                     rhs.display_tree(),
                 )
             }
+            VortexFuzzError::DomainValidationFailed(expected, idx, lhs, step, backtrace) => {
+                write!(
+                    f,
+                    "Domain validation failed:\n  Scalar: {expected}\n  Index: {idx}\n  Array: {}\n  Step: {step}\nBacktrace:\n{backtrace}",
+                    lhs.display_tree(),
+                )
+            }
             VortexFuzzError::VortexError(err, backtrace) => {
                 write!(f, "{err}\nBacktrace:\n{backtrace}")
             }
@@ -137,6 +146,7 @@ impl Error for VortexFuzzError {
             | VortexFuzzError::LengthMismatch(..)
             | VortexFuzzError::ScalarMismatch(..)
             | VortexFuzzError::MinMaxMismatch(..)
+            | VortexFuzzError::DomainValidationFailed(..)
             | VortexFuzzError::DTypeMismatch(..) => None,
         }
     }
diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs
@@ -31,6 +31,7 @@ use crate::IntoArray;
 use crate::ToCanonical;
 use crate::arrays::VarBinArray;
 use crate::arrays::VarBinViewArray;
+use crate::arrays::validator_for_ext_type;
 use crate::builders::ArrayBuilder;
 use crate::builders::DecimalBuilder;
 use crate::builders::FixedSizeListBuilder;
@@ -170,7 +171,6 @@ fn random_array_chunk(
 }
 
 /// Creates a random extension array.
-///
 /// If the `chunk_len` is specified, the length of the array will be equal to the chunk length.
 fn random_extension(
     u: &mut Unstructured,
@@ -179,24 +179,34 @@ fn random_extension(
 ) -> Result<ArrayRef> {
     use crate::builders::ExtensionBuilder;
 
+    // Get the validator for this extension type
+    let validator = validator_for_ext_type(ext_dtype);
+
     // Determine array length
     let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
 
     // Create builder for the extension array
     let mut builder = ExtensionBuilder::with_capacity(ext_dtype.clone(), array_length);
 
-    // Generate random values
+    // Generate random values, retrying if they don't pass validation
     for _ in 0..array_length {
-        // Wrap in extension scalar using Scalar::extension()
-        let ext_scalar = Scalar::extension(
-            ext_dtype.clone(),
-            random_scalar(u, ext_dtype.storage_dtype())?,
-        );
-
-        // Append to builder
-        builder
-            .append_scalar(&ext_scalar)
-            .vortex_expect("can append extension scalar");
+        // Retry loop to generate valid values
+        loop {
+            // Generate a random storage scalar
+            let storage_scalar = random_scalar(u, ext_dtype.storage_dtype())?;
+
+            // Wrap it in an extension scalar
+            let ext_scalar = Scalar::extension(ext_dtype.clone(), storage_scalar);
+
+            // Validate the scalar
+            if validator(&ext_scalar) {
+                // Valid value - append and break
+                builder
+                    .append_scalar(&ext_scalar)
+                    .vortex_expect("can append extension scalar");
+                break;
+            }
+        }
     }
 
     Ok(builder.finish())
diff --git a/vortex-array/src/arrays/extension/mod.rs b/vortex-array/src/arrays/extension/mod.rs
@@ -8,3 +8,6 @@ mod compute;
 
 mod vtable;
 pub use vtable::ExtensionVTable;
+
+mod validate;
+pub use validate::validator_for_ext_type;
diff --git a/vortex-array/src/arrays/extension/validate.rs b/vortex-array/src/arrays/extension/validate.rs
@@ -0,0 +1,138 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Domain validation for extension types.
+//!
+//! This module provides validators to ensure that extension array values
+//! are always within the valid domain for their type. For example, temporal
+//! arrays should not overflow when converted to Jiff types.
+
+use vortex_dtype::ExtDType;
+use vortex_dtype::PType;
+use vortex_dtype::datetime::TemporalMetadata;
+use vortex_dtype::datetime::is_temporal_ext_type;
+use vortex_error::VortexExpect;
+use vortex_scalar::Scalar;
+
+/// Type alias for a domain validator function.
+///
+/// A domain validator checks whether a scalar value is valid for a given extension type.
+/// For example, temporal extension types validate that values don't overflow when converted to Jiff types.
+pub type DomainValidator = Box<dyn Fn(&Scalar) -> bool + Send + Sync>;
+
+/// Creates a domain validator for the given extension type.
+///
+/// This function returns a validator that checks if scalar values are in the valid domain
+/// for the extension type. For temporal types (date, time, timestamp), it validates that
+/// the values can be successfully converted to Jiff types without overflow.
+///
+/// # Examples
+///
+/// ```
+/// use std::sync::Arc;
+/// use vortex_array::arrays::extension::validator_for_ext_type;
+/// use vortex_dtype::{ExtDType, ExtMetadata, DType, PType, Nullability};
+/// use vortex_dtype::datetime::{TemporalMetadata, TimeUnit, DATE_ID};
+/// use vortex_scalar::Scalar;
+///
+/// let metadata: ExtMetadata = TemporalMetadata::Date(TimeUnit::Days).into();
+/// let ext_dtype = ExtDType::new(
+///     DATE_ID.clone(),
+///     Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
+///     Some(metadata),
+/// );
+///
+/// let validator = validator_for_ext_type(&ext_dtype);
+///
+/// // Valid date value
+/// let valid_scalar = Scalar::extension(
+///     Arc::new(ext_dtype.clone()),
+///     Scalar::primitive(18000i32, Nullability::NonNullable),
+/// );
+/// assert!(validator(&valid_scalar));
+///
+/// // Null is always valid
+/// let null_scalar = Scalar::null(DType::Extension(Arc::new(ext_dtype)));
+/// assert!(validator(&null_scalar));
+/// ```
+pub fn validator_for_ext_type(ext_dtype: &ExtDType) -> DomainValidator {
+    if is_temporal_ext_type(ext_dtype.id()) {
+        // For temporal types, validate that the value can be converted to Jiff
+        let metadata = TemporalMetadata::try_from(ext_dtype)
+            .vortex_expect("temporal ext_dtype should have valid metadata");
+
+        Box::new(move |scalar: &Scalar| {
+            if scalar.is_null() {
+                return true;
+            }
+
+            // Extract the storage value and validate it can be converted to Jiff
+            let ext_scalar = scalar.as_extension();
+            let storage = ext_scalar.storage();
+            let primitive = storage.as_primitive();
+
+            // Get the i64 value from the primitive (temporal types use i32 or i64)
+            let value = match primitive.ptype() {
+                PType::I32 => primitive.typed_value::<i32>().map(|v| v as i64),
+                PType::I64 => primitive.typed_value::<i64>(),
+                _ => None,
+            };
+
+            value.map(|v| metadata.to_jiff(v).is_ok()).unwrap_or(false)
+        })
+    } else {
+        // Unknown extension type - accept all values
+        Box::new(|_| true)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use vortex_dtype::DType;
+    use vortex_dtype::ExtDType;
+    use vortex_dtype::ExtMetadata;
+    use vortex_dtype::Nullability;
+    use vortex_dtype::PType;
+    use vortex_dtype::datetime::DATE_ID;
+    use vortex_dtype::datetime::TemporalMetadata;
+    use vortex_dtype::datetime::TimeUnit;
+    use vortex_scalar::Scalar;
+
+    use super::*;
+
+    #[test]
+    fn test_temporal_validator_accepts_valid_values() {
+        let metadata: ExtMetadata = TemporalMetadata::Date(TimeUnit::Days).into();
+        let ext_dtype = ExtDType::new(
+            DATE_ID.clone(),
+            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
+            Some(metadata),
+        );
+
+        let validator = validator_for_ext_type(&ext_dtype);
+
+        // Valid date (days since epoch)
+        let valid_scalar = Scalar::extension(
+            Arc::new(ext_dtype.clone()),
+            Scalar::primitive(18000i32, Nullability::NonNullable),
+        );
+        assert!(validator(&valid_scalar));
+    }
+
+    #[test]
+    fn test_temporal_validator_accepts_null() {
+        let metadata: ExtMetadata = TemporalMetadata::Date(TimeUnit::Days).into();
+        let ext_dtype = ExtDType::new(
+            DATE_ID.clone(),
+            Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
+            Some(metadata),
+        );
+
+        let validator = validator_for_ext_type(&ext_dtype);
+
+        let null_scalar = Scalar::null(DType::Extension(Arc::new(ext_dtype)));
+        assert!(validator(&null_scalar));
+    }
+}
diff --git a/vortex-dtype/src/arbitrary/mod.rs b/vortex-dtype/src/arbitrary/mod.rs
@@ -10,13 +10,17 @@ use arbitrary::Unstructured;
 use crate::DType;
 use crate::DecimalDType;
 use crate::ExtDType;
-use crate::ExtID;
 use crate::FieldName;
 use crate::FieldNames;
 use crate::NativeDecimalType;
 use crate::Nullability;
 use crate::PType;
 use crate::StructFields;
+use crate::datetime::DATE_ID;
+use crate::datetime::TIME_ID;
+use crate::datetime::TIMESTAMP_ID;
+use crate::datetime::TemporalMetadata;
+use crate::datetime::TimeUnit;
 use crate::i256;
 mod decimal;
 
@@ -112,24 +116,65 @@ impl<'a> Arbitrary<'a> for ExtDType {
 }
 
 fn random_ext_dtype(u: &mut Unstructured<'_>, _depth: u8) -> Result<ExtDType> {
-    let id_str = u.choose(&[
-        "test.ext",
-        "example.currency",
-        "example.uuid",
-        "example.ipv4",
-    ])?;
-    let ext_id = ExtID::from(*id_str);
-
-    // Supports only base types for now
-    let storage_dtype = match u.int_in_range(1..=3)? {
-        // base types
-        1 => DType::Bool(u.arbitrary()?),
-        2 => DType::Primitive(u.arbitrary()?, u.arbitrary()?),
-        3 => DType::Decimal(u.arbitrary()?, u.arbitrary()?),
-        _ => unreachable!("int_in_range(1..=3) returned value out of range"),
-    };
-
-    Ok(ExtDType::new(ext_id, storage_dtype.into(), None))
+    let choice = u.int_in_range(0..=2)?;
+
+    match choice {
+        0 => {
+            // DATE: i32 (Days) or i64 (Milliseconds)
+            let (ptype, time_unit) = match u.int_in_range(0..=1)? {
+                0 => (PType::I32, TimeUnit::Days),
+                1 => (PType::I64, TimeUnit::Milliseconds),
+                _ => unreachable!(),
+            };
+
+            Ok(ExtDType::new(
+                DATE_ID.clone(),
+                DType::Primitive(ptype, u.arbitrary()?).into(),
+                Some(TemporalMetadata::Date(time_unit).into()),
+            ))
+        }
+        1 => {
+            // TIME: i32 for Seconds/Milliseconds, i64 for Microseconds/Nanoseconds
+            let (ptype, time_unit) = match u.int_in_range(0..=3)? {
+                0 => (PType::I32, TimeUnit::Seconds),
+                1 => (PType::I32, TimeUnit::Milliseconds),
+                2 => (PType::I64, TimeUnit::Microseconds),
+                3 => (PType::I64, TimeUnit::Nanoseconds),
+                _ => unreachable!(),
+            };
+
+            Ok(ExtDType::new(
+                TIME_ID.clone(),
+                DType::Primitive(ptype, u.arbitrary()?).into(),
+                Some(TemporalMetadata::Time(time_unit).into()),
+            ))
+        }
+        2 => {
+            // TIMESTAMP: always i64 with time unit and optional timezone
+            let time_unit = match u.int_in_range(0..=3)? {
+                0 => TimeUnit::Seconds,
+                1 => TimeUnit::Milliseconds,
+                2 => TimeUnit::Microseconds,
+                3 => TimeUnit::Nanoseconds,
+                _ => unreachable!(),
+            };
+
+            let time_zone = u
+                .arbitrary::<bool>()?
+                .then(|| {
+                    u.choose(&["UTC", "America/New_York", "Europe/London", "Asia/Tokyo"])
+                        .map(|s| s.to_string())
+                })
+                .transpose()?;
+
+            Ok(ExtDType::new(
+                TIMESTAMP_ID.clone(),
+                DType::Primitive(PType::I64, u.arbitrary()?).into(),
+                Some(TemporalMetadata::Timestamp(time_unit, time_zone).into()),
+            ))
+        }
+        _ => unreachable!(),
+    }
 }
 
 impl<'a> Arbitrary<'a> for StructFields {

Original file line number	Diff line number	Diff line change
`@@ -59,6 +59,8 @@ pub enum VortexFuzzError {`
`59`	`59`	`LengthMismatch(usize, usize, ArrayRef, ArrayRef, usize, Backtrace),`
`60`	`60`
`61`	`61`	`VortexError(VortexError, Backtrace),`
	`62`	`+`
	`63`	`+ DomainValidationFailed(Scalar, usize, ArrayRef, usize, Backtrace),`
`62`	`64`	`}`
`63`	`65`
`64`	`66`	`impl Debug for VortexFuzzError {`
`@@ -121,6 +123,13 @@ impl Display for VortexFuzzError {`
`121`	`123`	`rhs.display_tree(),`
`122`	`124`	`)`
`123`	`125`	`}`
	`126`	`+ VortexFuzzError::DomainValidationFailed(expected, idx, lhs, step, backtrace) => {`
	`127`	`+ write!(`
	`128`	`+ f,`
	`129`	`+ "Domain validation failed:\n Scalar: {expected}\n Index: {idx}\n Array: {}\n Step: {step}\nBacktrace:\n{backtrace}",`
	`130`	`+ lhs.display_tree(),`
	`131`	`+ )`
	`132`	`+ }`
`124`	`133`	`VortexFuzzError::VortexError(err, backtrace) => {`
`125`	`134`	`write!(f, "{err}\nBacktrace:\n{backtrace}")`
`126`	`135`	`}`
`@@ -137,6 +146,7 @@ impl Error for VortexFuzzError {`
`137`	`146`	`\| VortexFuzzError::LengthMismatch(..)`
`138`	`147`	`\| VortexFuzzError::ScalarMismatch(..)`
`139`	`148`	`\| VortexFuzzError::MinMaxMismatch(..)`
	`149`	`+ \| VortexFuzzError::DomainValidationFailed(..)`
`140`	`150`	`\| VortexFuzzError::DTypeMismatch(..) => None,`
`141`	`151`	`}`
`142`	`152`	`}`