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

vortex_array/extension/datetime/
time.rs
