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

vortex_array/extension/datetime/
date.rs
