vortex_datetime_parts/
serde.rs

1use vortex_array::arrays::TemporalArray;
2use vortex_array::serde::ArrayChildren;
3use vortex_array::vtable::{EncodeVTable, SerdeVTable, VisitorVTable};
4use vortex_array::{
5    Array, ArrayBufferVisitor, ArrayChildVisitor, Canonical, DeserializeMetadata, ProstMetadata,
6};
7use vortex_buffer::ByteBuffer;
8use vortex_dtype::{DType, Nullability, PType};
9use vortex_error::{VortexExpect, VortexResult, vortex_bail};
10
11use crate::{DateTimePartsArray, DateTimePartsEncoding, DateTimePartsVTable};
12
13#[derive(Clone, prost::Message)]
14#[repr(C)]
15pub struct DateTimePartsMetadata {
16    // Validity lives in the days array
17    // TODO(ngates): we should actually model this with a Tuple array when we have one.
18    #[prost(enumeration = "PType", tag = "1")]
19    days_ptype: i32,
20    #[prost(enumeration = "PType", tag = "2")]
21    seconds_ptype: i32,
22    #[prost(enumeration = "PType", tag = "3")]
23    subseconds_ptype: i32,
24}
25
26impl SerdeVTable<DateTimePartsVTable> for DateTimePartsVTable {
27    type Metadata = ProstMetadata<DateTimePartsMetadata>;
28
29    fn metadata(array: &DateTimePartsArray) -> VortexResult<Option<Self::Metadata>> {
30        Ok(Some(ProstMetadata(DateTimePartsMetadata {
31            days_ptype: PType::try_from(array.days().dtype()).vortex_expect("Must be a valid PType")
32                as i32,
33            seconds_ptype: PType::try_from(array.seconds().dtype())
34                .vortex_expect("Must be a valid PType") as i32,
35            subseconds_ptype: PType::try_from(array.subseconds().dtype())
36                .vortex_expect("Must be a valid PType") as i32,
37        })))
38    }
39
40    fn build(
41        _encoding: &DateTimePartsEncoding,
42        dtype: &DType,
43        len: usize,
44        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
45        _buffers: &[ByteBuffer],
46        children: &dyn ArrayChildren,
47    ) -> VortexResult<DateTimePartsArray> {
48        if children.len() != 3 {
49            vortex_bail!(
50                "Expected 3 children for datetime-parts encoding, found {}",
51                children.len()
52            )
53        }
54
55        let days = children.get(
56            0,
57            &DType::Primitive(metadata.days_ptype(), dtype.nullability()),
58            len,
59        )?;
60        let seconds = children.get(
61            1,
62            &DType::Primitive(metadata.seconds_ptype(), Nullability::NonNullable),
63            len,
64        )?;
65        let subseconds = children.get(
66            2,
67            &DType::Primitive(metadata.subseconds_ptype(), Nullability::NonNullable),
68            len,
69        )?;
70
71        DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)
72    }
73}
74
75impl EncodeVTable<DateTimePartsVTable> for DateTimePartsVTable {
76    fn encode(
77        _encoding: &DateTimePartsEncoding,
78        canonical: &Canonical,
79        _like: Option<&DateTimePartsArray>,
80    ) -> VortexResult<Option<DateTimePartsArray>> {
81        let ext_array = canonical.clone().into_extension()?;
82        let temporal = TemporalArray::try_from(ext_array)?;
83
84        Ok(Some(DateTimePartsArray::try_from(temporal)?))
85    }
86}
87
88impl VisitorVTable<DateTimePartsVTable> for DateTimePartsVTable {
89    fn visit_buffers(_array: &DateTimePartsArray, _visitor: &mut dyn ArrayBufferVisitor) {}
90
91    fn visit_children(array: &DateTimePartsArray, visitor: &mut dyn ArrayChildVisitor) {
92        visitor.visit_child("days", array.days());
93        visitor.visit_child("seconds", array.seconds());
94        visitor.visit_child("subseconds", array.subseconds());
95    }
96}
97
98#[cfg(test)]
99mod test {
100    use vortex_array::ProstMetadata;
101    use vortex_array::test_harness::check_metadata;
102    use vortex_dtype::PType;
103
104    use super::*;
105
106    #[cfg_attr(miri, ignore)]
107    #[test]
108    fn test_datetimeparts_metadata() {
109        check_metadata(
110            "datetimeparts.metadata",
111            ProstMetadata(DateTimePartsMetadata {
112                days_ptype: PType::I64 as i32,
113                seconds_ptype: PType::I64 as i32,
114                subseconds_ptype: PType::I64 as i32,
115            }),
116        );
117    }
118}
vortex_datetime_parts/serde.rs

vortex_datetime_parts/
serde.rs
