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
29    fn id(_encoding: &Self::Encoding) -> EncodingId {
30        EncodingId::new_ref("vortex.datetimeparts")
31    }
32
33    fn encoding(_array: &Self::Array) -> EncodingRef {
34        EncodingRef::new_ref(DateTimePartsEncoding.as_ref())
35    }
36}
37
38#[derive(Clone, Debug)]
39pub struct DateTimePartsArray {
40    dtype: DType,
41    days: ArrayRef,
42    seconds: ArrayRef,
43    subseconds: ArrayRef,
44    stats_set: ArrayStats,
45}
46
47#[derive(Clone, Debug)]
48pub struct DateTimePartsEncoding;
49
50impl DateTimePartsArray {
51    pub fn try_new(
52        dtype: DType,
53        days: ArrayRef,
54        seconds: ArrayRef,
55        subseconds: ArrayRef,
56    ) -> VortexResult<Self> {
57        if !days.dtype().is_int() || (dtype.is_nullable() != days.dtype().is_nullable()) {
58            vortex_bail!(
59                "Expected integer with nullability {}, got {}",
60                dtype.is_nullable(),
61                days.dtype()
62            );
63        }
64        if !seconds.dtype().is_int() || seconds.dtype().is_nullable() {
65            vortex_bail!(MismatchedTypes: "non-nullable integer", seconds.dtype());
66        }
67        if !subseconds.dtype().is_int() || subseconds.dtype().is_nullable() {
68            vortex_bail!(MismatchedTypes: "non-nullable integer", subseconds.dtype());
69        }
70
71        let length = days.len();
72        if length != seconds.len() || length != subseconds.len() {
73            vortex_bail!(
74                "Mismatched lengths {} {} {}",
75                days.len(),
76                seconds.len(),
77                subseconds.len()
78            );
79        }
80
81        Ok(Self {
82            dtype,
83            days,
84            seconds,
85            subseconds,
86            stats_set: Default::default(),
87        })
88    }
89
90    pub fn days(&self) -> &ArrayRef {
91        &self.days
92    }
93
94    pub fn seconds(&self) -> &ArrayRef {
95        &self.seconds
96    }
97
98    pub fn subseconds(&self) -> &ArrayRef {
99        &self.subseconds
100    }
101}
102
103impl ArrayVTable<DateTimePartsVTable> for DateTimePartsVTable {
104    fn len(array: &DateTimePartsArray) -> usize {
105        array.days.len()
106    }
107
108    fn dtype(array: &DateTimePartsArray) -> &DType {
109        &array.dtype
110    }
111
112    fn stats(array: &DateTimePartsArray) -> StatsSetRef<'_> {
113        array.stats_set.to_ref(array.as_ref())
114    }
115}
116
117impl ValidityChild<DateTimePartsVTable> for DateTimePartsVTable {
118    fn validity_child(array: &DateTimePartsArray) -> &dyn Array {
119        array.days()
120    }
121}