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

vortex_array/extension/datetime/
timestamp.rs
