vortex_datetime_parts/
array.rs

1use std::fmt::Debug;
2
3use vortex_array::stats::{ArrayStats, StatsSetRef};
4use vortex_array::vtable::{
5    ArrayVTable, NotSupported, VTable, ValidityChild, ValidityVTableFromChild,
6};
7use vortex_array::{Array, ArrayRef, EncodingId, EncodingRef, vtable};
8use vortex_dtype::DType;
9use vortex_error::{VortexResult, vortex_bail};
10
11vtable!(DateTimeParts);
12
13impl VTable for DateTimePartsVTable {
14    type Array = DateTimePartsArray;
15    type Encoding = DateTimePartsEncoding;
16
17    type ArrayVTable = Self;
18    type CanonicalVTable = Self;
19    type OperationsVTable = Self;
20    type ValidityVTable = ValidityVTableFromChild;
21    type VisitorVTable = Self;
22    type ComputeVTable = NotSupported;
23    type EncodeVTable = Self;
24    type SerdeVTable = Self;
25
26    fn id(_encoding: &Self::Encoding) -> EncodingId {
27        EncodingId::new_ref("vortex.datetimeparts")
28    }
29
30    fn encoding(_array: &Self::Array) -> EncodingRef {
31        EncodingRef::new_ref(DateTimePartsEncoding.as_ref())
32    }
33}
34
35#[derive(Clone, Debug)]
36pub struct DateTimePartsArray {
37    dtype: DType,
38    days: ArrayRef,
39    seconds: ArrayRef,
40    subseconds: ArrayRef,
41    stats_set: ArrayStats,
42}
43
44#[derive(Clone, Debug)]
45pub struct DateTimePartsEncoding;
46
47impl DateTimePartsArray {
48    pub fn try_new(
49        dtype: DType,
50        days: ArrayRef,
51        seconds: ArrayRef,
52        subseconds: ArrayRef,
53    ) -> VortexResult<Self> {
54        if !days.dtype().is_int() || (dtype.is_nullable() != days.dtype().is_nullable()) {
55            vortex_bail!(
56                "Expected integer with nullability {}, got {}",
57                dtype.is_nullable(),
58                days.dtype()
59            );
60        }
61        if !seconds.dtype().is_int() || seconds.dtype().is_nullable() {
62            vortex_bail!(MismatchedTypes: "non-nullable integer", seconds.dtype());
63        }
64        if !subseconds.dtype().is_int() || subseconds.dtype().is_nullable() {
65            vortex_bail!(MismatchedTypes: "non-nullable integer", subseconds.dtype());
66        }
67
68        let length = days.len();
69        if length != seconds.len() || length != subseconds.len() {
70            vortex_bail!(
71                "Mismatched lengths {} {} {}",
72                days.len(),
73                seconds.len(),
74                subseconds.len()
75            );
76        }
77
78        Ok(Self {
79            dtype,
80            days,
81            seconds,
82            subseconds,
83            stats_set: Default::default(),
84        })
85    }
86
87    pub fn days(&self) -> &ArrayRef {
88        &self.days
89    }
90
91    pub fn seconds(&self) -> &ArrayRef {
92        &self.seconds
93    }
94
95    pub fn subseconds(&self) -> &ArrayRef {
96        &self.subseconds
97    }
98}
99
100impl ArrayVTable<DateTimePartsVTable> for DateTimePartsVTable {
101    fn len(array: &DateTimePartsArray) -> usize {
102        array.days.len()
103    }
104
105    fn dtype(array: &DateTimePartsArray) -> &DType {
106        &array.dtype
107    }
108
109    fn stats(array: &DateTimePartsArray) -> StatsSetRef<'_> {
110        array.stats_set.to_ref(array.as_ref())
111    }
112}
113
114impl ValidityChild<DateTimePartsVTable> for DateTimePartsVTable {
115    fn validity_child(array: &DateTimePartsArray) -> &dyn Array {
116        array.days()
117    }
118}