vortex_datetime_parts/
array.rs

1use std::fmt::Debug;
2
3use vortex_array::arrays::StructArray;
4use vortex_array::compute::try_cast;
5use vortex_array::stats::{ArrayStats, StatsSetRef};
6use vortex_array::validity::Validity;
7use vortex_array::variants::ExtensionArrayTrait;
8use vortex_array::vtable::VTableRef;
9use vortex_array::{
10    Array, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl, ArrayVariantsImpl,
11    Encoding, RkyvMetadata,
12};
13use vortex_dtype::DType;
14use vortex_error::{VortexExpect as _, VortexResult, VortexUnwrap, vortex_bail};
15use vortex_mask::Mask;
16
17use crate::serde::DateTimePartsMetadata;
18
19#[derive(Clone, Debug)]
20pub struct DateTimePartsArray {
21    dtype: DType,
22    days: ArrayRef,
23    seconds: ArrayRef,
24    subseconds: ArrayRef,
25    stats_set: ArrayStats,
26}
27
28pub struct DateTimePartsEncoding;
29impl Encoding for DateTimePartsEncoding {
30    type Array = DateTimePartsArray;
31    type Metadata = RkyvMetadata<DateTimePartsMetadata>;
32}
33
34impl DateTimePartsArray {
35    pub fn try_new(
36        dtype: DType,
37        days: ArrayRef,
38        seconds: ArrayRef,
39        subseconds: ArrayRef,
40    ) -> VortexResult<Self> {
41        if !days.dtype().is_int() || (dtype.is_nullable() != days.dtype().is_nullable()) {
42            vortex_bail!(
43                "Expected integer with nullability {}, got {}",
44                dtype.is_nullable(),
45                days.dtype()
46            );
47        }
48        if !seconds.dtype().is_int() || seconds.dtype().is_nullable() {
49            vortex_bail!(MismatchedTypes: "non-nullable integer", seconds.dtype());
50        }
51        if !subseconds.dtype().is_int() || subseconds.dtype().is_nullable() {
52            vortex_bail!(MismatchedTypes: "non-nullable integer", subseconds.dtype());
53        }
54
55        let length = days.len();
56        if length != seconds.len() || length != subseconds.len() {
57            vortex_bail!(
58                "Mismatched lengths {} {} {}",
59                days.len(),
60                seconds.len(),
61                subseconds.len()
62            );
63        }
64
65        Ok(Self {
66            dtype,
67            days,
68            seconds,
69            subseconds,
70            stats_set: Default::default(),
71        })
72    }
73
74    pub fn days(&self) -> &ArrayRef {
75        &self.days
76    }
77
78    pub fn seconds(&self) -> &ArrayRef {
79        &self.seconds
80    }
81
82    pub fn subseconds(&self) -> &ArrayRef {
83        &self.subseconds
84    }
85}
86
87impl ArrayImpl for DateTimePartsArray {
88    type Encoding = DateTimePartsEncoding;
89
90    fn _len(&self) -> usize {
91        self.days.len()
92    }
93
94    fn _dtype(&self) -> &DType {
95        &self.dtype
96    }
97
98    fn _vtable(&self) -> VTableRef {
99        VTableRef::new_ref(&DateTimePartsEncoding)
100    }
101
102    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
103        Self::try_new(
104            self.dtype.clone(),
105            children[0].clone(),
106            children[1].clone(),
107            children[2].clone(),
108        )
109    }
110}
111
112impl ArrayVariantsImpl for DateTimePartsArray {
113    fn _as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
114        Some(self)
115    }
116}
117
118impl ExtensionArrayTrait for DateTimePartsArray {
119    fn storage_data(&self) -> ArrayRef {
120        // FIXME(ngates): this needs to be a tuple array so we can implement Compare
121        // we don't want to write validity twice, so we pull it up to the top
122        let days = try_cast(self.days(), &self.days().dtype().as_nonnullable()).vortex_unwrap();
123        StructArray::try_new(
124            vec!["days".into(), "seconds".into(), "subseconds".into()].into(),
125            [days, self.seconds().clone(), self.subseconds().clone()].into(),
126            self.len(),
127            Validity::copy_from_array(self).vortex_expect("Failed to copy validity"),
128        )
129        .vortex_expect("Failed to create struct array")
130        .into_array()
131    }
132}
133
134impl ArrayStatisticsImpl for DateTimePartsArray {
135    fn _stats_ref(&self) -> StatsSetRef<'_> {
136        self.stats_set.to_ref(self)
137    }
138}
139
140impl ArrayValidityImpl for DateTimePartsArray {
141    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
142        self.days().is_valid(index)
143    }
144
145    fn _all_valid(&self) -> VortexResult<bool> {
146        self.days().all_valid()
147    }
148
149    fn _all_invalid(&self) -> VortexResult<bool> {
150        self.days().all_invalid()
151    }
152
153    fn _validity_mask(&self) -> VortexResult<Mask> {
154        self.days().validity_mask()
155    }
156}
vortex_datetime_parts/array.rs

vortex_datetime_parts/
array.rs
