cast: support vortex.date -> vortex.timestamp extension casts (#28)

phillipleblanc · web-flow · commit 7e5b081518ac · 2026-02-27T13:04:05.000+09:00
diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs
@@ -1,40 +1,187 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_buffer::BufferMut;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_mask::AllOr;
+
 use crate::ArrayRef;
 use crate::IntoArray;
+use crate::ToCanonical;
 use crate::arrays::ExtensionArray;
 use crate::arrays::ExtensionVTable;
+use crate::arrays::PrimitiveArray;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
+use crate::dtype::PType;
+use crate::extension::datetime::AnyTemporal;
+use crate::extension::datetime::TemporalMetadata;
+use crate::extension::datetime::TimeUnit;
 use crate::scalar_fn::fns::cast::CastReduce;
+use crate::vtable::ValidityHelper;
 
 impl CastReduce for ExtensionVTable {
-    fn cast(array: &ExtensionArray, dtype: &DType) -> vortex_error::VortexResult<Option<ArrayRef>> {
-        if !array.dtype().eq_ignore_nullability(dtype) {
+    fn cast(array: &ExtensionArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
+        let DType::Extension(ext_dtype) = dtype else {
             return Ok(None);
+        };
+
+        if array.ext_dtype().eq_ignore_nullability(ext_dtype) {
+            let new_storage = match array
+                .storage()
+                .cast(ext_dtype.storage_dtype().clone())
+                .and_then(|a| a.to_canonical().map(|c| c.into_array()))
+            {
+                Ok(arr) => arr,
+                Err(e) => {
+                    tracing::warn!("Failed to cast storage array: {e}");
+                    return Ok(None);
+                }
+            };
+
+            return Ok(Some(
+                ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(),
+            ));
         }
 
-        let DType::Extension(ext_dtype) = dtype else {
-            unreachable!("Already verified we have an extension dtype");
-        };
+        if let Some(new_storage) = cast_temporal_date_to_timestamp(array, dtype)? {
+            return Ok(Some(
+                ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(),
+            ));
+        }
+
+        Ok(None)
+    }
+}
+
+fn cast_temporal_date_to_timestamp(
+    array: &ExtensionArray,
+    target_dtype: &DType,
+) -> VortexResult<Option<ArrayRef>> {
+    let DType::Extension(target_ext_dtype) = target_dtype else {
+        return Ok(None);
+    };
+
+    let Some(source_temporal) = array.ext_dtype().metadata_opt::<AnyTemporal>() else {
+        return Ok(None);
+    };
+    let Some(target_temporal) = target_ext_dtype.metadata_opt::<AnyTemporal>() else {
+        return Ok(None);
+    };
+
+    let (TemporalMetadata::Date(source_unit), TemporalMetadata::Timestamp(target_unit, _)) =
+        (source_temporal, target_temporal)
+    else {
+        return Ok(None);
+    };
+
+    let source_i64 = array
+        .storage()
+        .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?;
+    let source_i64 = source_i64.to_primitive();
 
-        let new_storage = match array
-            .storage()
-            .cast(ext_dtype.storage_dtype().clone())
-            .and_then(|a| a.to_canonical().map(|c| c.into_array()))
-        {
-            Ok(arr) => arr,
-            Err(e) => {
-                tracing::warn!("Failed to cast storage array: {e}");
-                return Ok(None);
+    let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit)?;
+
+    converted
+        .to_array()
+        .cast(target_ext_dtype.storage_dtype().clone())
+        .map(Some)
+}
+
+fn cast_date_values_to_timestamp(
+    values: &PrimitiveArray,
+    source_unit: TimeUnit,
+    target_unit: TimeUnit,
+) -> VortexResult<PrimitiveArray> {
+    let (multiply, divide) = date_to_timestamp_scale(source_unit, target_unit)?;
+
+    let input = values.as_slice::<i64>();
+    let mut output = BufferMut::with_capacity(input.len());
+    match values.validity_mask()?.bit_buffer() {
+        AllOr::All => {
+            for &value in input {
+                // SAFETY: output has sufficient capacity for all pushed values.
+                unsafe { output.push_unchecked(convert_temporal_value(value, multiply, divide)?) };
             }
-        };
+        }
+        AllOr::None => {
+            for _ in 0..input.len() {
+                // SAFETY: output has sufficient capacity for all pushed values.
+                unsafe { output.push_unchecked(0i64) };
+            }
+        }
+        AllOr::Some(bits) => {
+            for (&value, valid) in input.iter().zip(bits.iter()) {
+                if valid {
+                    // SAFETY: output has sufficient capacity for all pushed values.
+                    unsafe {
+                        output.push_unchecked(convert_temporal_value(value, multiply, divide)?)
+                    };
+                } else {
+                    // SAFETY: output has sufficient capacity for all pushed values.
+                    unsafe { output.push_unchecked(0i64) };
+                }
+            }
+        }
+    }
+
+    Ok(PrimitiveArray::new(
+        output.freeze(),
+        values.validity().clone(),
+    ))
+}
+
+fn date_to_timestamp_scale(
+    source_unit: TimeUnit,
+    target_unit: TimeUnit,
+) -> VortexResult<(i64, i64)> {
+    let source_ns = to_nanoseconds(source_unit)?;
+    let target_ns = to_nanoseconds(target_unit)?;
+
+    if source_ns >= target_ns {
+        let multiply = source_ns / target_ns;
+        return Ok((multiply, 1));
+    }
 
-        Ok(Some(
-            ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(),
-        ))
+    let divide = target_ns / source_ns;
+    Ok((1, divide))
+}
+
+fn to_nanoseconds(unit: TimeUnit) -> VortexResult<i64> {
+    match unit {
+        TimeUnit::Nanoseconds => Ok(1),
+        TimeUnit::Microseconds => Ok(1_000),
+        TimeUnit::Milliseconds => Ok(1_000_000),
+        TimeUnit::Seconds => Ok(1_000_000_000),
+        TimeUnit::Days => Ok(86_400_000_000_000),
+    }
+}
+
+fn convert_temporal_value(value: i64, multiply: i64, divide: i64) -> VortexResult<i64> {
+    let mut scaled = i128::from(value)
+        .checked_mul(i128::from(multiply))
+        .ok_or_else(|| {
+            vortex_error::vortex_err!(
+                Compute: "Date value {value} overflows while scaling to timestamp"
+            )
+        })?;
+
+    if divide != 1 {
+        let divisor = i128::from(divide);
+        if scaled % divisor != 0 {
+            vortex_bail!(
+                Compute: "Date value {value} cannot be represented exactly in target timestamp unit"
+            );
+        }
+        scaled /= divisor;
+    }
+
+    if scaled < i128::from(i64::MIN) || scaled > i128::from(i64::MAX) {
+        vortex_bail!(Compute: "Date value {value} overflows target timestamp range");
     }
+
+    Ok(scaled as i64)
 }
 
 #[cfg(test)]
@@ -45,11 +192,13 @@ mod tests {
     use vortex_buffer::buffer;
 
     use super::*;
+    use crate::Array;
     use crate::IntoArray;
     use crate::arrays::PrimitiveArray;
     use crate::builtins::ArrayBuiltins;
     use crate::compute::conformance::cast::test_cast_conformance;
     use crate::dtype::Nullability;
+    use crate::extension::datetime::Date;
     use crate::extension::datetime::TimeUnit;
     use crate::extension::datetime::Timestamp;
 
@@ -85,6 +234,57 @@ mod tests {
         assert_eq!(output.dtype(), &new_dtype);
     }
 
+    #[test]
+    fn cast_date_days_to_timestamp_nanoseconds() {
+        let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased();
+        let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased();
+
+        let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array());
+        let output = arr
+            .to_array()
+            .cast(DType::Extension(target_dtype.clone()))
+            .unwrap()
+            .to_extension();
+
+        assert_eq!(output.dtype(), &DType::Extension(target_dtype));
+
+        let storage = output.storage().to_primitive();
+        assert_eq!(
+            storage.as_slice::<i64>(),
+            &[0, 86_400_000_000_000, -86_400_000_000_000]
+        );
+    }
+
+    #[test]
+    fn cast_date_days_to_timestamp_seconds_nullable() {
+        let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased();
+        let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased();
+
+        let arr = ExtensionArray::new(
+            source_dtype,
+            PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(),
+        );
+
+        let output = arr
+            .to_array()
+            .cast(DType::Extension(target_dtype.clone()))
+            .unwrap()
+            .to_extension();
+
+        assert_eq!(output.dtype(), &DType::Extension(target_dtype));
+
+        let storage = output.storage().to_primitive();
+        assert_eq!(
+            storage.scalar_at(0).unwrap().as_primitive().as_::<i64>(),
+            Some(0)
+        );
+        assert!(storage.scalar_at(1).unwrap().is_null());
+        assert_eq!(
+            storage.scalar_at(2).unwrap().as_primitive().as_::<i64>(),
+            Some(172_800)
+        );
+    }
+
     #[test]
     fn cast_different_ext_dtype() {
         let original_dtype =
