vortex_datetime_parts/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::hash::Hash;
6
7use vortex_array::ArrayEq;
8use vortex_array::ArrayHash;
9use vortex_array::ArrayRef;
10use vortex_array::DeserializeMetadata;
11use vortex_array::DynArray;
12use vortex_array::ExecutionCtx;
13use vortex_array::ExecutionStep;
14use vortex_array::IntoArray;
15use vortex_array::Precision;
16use vortex_array::ProstMetadata;
17use vortex_array::SerializeMetadata;
18use vortex_array::buffer::BufferHandle;
19use vortex_array::dtype::DType;
20use vortex_array::dtype::Nullability;
21use vortex_array::dtype::PType;
22use vortex_array::serde::ArrayChildren;
23use vortex_array::stats::ArrayStats;
24use vortex_array::stats::StatsSetRef;
25use vortex_array::vtable;
26use vortex_array::vtable::ArrayId;
27use vortex_array::vtable::VTable;
28use vortex_array::vtable::ValidityChild;
29use vortex_array::vtable::ValidityVTableFromChild;
30use vortex_error::VortexExpect;
31use vortex_error::VortexResult;
32use vortex_error::vortex_bail;
33use vortex_error::vortex_ensure;
34use vortex_error::vortex_err;
35use vortex_error::vortex_panic;
36use vortex_session::VortexSession;
37
38use crate::canonical::decode_to_temporal;
39use crate::compute::kernel::PARENT_KERNELS;
40use crate::compute::rules::PARENT_RULES;
41
42vtable!(DateTimeParts);
43
44#[derive(Clone, prost::Message)]
45#[repr(C)]
46pub struct DateTimePartsMetadata {
47    // Validity lives in the days array
48    // TODO(ngates): we should actually model this with a Tuple array when we have one.
49    #[prost(enumeration = "PType", tag = "1")]
50    pub days_ptype: i32,
51    #[prost(enumeration = "PType", tag = "2")]
52    pub seconds_ptype: i32,
53    #[prost(enumeration = "PType", tag = "3")]
54    pub subseconds_ptype: i32,
55}
56
57impl DateTimePartsMetadata {
58    pub fn get_days_ptype(&self) -> VortexResult<PType> {
59        PType::try_from(self.days_ptype)
60            .map_err(|_| vortex_err!("Invalid PType {}", self.days_ptype))
61    }
62
63    pub fn get_seconds_ptype(&self) -> VortexResult<PType> {
64        PType::try_from(self.seconds_ptype)
65            .map_err(|_| vortex_err!("Invalid PType {}", self.seconds_ptype))
66    }
67
68    pub fn get_subseconds_ptype(&self) -> VortexResult<PType> {
69        PType::try_from(self.subseconds_ptype)
70            .map_err(|_| vortex_err!("Invalid PType {}", self.subseconds_ptype))
71    }
72}
73
74impl VTable for DateTimePartsVTable {
75    type Array = DateTimePartsArray;
76
77    type Metadata = ProstMetadata<DateTimePartsMetadata>;
78    type OperationsVTable = Self;
79    type ValidityVTable = ValidityVTableFromChild;
80
81    fn id(_array: &Self::Array) -> ArrayId {
82        Self::ID
83    }
84
85    fn len(array: &DateTimePartsArray) -> usize {
86        array.days.len()
87    }
88
89    fn dtype(array: &DateTimePartsArray) -> &DType {
90        &array.dtype
91    }
92
93    fn stats(array: &DateTimePartsArray) -> StatsSetRef<'_> {
94        array.stats_set.to_ref(array.as_ref())
95    }
96
97    fn array_hash<H: std::hash::Hasher>(
98        array: &DateTimePartsArray,
99        state: &mut H,
100        precision: Precision,
101    ) {
102        array.dtype.hash(state);
103        array.days.array_hash(state, precision);
104        array.seconds.array_hash(state, precision);
105        array.subseconds.array_hash(state, precision);
106    }
107
108    fn array_eq(
109        array: &DateTimePartsArray,
110        other: &DateTimePartsArray,
111        precision: Precision,
112    ) -> bool {
113        array.dtype == other.dtype
114            && array.days.array_eq(&other.days, precision)
115            && array.seconds.array_eq(&other.seconds, precision)
116            && array.subseconds.array_eq(&other.subseconds, precision)
117    }
118
119    fn nbuffers(_array: &DateTimePartsArray) -> usize {
120        0
121    }
122
123    fn buffer(_array: &DateTimePartsArray, idx: usize) -> BufferHandle {
124        vortex_panic!("DateTimePartsArray buffer index {idx} out of bounds")
125    }
126
127    fn buffer_name(_array: &DateTimePartsArray, idx: usize) -> Option<String> {
128        vortex_panic!("DateTimePartsArray buffer_name index {idx} out of bounds")
129    }
130
131    fn nchildren(_array: &DateTimePartsArray) -> usize {
132        3
133    }
134
135    fn child(array: &DateTimePartsArray, idx: usize) -> ArrayRef {
136        match idx {
137            0 => array.days().clone(),
138            1 => array.seconds().clone(),
139            2 => array.subseconds().clone(),
140            _ => vortex_panic!("DateTimePartsArray child index {idx} out of bounds"),
141        }
142    }
143
144    fn child_name(_array: &DateTimePartsArray, idx: usize) -> String {
145        match idx {
146            0 => "days".to_string(),
147            1 => "seconds".to_string(),
148            2 => "subseconds".to_string(),
149            _ => vortex_panic!("DateTimePartsArray child_name index {idx} out of bounds"),
150        }
151    }
152
153    fn metadata(array: &DateTimePartsArray) -> VortexResult<Self::Metadata> {
154        Ok(ProstMetadata(DateTimePartsMetadata {
155            days_ptype: PType::try_from(array.days().dtype())? as i32,
156            seconds_ptype: PType::try_from(array.seconds().dtype())? as i32,
157            subseconds_ptype: PType::try_from(array.subseconds().dtype())? as i32,
158        }))
159    }
160
161    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
162        Ok(Some(metadata.serialize()))
163    }
164
165    fn deserialize(
166        bytes: &[u8],
167        _dtype: &DType,
168        _len: usize,
169        _buffers: &[BufferHandle],
170        _session: &VortexSession,
171    ) -> VortexResult<Self::Metadata> {
172        Ok(ProstMetadata(
173            <ProstMetadata<DateTimePartsMetadata> as DeserializeMetadata>::deserialize(bytes)?,
174        ))
175    }
176
177    fn build(
178        dtype: &DType,
179        len: usize,
180        metadata: &Self::Metadata,
181        _buffers: &[BufferHandle],
182        children: &dyn ArrayChildren,
183    ) -> VortexResult<DateTimePartsArray> {
184        if children.len() != 3 {
185            vortex_bail!(
186                "Expected 3 children for datetime-parts encoding, found {}",
187                children.len()
188            )
189        }
190
191        let days = children.get(
192            0,
193            &DType::Primitive(metadata.0.get_days_ptype()?, dtype.nullability()),
194            len,
195        )?;
196        let seconds = children.get(
197            1,
198            &DType::Primitive(metadata.0.get_seconds_ptype()?, Nullability::NonNullable),
199            len,
200        )?;
201        let subseconds = children.get(
202            2,
203            &DType::Primitive(metadata.0.get_subseconds_ptype()?, Nullability::NonNullable),
204            len,
205        )?;
206
207        DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)
208    }
209
210    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
211        vortex_ensure!(
212            children.len() == 3,
213            "DateTimePartsArray expects exactly 3 children (days, seconds, subseconds), got {}",
214            children.len()
215        );
216
217        let mut children_iter = children.into_iter();
218        array.days = children_iter.next().vortex_expect("checked");
219        array.seconds = children_iter.next().vortex_expect("checked");
220        array.subseconds = children_iter.next().vortex_expect("checked");
221
222        Ok(())
223    }
224
225    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
226        Ok(ExecutionStep::Done(
227            decode_to_temporal(array, ctx)?.into_array(),
228        ))
229    }
230
231    fn reduce_parent(
232        array: &Self::Array,
233        parent: &ArrayRef,
234        child_idx: usize,
235    ) -> VortexResult<Option<ArrayRef>> {
236        PARENT_RULES.evaluate(array, parent, child_idx)
237    }
238
239    fn execute_parent(
240        array: &Self::Array,
241        parent: &ArrayRef,
242        child_idx: usize,
243        ctx: &mut ExecutionCtx,
244    ) -> VortexResult<Option<ArrayRef>> {
245        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
246    }
247}
248
249#[derive(Clone, Debug)]
250pub struct DateTimePartsArray {
251    dtype: DType,
252    days: ArrayRef,
253    seconds: ArrayRef,
254    subseconds: ArrayRef,
255    stats_set: ArrayStats,
256}
257
258#[derive(Clone, Debug)]
259pub struct DateTimePartsArrayParts {
260    pub dtype: DType,
261    pub days: ArrayRef,
262    pub seconds: ArrayRef,
263    pub subseconds: ArrayRef,
264}
265
266#[derive(Debug)]
267pub struct DateTimePartsVTable;
268
269impl DateTimePartsVTable {
270    pub const ID: ArrayId = ArrayId::new_ref("vortex.datetimeparts");
271}
272
273impl DateTimePartsArray {
274    pub fn try_new(
275        dtype: DType,
276        days: ArrayRef,
277        seconds: ArrayRef,
278        subseconds: ArrayRef,
279    ) -> VortexResult<Self> {
280        if !days.dtype().is_int() || (dtype.is_nullable() != days.dtype().is_nullable()) {
281            vortex_bail!(
282                "Expected integer with nullability {}, got {}",
283                dtype.is_nullable(),
284                days.dtype()
285            );
286        }
287        if !seconds.dtype().is_int() || seconds.dtype().is_nullable() {
288            vortex_bail!(MismatchedTypes: "non-nullable integer", seconds.dtype());
289        }
290        if !subseconds.dtype().is_int() || subseconds.dtype().is_nullable() {
291            vortex_bail!(MismatchedTypes: "non-nullable integer", subseconds.dtype());
292        }
293
294        let length = days.len();
295        if length != seconds.len() || length != subseconds.len() {
296            vortex_bail!(
297                "Mismatched lengths {} {} {}",
298                days.len(),
299                seconds.len(),
300                subseconds.len()
301            );
302        }
303
304        Ok(Self {
305            dtype,
306            days,
307            seconds,
308            subseconds,
309            stats_set: Default::default(),
310        })
311    }
312
313    pub(crate) unsafe fn new_unchecked(
314        dtype: DType,
315        days: ArrayRef,
316        seconds: ArrayRef,
317        subseconds: ArrayRef,
318    ) -> Self {
319        Self {
320            dtype,
321            days,
322            seconds,
323            subseconds,
324            stats_set: Default::default(),
325        }
326    }
327
328    pub fn into_parts(self) -> DateTimePartsArrayParts {
329        DateTimePartsArrayParts {
330            dtype: self.dtype,
331            days: self.days,
332            seconds: self.seconds,
333            subseconds: self.subseconds,
334        }
335    }
336
337    pub fn days(&self) -> &ArrayRef {
338        &self.days
339    }
340
341    pub fn seconds(&self) -> &ArrayRef {
342        &self.seconds
343    }
344
345    pub fn subseconds(&self) -> &ArrayRef {
346        &self.subseconds
347    }
348}
349
350impl ValidityChild<DateTimePartsVTable> for DateTimePartsVTable {
351    fn validity_child(array: &DateTimePartsArray) -> &ArrayRef {
352        array.days()
353    }
354}
vortex_datetime_parts/array.rs

vortex_datetime_parts/
array.rs
