diff --git a/src/compute/cast/mod.rs b/src/compute/cast/mod.rs index 029b6a3cd9e..71e21fa9a29 100644 --- a/src/compute/cast/mod.rs +++ b/src/compute/cast/mod.rs @@ -290,6 +290,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { (Timestamp(_, _), Date64) => true, (Int64, Duration(_)) => true, (Duration(_), Int64) => true, + (Interval(_), Interval(IntervalUnit::MonthDayNano)) => true, (_, _) => false, } } @@ -842,6 +843,13 @@ pub fn cast(array: &dyn Array, to_type: &DataType, options: CastOptions) -> Resu (Int64, Duration(_)) => primitive_to_same_primitive_dyn::(array, to_type), (Duration(_), Int64) => primitive_to_same_primitive_dyn::(array, to_type), + (Interval(IntervalUnit::DayTime), Interval(IntervalUnit::MonthDayNano)) => { + primitive_dyn!(array, days_ms_to_months_days_ns) + } + (Interval(IntervalUnit::YearMonth), Interval(IntervalUnit::MonthDayNano)) => { + primitive_dyn!(array, months_to_months_days_ns) + } + (_, _) => Err(ArrowError::NotYetImplemented(format!( "Casting from {:?} to {:?} not supported", from_type, to_type, diff --git a/src/compute/cast/primitive_to.rs b/src/compute/cast/primitive_to.rs index cdd1317e364..fea42348c2c 100644 --- a/src/compute/cast/primitive_to.rs +++ b/src/compute/cast/primitive_to.rs @@ -2,7 +2,9 @@ use std::hash::Hash; use num_traits::{AsPrimitive, Float, ToPrimitive}; +use crate::datatypes::IntervalUnit; use crate::error::Result; +use crate::types::{days_ms, months_days_ns}; use crate::{ array::*, bitmap::Bitmap, @@ -551,3 +553,31 @@ pub fn naive_timestamp_to_utf8( } } } + +#[inline] +fn days_ms_to_months_days_ns_scalar(from: days_ms) -> months_days_ns { + months_days_ns::new(0, from.days(), from.milliseconds() as i64 * 1000) +} + +/// Casts [`days_ms`]s to [`months_days_ns`]. This operation is infalible and lossless. +pub fn days_ms_to_months_days_ns(from: &PrimitiveArray) -> PrimitiveArray { + unary( + from, + days_ms_to_months_days_ns_scalar, + DataType::Interval(IntervalUnit::MonthDayNano), + ) +} + +#[inline] +fn months_to_months_days_ns_scalar(from: i32) -> months_days_ns { + months_days_ns::new(from, 0, 0) +} + +/// Casts months represented as [`i32`]s to [`months_days_ns`]. This operation is infalible and lossless. +pub fn months_to_months_days_ns(from: &PrimitiveArray) -> PrimitiveArray { + unary( + from, + months_to_months_days_ns_scalar, + DataType::Interval(IntervalUnit::MonthDayNano), + ) +} diff --git a/tests/it/compute/cast.rs b/tests/it/compute/cast.rs index 9c3d16d7491..501273c094a 100644 --- a/tests/it/compute/cast.rs +++ b/tests/it/compute/cast.rs @@ -1,7 +1,7 @@ use arrow2::array::*; use arrow2::compute::cast::{can_cast_types, cast, CastOptions}; use arrow2::datatypes::*; -use arrow2::types::NativeType; +use arrow2::types::{days_ms, months_days_ns, NativeType}; #[test] fn i32_to_f64() { @@ -512,6 +512,29 @@ fn date32_to_date64() { ); } +#[test] +fn days_ms_to_months_days_ns() { + test_primitive_to_primitive( + &[days_ms::new(1, 1), days_ms::new(1, 2)], + DataType::Interval(IntervalUnit::DayTime), + &[ + months_days_ns::new(0, 1, 1000), + months_days_ns::new(0, 1, 2000), + ], + DataType::Interval(IntervalUnit::MonthDayNano), + ); +} + +#[test] +fn months_to_months_days_ns() { + test_primitive_to_primitive( + &[1, 2], + DataType::Interval(IntervalUnit::YearMonth), + &[months_days_ns::new(1, 0, 0), months_days_ns::new(2, 0, 0)], + DataType::Interval(IntervalUnit::MonthDayNano), + ); +} + #[test] fn date64_to_date32() { test_primitive_to_primitive(