Skip to main content

vortex_array/extension/datetime/
matcher.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::sync::Arc;
6
7use vortex_error::VortexResult;
8
9use crate::dtype::extension::ExtDTypeRef;
10use crate::dtype::extension::ExtVTable;
11use crate::dtype::extension::Matcher;
12use crate::extension::datetime::Date;
13use crate::extension::datetime::Time;
14use crate::extension::datetime::TimeUnit;
15use crate::extension::datetime::Timestamp;
16
17/// Matcher for temporal extension data types.
18pub struct AnyTemporal;
19
20impl Matcher for AnyTemporal {
21    type Match<'a> = TemporalMetadata<'a>;
22
23    fn try_match<'a>(item: &'a ExtDTypeRef) -> Option<Self::Match<'a>> {
24        if let Some(opts) = item.metadata_opt::<Timestamp>() {
25            return Some(TemporalMetadata::Timestamp(&opts.unit, &opts.tz));
26        }
27        if let Some(opts) = item.metadata_opt::<Date>() {
28            return Some(TemporalMetadata::Date(opts));
29        }
30        if let Some(opts) = item.metadata_opt::<Time>() {
31            return Some(TemporalMetadata::Time(opts));
32        }
33        None
34    }
35}
36
37/// Metadata for temporal extension data types.
38#[derive(Debug, PartialEq, Eq)]
39pub enum TemporalMetadata<'a> {
40    /// Metadata for Timestamp dtypes, a tuple of time unit and optional timezone.
41    Timestamp(&'a TimeUnit, &'a Option<Arc<str>>),
42    /// Metadata for Date dtypes
43    Date(&'a <Date as ExtVTable>::Metadata),
44    /// Metadata for Time dtypes
45    Time(&'a <Time as ExtVTable>::Metadata),
46}
47
48// TODO(ngates): remove this logic in favor of having an ExtScalarVTable in vortex_array::scalar.
49//  Currently this is used largely to implement scalar display hacks.
50impl TemporalMetadata<'_> {
51    /// Get the time unit of the temporal dtype.
52    pub fn time_unit(&self) -> TimeUnit {
53        match self {
54            TemporalMetadata::Time(unit) => **unit,
55            TemporalMetadata::Date(unit) => **unit,
56            TemporalMetadata::Timestamp(unit, _tz) => **unit,
57        }
58    }
59
60    /// Convert a timestamp value to a Jiff value.
61    pub fn to_jiff(&self, v: i64) -> VortexResult<TemporalJiff> {
62        match self {
63            TemporalMetadata::Time(unit) => Ok(TemporalJiff::Time(
64                jiff::civil::Time::MIN.checked_add(unit.to_jiff_span(v)?)?,
65            )),
66            TemporalMetadata::Date(unit) => Ok(TemporalJiff::Date(
67                jiff::civil::Date::new(1970, 1, 1)?.checked_add(unit.to_jiff_span(v)?)?,
68            )),
69            TemporalMetadata::Timestamp(unit, tz) => match tz {
70                None => Ok(TemporalJiff::Unzoned(
71                    jiff::civil::DateTime::new(1970, 1, 1, 0, 0, 0, 0)?
72                        .checked_add(unit.to_jiff_span(v)?)?,
73                )),
74                Some(tz) => Ok(TemporalJiff::Zoned(
75                    jiff::Timestamp::UNIX_EPOCH
76                        .checked_add(unit.to_jiff_span(v)?)?
77                        .in_tz(tz.as_ref())?,
78                )),
79            },
80        }
81    }
82}
83
84/// A Jiff representation of a temporal value.
85pub enum TemporalJiff {
86    /// A time value.
87    Time(jiff::civil::Time),
88    /// A date value.
89    Date(jiff::civil::Date),
90    /// A zone-naive timestamp value.
91    Unzoned(jiff::civil::DateTime),
92    /// A zoned timestamp value.
93    Zoned(jiff::Zoned),
94}
95
96impl Display for TemporalJiff {
97    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98        match self {
99            TemporalJiff::Time(t) => write!(f, "{t}"),
100            TemporalJiff::Date(d) => write!(f, "{d}"),
101            TemporalJiff::Unzoned(dt) => write!(f, "{dt}"),
102            TemporalJiff::Zoned(z) => write!(f, "{z}"),
103        }
104    }
105}