datatypes/
timestamp.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_array::{
16    Array, ArrayRef, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray,
17    TimestampNanosecondArray, TimestampSecondArray,
18};
19use arrow_schema::DataType;
20use common_time::Timestamp;
21use common_time::timestamp::TimeUnit;
22use paste::paste;
23use serde::{Deserialize, Serialize};
24
25use crate::prelude::{Scalar, Value, ValueRef};
26use crate::scalars::ScalarRef;
27use crate::types::{
28    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
29    TimestampSecondType, WrapperType,
30};
31use crate::vectors::{
32    TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector,
33    TimestampSecondVector,
34};
35
36macro_rules! define_timestamp_with_unit {
37    ($unit: ident) => {
38        paste! {
39            #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
40            pub struct [<Timestamp $unit>](pub Timestamp);
41
42            impl [<Timestamp $unit>] {
43                pub fn new(val: i64) -> Self {
44                    Self(Timestamp::new(val, TimeUnit::$unit))
45                }
46            }
47
48            impl Default for [<Timestamp $unit>] {
49                fn default() -> Self {
50                    Self::new(0)
51                }
52            }
53
54            impl From<[<Timestamp $unit>]> for Value {
55                fn from(t: [<Timestamp $unit>]) -> Value {
56                    Value::Timestamp(t.0)
57                }
58            }
59
60            impl From<[<Timestamp $unit>]> for serde_json::Value {
61                fn from(t: [<Timestamp $unit>]) -> Self {
62                    t.0.into()
63                }
64            }
65
66            impl From<[<Timestamp $unit>]> for ValueRef<'static> {
67                fn from(t: [<Timestamp $unit>]) -> Self {
68                    ValueRef::Timestamp(t.0)
69                }
70            }
71
72            impl Scalar for [<Timestamp $unit>] {
73                type VectorType = [<Timestamp $unit Vector>];
74                type RefType<'a> = [<Timestamp $unit>];
75
76                fn as_scalar_ref(&self) -> Self::RefType<'_> {
77                    *self
78                }
79
80                fn upcast_gat<'short, 'long: 'short>(
81                    long: Self::RefType<'long>,
82                ) -> Self::RefType<'short> {
83                    long
84                }
85            }
86
87            impl<'a> ScalarRef<'a> for [<Timestamp $unit>] {
88                type ScalarType = [<Timestamp $unit>];
89
90                fn to_owned_scalar(&self) -> Self::ScalarType {
91                    *self
92                }
93            }
94
95            impl WrapperType for [<Timestamp $unit>] {
96                type LogicalType = [<Timestamp $unit Type>];
97                type Native = i64;
98
99                fn from_native(value: Self::Native) -> Self {
100                    Self::new(value)
101                }
102
103                fn into_native(self) -> Self::Native {
104                    self.0.into()
105                }
106            }
107
108            impl From<i64> for [<Timestamp $unit>] {
109                fn from(val: i64) -> Self {
110                    [<Timestamp $unit>]::from_native(val)
111                }
112            }
113
114            impl From<[<Timestamp $unit>]> for i64{
115                fn from(val: [<Timestamp $unit>]) -> Self {
116                    val.0.value()
117                }
118            }
119
120            impl TryFrom<Value> for Option<[<Timestamp $unit>]> {
121                type Error = $crate::error::Error;
122
123                #[inline]
124                fn try_from(from: Value) -> std::result::Result<Self, Self::Error> {
125                    match from {
126                        Value::Timestamp(v) if v.unit() == TimeUnit::$unit => {
127                            Ok(Some([<Timestamp $unit>](v)))
128                        },
129                        Value::Null => Ok(None),
130                        _ => $crate::error::TryFromValueSnafu {
131                            reason: format!("{:?} is not a {}", from, stringify!([<Timestamp $unit>])),
132                        }
133                        .fail(),
134                    }
135                }
136            }
137        }
138    };
139}
140
141define_timestamp_with_unit!(Second);
142define_timestamp_with_unit!(Millisecond);
143define_timestamp_with_unit!(Microsecond);
144define_timestamp_with_unit!(Nanosecond);
145
146/// Converts a timestamp array to a primitive array and the time unit.
147pub fn timestamp_array_to_primitive(
148    ts_array: &ArrayRef,
149) -> Option<(
150    PrimitiveArray<arrow_array::types::Int64Type>,
151    arrow::datatypes::TimeUnit,
152)> {
153    let DataType::Timestamp(unit, _) = ts_array.data_type() else {
154        return None;
155    };
156
157    let ts_primitive = match unit {
158        arrow_schema::TimeUnit::Second => ts_array
159            .as_any()
160            .downcast_ref::<TimestampSecondArray>()
161            .unwrap()
162            .reinterpret_cast::<arrow_array::types::Int64Type>(),
163        arrow_schema::TimeUnit::Millisecond => ts_array
164            .as_any()
165            .downcast_ref::<TimestampMillisecondArray>()
166            .unwrap()
167            .reinterpret_cast::<arrow_array::types::Int64Type>(),
168        arrow_schema::TimeUnit::Microsecond => ts_array
169            .as_any()
170            .downcast_ref::<TimestampMicrosecondArray>()
171            .unwrap()
172            .reinterpret_cast::<arrow_array::types::Int64Type>(),
173        arrow_schema::TimeUnit::Nanosecond => ts_array
174            .as_any()
175            .downcast_ref::<TimestampNanosecondArray>()
176            .unwrap()
177            .reinterpret_cast::<arrow_array::types::Int64Type>(),
178    };
179    Some((ts_primitive, *unit))
180}
181
182/// Appends non-null timestamps in the source array's native time unit.
183///
184/// Returns `None` for a non-timestamp array without changing `timestamps`.
185pub fn append_timestamps(ts_array: &ArrayRef, timestamps: &mut Vec<i64>) -> Option<()> {
186    let (values, _) = timestamp_array_to_primitive(ts_array)?;
187    if values.null_count() == 0 {
188        timestamps.extend_from_slice(values.values());
189    } else {
190        timestamps.extend(values.iter().flatten());
191    }
192    Some(())
193}
194
195#[cfg(test)]
196mod tests {
197    use std::sync::Arc;
198
199    use arrow_array::Int64Array;
200    use common_time::timezone::set_default_timezone;
201
202    use super::*;
203
204    #[test]
205    fn test_to_serde_json_value() {
206        set_default_timezone(Some("Asia/Shanghai")).unwrap();
207        let ts = TimestampSecond::new(123);
208        let val = serde_json::Value::from(ts);
209        match val {
210            serde_json::Value::String(s) => {
211                assert_eq!("1970-01-01 08:02:03+0800", s);
212            }
213            _ => unreachable!(),
214        }
215    }
216
217    #[test]
218    fn test_timestamp_scalar() {
219        let ts = TimestampSecond::new(123);
220        assert_eq!(ts, ts.as_scalar_ref());
221        assert_eq!(ts, ts.to_owned_scalar());
222        let ts = TimestampMillisecond::new(123);
223        assert_eq!(ts, ts.as_scalar_ref());
224        assert_eq!(ts, ts.to_owned_scalar());
225        let ts = TimestampMicrosecond::new(123);
226        assert_eq!(ts, ts.as_scalar_ref());
227        assert_eq!(ts, ts.to_owned_scalar());
228        let ts = TimestampNanosecond::new(123);
229        assert_eq!(ts, ts.as_scalar_ref());
230        assert_eq!(ts, ts.to_owned_scalar());
231    }
232
233    #[test]
234    fn test_append_timestamps() {
235        let cases = [
236            vec![Some(i64::MIN), Some(-1), Some(0), Some(i64::MAX)],
237            vec![Some(-1), None, Some(2), None],
238            vec![None, None],
239            vec![],
240        ];
241        for values in cases {
242            let arrays: [ArrayRef; 4] = [
243                Arc::new(TimestampSecondArray::from(values.clone())),
244                Arc::new(TimestampMillisecondArray::from(values.clone())),
245                Arc::new(TimestampMicrosecondArray::from(values.clone())),
246                Arc::new(TimestampNanosecondArray::from(values.clone())),
247            ];
248            let mut expected = vec![42];
249            expected.extend(values.iter().flatten().copied());
250            for array in arrays {
251                for _ in 0..2 {
252                    let (primitive, _) = timestamp_array_to_primitive(&array).unwrap();
253                    let mut reference = vec![42];
254                    reference.extend(primitive.iter().flatten());
255
256                    let mut timestamps = vec![42];
257                    assert_eq!(append_timestamps(&array, &mut timestamps), Some(()));
258                    assert_eq!(timestamps, expected);
259                    assert_eq!(timestamps, reference);
260                }
261            }
262        }
263    }
264
265    #[test]
266    fn test_append_timestamps_invalid_array_preserves_prefix() {
267        for values in [vec![], vec![Some(1), None]] {
268            let array: ArrayRef = Arc::new(Int64Array::from(values));
269            let mut timestamps = vec![42, -1];
270            assert_eq!(append_timestamps(&array, &mut timestamps), None);
271            assert_eq!(timestamps, vec![42, -1]);
272        }
273    }
274}
datatypes/timestamp.rs

datatypes/
timestamp.rs
