1use arrow::datatypes::{
18 DataType as ArrowDataType, Time32MillisecondType as ArrowTimeMillisecondType,
19 Time32SecondType as ArrowTimeSecondType, Time64MicrosecondType as ArrowTimeMicrosecondType,
20 Time64NanosecondType as ArrowTimeNanosecondType, TimeUnit as ArrowTimeUnit,
21};
22use common_time::time::Time;
23use common_time::timestamp::TimeUnit;
24use enum_dispatch::enum_dispatch;
25use paste::paste;
26use serde::{Deserialize, Serialize};
27use snafu::OptionExt;
28
29use crate::error;
30use crate::prelude::{
31 ConcreteDataType, DataType, LogicalTypeId, MutableVector, ScalarVectorBuilder, Value, ValueRef,
32 Vector,
33};
34use crate::time::{TimeMicrosecond, TimeMillisecond, TimeNanosecond, TimeSecond};
35use crate::types::LogicalPrimitiveType;
36use crate::vectors::{
37 PrimitiveVector, TimeMicrosecondVector, TimeMicrosecondVectorBuilder, TimeMillisecondVector,
38 TimeMillisecondVectorBuilder, TimeNanosecondVector, TimeNanosecondVectorBuilder,
39 TimeSecondVector, TimeSecondVectorBuilder,
40};
41
42const SECOND_VARIATION: u64 = 0;
43const MILLISECOND_VARIATION: u64 = 3;
44const MICROSECOND_VARIATION: u64 = 6;
45const NANOSECOND_VARIATION: u64 = 9;
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
48#[enum_dispatch(DataType)]
49pub enum TimeType {
50 Second(TimeSecondType),
51 Millisecond(TimeMillisecondType),
52 Microsecond(TimeMicrosecondType),
53 Nanosecond(TimeNanosecondType),
54}
55
56impl TimeType {
57 pub fn from_unit(unit: TimeUnit) -> Self {
59 match unit {
60 TimeUnit::Second => Self::Second(TimeSecondType),
61 TimeUnit::Millisecond => Self::Millisecond(TimeMillisecondType),
62 TimeUnit::Microsecond => Self::Microsecond(TimeMicrosecondType),
63 TimeUnit::Nanosecond => Self::Nanosecond(TimeNanosecondType),
64 }
65 }
66
67 pub fn unit(&self) -> TimeUnit {
69 match self {
70 TimeType::Second(_) => TimeUnit::Second,
71 TimeType::Millisecond(_) => TimeUnit::Millisecond,
72 TimeType::Microsecond(_) => TimeUnit::Microsecond,
73 TimeType::Nanosecond(_) => TimeUnit::Nanosecond,
74 }
75 }
76
77 pub fn precision(&self) -> u64 {
79 match self {
80 TimeType::Second(_) => SECOND_VARIATION,
81 TimeType::Millisecond(_) => MILLISECOND_VARIATION,
82 TimeType::Microsecond(_) => MICROSECOND_VARIATION,
83 TimeType::Nanosecond(_) => NANOSECOND_VARIATION,
84 }
85 }
86}
87
88macro_rules! impl_data_type_for_time {
89 ($unit: ident,$arrow_type: ident, $type: ty, $TargetType: ident) => {
90 paste! {
91 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
92 pub struct [<Time $unit Type>];
93
94 impl DataType for [<Time $unit Type>] {
95 fn name(&self) -> String {
96 stringify!([<Time $unit>]).to_string()
97 }
98
99 fn logical_type_id(&self) -> LogicalTypeId {
100 LogicalTypeId::[<Time $unit>]
101 }
102
103 fn default_value(&self) -> Value {
104 Value::Time(Time::new(0, TimeUnit::$unit))
105 }
106
107 fn as_arrow_type(&self) -> ArrowDataType {
108 ArrowDataType::$arrow_type(ArrowTimeUnit::$unit)
109 }
110
111 fn create_mutable_vector(&self, capacity: usize) -> Box<dyn MutableVector> {
112 Box::new([<Time $unit Vector Builder>]::with_capacity(capacity))
113 }
114
115 fn try_cast(&self, from: Value) -> Option<Value> {
116 match from {
117 Value::$TargetType(v) => Some(Value::Time(Time::new(v as i64, TimeUnit::$unit))),
118 Value::Time(v) => v.convert_to(TimeUnit::$unit).map(Value::Time),
119 _ => None,
120 }
121 }
122 }
123
124 impl LogicalPrimitiveType for [<Time $unit Type>] {
125 type ArrowPrimitive = [<Arrow Time $unit Type>];
126 type Native = $type;
127 type Wrapper = [<Time $unit>];
128 type LargestType = Self;
129
130 fn build_data_type() -> ConcreteDataType {
131 ConcreteDataType::Time(TimeType::$unit(
132 [<Time $unit Type>]::default(),
133 ))
134 }
135
136 fn type_name() -> &'static str {
137 stringify!([<Time $unit Type>])
138 }
139
140 fn cast_vector(vector: &dyn Vector) -> crate::Result<&PrimitiveVector<Self>> {
141 vector
142 .as_any()
143 .downcast_ref::<[<Time $unit Vector>]>()
144 .with_context(|| error::CastTypeSnafu {
145 msg: format!(
146 "Failed to cast {} to {}",
147 vector.vector_type_name(), stringify!([<Time $unit Vector>])
148 ),
149 })
150 }
151
152 fn cast_value_ref(value: ValueRef) -> crate::Result<Option<Self::Wrapper>> {
153 match value {
154 ValueRef::Null => Ok(None),
155 ValueRef::Int64(v) =>{
156 Ok(Some([<Time $unit>]::from(v)))
157 }
158 ValueRef::Time(t) => match t.unit() {
159 TimeUnit::$unit => Ok(Some([<Time $unit>](t))),
160 other => error::CastTypeSnafu {
161 msg: format!(
162 "Failed to cast Time value with different unit {:?} to {}",
163 other, stringify!([<Time $unit>])
164 ),
165 }
166 .fail(),
167 },
168 other => error::CastTypeSnafu {
169 msg: format!("Failed to cast value {:?} to {}", other, stringify!([<Time $unit>])),
170 }
171 .fail(),
172 }
173 }
174 }
175 }
176 }
177}
178
179impl_data_type_for_time!(Second, Time32, i32, Int32);
180impl_data_type_for_time!(Millisecond, Time32, i32, Int32);
181impl_data_type_for_time!(Nanosecond, Time64, i64, Int64);
182impl_data_type_for_time!(Microsecond, Time64, i64, Int64);
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187
188 #[test]
189 fn test_time_type_unit() {
190 assert_eq!(TimeUnit::Second, TimeType::Second(TimeSecondType).unit());
191 assert_eq!(
192 TimeUnit::Millisecond,
193 TimeType::Millisecond(TimeMillisecondType).unit()
194 );
195 assert_eq!(
196 TimeUnit::Microsecond,
197 TimeType::Microsecond(TimeMicrosecondType).unit()
198 );
199 assert_eq!(
200 TimeUnit::Nanosecond,
201 TimeType::Nanosecond(TimeNanosecondType).unit()
202 );
203 }
204
205 #[test]
206 fn test_as_arrow_datatype() {
207 assert_eq!(
208 ArrowDataType::Time32(ArrowTimeUnit::Second),
209 TimeSecondType.as_arrow_type()
210 );
211 assert_eq!(
212 ArrowDataType::Time32(ArrowTimeUnit::Millisecond),
213 TimeMillisecondType.as_arrow_type()
214 );
215 assert_eq!(
216 ArrowDataType::Time64(ArrowTimeUnit::Microsecond),
217 TimeMicrosecondType.as_arrow_type()
218 );
219 assert_eq!(
220 ArrowDataType::Time64(ArrowTimeUnit::Nanosecond),
221 TimeNanosecondType.as_arrow_type()
222 );
223 }
224
225 #[test]
226 fn test_time_cast() {
227 let val = Value::Int32(1000);
229 let time = ConcreteDataType::time_second_datatype()
230 .try_cast(val)
231 .unwrap();
232 assert_eq!(time, Value::Time(Time::new_second(1000)));
233
234 let val = Value::Int32(2000);
236 let time = ConcreteDataType::time_millisecond_datatype()
237 .try_cast(val)
238 .unwrap();
239 assert_eq!(time, Value::Time(Time::new_millisecond(2000)));
240
241 let val = Value::Int64(3000);
243 let time = ConcreteDataType::time_microsecond_datatype()
244 .try_cast(val)
245 .unwrap();
246 assert_eq!(time, Value::Time(Time::new_microsecond(3000)));
247
248 let val = Value::Int64(4000);
250 let time = ConcreteDataType::time_nanosecond_datatype()
251 .try_cast(val)
252 .unwrap();
253 assert_eq!(time, Value::Time(Time::new_nanosecond(4000)));
254
255 let val = Value::Int64(123);
258 let time = ConcreteDataType::time_second_datatype().try_cast(val);
259 assert_eq!(time, None);
260
261 let val = Value::Int32(123);
262 let time = ConcreteDataType::time_microsecond_datatype().try_cast(val);
263 assert_eq!(time, None);
264
265 let second = Value::Time(Time::new_second(2023));
267 let microsecond = ConcreteDataType::time_microsecond_datatype()
268 .try_cast(second)
269 .unwrap();
270 assert_eq!(
271 microsecond,
272 Value::Time(Time::new_microsecond(2023 * 1000000))
273 );
274
275 let second = Value::Time(Time::new_second(i64::MAX));
277 let microsecond = ConcreteDataType::time_microsecond_datatype().try_cast(second);
278 assert_eq!(microsecond, None);
279 }
280}