datatypes/types/
time_type.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! TimeType represents the elapsed time since midnight in the unit of `TimeUnit`.
16
17use 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    /// Creates time type from `TimeUnit`.
58    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    /// Returns the time type's `TimeUnit`.
68    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    /// Returns the time type's precision.
78    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        // Int32 -> TimeSecondType
228        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        // Int32 -> TimeMillisecondType
235        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        // Int64 -> TimeMicrosecondType
242        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        // Int64 -> TimeNanosecondType
249        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        // Other situations will return None, such as Int64 -> TimeSecondType or
256        // Int32 -> TimeMicrosecondType etc.
257        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        // TimeSecond -> TimeMicroSecond
266        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        // test overflow
276        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}