datafusion_comet_spark_expr/
timezone.rs1use arrow::error::ArrowError;
20use chrono::{
21 format::{parse, Parsed, StrftimeItems},
22 offset::TimeZone,
23 FixedOffset, LocalResult, NaiveDate, NaiveDateTime, Offset,
24};
25use std::str::FromStr;
26
27fn parse_fixed_offset(tz: &str) -> Result<FixedOffset, ArrowError> {
29 let mut parsed = Parsed::new();
30
31 if let Ok(fixed_offset) =
32 parse(&mut parsed, tz, StrftimeItems::new("%:z")).and_then(|_| parsed.to_fixed_offset())
33 {
34 return Ok(fixed_offset);
35 }
36
37 if let Ok(fixed_offset) =
38 parse(&mut parsed, tz, StrftimeItems::new("%#z")).and_then(|_| parsed.to_fixed_offset())
39 {
40 return Ok(fixed_offset);
41 }
42
43 Err(ArrowError::ParseError(format!(
44 "Invalid timezone \"{tz}\": Expected format [+-]XX:XX, [+-]XX, or [+-]XXXX"
45 )))
46}
47
48#[derive(Debug, Copy, Clone)]
50pub struct TzOffset {
51 tz: Tz,
52 offset: FixedOffset,
53}
54
55impl std::fmt::Display for TzOffset {
56 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57 self.offset.fmt(f)
58 }
59}
60
61impl Offset for TzOffset {
62 fn fix(&self) -> FixedOffset {
63 self.offset
64 }
65}
66
67#[derive(Debug, Copy, Clone)]
69pub struct Tz(TzInner);
70
71#[derive(Debug, Copy, Clone)]
72enum TzInner {
73 Timezone(chrono_tz::Tz),
74 Offset(FixedOffset),
75}
76
77impl FromStr for Tz {
78 type Err = ArrowError;
79
80 fn from_str(tz: &str) -> Result<Self, Self::Err> {
81 if tz.starts_with('+') || tz.starts_with('-') {
82 Ok(Self(TzInner::Offset(parse_fixed_offset(tz)?)))
83 } else {
84 Ok(Self(TzInner::Timezone(tz.parse().map_err(|e| {
85 ArrowError::ParseError(format!("Invalid timezone \"{tz}\": {e}"))
86 })?)))
87 }
88 }
89}
90
91macro_rules! tz {
92 ($s:ident, $tz:ident, $b:block) => {
93 match $s.0 {
94 TzInner::Timezone($tz) => $b,
95 TzInner::Offset($tz) => $b,
96 }
97 };
98}
99
100impl TimeZone for Tz {
101 type Offset = TzOffset;
102
103 fn from_offset(offset: &Self::Offset) -> Self {
104 offset.tz
105 }
106
107 fn offset_from_local_date(&self, local: &NaiveDate) -> LocalResult<Self::Offset> {
108 tz!(self, tz, {
109 tz.offset_from_local_date(local).map(|x| TzOffset {
110 tz: *self,
111 offset: x.fix(),
112 })
113 })
114 }
115
116 fn offset_from_local_datetime(&self, local: &NaiveDateTime) -> LocalResult<Self::Offset> {
117 tz!(self, tz, {
118 tz.offset_from_local_datetime(local).map(|x| TzOffset {
119 tz: *self,
120 offset: x.fix(),
121 })
122 })
123 }
124
125 fn offset_from_utc_date(&self, utc: &NaiveDate) -> Self::Offset {
126 tz!(self, tz, {
127 TzOffset {
128 tz: *self,
129 offset: tz.offset_from_utc_date(utc).fix(),
130 }
131 })
132 }
133
134 fn offset_from_utc_datetime(&self, utc: &NaiveDateTime) -> Self::Offset {
135 tz!(self, tz, {
136 TzOffset {
137 tz: *self,
138 offset: tz.offset_from_utc_datetime(utc).fix(),
139 }
140 })
141 }
142}