datatypes/types/
duration_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
15use arrow::datatypes::{
16    DataType as ArrowDataType, DurationMicrosecondType as ArrowDurationMicrosecondType,
17    DurationMillisecondType as ArrowDurationMillisecondType,
18    DurationNanosecondType as ArrowDurationNanosecondType,
19    DurationSecondType as ArrowDurationSecondType, TimeUnit as ArrowTimeUnit,
20};
21use common_time::duration::Duration;
22use common_time::timestamp::TimeUnit;
23use enum_dispatch::enum_dispatch;
24use paste::paste;
25use serde::{Deserialize, Serialize};
26use snafu::OptionExt;
27
28use crate::data_type::DataType;
29use crate::duration::{
30    DurationMicrosecond, DurationMillisecond, DurationNanosecond, DurationSecond,
31};
32use crate::error;
33use crate::prelude::{
34    ConcreteDataType, LogicalTypeId, MutableVector, ScalarVectorBuilder, Value, ValueRef, Vector,
35};
36use crate::types::LogicalPrimitiveType;
37use crate::vectors::{
38    DurationMicrosecondVector, DurationMicrosecondVectorBuilder, DurationMillisecondVector,
39    DurationMillisecondVectorBuilder, DurationNanosecondVector, DurationNanosecondVectorBuilder,
40    DurationSecondVector, DurationSecondVectorBuilder, PrimitiveVector,
41};
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
44#[enum_dispatch(DataType)]
45pub enum DurationType {
46    Second(DurationSecondType),
47    Millisecond(DurationMillisecondType),
48    Microsecond(DurationMicrosecondType),
49    Nanosecond(DurationNanosecondType),
50}
51
52impl DurationType {
53    /// Creates time type from `TimeUnit`.
54    pub fn from_unit(unit: TimeUnit) -> Self {
55        match unit {
56            TimeUnit::Second => Self::Second(DurationSecondType),
57            TimeUnit::Millisecond => Self::Millisecond(DurationMillisecondType),
58            TimeUnit::Microsecond => Self::Microsecond(DurationMicrosecondType),
59            TimeUnit::Nanosecond => Self::Nanosecond(DurationNanosecondType),
60        }
61    }
62
63    /// Returns the [`TimeUnit`] of this type.
64    pub fn unit(&self) -> TimeUnit {
65        match self {
66            DurationType::Second(_) => TimeUnit::Second,
67            DurationType::Millisecond(_) => TimeUnit::Millisecond,
68            DurationType::Microsecond(_) => TimeUnit::Microsecond,
69            DurationType::Nanosecond(_) => TimeUnit::Nanosecond,
70        }
71    }
72}
73
74macro_rules! impl_data_type_for_duration {
75    ($unit: ident) => {
76        paste! {
77            #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
78            pub struct [<Duration $unit Type>];
79
80            impl DataType for [<Duration $unit Type>] {
81                fn name(&self) -> String {
82                    stringify!([<Duration $unit>]).to_string()
83                }
84
85                fn logical_type_id(&self) -> LogicalTypeId {
86                    LogicalTypeId::[<Duration $unit>]
87                }
88
89                fn default_value(&self) -> Value {
90                    Value::Duration(Duration::new(0, TimeUnit::$unit))
91                }
92
93                fn as_arrow_type(&self) -> ArrowDataType {
94                    ArrowDataType::Duration(ArrowTimeUnit::$unit)
95                }
96
97                fn create_mutable_vector(&self, capacity: usize) -> Box<dyn MutableVector> {
98                    Box::new([<Duration $unit Vector Builder>]::with_capacity(capacity))
99                }
100
101
102                fn try_cast(&self, _: Value) -> Option<Value> {
103                    // TODO(QuenKar): Implement casting for duration types.
104                    None
105                }
106            }
107
108            impl LogicalPrimitiveType for [<Duration $unit Type>] {
109                type ArrowPrimitive = [<Arrow Duration $unit Type>];
110                type Native = i64;
111                type Wrapper = [<Duration $unit>];
112                type LargestType = Self;
113
114                fn build_data_type() -> ConcreteDataType {
115                    ConcreteDataType::Duration(DurationType::$unit(
116                        [<Duration $unit Type>]::default(),
117                    ))
118                }
119
120                fn type_name() -> &'static str {
121                    stringify!([<Duration $unit Type>])
122                }
123
124                fn cast_vector(vector: &dyn Vector) -> crate::Result<&PrimitiveVector<Self>> {
125                    vector
126                        .as_any()
127                        .downcast_ref::<[<Duration $unit Vector>]>()
128                        .with_context(|| error::CastTypeSnafu {
129                            msg: format!(
130                                "Failed to cast {} to {}",
131                                vector.vector_type_name(), stringify!([<Duration $unit Vector>])
132                            ),
133                        })
134                }
135
136                fn cast_value_ref(value: ValueRef) -> crate::Result<Option<Self::Wrapper>> {
137                    match value {
138                        ValueRef::Null => Ok(None),
139                        ValueRef::Duration(t) => match t.unit() {
140                            TimeUnit::$unit => Ok(Some([<Duration $unit>](t))),
141                            other => error::CastTypeSnafu {
142                                msg: format!(
143                                    "Failed to cast Duration value with different unit {:?} to {}",
144                                    other, stringify!([<Duration $unit>])
145                                ),
146                            }
147                            .fail(),
148                        },
149                        other => error::CastTypeSnafu {
150                            msg: format!("Failed to cast value {:?} to {}", other, stringify!([<Duration $unit>])),
151                        }
152                        .fail(),
153                    }
154                }
155            }
156        }
157    }
158}
159
160impl_data_type_for_duration!(Second);
161impl_data_type_for_duration!(Millisecond);
162impl_data_type_for_duration!(Microsecond);
163impl_data_type_for_duration!(Nanosecond);
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn test_duration_type_unit() {
171        assert_eq!(
172            TimeUnit::Second,
173            DurationType::Second(DurationSecondType).unit()
174        );
175
176        assert_eq!(
177            TimeUnit::Millisecond,
178            DurationType::Millisecond(DurationMillisecondType).unit()
179        );
180
181        assert_eq!(
182            TimeUnit::Microsecond,
183            DurationType::Microsecond(DurationMicrosecondType).unit()
184        );
185
186        assert_eq!(
187            TimeUnit::Nanosecond,
188            DurationType::Nanosecond(DurationNanosecondType).unit()
189        );
190    }
191
192    #[test]
193    fn test_from_unit() {
194        assert_eq!(
195            DurationType::Second(DurationSecondType),
196            DurationType::from_unit(TimeUnit::Second)
197        );
198
199        assert_eq!(
200            DurationType::Millisecond(DurationMillisecondType),
201            DurationType::from_unit(TimeUnit::Millisecond)
202        );
203
204        assert_eq!(
205            DurationType::Microsecond(DurationMicrosecondType),
206            DurationType::from_unit(TimeUnit::Microsecond)
207        );
208
209        assert_eq!(
210            DurationType::Nanosecond(DurationNanosecondType),
211            DurationType::from_unit(TimeUnit::Nanosecond)
212        );
213    }
214}
datatypes/types/duration_type.rs

datatypes/types/
duration_type.rs
