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

vortex_array/extension/datetime/
time.rs
