vortex_array/extension/datetime/
timestamp.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Temporal extension data types.
5
6use std::fmt;
7use std::sync::Arc;
8
9use jiff::Span;
10use vortex_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_ensure;
14use vortex_error::vortex_err;
15use vortex_error::vortex_panic;
16
17use crate::dtype::DType;
18use crate::dtype::Nullability;
19use crate::dtype::PType;
20use crate::dtype::extension::ExtDType;
21use crate::dtype::extension::ExtId;
22use crate::dtype::extension::ExtVTable;
23use crate::extension::datetime::TimeUnit;
24use crate::scalar::ScalarValue;
25
26/// Timestamp DType.
27#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
28pub struct Timestamp;
29
30impl Timestamp {
31    /// Creates a new Timestamp extension =dtype with the given time unit and nullability.
32    pub fn new(time_unit: TimeUnit, nullability: Nullability) -> ExtDType<Self> {
33        Self::new_with_tz(time_unit, None, nullability)
34    }
35
36    /// Creates a new Timestamp extension dtype with the given time unit, timezone, and nullability.
37    pub fn new_with_tz(
38        time_unit: TimeUnit,
39        timezone: Option<Arc<str>>,
40        nullability: Nullability,
41    ) -> ExtDType<Self> {
42        ExtDType::try_new(
43            TimestampOptions {
44                unit: time_unit,
45                tz: timezone,
46            },
47            DType::Primitive(PType::I64, nullability),
48        )
49        .vortex_expect("failed to create timestamp dtype")
50    }
51
52    /// Creates a new `Timestamp` extension dtype with the given options and nullability.
53    pub fn new_with_options(options: TimestampOptions, nullability: Nullability) -> ExtDType<Self> {
54        ExtDType::try_new(options, DType::Primitive(PType::I64, nullability))
55            .vortex_expect("failed to create timestamp dtype")
56    }
57}
58
59/// Options for the Timestamp DType.
60#[derive(Clone, Debug, PartialEq, Eq, Hash)]
61pub struct TimestampOptions {
62    /// The time unit of the timestamp.
63    pub unit: TimeUnit,
64    /// The timezone of the timestamp, if any.
65    pub tz: Option<Arc<str>>,
66}
67
68impl fmt::Display for TimestampOptions {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        match &self.tz {
71            Some(tz) => write!(f, "{}, tz={}", self.unit, tz),
72            None => write!(f, "{}", self.unit),
73        }
74    }
75}
76
77/// Unpacked value of a [`Timestamp`] extension scalar.
78///
79/// Each variant carries the raw storage value and an optional timezone.
80pub enum TimestampValue<'a> {
81    /// Seconds since the Unix epoch.
82    Seconds(i64, Option<&'a Arc<str>>),
83    /// Milliseconds since the Unix epoch.
84    Milliseconds(i64, Option<&'a Arc<str>>),
85    /// Microseconds since the Unix epoch.
86    Microseconds(i64, Option<&'a Arc<str>>),
87    /// Nanoseconds since the Unix epoch.
88    Nanoseconds(i64, Option<&'a Arc<str>>),
89}
90
91impl fmt::Display for TimestampValue<'_> {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        let (span, tz) = match self {
94            TimestampValue::Seconds(v, tz) => (Span::new().seconds(*v), *tz),
95            TimestampValue::Milliseconds(v, tz) => (Span::new().milliseconds(*v), *tz),
96            TimestampValue::Microseconds(v, tz) => (Span::new().microseconds(*v), *tz),
97            TimestampValue::Nanoseconds(v, tz) => (Span::new().nanoseconds(*v), *tz),
98        };
99        let ts = jiff::Timestamp::UNIX_EPOCH + span;
100
101        match tz {
102            None => write!(f, "{ts}"),
103            Some(tz) => {
104                let adjusted_ts = ts.in_tz(tz.as_ref()).vortex_expect("unknown timezone");
105                write!(f, "{adjusted_ts}",)
106            }
107        }
108    }
109}
110
111impl ExtVTable for Timestamp {
112    type Metadata = TimestampOptions;
113
114    type NativeValue<'a> = TimestampValue<'a>;
115
116    fn id(&self) -> ExtId {
117        ExtId::new("vortex.timestamp")
118    }
119
120    // NOTE(ngates): unfortunately we're stuck with this hand-rolled serialization format for
121    //  backwards compatibility.
122    fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
123        let mut bytes = Vec::with_capacity(4);
124        let unit_tag: u8 = metadata.unit.into();
125
126        bytes.push(unit_tag);
127
128        // Encode time_zone as u16 length followed by utf8 bytes.
129        match &metadata.tz {
130            None => bytes.extend_from_slice(0u16.to_le_bytes().as_slice()),
131            Some(tz) => {
132                let tz_bytes = tz.as_bytes();
133                let tz_len = u16::try_from(tz_bytes.len())
134                    .unwrap_or_else(|err| vortex_panic!("tz did not fit in u16: {}", err));
135                bytes.extend_from_slice(tz_len.to_le_bytes().as_slice());
136                bytes.extend_from_slice(tz_bytes);
137            }
138        }
139
140        Ok(bytes)
141    }
142
143    fn deserialize_metadata(&self, data: &[u8]) -> VortexResult<Self::Metadata> {
144        vortex_ensure!(
145            data.len() >= 3,
146            "Timestamp metadata must have at least 3 bytes, got {}",
147            data.len()
148        );
149
150        let tag = data[0];
151        let time_unit = TimeUnit::try_from(tag)?;
152        let tz_len_bytes: [u8; 2] = data[1..3]
153            .try_into()
154            .ok()
155            .vortex_expect("Verified to have two bytes");
156        let tz_len = u16::from_le_bytes(tz_len_bytes) as usize;
157        if tz_len == 0 {
158            return Ok(TimestampOptions {
159                unit: time_unit,
160                tz: None,
161            });
162        }
163
164        // Attempt to load from len-prefixed bytes
165        vortex_ensure!(
166            data.len() >= 3 + tz_len,
167            "Timestamp metadata is truncated: declared timezone length {} but only {} bytes available",
168            tz_len,
169            data.len() - 3
170        );
171        let tz_bytes = &data[3..3 + tz_len];
172        let tz: Arc<str> = str::from_utf8(tz_bytes)
173            .map_err(|e| vortex_err!("timezone is not valid utf8 string: {e}"))?
174            .to_string()
175            .into();
176        Ok(TimestampOptions {
177            unit: time_unit,
178            tz: Some(tz),
179        })
180    }
181
182    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
183        vortex_ensure!(
184            matches!(ext_dtype.storage_dtype(), DType::Primitive(PType::I64, _)),
185            "Timestamp storage dtype must be i64"
186        );
187        Ok(())
188    }
189
190    fn unpack_native<'a>(
191        ext_dtype: &'a ExtDType<Self>,
192        storage_value: &'a ScalarValue,
193    ) -> VortexResult<Self::NativeValue<'a>> {
194        let metadata = ext_dtype.metadata();
195        let ts_value = storage_value.as_primitive().cast::<i64>()?;
196        let tz = metadata.tz.as_ref();
197
198        let (span, value) = match metadata.unit {
199            TimeUnit::Nanoseconds => (
200                Span::new().nanoseconds(ts_value),
201                TimestampValue::Nanoseconds(ts_value, tz),
202            ),
203            TimeUnit::Microseconds => (
204                Span::new().microseconds(ts_value),
205                TimestampValue::Microseconds(ts_value, tz),
206            ),
207            TimeUnit::Milliseconds => (
208                Span::new().milliseconds(ts_value),
209                TimestampValue::Milliseconds(ts_value, tz),
210            ),
211            TimeUnit::Seconds => (
212                Span::new().seconds(ts_value),
213                TimestampValue::Seconds(ts_value, tz),
214            ),
215            TimeUnit::Days => vortex_bail!("Timestamp does not support Days time unit"),
216        };
217
218        // Validate the storage value is within the valid range for Timestamp.
219        let ts = jiff::Timestamp::UNIX_EPOCH
220            .checked_add(span)
221            .map_err(|e| vortex_err!("Invalid timestamp scalar: {}", e))?;
222
223        if let Some(tz) = tz {
224            ts.in_tz(tz.as_ref())
225                .map_err(|e| vortex_err!("Invalid timezone for timestamp scalar: {}", e))?;
226        }
227
228        Ok(value)
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use std::sync::Arc;
235
236    use vortex_error::VortexResult;
237
238    use crate::dtype::DType;
239    use crate::dtype::Nullability::Nullable;
240    use crate::extension::datetime::TimeUnit;
241    use crate::extension::datetime::Timestamp;
242    use crate::scalar::PValue;
243    use crate::scalar::Scalar;
244    use crate::scalar::ScalarValue;
245
246    #[test]
247    fn validate_timestamp_scalar() -> VortexResult<()> {
248        let dtype = DType::Extension(Timestamp::new(TimeUnit::Seconds, Nullable).erased());
249        Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(0))))?;
250
251        Ok(())
252    }
253
254    #[cfg_attr(miri, ignore)]
255    #[test]
256    fn reject_timestamp_with_invalid_timezone() {
257        let dtype = DType::Extension(
258            Timestamp::new_with_tz(
259                TimeUnit::Seconds,
260                Some(Arc::from("Not/A/Timezone")),
261                Nullable,
262            )
263            .erased(),
264        );
265        let result = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(0))));
266        assert!(result.is_err());
267    }
268
269    #[cfg_attr(miri, ignore)]
270    #[test]
271    fn display_timestamp_scalar() {
272        // Local (no timezone) timestamp.
273        let local_dtype = DType::Extension(Timestamp::new(TimeUnit::Seconds, Nullable).erased());
274        let scalar = Scalar::new(local_dtype, Some(ScalarValue::Primitive(PValue::I64(0))));
275        assert_eq!(format!("{}", scalar.as_extension()), "1970-01-01T00:00:00Z");
276
277        // Zoned timestamp.
278        let zoned_dtype = DType::Extension(
279            Timestamp::new_with_tz(
280                TimeUnit::Seconds,
281                Some(Arc::from("America/New_York")),
282                Nullable,
283            )
284            .erased(),
285        );
286        let scalar = Scalar::new(zoned_dtype, Some(ScalarValue::Primitive(PValue::I64(0))));
287        assert_eq!(
288            format!("{}", scalar.as_extension()),
289            "1969-12-31T19:00:00-05:00[America/New_York]"
290        );
291    }
292
293    #[test]
294    fn deserialize_empty_metadata_returns_error() {
295        use crate::dtype::extension::ExtVTable;
296
297        let vtable = Timestamp;
298        assert!(vtable.deserialize_metadata(&[]).is_err());
299    }
300
301    #[test]
302    fn deserialize_too_short_metadata_returns_error() {
303        use crate::dtype::extension::ExtVTable;
304
305        let vtable = Timestamp;
306        // Only 2 bytes - too short for the required 3-byte header.
307        assert!(vtable.deserialize_metadata(&[0x00, 0x01]).is_err());
308    }
309
310    #[test]
311    fn deserialize_truncated_timezone_returns_error() {
312        use crate::dtype::extension::ExtVTable;
313
314        let vtable = Timestamp;
315        // Valid tag (0x00 = Nanoseconds), tz_len = 10 (little-endian: [0x0A, 0x00]),
316        // but only 3 bytes of timezone data instead of the declared 10.
317        let data = [0x00u8, 0x0A, 0x00, b'U', b'T', b'C'];
318        assert!(vtable.deserialize_metadata(&data).is_err());
319    }
320}
vortex_array/extension/datetime/timestamp.rs

vortex_array/extension/datetime/
timestamp.rs
