feat(type): support timestamp and timestamptz (#769)

* feat(type):support timestamp and timestamptz

Signed-off-by: skyfan <3374614481@qq.com>

* fix:fmt

Signed-off-by: skyfan <3374614481@qq.com>

* fix:simplify time zone conversion

Signed-off-by: skyfan <3374614481@qq.com>

* fix: use slt to test timestamp

Signed-off-by: skyfan <3374614481@qq.com>

* fix:Switch from CRLF to LF

Signed-off-by: skyfan <3374614481@qq.com>

* small modification

Signed-off-by: skyfan <3374614481@qq.com>

Signed-off-by: skyfan <3374614481@qq.com>

Cargo.toml

@@ -62,7 +62,7 @@ smallvec = { version = "1", features = ["serde"] }
 sqllogictest = "0.9"
 sqlparser = { version = "0.30", features = ["serde"] }
 thiserror = "1"
-tikv-jemallocator = { version = "0.5", optional = true, features=["disable_initial_exec_tls"] }
+tikv-jemallocator = { version = "0.5", optional = true, features = ["disable_initial_exec_tls"] }
 tokio = { version = "1", features = ["full"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "parking_lot"] }
@@ -89,14 +89,14 @@ lto = 'thin'
 
 [workspace]
 members = [
-  "proto",
-  "tests/sqllogictest",
-  "tests/sqlplannertest"
+    "proto",
+    "tests/sqllogictest",
+    "tests/sqlplannertest"
 ]
 
 [patch.crates-io]
 risinglight_proto = { path = "proto" }
 
 [lib]
 name = "risinglight"
-crate-type = ["cdylib",  "rlib"]
+crate-type = ["cdylib", "rlib"]

src/array/data_chunk.rs

@@ -263,6 +263,8 @@ pub fn datachunk_to_sqllogictest_string(chunk: &Chunk) -> Vec<Vec<String>> {
                     DataValue::Blob(s) => s.to_string(),
                     DataValue::Decimal(v) => v.to_string(),
                     DataValue::Date(v) => v.to_string(),
+                    DataValue::Timestamp(v) => v.to_string(),
+                    DataValue::TimestampTz(v) => v.to_string(),
                     DataValue::Interval(v) => v.to_string(),
                 };
                 row_vec.push(s);

src/array/mod.rs

@@ -11,7 +11,8 @@ use rust_decimal::prelude::FromStr;
 use rust_decimal::Decimal;
 
 use crate::types::{
-    Blob, ConvertError, DataType, DataTypeKind, DataValue, Date, Interval, F32, F64,
+    Blob, ConvertError, DataType, DataTypeKind, DataValue, Date, Interval, Timestamp, TimestampTz,
+    F32, F64,
 };
 
 mod data_chunk;
@@ -26,9 +27,11 @@ pub use self::primitive_array::*;
 pub use self::utf8_array::*;
 
 mod internal_ext;
+
 pub use internal_ext::*;
 
 mod shuffle_ext;
+
 pub use shuffle_ext::*;
 
 /// A trait over all array builders.
@@ -181,6 +184,8 @@ pub type F32Array = PrimitiveArray<F32>;
 pub type F64Array = PrimitiveArray<F64>;
 pub type DecimalArray = PrimitiveArray<Decimal>;
 pub type DateArray = PrimitiveArray<Date>;
+pub type TimestampArray = PrimitiveArray<Timestamp>;
+pub type TimestampTzArray = PrimitiveArray<TimestampTz>;
 pub type IntervalArray = PrimitiveArray<Interval>;
 
 /// Embeds all types of arrays in `array` module.
@@ -197,6 +202,8 @@ pub enum ArrayImpl {
     Blob(Arc<BlobArray>),
     Decimal(Arc<DecimalArray>),
     Date(Arc<DateArray>),
+    Timestamp(Arc<TimestampArray>),
+    TimestampTz(Arc<TimestampTzArray>),
     Interval(Arc<IntervalArray>),
 }
 
@@ -209,6 +216,8 @@ pub type F32ArrayBuilder = PrimitiveArrayBuilder<F32>;
 pub type F64ArrayBuilder = PrimitiveArrayBuilder<F64>;
 pub type DecimalArrayBuilder = PrimitiveArrayBuilder<Decimal>;
 pub type DateArrayBuilder = PrimitiveArrayBuilder<Date>;
+pub type TimestampArrayBuilder = PrimitiveArrayBuilder<Timestamp>;
+pub type TimestampTzArrayBuilder = PrimitiveArrayBuilder<TimestampTz>;
 pub type IntervalArrayBuilder = PrimitiveArrayBuilder<Interval>;
 
 /// Embeds all types of array builders in `array` module.
@@ -224,6 +233,8 @@ pub enum ArrayBuilderImpl {
     Blob(BlobArrayBuilder),
     Decimal(DecimalArrayBuilder),
     Date(DateArrayBuilder),
+    Timestamp(TimestampArrayBuilder),
+    TimestampTz(TimestampTzArrayBuilder),
     Interval(IntervalArrayBuilder),
 }
 
@@ -248,6 +259,8 @@ macro_rules! for_all_variants {
             { Float64, F64, float64, F64Array, F64ArrayBuilder, Float64, Float64 },
             { Decimal, Decimal, decimal, DecimalArray, DecimalArrayBuilder, Decimal, Decimal(_, _) },
             { Date, Date, date, DateArray, DateArrayBuilder, Date, Date },
+            { Timestamp, Timestamp, timestamp, TimestampArray, TimestampArrayBuilder, Timestamp, Timestamp },
+            { TimestampTz, TimestampTz, timestamp_tz, TimestampTzArray, TimestampTzArrayBuilder, TimestampTz, TimestampTz },
             { Interval, Interval, interval, IntervalArray, IntervalArrayBuilder, Interval, Interval },
             { Utf8, str, utf8, Utf8Array, Utf8ArrayBuilder, String, String },
             { Blob, BlobRef, blob, BlobArray, BlobArrayBuilder, Blob, Blob }
@@ -267,6 +280,8 @@ macro_rules! for_all_variants_without_null {
             { Float64, F64, float64, F64Array, F64ArrayBuilder, Float64, Float64 },
             { Decimal, Decimal, decimal, DecimalArray, DecimalArrayBuilder, Decimal, Decimal(_, _) },
             { Date, Date, date, DateArray, DateArrayBuilder, Date, Date },
+            { Timestamp, Timestamp, timestamp, TimestampArray, TimestampArrayBuilder, Timestamp, Timestamp },
+            { TimestampTz, TimestampTz, timestamp_tz, TimestampTzArray, TimestampTzArrayBuilder, TimestampTz, TimestampTz },
             { Interval, Interval, interval, IntervalArray, IntervalArrayBuilder, Interval, Interval },
             { Utf8, str, utf8, Utf8Array, Utf8ArrayBuilder, String, String },
             { Blob, BlobRef, blob, BlobArray, BlobArrayBuilder, Blob, Blob }
@@ -465,6 +480,8 @@ impl ArrayBuilderImpl {
             Self::Blob(a) if null => a.push(None),
             Self::Decimal(a) if null => a.push(None),
             Self::Date(a) if null => a.push(None),
+            Self::Timestamp(a) if null => a.push(None),
+            Self::TimestampTz(a) if null => a.push(None),
             Self::Interval(a) if null => a.push(None),
             Self::Bool(a) => a.push(Some(
                 &s.parse::<bool>()
@@ -497,6 +514,14 @@ impl ArrayBuilderImpl {
             Self::Date(a) => a.push(Some(
                 &Date::from_str(s).map_err(|e| ConvertError::ParseDate(s.to_string(), e))?,
             )),
+            Self::Timestamp(a) => a
+                .push(Some(&Timestamp::from_str(s).map_err(|e| {
+                    ConvertError::ParseTimestamp(s.to_string(), e)
+                })?)),
+            Self::TimestampTz(a) => a
+                .push(Some(&TimestampTz::from_str(s).map_err(|e| {
+                    ConvertError::ParseTimestampTz(s.to_string(), e)
+                })?)),
             Self::Interval(a) => a.push(Some(
                 &Interval::from_str(s)
                     .map_err(|e| ConvertError::ParseInterval(s.to_string(), e))?,
@@ -614,6 +639,8 @@ impl From<&DataValue> for ArrayImpl {
             DataValue::Blob(v) => Self::new_blob([Some(v)].into_iter().collect()),
             &DataValue::Decimal(v) => Self::new_decimal([v].into_iter().collect()),
             &DataValue::Date(v) => Self::new_date([v].into_iter().collect()),
+            &DataValue::Timestamp(v) => Self::new_timestamp([v].into_iter().collect()),
+            &DataValue::TimestampTz(v) => Self::new_timestamp_tz([v].into_iter().collect()),
             &DataValue::Interval(v) => Self::new_interval([v].into_iter().collect()),
         }
     }

src/array/ops.rs

@@ -11,7 +11,8 @@ use super::*;
 use crate::for_all_variants;
 use crate::parser::{BinaryOperator, UnaryOperator};
 use crate::types::{
-    Blob, ConvertError, DataTypeKind, DataValue, Date, DateTimeField, Interval, F64,
+    Blob, ConvertError, DataTypeKind, DataValue, Date, DateTimeField, Interval, Timestamp,
+    TimestampTz, F64,
 };
 
 type A = ArrayImpl;
@@ -216,7 +217,7 @@ impl ArrayImpl {
                 return Err(ConvertError::NoUnaryOp(
                     "extract".into(),
                     self.type_string(),
-                ))
+                ));
             }
         })
     }
@@ -276,7 +277,13 @@ impl ArrayImpl {
                 Type::Decimal(_, _) => {
                     Self::new_decimal(unary_op(a.as_ref(), |&b| Decimal::from(b as u8)))
                 }
-                Type::Null | Type::Date | Type::Interval | Type::Blob | Type::Struct(_) => {
+                Type::Null
+                | Type::Date
+                | Type::Timestamp
+                | Type::TimestampTz
+                | Type::Interval
+                | Type::Blob
+                | Type::Struct(_) => {
                     return Err(ConvertError::NoCast("BOOLEAN", data_type.clone()));
                 }
             },
@@ -290,7 +297,13 @@ impl ArrayImpl {
                 Type::Decimal(_, _) => {
                     Self::new_decimal(unary_op(a.as_ref(), |&i| Decimal::from(i)))
                 }
-                Type::Null | Type::Date | Type::Interval | Type::Blob | Type::Struct(_) => {
+                Type::Null
+                | Type::Date
+                | Type::Timestamp
+                | Type::TimestampTz
+                | Type::Interval
+                | Type::Blob
+                | Type::Struct(_) => {
                     return Err(ConvertError::NoCast("SMALLINT", data_type.clone()));
                 }
             },
@@ -307,7 +320,13 @@ impl ArrayImpl {
                 Type::Decimal(_, _) => {
                     Self::new_decimal(unary_op(a.as_ref(), |&i| Decimal::from(i)))
                 }
-                Type::Null | Type::Date | Type::Interval | Type::Blob | Type::Struct(_) => {
+                Type::Null
+                | Type::Date
+                | Type::Timestamp
+                | Type::TimestampTz
+                | Type::Interval
+                | Type::Blob
+                | Type::Struct(_) => {
                     return Err(ConvertError::NoCast("INT", data_type.clone()));
                 }
             },
@@ -327,7 +346,13 @@ impl ArrayImpl {
                 Type::Decimal(_, _) => {
                     Self::new_decimal(unary_op(a.as_ref(), |&i| Decimal::from(i)))
                 }
-                Type::Null | Type::Date | Type::Interval | Type::Blob | Type::Struct(_) => {
+                Type::Null
+                | Type::Date
+                | Type::Timestamp
+                | Type::TimestampTz
+                | Type::Interval
+                | Type::Blob
+                | Type::Struct(_) => {
                     return Err(ConvertError::NoCast("BIGINT", data_type.clone()));
                 }
             },
@@ -350,7 +375,13 @@ impl ArrayImpl {
                 Type::Decimal(_, _) => Self::new_decimal(unary_op(a.as_ref(), |&f| {
                     Decimal::from_f64_retain(f.0).unwrap()
                 })),
-                Type::Null | Type::Date | Type::Interval | Type::Blob | Type::Struct(_) => {
+                Type::Null
+                | Type::Date
+                | Type::Timestamp
+                | Type::TimestampTz
+                | Type::Interval
+                | Type::Blob
+                | Type::Struct(_) => {
                     return Err(ConvertError::NoCast("DOUBLE", data_type.clone()));
                 }
             },
@@ -382,6 +413,14 @@ impl ArrayImpl {
                 Type::Date => Self::new_date(try_unary_op(a.as_ref(), |s| {
                     Date::from_str(s).map_err(|e| ConvertError::ParseDate(s.to_string(), e))
                 })?),
+                Type::Timestamp => Self::new_timestamp(try_unary_op(a.as_ref(), |s| {
+                    Timestamp::from_str(s)
+                        .map_err(|e| ConvertError::ParseTimestamp(s.to_string(), e))
+                })?),
+                Type::TimestampTz => Self::new_timestamp_tz(try_unary_op(a.as_ref(), |s| {
+                    TimestampTz::from_str(s)
+                        .map_err(|e| ConvertError::ParseTimestampTz(s.to_string(), e))
+                })?),
                 Type::Interval => Self::new_interval(try_unary_op(a.as_ref(), |s| {
                     Interval::from_str(s).map_err(|e| ConvertError::ParseInterval(s.to_string(), e))
                 })?),
@@ -414,7 +453,13 @@ impl ArrayImpl {
                 })?),
                 Type::String => Self::new_utf8(Utf8Array::from_iter_display(a.iter())),
                 Type::Decimal(_, _) => self.clone(),
-                Type::Null | Type::Blob | Type::Date | Type::Interval | Type::Struct(_) => {
+                Type::Null
+                | Type::Blob
+                | Type::Date
+                | Type::Timestamp
+                | Type::TimestampTz
+                | Type::Interval
+                | Type::Struct(_) => {
                     return Err(ConvertError::NoCast("DOUBLE", data_type.clone()));
                 }
             },
@@ -423,6 +468,21 @@ impl ArrayImpl {
                 Type::String => Self::new_utf8(Utf8Array::from_iter_display(a.iter())),
                 _ => return Err(ConvertError::NoCast("DATE", data_type.clone())),
             },
+            Self::Timestamp(a) => match data_type {
+                Type::Timestamp => self.clone(),
+                Type::String => Self::new_utf8(Utf8Array::from_iter_display(a.iter())),
+                _ => return Err(ConvertError::NoCast("TIMESTAMP", data_type.clone())),
+            },
+            Self::TimestampTz(a) => match data_type {
+                Type::TimestampTz => self.clone(),
+                Type::String => Self::new_utf8(Utf8Array::from_iter_display(a.iter())),
+                _ => {
+                    return Err(ConvertError::NoCast(
+                        "TIMESTAMP WITH TIME ZONE",
+                        data_type.clone(),
+                    ))
+                }
+            },
             Self::Interval(a) => match data_type {
                 Type::Interval => self.clone(),
                 Type::String => Self::new_utf8(Utf8Array::from_iter_display(a.iter())),

src/python_extension/mod.rs

@@ -67,6 +67,8 @@ pub fn datachunk_to_python_list(py: Python, chunk: &Chunk) -> Vec<Vec<PyObject>>
                     DataValue::Blob(s) => s.to_string().to_object(py),
                     DataValue::Decimal(v) => v.to_string().to_object(py),
                     DataValue::Date(v) => v.to_string().to_object(py),
+                    DataValue::Timestamp(v) => v.to_string().to_object(py),
+                    DataValue::TimestampTz(v) => v.to_string().to_object(py),
                     DataValue::Interval(v) => v.to_string().to_object(py),
                 };
                 row_vec.push(s);

src/storage/secondary/column/column_builder.rs

@@ -11,7 +11,9 @@ use super::primitive_column_builder::{
 };
 use super::{BoolColumnBuilder, ColumnBuilder};
 use crate::array::ArrayImpl;
-use crate::storage::secondary::column::IntervalColumnBuilder;
+use crate::storage::secondary::column::{
+    IntervalColumnBuilder, TimestampColumnBuilder, TimestampTzColumnBuilder,
+};
 use crate::types::{DataType, DataTypeKind};
 
 /// [`ColumnBuilder`] of all types
@@ -24,6 +26,8 @@ pub enum ColumnBuilderImpl {
     Utf8(CharColumnBuilder),
     Decimal(DecimalColumnBuilder),
     Date(DateColumnBuilder),
+    Timestamp(TimestampColumnBuilder),
+    TimestampTz(TimestampTzColumnBuilder),
     Interval(IntervalColumnBuilder),
     Blob(BlobColumnBuilder),
 }
@@ -41,6 +45,10 @@ impl ColumnBuilderImpl {
             String => Self::Utf8(CharColumnBuilder::new(datatype.nullable, None, options)),
             Decimal(_, _) => Self::Decimal(DecimalColumnBuilder::new(datatype.nullable, options)),
             Date => Self::Date(DateColumnBuilder::new(datatype.nullable, options)),
+            Timestamp => Self::Timestamp(TimestampColumnBuilder::new(datatype.nullable, options)),
+            TimestampTz => {
+                Self::TimestampTz(TimestampTzColumnBuilder::new(datatype.nullable, options))
+            }
             Interval => Self::Interval(IntervalColumnBuilder::new(datatype.nullable, options)),
             Blob => Self::Blob(BlobColumnBuilder::new(datatype.nullable, options)),
             Struct(_) => todo!("struct column builder"),
@@ -57,6 +65,8 @@ impl ColumnBuilderImpl {
             (Self::Utf8(builder), ArrayImpl::Utf8(array)) => builder.append(array),
             (Self::Decimal(builder), ArrayImpl::Decimal(array)) => builder.append(array),
             (Self::Date(builder), ArrayImpl::Date(array)) => builder.append(array),
+            (Self::Timestamp(builder), ArrayImpl::Timestamp(array)) => builder.append(array),
+            (Self::TimestampTz(builder), ArrayImpl::TimestampTz(array)) => builder.append(array),
             (Self::Interval(builder), ArrayImpl::Interval(array)) => builder.append(array),
             (Self::Blob(builder), ArrayImpl::Blob(array)) => builder.append(array),
             _ => todo!(),
@@ -73,6 +83,8 @@ impl ColumnBuilderImpl {
             Self::Utf8(builder) => builder.finish(),
             Self::Decimal(builder) => builder.finish(),
             Self::Date(builder) => builder.finish(),
+            Self::Timestamp(builder) => builder.finish(),
+            Self::TimestampTz(builder) => builder.finish(),
             Self::Interval(builder) => builder.finish(),
             Self::Blob(builder) => builder.finish(),
         }