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_ref("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!(data.len() >= 3);
145
146        let tag = data[0];
147        let time_unit = TimeUnit::try_from(tag)?;
148        let tz_len_bytes: [u8; 2] = data[1..3]
149            .try_into()
150            .ok()
151            .vortex_expect("Verified to have two bytes");
152        let tz_len = u16::from_le_bytes(tz_len_bytes) as usize;
153        if tz_len == 0 {
154            return Ok(TimestampOptions {
155                unit: time_unit,
156                tz: None,
157            });
158        }
159
160        // Attempt to load from len-prefixed bytes
161        let tz_bytes = &data[3..][..tz_len];
162        let tz: Arc<str> = str::from_utf8(tz_bytes)
163            .map_err(|e| vortex_err!("timezone is not valid utf8 string: {e}"))?
164            .to_string()
165            .into();
166        Ok(TimestampOptions {
167            unit: time_unit,
168            tz: Some(tz),
169        })
170    }
171
172    fn validate_dtype(&self, ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
173        vortex_ensure!(
174            matches!(ext_dtype.storage_dtype(), DType::Primitive(PType::I64, _)),
175            "Timestamp storage dtype must be i64"
176        );
177        Ok(())
178    }
179
180    fn unpack_native<'a>(
181        &self,
182        ext_dtype: &'a ExtDType<Self>,
183        storage_value: &'a ScalarValue,
184    ) -> VortexResult<Self::NativeValue<'a>> {
185        let metadata = ext_dtype.metadata();
186        let ts_value = storage_value.as_primitive().cast::<i64>()?;
187        let tz = metadata.tz.as_ref();
188
189        let (span, value) = match metadata.unit {
190            TimeUnit::Nanoseconds => (
191                Span::new().nanoseconds(ts_value),
192                TimestampValue::Nanoseconds(ts_value, tz),
193            ),
194            TimeUnit::Microseconds => (
195                Span::new().microseconds(ts_value),
196                TimestampValue::Microseconds(ts_value, tz),
197            ),
198            TimeUnit::Milliseconds => (
199                Span::new().milliseconds(ts_value),
200                TimestampValue::Milliseconds(ts_value, tz),
201            ),
202            TimeUnit::Seconds => (
203                Span::new().seconds(ts_value),
204                TimestampValue::Seconds(ts_value, tz),
205            ),
206            TimeUnit::Days => vortex_bail!("Timestamp does not support Days time unit"),
207        };
208
209        // Validate the storage value is within the valid range for Timestamp.
210        let ts = jiff::Timestamp::UNIX_EPOCH
211            .checked_add(span)
212            .map_err(|e| vortex_err!("Invalid timestamp scalar: {}", e))?;
213
214        if let Some(tz) = tz {
215            ts.in_tz(tz.as_ref())
216                .map_err(|e| vortex_err!("Invalid timezone for timestamp scalar: {}", e))?;
217        }
218
219        Ok(value)
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use std::sync::Arc;
226
227    use vortex_error::VortexResult;
228
229    use crate::dtype::DType;
230    use crate::dtype::Nullability::Nullable;
231    use crate::extension::datetime::TimeUnit;
232    use crate::extension::datetime::Timestamp;
233    use crate::scalar::PValue;
234    use crate::scalar::Scalar;
235    use crate::scalar::ScalarValue;
236
237    #[test]
238    fn validate_timestamp_scalar() -> VortexResult<()> {
239        let dtype = DType::Extension(Timestamp::new(TimeUnit::Seconds, Nullable).erased());
240        Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(0))))?;
241
242        Ok(())
243    }
244
245    #[cfg_attr(miri, ignore)]
246    #[test]
247    fn reject_timestamp_with_invalid_timezone() {
248        let dtype = DType::Extension(
249            Timestamp::new_with_tz(
250                TimeUnit::Seconds,
251                Some(Arc::from("Not/A/Timezone")),
252                Nullable,
253            )
254            .erased(),
255        );
256        let result = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(0))));
257        assert!(result.is_err());
258    }
259
260    #[cfg_attr(miri, ignore)]
261    #[test]
262    fn display_timestamp_scalar() {
263        // Local (no timezone) timestamp.
264        let local_dtype = DType::Extension(Timestamp::new(TimeUnit::Seconds, Nullable).erased());
265        let scalar = Scalar::new(local_dtype, Some(ScalarValue::Primitive(PValue::I64(0))));
266        assert_eq!(format!("{}", scalar.as_extension()), "1970-01-01T00:00:00Z");
267
268        // Zoned timestamp.
269        let zoned_dtype = DType::Extension(
270            Timestamp::new_with_tz(
271                TimeUnit::Seconds,
272                Some(Arc::from("America/New_York")),
273                Nullable,
274            )
275            .erased(),
276        );
277        let scalar = Scalar::new(zoned_dtype, Some(ScalarValue::Primitive(PValue::I64(0))));
278        assert_eq!(
279            format!("{}", scalar.as_extension()),
280            "1969-12-31T19:00:00-05:00[America/New_York]"
281        );
282    }
283}
vortex_array/extension/datetime/timestamp.rs

vortex_array/extension/datetime/
timestamp.rs
