vortex_datetime_parts/compute/
cast.rs

1use vortex_array::compute::{CastKernel, CastKernelAdapter, cast};
2use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
3use vortex_dtype::DType;
4use vortex_error::{VortexResult, vortex_bail};
5
6use crate::{DateTimePartsArray, DateTimePartsVTable};
7
8impl CastKernel for DateTimePartsVTable {
9    fn cast(&self, array: &DateTimePartsArray, dtype: &DType) -> VortexResult<ArrayRef> {
10        if !array.dtype().eq_ignore_nullability(dtype) {
11            vortex_bail!("cannot cast from {} to {}", array.dtype(), dtype);
12        };
13
14        Ok(DateTimePartsArray::try_new(
15            dtype.clone(),
16            cast(
17                array.days().as_ref(),
18                &array.days().dtype().with_nullability(dtype.nullability()),
19            )?,
20            array.seconds().clone(),
21            array.subseconds().clone(),
22        )?
23        .into_array())
24    }
25}
26
27register_kernel!(CastKernelAdapter(DateTimePartsVTable).lift());
28
29#[cfg(test)]
30mod tests {
31    use rstest::rstest;
32    use vortex_array::arrays::{PrimitiveArray, TemporalArray};
33    use vortex_array::compute::cast;
34    use vortex_array::validity::Validity;
35    use vortex_array::{Array, ArrayRef, IntoArray};
36    use vortex_buffer::buffer;
37    use vortex_dtype::datetime::TimeUnit;
38    use vortex_dtype::{DType, Nullability};
39
40    use crate::DateTimePartsArray;
41
42    fn date_time_array(validity: Validity) -> ArrayRef {
43        DateTimePartsArray::try_from(TemporalArray::new_timestamp(
44            PrimitiveArray::new(
45                buffer![
46                    86_400i64,            // element with only day component
47                    86_400i64 + 1000,     // element with day + second components
48                    86_400i64 + 1000 + 1, // element with day + second + sub-second components
49                ],
50                validity,
51            )
52            .into_array(),
53            TimeUnit::Ms,
54            Some("UTC".to_string()),
55        ))
56        .unwrap()
57        .into_array()
58    }
59
60    #[rstest]
61    #[case(Validity::NonNullable, Nullability::Nullable)]
62    #[case(Validity::AllValid, Nullability::Nullable)]
63    #[case(Validity::AllInvalid, Nullability::Nullable)]
64    #[case(Validity::from_iter([true, false, true]), Nullability::Nullable)]
65    #[case(Validity::NonNullable, Nullability::NonNullable)]
66    #[case(Validity::AllValid, Nullability::NonNullable)]
67    #[case(Validity::from_iter([true, true, true]), Nullability::Nullable)]
68    fn test_cast_to_compatibile_nullability(
69        #[case] validity: Validity,
70        #[case] cast_to_nullability: Nullability,
71    ) {
72        let array = date_time_array(validity);
73        let new_dtype = array.dtype().with_nullability(cast_to_nullability);
74        let result = cast(&array, &new_dtype);
75        assert!(result.is_ok(), "{result:?}");
76        assert_eq!(result.unwrap().dtype(), &new_dtype);
77    }
78
79    #[rstest]
80    #[case(Validity::AllInvalid)]
81    #[case(Validity::from_iter([true, false, true]))]
82    fn test_bad_cast_fails(#[case] validity: Validity) {
83        let array = date_time_array(validity);
84        let result = cast(&array, &DType::Bool(Nullability::NonNullable));
85        assert!(
86            result
87                .as_ref()
88                .is_err_and(|err| err.to_string().contains("cannot cast from")),
89            "{result:?}"
90        );
91
92        let result = cast(
93            &array,
94            &array.dtype().with_nullability(Nullability::NonNullable),
95        );
96        assert!(
97            result.as_ref().is_err_and(|err| err
98                .to_string()
99                .contains("invalid cast from nullable to non-nullable")),
100            "{result:?}"
101        );
102    }
103}
vortex_datetime_parts/compute/cast.rs

vortex_datetime_parts/compute/
cast.rs
