datatypes/types/
interval_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, IntervalDayTime as ArrowIntervalDayTime,
17    IntervalDayTimeType as ArrowIntervalDayTimeType,
18    IntervalMonthDayNano as ArrowIntervalMonthDayNano,
19    IntervalMonthDayNanoType as ArrowIntervalMonthDayNanoType, IntervalUnit as ArrowIntervalUnit,
20    IntervalYearMonthType as ArrowIntervalYearMonthType,
21};
22use common_time::interval::{
23    IntervalDayTime, IntervalMonthDayNano, IntervalUnit, IntervalYearMonth,
24};
25use enum_dispatch::enum_dispatch;
26use paste::paste;
27use serde::{Deserialize, Serialize};
28use snafu::OptionExt;
29
30use crate::data_type::ConcreteDataType;
31use crate::error;
32use crate::prelude::{
33    DataType, LogicalTypeId, MutableVector, ScalarVectorBuilder, Value, ValueRef, Vector,
34};
35use crate::types::LogicalPrimitiveType;
36use crate::vectors::{
37    IntervalDayTimeVector, IntervalDayTimeVectorBuilder, IntervalMonthDayNanoVector,
38    IntervalMonthDayNanoVectorBuilder, IntervalYearMonthVector, IntervalYearMonthVectorBuilder,
39    PrimitiveVector,
40};
41
42/// The "calendar" interval is a type of time interval that does not
43/// have a precise duration without taking into account a specific
44/// base timestamp.
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
46#[enum_dispatch(DataType)]
47pub enum IntervalType {
48    YearMonth(IntervalYearMonthType),
49    DayTime(IntervalDayTimeType),
50    MonthDayNano(IntervalMonthDayNanoType),
51}
52
53impl IntervalType {
54    /// Returns the unit of the interval.
55    pub fn unit(&self) -> IntervalUnit {
56        match self {
57            IntervalType::YearMonth(_) => IntervalUnit::YearMonth,
58            IntervalType::DayTime(_) => IntervalUnit::DayTime,
59            IntervalType::MonthDayNano(_) => IntervalUnit::MonthDayNano,
60        }
61    }
62}
63
64macro_rules! impl_data_type_for_interval {
65    ($unit: ident, $type: ty) => {
66        paste! {
67            #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
68            pub struct [<Interval $unit Type>];
69
70            impl DataType for [<Interval $unit Type>] {
71                fn name(&self) -> String {
72                    stringify!([<Interval $unit>]).to_string()
73                }
74
75                fn logical_type_id(&self) -> LogicalTypeId {
76                    LogicalTypeId::[<Interval $unit>]
77                }
78
79                fn default_value(&self) -> Value {
80                    Value::[<Interval $unit>]([<Interval $unit>]::default())
81                }
82
83                fn as_arrow_type(&self) -> ArrowDataType {
84                    ArrowDataType::Interval(ArrowIntervalUnit::$unit)
85                }
86
87                fn create_mutable_vector(&self, capacity: usize) -> Box<dyn MutableVector> {
88                    Box::new([<Interval $unit Vector Builder>]::with_capacity(capacity))
89                }
90
91
92                fn try_cast(&self, _: Value) -> Option<Value> {
93                    // TODO(QuenKar): Implement casting for interval types.
94                    None
95                }
96            }
97
98            impl LogicalPrimitiveType for [<Interval $unit Type>] {
99                type ArrowPrimitive = [<Arrow Interval $unit Type>];
100                type Native = $type;
101                type Wrapper = [<Interval $unit>];
102                type LargestType = Self;
103
104                fn build_data_type() -> ConcreteDataType {
105                    ConcreteDataType::Interval(IntervalType::$unit(
106                        [<Interval $unit Type>]::default(),
107                    ))
108                }
109
110                fn type_name() -> &'static str {
111                    stringify!([<Interval $unit Type>])
112                }
113
114                fn cast_vector(vector: &dyn Vector) -> crate::Result<&PrimitiveVector<Self>> {
115                    vector
116                        .as_any()
117                        .downcast_ref::<[<Interval $unit Vector>]>()
118                        .with_context(|| error::CastTypeSnafu {
119                            msg: format!(
120                                "Failed to cast {} to {}",
121                                vector.vector_type_name(), stringify!([<Interval $unit Vector>])
122                            ),
123                        })
124                }
125
126                fn cast_value_ref(value: ValueRef) -> crate::Result<Option<Self::Wrapper>> {
127                    match value {
128                        ValueRef::Null => Ok(None),
129                        ValueRef::[<Interval $unit>](t) => Ok(Some(t)),
130                        other => error::CastTypeSnafu {
131                            msg: format!("Failed to cast value {:?} to {}", other, stringify!([<Interval $unit>])),
132                        }
133                        .fail(),
134                    }
135                }
136            }
137        }
138    }
139}
140
141impl_data_type_for_interval!(YearMonth, i32);
142impl_data_type_for_interval!(DayTime, ArrowIntervalDayTime);
143impl_data_type_for_interval!(MonthDayNano, ArrowIntervalMonthDayNano);
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn test_interval_type_unit() {
151        assert_eq!(
152            IntervalUnit::DayTime,
153            IntervalType::DayTime(IntervalDayTimeType).unit()
154        );
155        assert_eq!(
156            IntervalUnit::MonthDayNano,
157            IntervalType::MonthDayNano(IntervalMonthDayNanoType).unit()
158        );
159        assert_eq!(
160            IntervalUnit::YearMonth,
161            IntervalType::YearMonth(IntervalYearMonthType).unit()
162        );
163    }
164
165    #[test]
166    fn test_interval_as_arrow_type() {
167        assert_eq!(
168            ArrowDataType::Interval(ArrowIntervalUnit::DayTime),
169            IntervalType::DayTime(IntervalDayTimeType).as_arrow_type()
170        );
171        assert_eq!(
172            ArrowDataType::Interval(ArrowIntervalUnit::MonthDayNano),
173            IntervalType::MonthDayNano(IntervalMonthDayNanoType).as_arrow_type()
174        );
175        assert_eq!(
176            ArrowDataType::Interval(ArrowIntervalUnit::YearMonth),
177            IntervalType::YearMonth(IntervalYearMonthType).as_arrow_type()
178        );
179    }
180}
datatypes/types/interval_type.rs

datatypes/types/
interval_type.rs
