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

vortex_datetime_parts/
array.rs
