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_ensure_eq;
12use vortex_error::vortex_err;
13use vortex_session::registry::CachedId;
14
15use crate::dtype::DType;
16use crate::dtype::Nullability;
17use crate::dtype::PType;
18use crate::dtype::extension::ExtDType;
19use crate::dtype::extension::ExtId;
20use crate::dtype::extension::ExtVTable;
21use crate::extension::datetime::TimeUnit;
22use crate::scalar::ScalarValue;
23
24/// The Unix epoch date (1970-01-01).
25const EPOCH: jiff::civil::Date = jiff::civil::Date::constant(1970, 1, 1);
26
27/// Date DType.
28#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
29pub struct Date;
30
31fn date_ptype(time_unit: &TimeUnit) -> Option<PType> {
32    match time_unit {
33        TimeUnit::Nanoseconds => None,
34        TimeUnit::Microseconds => None,
35        TimeUnit::Milliseconds => Some(PType::I64),
36        TimeUnit::Seconds => None,
37        TimeUnit::Days => Some(PType::I32),
38    }
39}
40
41impl Date {
42    /// Creates a new Date extension dtype with the given time unit and nullability.
43    ///
44    /// Note that only Milliseconds and Days time units are supported for Date.
45    pub fn try_new(time_unit: TimeUnit, nullability: Nullability) -> VortexResult<ExtDType<Self>> {
46        let ptype = date_ptype(&time_unit)
47            .ok_or_else(|| vortex_err!("Date type does not support time unit {}", time_unit))?;
48        ExtDType::try_new(time_unit, DType::Primitive(ptype, nullability))
49    }
50
51    /// Creates a new Date extension dtype with the given time unit and nullability.
52    ///
53    /// # Panics
54    ///
55    /// Panics if the `time_unit` is not supported by date types.
56    pub fn new(time_unit: TimeUnit, nullability: Nullability) -> ExtDType<Self> {
57        Self::try_new(time_unit, nullability).vortex_expect("failed to create date dtype")
58    }
59}
60
61/// Unpacked value of a [`Date`] extension scalar.
62pub enum DateValue {
63    /// Days since the Unix epoch.
64    Days(i32),
65    /// Milliseconds since the Unix epoch.
66    Milliseconds(i64),
67}
68
69impl fmt::Display for DateValue {
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        let date = match self {
72            DateValue::Days(days) => EPOCH + Span::new().days(*days),
73            DateValue::Milliseconds(ms) => EPOCH + Span::new().milliseconds(*ms),
74        };
75        write!(f, "{}", date)
76    }
77}
78
79impl ExtVTable for Date {
80    type Metadata = TimeUnit;
81    type NativeValue<'a> = DateValue;
82
83    fn id(&self) -> ExtId {
84        static ID: CachedId = CachedId::new("vortex.date");
85        *ID
86    }
87
88    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
89        Ok(vec![u8::from(*metadata)])
90    }
91
92    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
93        vortex_ensure!(!metadata.is_empty(), "Date metadata must not be empty");
94        let tag = metadata[0];
95        TimeUnit::try_from(tag)
96    }
97
98    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
99        let metadata = ext_dtype.metadata();
100        let ptype = date_ptype(metadata)
101            .ok_or_else(|| vortex_err!("Date type does not support time unit {}", metadata))?;
102
103        vortex_ensure_eq!(
104            ext_dtype.storage_dtype().as_ptype(),
105            ptype,
106            "Date storage dtype for {} has the wrong ptype",
107            metadata
108        );
109
110        Ok(())
111    }
112
113    fn unpack_native<'a>(
114        ext_dtype: &'a ExtDType<Self>,
115        storage_value: &'a ScalarValue,
116    ) -> VortexResult<Self::NativeValue<'a>> {
117        let metadata = ext_dtype.metadata();
118        match metadata {
119            TimeUnit::Milliseconds => Ok(DateValue::Milliseconds(
120                storage_value.as_primitive().cast::<i64>()?,
121            )),
122            TimeUnit::Days => Ok(DateValue::Days(storage_value.as_primitive().cast::<i32>()?)),
123            _ => vortex_bail!("Date type does not support time unit {}", metadata),
124        }
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use vortex_error::VortexResult;
131
132    use crate::dtype::DType;
133    use crate::dtype::Nullability::Nullable;
134    use crate::extension::datetime::Date;
135    use crate::extension::datetime::TimeUnit;
136    use crate::scalar::PValue;
137    use crate::scalar::Scalar;
138    use crate::scalar::ScalarValue;
139
140    #[test]
141    fn validate_date_scalar() -> VortexResult<()> {
142        let days_dtype = DType::Extension(Date::new(TimeUnit::Days, Nullable).erased());
143        Scalar::try_new(days_dtype, Some(ScalarValue::Primitive(PValue::I32(0))))?;
144
145        let ms_dtype = DType::Extension(Date::new(TimeUnit::Milliseconds, Nullable).erased());
146        Scalar::try_new(
147            ms_dtype,
148            Some(ScalarValue::Primitive(PValue::I64(86_400_000))),
149        )?;
150
151        Ok(())
152    }
153
154    #[test]
155    fn reject_date_with_overflowing_value() {
156        // Days storage is `I32`, so an `I64` value that overflows `i32` should fail the cast.
157        let dtype = DType::Extension(Date::new(TimeUnit::Days, Nullable).erased());
158        let result = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(i64::MAX))));
159        assert!(result.is_err());
160    }
161
162    #[test]
163    fn display_date_scalar() {
164        let dtype = DType::Extension(Date::new(TimeUnit::Days, Nullable).erased());
165
166        let scalar = Scalar::new(dtype.clone(), Some(ScalarValue::Primitive(PValue::I32(0))));
167        assert_eq!(format!("{}", scalar.as_extension()), "1970-01-01");
168
169        let scalar = Scalar::new(dtype, Some(ScalarValue::Primitive(PValue::I32(365))));
170        assert_eq!(format!("{}", scalar.as_extension()), "1971-01-01");
171    }
172
173    #[test]
174    fn deserialize_empty_metadata_returns_error() {
175        use crate::dtype::extension::ExtVTable;
176
177        let vtable = Date;
178        assert!(vtable.deserialize_metadata(&[]).is_err());
179    }
180
181    #[test]
182    fn deserialize_invalid_tag_returns_error() {
183        use crate::dtype::extension::ExtVTable;
184
185        let vtable = Date;
186        // 0xFF is not a valid TimeUnit tag.
187        assert!(vtable.deserialize_metadata(&[0xFF]).is_err());
188    }
189}
vortex_array/extension/datetime/date.rs

vortex_array/extension/datetime/
date.rs
