Change extension dtype nullability without validating again

An extension type cannot restrict whether its values can be null, so
`ExtVTable::validate_dtype` must not depend on the top-level nullability
of the storage dtype. `ExtDType::with_nullability` now relies on that
contract and rebuilds the dtype without validating again, so nullability
changes in `take`, `mask` and `min_max` cannot panic.

Also fixes `Debug for ExtDTypeRef`, which printed the metadata before the
struct and then `metadata: Ok(())`.

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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

@@ -180,11 +180,9 @@ impl Hash for ExtDTypeRef {
 
 impl fmt::Debug for ExtDTypeRef {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let metadata = self.0.metadata_debug(f);
-
         f.debug_struct("ExtDType")
             .field("id", &self.id())
-            .field("metadata", &metadata)
+            .field("metadata", &MetadataDebug(&*self.0))
             .field("storage_dtype", &self.storage_dtype())
             .finish()
     }
@@ -206,6 +204,13 @@ impl fmt::Display for ExtDTypeRef {
 
 // Private formatting helpers for Display and Debug impls.
 
+struct MetadataDebug<'a>(&'a dyn DynExtDType);
+impl fmt::Debug for MetadataDebug<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.metadata_debug(f)
+    }
+}
+
 struct MetadataDisplay<'a>(&'a dyn DynExtDType);
 impl fmt::Display for MetadataDisplay<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

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

@@ -55,3 +55,6 @@ mod sealed {
     /// This can be the **only** implementor for [`super::typed::DynExtDType`].
     impl<V: ExtVTable> Sealed for ExtDType<V> {}
 }
+
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,115 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Regression tests for extension dtype nullability changes and debug output.
+
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use rstest::rstest;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+use vortex_session::registry::CachedId;
+
+use super::ExtDType;
+use super::ExtId;
+use super::ExtVTable;
+use crate::dtype::DType;
+use crate::dtype::Nullability;
+use crate::dtype::PType;
+use crate::extension::datetime::TimeUnit;
+use crate::extension::datetime::Timestamp;
+use crate::scalar::ScalarValue;
+
+/// The number of times [`CountingValidator::validate_dtype`] ran.
+static VALIDATIONS: AtomicUsize = AtomicUsize::new(0);
+
+/// An extension type that counts how often its dtype is validated.
+#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
+struct CountingValidator;
+
+impl ExtVTable for CountingValidator {
+    type Metadata = u32;
+    type NativeValue<'a> = &'a ScalarValue;
+
+    fn id(&self) -> ExtId {
+        static ID: CachedId = CachedId::new("test.counting_validator");
+        *ID
+    }
+
+    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
+        Ok(metadata.to_le_bytes().to_vec())
+    }
+
+    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
+        let bytes = metadata
+            .try_into()
+            .map_err(|_| vortex_err!("Expected four metadata bytes, got {}", metadata.len()))?;
+
+        Ok(u32::from_le_bytes(bytes))
+    }
+
+    fn validate_dtype(_ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
+        VALIDATIONS.fetch_add(1, Ordering::Relaxed);
+        Ok(())
+    }
+
+    fn unpack_native<'a>(
+        _ext_dtype: &'a ExtDType<Self>,
+        storage_value: &'a ScalarValue,
+    ) -> VortexResult<Self::NativeValue<'a>> {
+        Ok(storage_value)
+    }
+}
+
+// One test owns `VALIDATIONS`, so test threads cannot race on the counter.
+#[test]
+fn test_nullability_changes_do_not_validate_again() -> VortexResult<()> {
+    for initial in [Nullability::NonNullable, Nullability::Nullable] {
+        VALIDATIONS.store(0, Ordering::Relaxed);
+        let dtype =
+            ExtDType::<CountingValidator>::try_new(42, DType::Primitive(PType::I64, initial))?;
+        assert_eq!(VALIDATIONS.load(Ordering::Relaxed), 1);
+
+        for target in [Nullability::Nullable, Nullability::NonNullable] {
+            let changed = dtype.with_nullability(target);
+            assert_eq!(
+                changed.storage_dtype(),
+                &DType::Primitive(PType::I64, target)
+            );
+            assert_eq!(changed.metadata::<CountingValidator>(), &42);
+            assert_eq!(changed.id(), dtype.id());
+        }
+        assert_eq!(VALIDATIONS.load(Ordering::Relaxed), 1);
+    }
+
+    Ok(())
+}
+
+#[rstest]
+#[case::compact(false)]
+#[case::pretty(true)]
+fn test_debug_preserves_metadata(#[case] pretty: bool) {
+    let dtype = Timestamp::new_with_tz(
+        TimeUnit::Milliseconds,
+        Some("UTC".into()),
+        Nullability::Nullable,
+    )
+    .erased();
+    let output = if pretty {
+        format!("{dtype:#?}")
+    } else {
+        format!("{dtype:?}")
+    };
+
+    assert!(output.starts_with("ExtDType {"), "{output}");
+    assert!(output.contains("metadata: TimestampOptions {"), "{output}");
+    assert!(output.contains("unit: Milliseconds"), "{output}");
+    assert_eq!(
+        output.contains("metadata: TimestampOptions {\n"),
+        pretty,
+        "{output}"
+    );
+    assert_eq!(output.matches("UTC").count(), 1, "{output}");
+    assert!(!output.contains("Ok(())"), "{output}");
+}

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

@@ -12,7 +12,6 @@ use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
 
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use crate::dtype::DType;
@@ -88,14 +87,16 @@ impl<V: ExtVTable> ExtDType<V> {
     }
 
     /// Returns a new [`ExtDTypeRef`] with the given nullability.
+    ///
+    /// This does not validate the dtype again, because [`ExtVTable::validate_dtype`] must not
+    /// depend on the top-level nullability of the storage dtype.
     pub fn with_nullability(&self, nullability: Nullability) -> ExtDTypeRef {
-        let storage_dtype = self.storage_dtype.with_nullability(nullability);
-        ExtDType::<V>::try_with_vtable(self.vtable.clone(), self.metadata.clone(), storage_dtype)
-            .vortex_expect(
-                "Extension DType should not fail validation with the same storage type \
-                 but different nullability",
-            )
-            .erased()
+        Self {
+            vtable: self.vtable.clone(),
+            metadata: self.metadata.clone(),
+            storage_dtype: self.storage_dtype.with_nullability(nullability),
+        }
+        .erased()
     }
 
     /// Serializes the metadata into a byte vector.

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

@@ -46,6 +46,13 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
     ///
     /// This is called when constructing an [`ExtDType`] and should check both storage dtype and
     /// extension metadata.
+    ///
+    /// The result **must not** depend on the top-level nullability of the storage dtype. An
+    /// extension type cannot restrict whether its values can be null, and
+    /// [`ExtDType::with_nullability`] changes the nullability without validating again. Nested
+    /// nullability belongs to the extension type, so the validator can check it. A validator that
+    /// breaks this rule does not cause undefined behavior, but `with_nullability` can then return a
+    /// dtype that the validator rejects.
     fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()>;
 
     // Methods related to the extension scalar values.