vortex-datetime-parts 0.69.0

Vortex physical encoding that compresses temporal components individually
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use vortex_array::ArrayRef;
use vortex_array::ArrayView;
use vortex_array::IntoArray;
use vortex_array::builtins::ArrayBuiltins;
use vortex_array::dtype::DType;
use vortex_array::scalar_fn::fns::cast::CastReduce;
use vortex_error::VortexResult;

use crate::DateTimeParts;
use crate::array::DateTimePartsArrayExt;
impl CastReduce for DateTimeParts {
    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
        if !array.dtype().eq_ignore_nullability(dtype) {
            return Ok(None);
        };

        Ok(Some(
            DateTimeParts::try_new(
                dtype.clone(),
                array
                    .days()
                    .cast(array.days().dtype().with_nullability(dtype.nullability()))?,
                array.seconds().clone(),
                array.subseconds().clone(),
            )?
            .into_array(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;
    use vortex_array::ArrayRef;
    use vortex_array::IntoArray;
    use vortex_array::arrays::PrimitiveArray;
    use vortex_array::arrays::TemporalArray;
    use vortex_array::builtins::ArrayBuiltins;
    use vortex_array::dtype::DType;
    use vortex_array::dtype::Nullability;
    use vortex_array::extension::datetime::TimeUnit;
    use vortex_array::validity::Validity;
    use vortex_buffer::buffer;

    use crate::DateTimeParts;
    use crate::DateTimePartsArray;

    fn date_time_array(validity: Validity) -> ArrayRef {
        DateTimeParts::try_from_temporal(TemporalArray::new_timestamp(
            PrimitiveArray::new(
                buffer![
                    86_400i64,            // element with only day component
                    86_400i64 + 1000,     // element with day + second components
                    86_400i64 + 1000 + 1, // element with day + second + sub-second components
                ],
                validity,
            )
            .into_array(),
            TimeUnit::Milliseconds,
            Some("UTC".into()),
        ))
        .unwrap()
        .into_array()
    }

    #[rstest]
    #[case(Validity::NonNullable, Nullability::Nullable)]
    #[case(Validity::AllValid, Nullability::Nullable)]
    #[case(Validity::AllInvalid, Nullability::Nullable)]
    #[case(Validity::from_iter([true, false, true]), Nullability::Nullable)]
    #[case(Validity::NonNullable, Nullability::NonNullable)]
    #[case(Validity::AllValid, Nullability::NonNullable)]
    #[case(Validity::from_iter([true, true, true]), Nullability::Nullable)]
    fn test_cast_to_compatible_nullability(
        #[case] validity: Validity,
        #[case] cast_to_nullability: Nullability,
    ) {
        let array = date_time_array(validity);
        let new_dtype = array.dtype().with_nullability(cast_to_nullability);
        let result = array.cast(new_dtype.clone());
        assert!(result.is_ok(), "{result:?}");
        assert_eq!(result.unwrap().dtype(), &new_dtype);
    }

    #[rstest]
    #[case(Validity::AllInvalid)]
    #[case(Validity::from_iter([true, false, true]))]
    fn test_bad_cast_fails(#[case] validity: Validity) {
        let array = date_time_array(validity);
        // Cast to incompatible type - force evaluation via to_canonical
        let result = array
            .cast(DType::Bool(Nullability::NonNullable))
            .and_then(|a| a.to_canonical().map(|c| c.into_array()));
        assert!(result.is_err(), "Expected error, got: {result:?}");

        // Cast nullable with nulls to non-nullable - force evaluation via to_canonical
        let result = array
            .cast(array.dtype().with_nullability(Nullability::NonNullable))
            .and_then(|a| a.to_canonical().map(|c| c.into_array()));
        assert!(result.is_err(), "Expected error, got: {result:?}");
    }

    #[rstest]
    #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp(
        buffer![
            0i64,
            86_400_000,  // 1 day in ms
            172_800_000, // 2 days in ms
            259_200_000, // 3 days in ms
            345_600_000, // 4 days in ms
        ].into_array(),
        TimeUnit::Milliseconds,
        Some("UTC".into())
    )).unwrap())]
    #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp(
        PrimitiveArray::from_option_iter([
            Some(0i64),
            None,
            Some(172_800_000), // 2 days in ms
            Some(259_200_000), // 3 days in ms
            None,
        ]).into_array(),
        TimeUnit::Milliseconds,
        Some("UTC".into())
    )).unwrap())]
    #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp(
        buffer![86_400_000_000_000i64].into_array(), // 1 day in ns
        TimeUnit::Nanoseconds,
        Some("UTC".into())
    )).unwrap())]
    fn test_cast_datetime_parts_conformance(#[case] array: DateTimePartsArray) {
        use vortex_array::compute::conformance::cast::test_cast_conformance;
        test_cast_conformance(&array.into_array());
    }
}