diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index e080eb7f04c..be8f9ee29d9 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/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> { 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 = 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> { + Ok(Vec::new()) + } + + fn deserialize_metadata(&self, _metadata: &[u8]) -> VortexResult { + Ok(EmptyMetadata) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> 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, target: &DType) -> bool { + let Some(target) = target.as_extension_opt() else { + return false; + }; + let Some(options) = target.metadata_opt::() 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, + storage_value: &'a ScalarValue, + ) -> VortexResult> { + 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::::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::(&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(); diff --git a/vortex-array/src/dtype/coercion.rs b/vortex-array/src/dtype/coercion.rs index 5021c1ae013..9b088f2ba5c 100644 --- a/vortex-array/src/dtype/coercion.rs +++ b/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? diff --git a/vortex-array/src/dtype/extension/erased.rs b/vortex-array/src/dtype/extension/erased.rs index a80202ba4e1..2e18c3fbe17 100644 --- a/vortex-array/src/dtype/extension/erased.rs +++ b/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. diff --git a/vortex-array/src/dtype/extension/typed.rs b/vortex-array/src/dtype/extension/typed.rs index 5b176bbb34d..50adde5cfd9 100644 --- a/vortex-array/src/dtype/extension/typed.rs +++ b/vortex-array/src/dtype/extension/typed.rs @@ -108,14 +108,12 @@ impl ExtDType { 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; } @@ -221,12 +219,12 @@ impl DynExtDType for ExtDType { } } - 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 { diff --git a/vortex-array/src/dtype/extension/vtable.rs b/vortex-array/src/dtype/extension/vtable.rs index c89f296a22b..4425d5d0f02 100644 --- a/vortex-array/src/dtype/extension/vtable.rs +++ b/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) -> 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, other: &DType) -> bool { - let _ = (ext_dtype, other); + /// Returns whether `source` can be coerced into this target. + fn can_coerce_from(target: &ExtDType, 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, 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, target: &DType) -> bool { + let _ = (source, target); false } diff --git a/vortex-array/src/extension/datetime/date.rs b/vortex-array/src/extension/datetime/date.rs index 9eda23dd2d1..cf412565470 100644 --- a/vortex-array/src/extension/datetime/date.rs +++ b/vortex-array/src/extension/datetime/date.rs @@ -109,16 +109,16 @@ impl ExtVTable for Date { Ok(()) } - fn can_coerce_from(ext_dtype: &ExtDType, other: &DType) -> bool { - let DType::Extension(other_ext) = other else { + fn can_coerce_from(target: &ExtDType, source: &DType) -> bool { + let DType::Extension(source_ext) = source else { return false; }; - let Some(other_unit) = other_ext.metadata_opt::() else { + let Some(source_unit) = source_ext.metadata_opt::() else { return false; }; - let our_unit = ext_dtype.metadata(); - // We can coerce from other if our unit is finer (<=) and nullability is compatible. - our_unit <= other_unit && (ext_dtype.storage_dtype().is_nullable() || !other.is_nullable()) + + target.metadata() <= source_unit + && (target.storage_dtype().is_nullable() || !source.is_nullable()) } fn least_supertype(ext_dtype: &ExtDType, other: &DType) -> Option { diff --git a/vortex-array/src/extension/datetime/time.rs b/vortex-array/src/extension/datetime/time.rs index 88455a121b5..b122defeaf4 100644 --- a/vortex-array/src/extension/datetime/time.rs +++ b/vortex-array/src/extension/datetime/time.rs @@ -110,15 +110,16 @@ impl ExtVTable for Time { Ok(()) } - fn can_coerce_from(ext_dtype: &ExtDType, other: &DType) -> bool { - let DType::Extension(other_ext) = other else { + fn can_coerce_from(target: &ExtDType, source: &DType) -> bool { + let DType::Extension(source_ext) = source else { return false; }; - let Some(other_unit) = other_ext.metadata_opt::