vortex_array/extension/datetime/
time.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/// Time DType.
24#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
25pub struct Time;
26
27fn time_ptype(time_unit: &TimeUnit) -> Option<PType> {
28    Some(match time_unit {
29        TimeUnit::Nanoseconds | TimeUnit::Microseconds => PType::I64,
30        TimeUnit::Milliseconds | TimeUnit::Seconds => PType::I32,
31        TimeUnit::Days => return None,
32    })
33}
34
35impl Time {
36    /// Creates a new Time extension dtype with the given time unit and nullability.
37    ///
38    /// Note that Days units are not supported for Time.
39    pub fn try_new(time_unit: TimeUnit, nullability: Nullability) -> VortexResult<ExtDType<Self>> {
40        let ptype = time_ptype(&time_unit)
41            .ok_or_else(|| vortex_err!("Time type does not support time unit {}", time_unit))?;
42        ExtDType::try_new(time_unit, DType::Primitive(ptype, nullability))
43    }
44
45    /// Creates a new Time extension dtype with the given time unit and nullability.
46    pub fn new(time_unit: TimeUnit, nullability: Nullability) -> ExtDType<Self> {
47        Self::try_new(time_unit, nullability).vortex_expect("failed to create time dtype")
48    }
49}
50
51/// Unpacked value of a [`Time`] extension scalar.
52pub enum TimeValue {
53    /// Seconds since midnight.
54    Seconds(i32),
55    /// Milliseconds since midnight.
56    Milliseconds(i32),
57    /// Microseconds since midnight.
58    Microseconds(i64),
59    /// Nanoseconds since midnight.
60    Nanoseconds(i64),
61}
62
63impl fmt::Display for TimeValue {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        let min = jiff::civil::Time::MIN;
66
67        let time = match self {
68            TimeValue::Seconds(s) => min + Span::new().seconds(*s),
69            TimeValue::Milliseconds(ms) => min + Span::new().milliseconds(*ms),
70            TimeValue::Microseconds(us) => min + Span::new().microseconds(*us),
71            TimeValue::Nanoseconds(ns) => min + Span::new().nanoseconds(*ns),
72        };
73
74        write!(f, "{}", time)
75    }
76}
77
78impl ExtVTable for Time {
79    type Metadata = TimeUnit;
80
81    type NativeValue<'a> = TimeValue;
82
83    fn id(&self) -> ExtId {
84        static ID: CachedId = CachedId::new("vortex.time");
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, data: &[u8]) -> VortexResult<Self::Metadata> {
93        vortex_ensure!(!data.is_empty(), "Time metadata must not be empty");
94        let tag = data[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 = time_ptype(metadata)
101            .ok_or_else(|| vortex_err!("Time type does not support time unit {}", metadata))?;
102
103        vortex_ensure!(
104            ext_dtype.storage_dtype().as_ptype() == ptype,
105            "Time storage dtype for {} must be {}",
106            metadata,
107            ptype
108        );
109
110        Ok(())
111    }
112
113    fn can_coerce_from(ext_dtype: &ExtDType<Self>, other: &DType) -> bool {
114        let DType::Extension(other_ext) = other else {
115            return false;
116        };
117        let Some(other_unit) = other_ext.metadata_opt::<Time>() else {
118            return false;
119        };
120        let our_unit = ext_dtype.metadata();
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::<Time>()?;
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(Time::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 length_of_time = storage_value.as_primitive().cast::<i64>()?;
140
141        let (span, value) = match *ext_dtype.metadata() {
142            TimeUnit::Seconds => {
143                let v = i32::try_from(length_of_time)
144                    .map_err(|e| vortex_err!("Time seconds value out of i32 range: {e}"))?;
145                (Span::new().seconds(v), TimeValue::Seconds(v))
146            }
147            TimeUnit::Milliseconds => {
148                let v = i32::try_from(length_of_time)
149                    .map_err(|e| vortex_err!("Time milliseconds value out of i32 range: {e}"))?;
150                (Span::new().milliseconds(v), TimeValue::Milliseconds(v))
151            }
152            TimeUnit::Microseconds => (
153                Span::new().microseconds(length_of_time),
154                TimeValue::Microseconds(length_of_time),
155            ),
156            TimeUnit::Nanoseconds => (
157                Span::new().nanoseconds(length_of_time),
158                TimeValue::Nanoseconds(length_of_time),
159            ),
160            d @ TimeUnit::Days => vortex_bail!("Time type does not support time unit {d}"),
161        };
162
163        // Validate the storage value is within the valid range for Time.
164        jiff::civil::Time::MIN
165            .checked_add(span)
166            .map_err(|e| vortex_err!("Invalid time scalar: {}", e))?;
167
168        Ok(value)
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use vortex_error::VortexResult;
175
176    use crate::dtype::DType;
177    use crate::dtype::Nullability::Nullable;
178    use crate::extension::datetime::Time;
179    use crate::extension::datetime::TimeUnit;
180    use crate::scalar::PValue;
181    use crate::scalar::Scalar;
182    use crate::scalar::ScalarValue;
183
184    #[test]
185    fn validate_time_scalar() -> VortexResult<()> {
186        // 3661 seconds = 1 hour, 1 minute, 1 second.
187        let dtype = DType::Extension(Time::new(TimeUnit::Seconds, Nullable).erased());
188        Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I32(3661))))?;
189
190        Ok(())
191    }
192
193    #[test]
194    fn reject_time_out_of_range() {
195        // 86400 seconds = exactly 24 hours, which exceeds the valid `jiff::civil::Time` range.
196        let dtype = DType::Extension(Time::new(TimeUnit::Seconds, Nullable).erased());
197        let result = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I32(86400))));
198        assert!(result.is_err());
199    }
200
201    #[test]
202    fn display_time_scalar() {
203        let dtype = DType::Extension(Time::new(TimeUnit::Seconds, Nullable).erased());
204
205        let scalar = Scalar::new(
206            dtype.clone(),
207            Some(ScalarValue::Primitive(PValue::I32(3661))),
208        );
209        assert_eq!(format!("{}", scalar.as_extension()), "01:01:01");
210
211        let scalar = Scalar::new(dtype, Some(ScalarValue::Primitive(PValue::I32(0))));
212        assert_eq!(format!("{}", scalar.as_extension()), "00:00:00");
213    }
214
215    #[test]
216    fn least_supertype_time_units() {
217        use crate::dtype::Nullability::NonNullable;
218
219        let secs = DType::Extension(Time::new(TimeUnit::Seconds, NonNullable).erased());
220        let ns = DType::Extension(Time::new(TimeUnit::Nanoseconds, NonNullable).erased());
221        let expected = DType::Extension(Time::new(TimeUnit::Nanoseconds, NonNullable).erased());
222        assert_eq!(secs.least_supertype(&ns).unwrap(), expected);
223        assert_eq!(ns.least_supertype(&secs).unwrap(), expected);
224    }
225
226    #[test]
227    fn deserialize_empty_metadata_returns_error() {
228        use crate::dtype::extension::ExtVTable;
229
230        let vtable = Time;
231        assert!(vtable.deserialize_metadata(&[]).is_err());
232    }
233
234    #[test]
235    fn deserialize_invalid_tag_returns_error() {
236        use crate::dtype::extension::ExtVTable;
237
238        let vtable = Time;
239        // 0xFF is not a valid TimeUnit tag.
240        assert!(vtable.deserialize_metadata(&[0xFF]).is_err());
241    }
242}
vortex_array/extension/datetime/time.rs

vortex_array/extension/datetime/
time.rs
