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::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hasher;
8
9use prost::Message;
10use vortex_array::AnyCanonical;
11use vortex_array::Array;
12use vortex_array::ArrayEq;
13use vortex_array::ArrayHash;
14use vortex_array::ArrayId;
15use vortex_array::ArrayParts;
16use vortex_array::ArrayRef;
17use vortex_array::ArrayView;
18use vortex_array::EqMode;
19use vortex_array::ExecutionCtx;
20use vortex_array::ExecutionResult;
21use vortex_array::IntoArray;
22use vortex_array::array_slots;
23use vortex_array::arrays::Primitive;
24use vortex_array::arrays::TemporalArray;
25use vortex_array::buffer::BufferHandle;
26use vortex_array::dtype::DType;
27use vortex_array::dtype::Nullability;
28use vortex_array::dtype::PType;
29use vortex_array::require_child;
30use vortex_array::serde::ArrayChildren;
31use vortex_array::smallvec::smallvec;
32use vortex_array::vtable::VTable;
33use vortex_array::vtable::ValidityChild;
34use vortex_array::vtable::ValidityVTableFromChild;
35use vortex_error::VortexResult;
36use vortex_error::vortex_bail;
37use vortex_error::vortex_ensure;
38use vortex_error::vortex_err;
39use vortex_error::vortex_panic;
40use vortex_session::VortexSession;
41use vortex_session::registry::CachedId;
42
43use crate::TemporalParts;
44use crate::canonical::decode_to_temporal;
45use crate::compute::rules::PARENT_RULES;
46use crate::split_temporal;
47
48/// A [`DateTimeParts`]-encoded Vortex array.
49pub type DateTimePartsArray = Array<DateTimeParts>;
50
51impl ArrayHash for DateTimePartsData {
52    fn array_hash<H: Hasher>(&self, _state: &mut H, _accuracy: EqMode) {}
53}
54
55impl ArrayEq for DateTimePartsData {
56    fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool {
57        true
58    }
59}
60
61#[derive(Clone, prost::Message)]
62#[repr(C)]
63pub struct DateTimePartsMetadata {
64    // Validity lives in the days array
65    // TODO(ngates): we should actually model this with a Tuple array when we have one.
66    #[prost(enumeration = "PType", tag = "1")]
67    pub days_ptype: i32,
68    #[prost(enumeration = "PType", tag = "2")]
69    pub seconds_ptype: i32,
70    #[prost(enumeration = "PType", tag = "3")]
71    pub subseconds_ptype: i32,
72}
73
74impl DateTimePartsMetadata {
75    pub fn get_days_ptype(&self) -> VortexResult<PType> {
76        PType::try_from(self.days_ptype)
77            .map_err(|_| vortex_err!("Invalid PType {}", self.days_ptype))
78    }
79
80    pub fn get_seconds_ptype(&self) -> VortexResult<PType> {
81        PType::try_from(self.seconds_ptype)
82            .map_err(|_| vortex_err!("Invalid PType {}", self.seconds_ptype))
83    }
84
85    pub fn get_subseconds_ptype(&self) -> VortexResult<PType> {
86        PType::try_from(self.subseconds_ptype)
87            .map_err(|_| vortex_err!("Invalid PType {}", self.subseconds_ptype))
88    }
89}
90
91impl VTable for DateTimeParts {
92    type TypedArrayData = DateTimePartsData;
93
94    type OperationsVTable = Self;
95    type ValidityVTable = ValidityVTableFromChild;
96
97    fn id(&self) -> ArrayId {
98        static ID: CachedId = CachedId::new("vortex.datetimeparts");
99        *ID
100    }
101
102    fn validate(
103        &self,
104        _data: &Self::TypedArrayData,
105        dtype: &DType,
106        len: usize,
107        slots: &[Option<ArrayRef>],
108    ) -> VortexResult<()> {
109        let slots = DateTimePartsSlotsView::from_slots(slots);
110        DateTimePartsData::validate(dtype, slots.days, slots.seconds, slots.subseconds, len)
111    }
112
113    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
114        0
115    }
116
117    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
118        vortex_panic!("DateTimePartsArray buffer index {idx} out of bounds")
119    }
120
121    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
122        vortex_panic!("DateTimePartsArray buffer_name index {idx} out of bounds")
123    }
124
125    fn with_buffers(
126        &self,
127        array: ArrayView<'_, Self>,
128        buffers: &[BufferHandle],
129    ) -> VortexResult<ArrayParts<Self>> {
130        vortex_array::vtable::with_empty_buffers(self, array, buffers)
131    }
132
133    fn serialize(
134        array: ArrayView<'_, Self>,
135        _session: &VortexSession,
136    ) -> VortexResult<Option<Vec<u8>>> {
137        Ok(Some(
138            DateTimePartsMetadata {
139                days_ptype: PType::try_from(array.days().dtype())? as i32,
140                seconds_ptype: PType::try_from(array.seconds().dtype())? as i32,
141                subseconds_ptype: PType::try_from(array.subseconds().dtype())? as i32,
142            }
143            .encode_to_vec(),
144        ))
145    }
146
147    fn deserialize(
148        &self,
149        dtype: &DType,
150        len: usize,
151        metadata: &[u8],
152        _buffers: &[BufferHandle],
153        children: &dyn ArrayChildren,
154        _session: &VortexSession,
155    ) -> VortexResult<ArrayParts<Self>> {
156        let metadata = DateTimePartsMetadata::decode(metadata)?;
157        if children.len() != 3 {
158            vortex_bail!(
159                "Expected 3 children for datetime-parts encoding, found {}",
160                children.len()
161            )
162        }
163
164        let days = children.get(
165            0,
166            &DType::Primitive(metadata.get_days_ptype()?, dtype.nullability()),
167            len,
168        )?;
169        let seconds = children.get(
170            1,
171            &DType::Primitive(metadata.get_seconds_ptype()?, Nullability::NonNullable),
172            len,
173        )?;
174        let subseconds = children.get(
175            2,
176            &DType::Primitive(metadata.get_subseconds_ptype()?, Nullability::NonNullable),
177            len,
178        )?;
179
180        let slots = smallvec![Some(days), Some(seconds), Some(subseconds)];
181        let data = DateTimePartsData {};
182        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
183    }
184
185    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
186        DateTimePartsSlots::NAMES[idx].to_string()
187    }
188
189    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
190        let array = require_child!(array, array.days(), DateTimePartsSlots::DAYS => Primitive);
191        let array =
192            require_child!(array, array.seconds(), DateTimePartsSlots::SECONDS => AnyCanonical);
193        let array = require_child!(array, array.subseconds(), DateTimePartsSlots::SUBSECONDS => AnyCanonical);
194
195        let dtype = array.dtype().clone();
196        let parts = array.into_parts();
197
198        Ok(ExecutionResult::done(
199            decode_to_temporal(parts, &dtype, ctx)?.into_array(),
200        ))
201    }
202
203    fn reduce_parent(
204        array: ArrayView<'_, Self>,
205        parent: &ArrayRef,
206        child_idx: usize,
207    ) -> VortexResult<Option<ArrayRef>> {
208        PARENT_RULES.evaluate(array, parent, child_idx)
209    }
210}
211
212#[array_slots(DateTimeParts)]
213pub struct DateTimePartsSlots {
214    /// The days component of the datetime, stored as an integer array.
215    pub days: ArrayRef,
216    /// The seconds component of the datetime (within the day).
217    pub seconds: ArrayRef,
218    /// The sub-second component of the datetime.
219    pub subseconds: ArrayRef,
220}
221
222#[derive(Clone, Debug)]
223pub struct DateTimePartsData {}
224
225pub struct DateTimePartsParts {
226    pub days: ArrayRef,
227    pub seconds: ArrayRef,
228    pub subseconds: ArrayRef,
229}
230
231pub trait DateTimePartsOwnedExt {
232    fn into_parts(self) -> DateTimePartsParts;
233}
234
235impl DateTimePartsOwnedExt for Array<DateTimeParts> {
236    fn into_parts(self) -> DateTimePartsParts {
237        match self.try_into_parts() {
238            Ok(parts) => {
239                let slots = DateTimePartsSlots::from_slots(parts.slots);
240                DateTimePartsParts {
241                    days: slots.days,
242                    seconds: slots.seconds,
243                    subseconds: slots.subseconds,
244                }
245            }
246            Err(array) => {
247                let view = DateTimePartsSlotsView::from_slots(array.as_ref().slots());
248                DateTimePartsParts {
249                    days: view.days.clone(),
250                    seconds: view.seconds.clone(),
251                    subseconds: view.subseconds.clone(),
252                }
253            }
254        }
255    }
256}
257
258impl Display for DateTimePartsData {
259    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
260        Ok(())
261    }
262}
263
264#[derive(Clone, Debug)]
265pub struct DateTimeParts;
266
267impl DateTimeParts {
268    /// Construct a new [`DateTimePartsArray`] from its components.
269    pub fn try_new(
270        dtype: DType,
271        days: ArrayRef,
272        seconds: ArrayRef,
273        subseconds: ArrayRef,
274    ) -> VortexResult<DateTimePartsArray> {
275        let len = days.len();
276        DateTimePartsData::validate(&dtype, &days, &seconds, &subseconds, len)?;
277        let slots = smallvec![Some(days), Some(seconds), Some(subseconds)];
278        let data = DateTimePartsData {};
279        Ok(unsafe {
280            Array::from_parts_unchecked(
281                ArrayParts::new(DateTimeParts, dtype, len, data).with_slots(slots),
282            )
283        })
284    }
285
286    /// Construct a [`DateTimePartsArray`] from a [`TemporalArray`].
287    pub fn try_from_temporal(
288        temporal: TemporalArray,
289        ctx: &mut ExecutionCtx,
290    ) -> VortexResult<DateTimePartsArray> {
291        let dtype = temporal.dtype().clone();
292        let TemporalParts {
293            days,
294            seconds,
295            subseconds,
296        } = split_temporal(temporal, ctx)?;
297        Self::try_new(dtype, days, seconds, subseconds)
298    }
299}
300
301impl DateTimePartsData {
302    pub fn validate(
303        dtype: &DType,
304        days: &ArrayRef,
305        seconds: &ArrayRef,
306        subseconds: &ArrayRef,
307        len: usize,
308    ) -> VortexResult<()> {
309        vortex_ensure!(days.len() == len, "expected len {len}, got {}", days.len());
310
311        if !days.dtype().is_int() || (dtype.is_nullable() != days.dtype().is_nullable()) {
312            vortex_bail!(
313                "Expected integer with nullability {}, got {}",
314                dtype.is_nullable(),
315                days.dtype()
316            );
317        }
318        if !seconds.dtype().is_int() || seconds.dtype().is_nullable() {
319            vortex_bail!(MismatchedTypes: "non-nullable integer", seconds.dtype());
320        }
321        if !subseconds.dtype().is_int() || subseconds.dtype().is_nullable() {
322            vortex_bail!(MismatchedTypes: "non-nullable integer", subseconds.dtype());
323        }
324
325        if len != seconds.len() || len != subseconds.len() {
326            vortex_bail!(
327                "Mismatched lengths {} {} {}",
328                days.len(),
329                seconds.len(),
330                subseconds.len()
331            );
332        }
333
334        Ok(())
335    }
336}
337
338impl ValidityChild<DateTimeParts> for DateTimeParts {
339    fn validity_child(array: ArrayView<'_, DateTimeParts>) -> ArrayRef {
340        array.days().clone()
341    }
342}
vortex_datetime_parts/array.rs

vortex_datetime_parts/
array.rs
