Honor source extension coercions

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

@@ -16,9 +16,25 @@ use crate::scalar_fn::fns::cast::CastReduce;
 impl CastReduce for Extension {
     fn cast(array: ArrayView<'_, Extension>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
         if !array.dtype().eq_ignore_nullability(dtype) {
-            // Target is not the same extension type.
-            // Delegate to the storage array's cast.
-            return Ok(Some(array.storage_array().cast(dtype.clone())?));
+            let DType::Extension(target_ext_dtype) = dtype else {
+                return Ok(Some(array.storage_array().cast(dtype.clone())?));
+            };
+
+            let source_ext_dtype = array.dtype().as_extension();
+
+            // `can_coerce_from` may require an extension-specific value conversion. This generic
+            // cast only supports `can_coerce_to`, where casting the storage is sufficient.
+            if !source_ext_dtype.can_coerce_to(dtype) {
+                return Ok(None);
+            }
+
+            let target_storage = array
+                .storage_array()
+                .cast(target_ext_dtype.storage_dtype().clone())?;
+
+            return Ok(Some(
+                ExtensionArray::new(target_ext_dtype.clone(), target_storage).into_array(),
+            ));
         }
 
         let DType::Extension(ext_dtype) = dtype else {
@@ -49,9 +65,11 @@ mod tests {
     use rstest::rstest;
     use vortex_buffer::Buffer;
     use vortex_buffer::buffer;
+    use vortex_error::vortex_ensure;
     use vortex_session::VortexSession;
 
     use super::*;
+    use crate::EmptyMetadata;
     use crate::IntoArray;
     use crate::arrays::PrimitiveArray;
     use crate::assert_arrays_eq;
@@ -60,12 +78,68 @@ mod tests {
     use crate::dtype::DType;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
+    use crate::dtype::extension::ExtDType;
+    use crate::dtype::extension::ExtId;
+    use crate::dtype::extension::ExtVTable;
     use crate::executor::VortexSessionExecute;
     use crate::extension::datetime::TimeUnit;
     use crate::extension::datetime::Timestamp;
+    use crate::scalar::ScalarValue;
 
     static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
 
+    #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+    struct MillisecondTimestamp;
+
+    impl ExtVTable for MillisecondTimestamp {
+        type Metadata = EmptyMetadata;
+        type NativeValue<'a> = &'a ScalarValue;
+
+        #[expect(clippy::disallowed_methods, reason = "test-only extension ID")]
+        fn id(&self) -> ExtId {
+            ExtId::new("vortex.test.millisecond_timestamp")
+        }
+
+        fn serialize_metadata(&self, _metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
+            Ok(Vec::new())
+        }
+
+        fn deserialize_metadata(&self, _metadata: &[u8]) -> VortexResult<Self::Metadata> {
+            Ok(EmptyMetadata)
+        }
+
+        fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
+            vortex_ensure!(
+                matches!(ext_dtype.storage_dtype(), DType::Primitive(PType::I64, _)),
+                "MillisecondTimestamp storage must be i64, got {}",
+                ext_dtype.storage_dtype(),
+            );
+            Ok(())
+        }
+
+        fn can_coerce_to(source: &ExtDType<Self>, target: &DType) -> bool {
+            let Some(target) = target.as_extension_opt() else {
+                return false;
+            };
+            let Some(options) = target.metadata_opt::<Timestamp>() else {
+                return false;
+            };
+
+            options.unit == TimeUnit::Milliseconds
+                && options.tz.is_none()
+                && target
+                    .storage_dtype()
+                    .can_coerce_from(source.storage_dtype())
+        }
+
+        fn unpack_native<'a>(
+            _ext_dtype: &'a ExtDType<Self>,
+            storage_value: &'a ScalarValue,
+        ) -> VortexResult<Self::NativeValue<'a>> {
+            Ok(storage_value)
+        }
+    }
+
     #[test]
     fn cast_same_ext_dtype() {
         let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased();
@@ -119,6 +193,34 @@ mod tests {
         assert!(result.is_err());
     }
 
+    #[test]
+    fn cast_uses_source_extension_coercion() -> VortexResult<()> {
+        let source_dtype = ExtDType::<MillisecondTimestamp>::try_new(
+            EmptyMetadata,
+            DType::Primitive(PType::I64, Nullability::NonNullable),
+        )?
+        .erased();
+        let target_dtype =
+            Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased();
+        let source = ExtensionArray::new(source_dtype, buffer![1i64].into_array()).into_array();
+        let target = DType::Extension(target_dtype);
+        let incompatible_target = DType::Extension(
+            Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased(),
+        );
+
+        assert!(target.can_coerce_from(source.dtype()));
+        assert!(source.dtype().can_coerce_to(&target));
+        assert!(!source.dtype().can_coerce_from(&target));
+        assert!(!target.can_coerce_to(source.dtype()));
+        assert!(!incompatible_target.can_coerce_from(source.dtype()));
+        let result = source
+            .cast(target.clone())?
+            .execute::<ExtensionArray>(&mut SESSION.create_execution_ctx())?;
+
+        assert_eq!(result.dtype(), &target);
+        Ok(())
+    }
+
     #[test]
     fn cast_timestamp_to_i64() -> VortexResult<()> {
         let mut ctx = SESSION.create_execution_ctx();

vortex-array/src/dtype/coercion.rs

@@ -198,93 +198,100 @@ impl DType {
             .try_fold(types[0].clone(), |acc, t| acc.least_supertype(t))
     }
 
-    /// Is there any implicit coercion path from `other` to `self`?
-    pub fn can_coerce_from(&self, other: &DType) -> bool {
+    /// Returns whether `source` has an implicit coercion path to this target dtype.
+    pub fn can_coerce_from(&self, source: &DType) -> bool {
         if let (
             DType::FixedSizeList(target_elem, target_size, _),
             DType::FixedSizeList(source_elem, source_size, _),
-        ) = (self, other)
+        ) = (self, source)
         {
             return target_size == source_size
-                && (self.is_nullable() || !other.is_nullable())
+                && (self.is_nullable() || !source.is_nullable())
                 && target_elem.can_coerce_from(source_elem);
         }
 
-        if let (DType::List(target_elem, _), DType::List(source_elem, _)) = (self, other) {
-            return (self.is_nullable() || !other.is_nullable())
+        if let (DType::List(target_elem, _), DType::List(source_elem, _)) = (self, source) {
+            return (self.is_nullable() || !source.is_nullable())
                 && target_elem.can_coerce_from(source_elem);
         }
 
-        if let (DType::Map(target, _), DType::Map(source, _)) = (self, other) {
-            return (self.is_nullable() || !other.is_nullable())
-                && (!target.keys_sorted() || source.keys_sorted())
-                && target.key_dtype().can_coerce_from(&source.key_dtype())
-                && target.value_dtype().can_coerce_from(&source.value_dtype());
+        if let (DType::Map(target, _), DType::Map(source_map, _)) = (self, source) {
+            return (self.is_nullable() || !source.is_nullable())
+                && (!target.keys_sorted() || source_map.keys_sorted())
+                && target.key_dtype().can_coerce_from(&source_map.key_dtype())
+                && target
+                    .value_dtype()
+                    .can_coerce_from(&source_map.value_dtype());
         }
 
-        if let (DType::Struct(target, _), DType::Struct(source, _)) = (self, other) {
-            return (self.is_nullable() || !other.is_nullable())
-                && target.nfields() == source.nfields()
-                && target.names() == source.names()
+        if let (DType::Struct(target, _), DType::Struct(source_struct, _)) = (self, source) {
+            return (self.is_nullable() || !source.is_nullable())
+                && target.nfields() == source_struct.nfields()
+                && target.names() == source_struct.names()
                 && target
                     .fields()
-                    .zip(source.fields())
+                    .zip(source_struct.fields())
                     .all(|(target, source)| target.can_coerce_from(&source));
         }
 
         // Same type (ignoring nullability): check nullability compatibility
-        if self.eq_ignore_nullability(other) {
-            return self.is_nullable() || !other.is_nullable();
+        if self.eq_ignore_nullability(source) {
+            return self.is_nullable() || !source.is_nullable();
         }
 
         // Null → nullable target
-        if matches!(other, DType::Null) {
+        if matches!(source, DType::Null) {
             return self.is_nullable();
         }
 
         // Bool → numeric
-        if other.is_boolean() && self.is_numeric() {
-            return self.is_nullable() || !other.is_nullable();
+        if source.is_boolean() && self.is_numeric() {
+            return self.is_nullable() || !source.is_nullable();
         }
 
         // Primitive widening: true if least_supertype(source, target) == target
-        if let (DType::Primitive(..), DType::Primitive(..)) = (self, other) {
-            return other
+        if let (DType::Primitive(..), DType::Primitive(..)) = (self, source) {
+            return source
                 .least_supertype(self)
                 .is_some_and(|st| st.eq_ignore_nullability(self))
-                && (self.is_nullable() || !other.is_nullable());
+                && (self.is_nullable() || !source.is_nullable());
         }
 
         // Decimal widening
-        if let (DType::Decimal(target, _), DType::Decimal(source, _)) = (self, other) {
+        if let (DType::Decimal(target, _), DType::Decimal(source_decimal, _)) = (self, source) {
             let target_integral = target.precision() as i16 - target.scale() as i16;
-            let source_integral = source.precision() as i16 - source.scale() as i16;
+            let source_integral = source_decimal.precision() as i16 - source_decimal.scale() as i16;
             return target_integral >= source_integral
-                && target.scale() >= source.scale()
-                && (self.is_nullable() || !other.is_nullable());
+                && target.scale() >= source_decimal.scale()
+                && (self.is_nullable() || !source.is_nullable());
         }
 
         // Integer → Decimal
-        if let (DType::Decimal(dec, _), DType::Primitive(p, _)) = (self, other)
+        if let (DType::Decimal(dec, _), DType::Primitive(p, _)) = (self, source)
             && p.is_int()
         {
             let needed = integer_decimal_precision(*p);
             let integral_digits = dec.precision() as i16 - dec.scale() as i16;
             return integral_digits >= needed as i16
-                && (self.is_nullable() || !other.is_nullable());
+                && (self.is_nullable() || !source.is_nullable());
         }
 
-        // Extension: delegate to vtable
-        if let DType::Extension(ext) = self {
-            return ext.can_coerce_from(other);
+        if let DType::Extension(target_ext) = self
+            && target_ext.can_coerce_from(source)
+        {
+            return true;
+        }
+
+        if let DType::Extension(source_ext) = source {
+            return source_ext.can_coerce_to(self);
         }
 
         false
     }
 
-    /// Convenience — is there a path from `self` to `other`?
-    pub fn can_coerce_to(&self, other: &DType) -> bool {
-        other.can_coerce_from(self)
+    /// Returns whether this source dtype has an implicit coercion path to `target`.
+    pub fn can_coerce_to(&self, target: &DType) -> bool {
+        target.can_coerce_from(self)
     }
 
     /// Are all types in the slice mutually coercible to a common type?

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

@@ -108,14 +108,12 @@ impl ExtDTypeRef {
         self.0.validate_scalar_value(storage_value)
     }
 
-    /// Can a value of `other` be implicitly coerced into this extension type?
-    pub fn can_coerce_from(&self, other: &DType) -> bool {
-        self.0.can_coerce_from(other)
+    pub fn can_coerce_from(&self, source: &DType) -> bool {
+        self.0.can_coerce_from(source)
     }
 
-    /// Can this extension type be implicitly coerced into `other`?
-    pub fn can_coerce_to(&self, other: &DType) -> bool {
-        self.0.can_coerce_to(other)
+    pub fn can_coerce_to(&self, target: &DType) -> bool {
+        self.0.can_coerce_to(target)
     }
 
     /// Compute the least supertype of this extension type and another type.

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

@@ -108,14 +108,12 @@ impl<V: ExtVTable> ExtDType<V> {
         V::validate_scalar_value(self, storage_value)
     }
 
-    /// Can a value of `other` be implicitly coerced into this extension type?
-    pub fn can_coerce_from(&self, other: &DType) -> bool {
-        V::can_coerce_from(self, other)
+    pub fn can_coerce_from(&self, source: &DType) -> bool {
+        V::can_coerce_from(self, source)
     }
 
-    /// Can this extension type be implicitly coerced into `other`?
-    pub fn can_coerce_to(&self, other: &DType) -> bool {
-        V::can_coerce_to(self, other)
+    pub fn can_coerce_to(&self, target: &DType) -> bool {
+        V::can_coerce_to(self, target)
     }
 
     /// Compute the least supertype of this extension type and another type.
@@ -148,8 +146,8 @@ pub(super) trait DynExtDType: 'static + Send + Sync + super::sealed::Sealed {
     fn validate_scalar_value(&self, storage_value: &ScalarValue) -> VortexResult<()>;
     fn value_display(&self, f: &mut fmt::Formatter<'_>, storage_value: &ScalarValue)
     -> fmt::Result;
-    fn can_coerce_from(&self, other: &DType) -> bool;
-    fn can_coerce_to(&self, other: &DType) -> bool;
+    fn can_coerce_from(&self, source: &DType) -> bool;
+    fn can_coerce_to(&self, target: &DType) -> bool;
     fn least_supertype(&self, other: &DType) -> Option<DType>;
 }
 
@@ -221,12 +219,12 @@ impl<V: ExtVTable> DynExtDType for ExtDType<V> {
         }
     }
 
-    fn can_coerce_from(&self, other: &DType) -> bool {
-        self.can_coerce_from(other)
+    fn can_coerce_from(&self, source: &DType) -> bool {
+        self.can_coerce_from(source)
     }
 
-    fn can_coerce_to(&self, other: &DType) -> bool {
-        self.can_coerce_to(other)
+    fn can_coerce_to(&self, target: &DType) -> bool {
+        self.can_coerce_to(target)
     }
 
     fn least_supertype(&self, other: &DType) -> Option<DType> {

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

@@ -22,6 +22,10 @@ use crate::scalar::ScalarValue;
 /// [`DType`] plus metadata. Implementations should keep [`validate_dtype`](Self::validate_dtype)
 /// strict enough that every valid storage scalar can be interpreted by
 /// [`unpack_native`](Self::unpack_native).
+///
+/// The target defines a coercion with [`can_coerce_from`](Self::can_coerce_from), and the source
+/// defines one with [`can_coerce_to`](Self::can_coerce_to). Implementors only need to override one
+/// of these methods.
 pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
     /// Associated type containing the deserialized metadata for this extension type.
     type Metadata: 'static + Send + Sync + Clone + Debug + Display + Eq + Hash;
@@ -48,22 +52,15 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
     /// extension metadata.
     fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()>;
 
-    /// Can a value of `other` be implicitly widened into this type? (e.g. GeographyType might
-    /// accept Point, LineString, etc.)
-    ///
-    /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to`. We have both
-    /// so that either side of the coercion can provide the logic.
-    fn can_coerce_from(ext_dtype: &ExtDType<Self>, other: &DType) -> bool {
-        let _ = (ext_dtype, other);
+    /// Returns whether `source` can be coerced into this target.
+    fn can_coerce_from(target: &ExtDType<Self>, source: &DType) -> bool {
+        let _ = (target, source);
         false
     }
 
-    /// Can this type be implicitly widened into `other`?
-    ///
-    /// Implementors only need to override one of `can_coerce_from` or `can_coerce_to`. We have both
-    /// so that either side of the coercion can provide the logic.
-    fn can_coerce_to(ext_dtype: &ExtDType<Self>, other: &DType) -> bool {
-        let _ = (ext_dtype, other);
+    /// Returns whether this source can be coerced into `target` by casting its storage.
+    fn can_coerce_to(source: &ExtDType<Self>, target: &DType) -> bool {
+        let _ = (source, target);
         false
     }