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("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        vortex_ensure!(!metadata.is_empty(), "Date metadata must not be empty");
91        let tag = metadata[0];
92        TimeUnit::try_from(tag)
93    }
94
95    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
96        let metadata = ext_dtype.metadata();
97        let ptype = date_ptype(metadata)
98            .ok_or_else(|| vortex_err!("Date type does not support time unit {}", metadata))?;
99
100        vortex_ensure!(
101            ext_dtype.storage_dtype().as_ptype() == ptype,
102            "Date storage dtype for {} must be {}",
103            metadata,
104            ptype
105        );
106
107        Ok(())
108    }
109
110    fn can_coerce_from(ext_dtype: &ExtDType<Self>, other: &DType) -> bool {
111        let DType::Extension(other_ext) = other else {
112            return false;
113        };
114        let Some(other_unit) = other_ext.metadata_opt::<Date>() else {
115            return false;
116        };
117        let our_unit = ext_dtype.metadata();
118        // We can coerce from other if our unit is finer (<=) and nullability is compatible.
119        our_unit <= other_unit && (ext_dtype.storage_dtype().is_nullable() || !other.is_nullable())
120    }
121
122    fn least_supertype(ext_dtype: &ExtDType<Self>, other: &DType) -> Option<DType> {
123        let DType::Extension(other_ext) = other else {
124            return None;
125        };
126        let other_unit = other_ext.metadata_opt::<Date>()?;
127        let our_unit = ext_dtype.metadata();
128        let finest = (*our_unit).min(*other_unit);
129        let union_null = ext_dtype.storage_dtype().nullability() | other.nullability();
130        Some(DType::Extension(Date::new(finest, union_null).erased()))
131    }
132
133    fn unpack_native<'a>(
134        ext_dtype: &'a ExtDType<Self>,
135        storage_value: &'a ScalarValue,
136    ) -> VortexResult<Self::NativeValue<'a>> {
137        let metadata = ext_dtype.metadata();
138        match metadata {
139            TimeUnit::Milliseconds => Ok(DateValue::Milliseconds(
140                storage_value.as_primitive().cast::<i64>()?,
141            )),
142            TimeUnit::Days => Ok(DateValue::Days(storage_value.as_primitive().cast::<i32>()?)),
143            _ => vortex_bail!("Date type does not support time unit {}", metadata),
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use vortex_error::VortexResult;
151
152    use crate::dtype::DType;
153    use crate::dtype::Nullability::Nullable;
154    use crate::extension::datetime::Date;
155    use crate::extension::datetime::TimeUnit;
156    use crate::scalar::PValue;
157    use crate::scalar::Scalar;
158    use crate::scalar::ScalarValue;
159
160    #[test]
161    fn validate_date_scalar() -> VortexResult<()> {
162        let days_dtype = DType::Extension(Date::new(TimeUnit::Days, Nullable).erased());
163        Scalar::try_new(days_dtype, Some(ScalarValue::Primitive(PValue::I32(0))))?;
164
165        let ms_dtype = DType::Extension(Date::new(TimeUnit::Milliseconds, Nullable).erased());
166        Scalar::try_new(
167            ms_dtype,
168            Some(ScalarValue::Primitive(PValue::I64(86_400_000))),
169        )?;
170
171        Ok(())
172    }
173
174    #[test]
175    fn reject_date_with_overflowing_value() {
176        // Days storage is `I32`, so an `I64` value that overflows `i32` should fail the cast.
177        let dtype = DType::Extension(Date::new(TimeUnit::Days, Nullable).erased());
178        let result = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(i64::MAX))));
179        assert!(result.is_err());
180    }
181
182    #[test]
183    fn display_date_scalar() {
184        let dtype = DType::Extension(Date::new(TimeUnit::Days, Nullable).erased());
185
186        let scalar = Scalar::new(dtype.clone(), Some(ScalarValue::Primitive(PValue::I32(0))));
187        assert_eq!(format!("{}", scalar.as_extension()), "1970-01-01");
188
189        let scalar = Scalar::new(dtype, Some(ScalarValue::Primitive(PValue::I32(365))));
190        assert_eq!(format!("{}", scalar.as_extension()), "1971-01-01");
191    }
192
193    #[test]
194    fn least_supertype_date_units() {
195        use crate::dtype::Nullability::NonNullable;
196
197        let days = DType::Extension(Date::new(TimeUnit::Days, NonNullable).erased());
198        let ms = DType::Extension(Date::new(TimeUnit::Milliseconds, NonNullable).erased());
199        let expected = DType::Extension(Date::new(TimeUnit::Milliseconds, NonNullable).erased());
200        assert_eq!(days.least_supertype(&ms).unwrap(), expected);
201        assert_eq!(ms.least_supertype(&days).unwrap(), expected);
202    }
203
204    #[test]
205    fn can_coerce_from_date() {
206        use crate::dtype::Nullability::NonNullable;
207
208        let days = DType::Extension(Date::new(TimeUnit::Days, NonNullable).erased());
209        let ms = DType::Extension(Date::new(TimeUnit::Milliseconds, NonNullable).erased());
210        assert!(ms.can_coerce_from(&days));
211        assert!(!days.can_coerce_from(&ms));
212    }
213
214    #[test]
215    fn deserialize_empty_metadata_returns_error() {
216        use crate::dtype::extension::ExtVTable;
217
218        let vtable = Date;
219        assert!(vtable.deserialize_metadata(&[]).is_err());
220    }
221
222    #[test]
223    fn deserialize_invalid_tag_returns_error() {
224        use crate::dtype::extension::ExtVTable;
225
226        let vtable = Date;
227        // 0xFF is not a valid TimeUnit tag.
228        assert!(vtable.deserialize_metadata(&[0xFF]).is_err());
229    }
230}
vortex_array/extension/datetime/date.rs

vortex_array/extension/datetime/
date.rs
