vortex_array/extension/datetime/
date.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5
6use jiff::Span;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_ensure;
11use vortex_error::vortex_err;
12
13use crate::dtype::DType;
14use crate::dtype::Nullability;
15use crate::dtype::PType;
16use crate::dtype::extension::ExtDType;
17use crate::dtype::extension::ExtId;
18use crate::dtype::extension::ExtVTable;
19use crate::extension::datetime::TimeUnit;
20use crate::scalar::ScalarValue;
21
22/// The Unix epoch date (1970-01-01).
23const EPOCH: jiff::civil::Date = jiff::civil::Date::constant(1970, 1, 1);
24
25/// Date DType.
26#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
27pub struct Date;
28
29fn date_ptype(time_unit: &TimeUnit) -> Option<PType> {
30    match time_unit {
31        TimeUnit::Nanoseconds => None,
32        TimeUnit::Microseconds => None,
33        TimeUnit::Milliseconds => Some(PType::I64),
34        TimeUnit::Seconds => None,
35        TimeUnit::Days => Some(PType::I32),
36    }
37}
38
39impl Date {
40    /// Creates a new Date extension dtype with the given time unit and nullability.
41    ///
42    /// Note that only Milliseconds and Days time units are supported for Date.
43    pub fn try_new(time_unit: TimeUnit, nullability: Nullability) -> VortexResult<ExtDType<Self>> {
44        let ptype = date_ptype(&time_unit)
45            .ok_or_else(|| vortex_err!("Date type does not support time unit {}", time_unit))?;
46        ExtDType::try_new(time_unit, DType::Primitive(ptype, nullability))
47    }
48
49    /// Creates a new Date extension dtype with the given time unit and nullability.
50    ///
51    /// # Panics
52    ///
53    /// Panics if the `time_unit` is not supported by date types.
54    pub fn new(time_unit: TimeUnit, nullability: Nullability) -> ExtDType<Self> {
55        Self::try_new(time_unit, nullability).vortex_expect("failed to create date dtype")
56    }
57}
58
59/// Unpacked value of a [`Date`] extension scalar.
60pub enum DateValue {
61    /// Days since the Unix epoch.
62    Days(i32),
63    /// Milliseconds since the Unix epoch.
64    Milliseconds(i64),
65}
66
67impl fmt::Display for DateValue {
68    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69        let date = match self {
70            DateValue::Days(days) => EPOCH + Span::new().days(*days),
71            DateValue::Milliseconds(ms) => EPOCH + Span::new().milliseconds(*ms),
72        };
73        write!(f, "{}", date)
74    }
75}
76
77impl ExtVTable for Date {
78    type Metadata = TimeUnit;
79    type NativeValue<'a> = DateValue;
80
81    fn id(&self) -> ExtId {
82        ExtId::new_ref("vortex.date")
83    }
84
85    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
86        Ok(vec![u8::from(*metadata)])
87    }
88
89    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
90        let tag = metadata[0];
91        TimeUnit::try_from(tag)
92    }
93
94    fn validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &DType) -> VortexResult<()> {
95        let ptype = date_ptype(metadata)
96            .ok_or_else(|| vortex_err!("Date type does not support time unit {}", metadata))?;
97
98        vortex_ensure!(
99            storage_dtype.as_ptype() == ptype,
100            "Date storage dtype for {} must be {}",
101            metadata,
102            ptype
103        );
104
105        Ok(())
106    }
107
108    fn unpack_native(
109        &self,
110        metadata: &Self::Metadata,
111        _storage_dtype: &DType,
112        storage_value: &ScalarValue,
113    ) -> VortexResult<Self::NativeValue<'_>> {
114        match metadata {
115            TimeUnit::Milliseconds => Ok(DateValue::Milliseconds(
116                storage_value.as_primitive().cast::<i64>()?,
117            )),
118            TimeUnit::Days => Ok(DateValue::Days(storage_value.as_primitive().cast::<i32>()?)),
119            _ => vortex_bail!("Date type does not support time unit {}", metadata),
120        }
121    }
122}
vortex_array/extension/datetime/date.rs

vortex_array/extension/datetime/
date.rs
