vortex_datetime_parts/
array.rs

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