feat: support DataType TimestampTz (#18892)

* feat: support DataType `TimestampTimezone`

* chore: add `to_timestamp_tz` from timestamp

* chore: add interval function for timestamp_tz

* chore: add meta test of TimestampTz

* chore: fix bug on codex

* chore: codefmt

* chore: fix e2e test

* chore: remove timestamp_tz on old datatype

Author: KKould

Cargo.lock

@@ -29,6 +29,7 @@ either = { workspace = true }
 ethnum = { workspace = true }
 foreign_vec = { workspace = true }
 hex = { workspace = true }
+jiff = { workspace = true }
 log = { workspace = true }
 match-template = { workspace = true }
 num-traits = { workspace = true }

src/common/column/Cargo.toml

@@ -87,6 +87,8 @@ pub enum PrimitiveType {
     DaysMs,
     /// months_days_micros(i32, i32, i64)
     MonthDayMicros,
+    /// timestamp_tz(i32, i64)
+    TimestampTz,
 }
 
 mod private {
@@ -115,5 +117,6 @@ mod private {
     impl Sealed for OrderedFloat<f64> {}
     impl Sealed for super::days_ms {}
     impl Sealed for super::months_days_micros {}
+    impl Sealed for super::timestamp_tz {}
     impl Sealed for View {}
 }

src/common/column/src/types/mod.rs

@@ -15,6 +15,9 @@
 
 use std::cmp::Ordering;
 use std::convert::TryFrom;
+use std::fmt::Debug;
+use std::fmt::Display;
+use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::ops::Add;
@@ -27,12 +30,17 @@ use borsh::BorshSerialize;
 use bytemuck::Pod;
 use bytemuck::Zeroable;
 use databend_common_base::base::OrderedFloat;
+use jiff::fmt::strtime;
+use jiff::tz;
+use jiff::Timestamp;
 use log::error;
 use serde_derive::Deserialize;
 use serde_derive::Serialize;
 
 use super::PrimitiveType;
 
+pub const TIMESTAMP_TIMEZONE_FORMAT: &str = "%Y-%m-%d %H:%M:%S%.6f %z";
+
 /// Sealed trait implemented by all physical types that can be allocated,
 /// serialized and deserialized by this crate.
 /// All O(N) allocations in this crate are done for this trait alone.
@@ -71,6 +79,10 @@ pub trait NativeType:
 
     /// From bytes in big endian
     fn from_be_bytes(bytes: Self::Bytes) -> Self;
+
+    fn size_of() -> usize {
+        std::mem::size_of::<Self>()
+    }
 }
 
 macro_rules! native_type {
@@ -428,6 +440,138 @@ impl Neg for months_days_micros {
     }
 }
 
+/// The in-memory representation of the MonthDayNano variant of the "Interval" logical type.
+#[derive(
+    Debug,
+    Copy,
+    Clone,
+    Default,
+    Eq,
+    Zeroable,
+    Pod,
+    Serialize,
+    Deserialize,
+    BorshSerialize,
+    BorshDeserialize,
+)]
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub struct timestamp_tz(pub i128);
+
+impl Hash for timestamp_tz {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.total_micros().hash(state)
+    }
+}
+impl PartialEq for timestamp_tz {
+    fn eq(&self, other: &Self) -> bool {
+        self.total_micros() == other.total_micros()
+    }
+}
+impl PartialOrd for timestamp_tz {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for timestamp_tz {
+    fn cmp(&self, other: &Self) -> Ordering {
+        let total_micros = self.total_micros();
+        let other_micros = other.total_micros();
+        total_micros.cmp(&other_micros)
+    }
+}
+
+impl timestamp_tz {
+    pub const MICROS_PER_SECOND: i64 = 1_000_000;
+
+    pub fn new(timestamp: i64, offset: i32) -> Self {
+        let ts = timestamp as u64 as i128; // <- 中间加一次 u64 屏蔽符号位
+        let off = (offset as i128) << 64;
+        Self(off | ts)
+    }
+
+    #[inline]
+    pub fn timestamp(&self) -> i64 {
+        self.0 as u64 as i64
+    }
+
+    #[inline]
+    pub fn seconds_offset(&self) -> i32 {
+        (self.0 >> 64) as i32
+    }
+
+    #[inline]
+    pub fn micros_offset(&self) -> Option<i64> {
+        (self.seconds_offset() as i64).checked_mul(Self::MICROS_PER_SECOND)
+    }
+
+    #[inline]
+    pub fn hours_offset(&self) -> i8 {
+        (self.seconds_offset() / 3600) as i8
+    }
+
+    #[inline]
+    pub fn total_micros(&self) -> i64 {
+        self.try_total_micros().unwrap_or_else(|| {
+            error!(
+                "interval is out of range: timestamp={}, offset={}",
+                self.timestamp(),
+                self.seconds_offset()
+            );
+            0
+        })
+    }
+
+    #[inline]
+    pub fn try_total_micros(&self) -> Option<i64> {
+        let offset_micros = self.micros_offset()?;
+        self.timestamp().checked_sub(offset_micros)
+    }
+}
+
+impl Display for timestamp_tz {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let timestamp = Timestamp::from_microsecond(self.timestamp()).unwrap();
+
+        let offset = tz::Offset::from_seconds(self.seconds_offset()).unwrap();
+        let string = strtime::format(
+            TIMESTAMP_TIMEZONE_FORMAT,
+            &timestamp.to_zoned(offset.to_time_zone()),
+        )
+        .unwrap();
+        write!(f, "{}", string)
+    }
+}
+
+impl NativeType for timestamp_tz {
+    const PRIMITIVE: PrimitiveType = PrimitiveType::TimestampTz;
+    type Bytes = [u8; 16];
+    #[inline]
+    fn to_le_bytes(&self) -> Self::Bytes {
+        self.0.to_le_bytes()
+    }
+
+    #[inline]
+    fn to_be_bytes(&self) -> Self::Bytes {
+        self.0.to_be_bytes()
+    }
+
+    #[inline]
+    fn from_le_bytes(bytes: Self::Bytes) -> Self {
+        let mut buf16 = [0u8; 16];
+        buf16.copy_from_slice(&bytes);
+        Self(i128::from_le_bytes(buf16))
+    }
+
+    #[inline]
+    fn from_be_bytes(bytes: Self::Bytes) -> Self {
+        let mut buf16 = [0u8; 16];
+        buf16.copy_from_slice(&bytes);
+        Self(i128::from_be_bytes(buf16))
+    }
+}
+
 /// Type representation of the Float16 physical type
 #[derive(Copy, Clone, Default, Zeroable, Pod)]
 #[allow(non_camel_case_types)]

src/common/column/src/types/native.rs

@@ -20,6 +20,7 @@ geozero = { workspace = true }
 gimli = { workspace = true }
 http = { workspace = true }
 hyper = { workspace = true }
+jiff = { workspace = true }
 libc = { workspace = true }
 object = { workspace = true }
 once_cell = { workspace = true }

src/common/exception/Cargo.toml

@@ -177,7 +177,7 @@ build_exceptions! {
     StagePermissionDenied(2506),
 }
 
-// Data Format and Parsing Errors [1046, 1057, 1060, 1064, 1072, 1074-1081, 1090, 1201-1202, 2507-2509]
+// Data Format and Parsing Errors [1046, 1057, 1060, 1064, 1072, 1074-1081, 1090, 1201-1202, 2507-2509, 2513]
 build_exceptions! {
     /// Bad bytes
     BadBytes(1046),
@@ -215,6 +215,8 @@ build_exceptions! {
     IllegalFileFormat(2508),
     /// File format already exists
     FileFormatAlreadyExists(2509),
+    /// Jiff error
+    JiffError(2513),
 }
 
 // Cluster and Resource Management Errors [1035, 1045, 1082, 1101, 2401-2410, 4013]

src/common/exception/src/exception_code.rs

@@ -228,6 +228,12 @@ impl From<ParseError> for ErrorCode {
     }
 }
 
+impl From<jiff::Error> for ErrorCode {
+    fn from(error: jiff::Error) -> Self {
+        ErrorCode::JiffError(error.to_string())
+    }
+}
+
 impl From<GeozeroError> for ErrorCode {
     fn from(value: GeozeroError) -> Self {
         ErrorCode::GeometryError(value.to_string())

src/common/exception/src/exception_into.rs

@@ -35,5 +35,8 @@ mod fixed_list;
 pub use fixed_list::*;
 mod interval;
 mod map;
+mod timestamp_tz;
+
 pub use interval::*;
 pub use map::*;
+pub use timestamp_tz::*;

src/common/native/src/read/array/mod.rs

@@ -0,0 +1,128 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::io::BufReader;
+use std::io::Cursor;
+
+use databend_common_column::buffer::Buffer;
+use databend_common_column::error::Result;
+use databend_common_column::types::timestamp_tz;
+use databend_common_expression::types::timestamp_tz::TimestampTzType;
+use databend_common_expression::types::ArgType;
+use databend_common_expression::Column;
+use databend_common_expression::TableDataType;
+
+use crate::compression::integer::decompress_integer;
+use crate::nested::InitNested;
+use crate::nested::NestedState;
+use crate::read::read_basic::read_nested;
+use crate::read::NativeReadBuf;
+use crate::read::PageIterator;
+use crate::PageMeta;
+
+#[derive(Debug)]
+pub struct TimestampTzNestedIter<I>
+where I: Iterator<Item = Result<(u64, Vec<u8>)>> + PageIterator + Send + Sync
+{
+    iter: I,
+    data_type: TableDataType,
+    init: Vec<InitNested>,
+    scratch: Vec<u8>,
+}
+
+impl<I> TimestampTzNestedIter<I>
+where I: Iterator<Item = Result<(u64, Vec<u8>)>> + PageIterator + Send + Sync
+{
+    pub fn new(iter: I, data_type: TableDataType, init: Vec<InitNested>) -> Self {
+        Self {
+            iter,
+            data_type,
+            init,
+            scratch: vec![],
+        }
+    }
+}
+
+impl<I> TimestampTzNestedIter<I>
+where I: Iterator<Item = Result<(u64, Vec<u8>)>> + PageIterator + Send + Sync
+{
+    fn deserialize(&mut self, num_values: u64, buffer: Vec<u8>) -> Result<(NestedState, Column)> {
+        let mut reader = BufReader::with_capacity(buffer.len(), Cursor::new(buffer));
+        let (nested, validity) = read_nested(&mut reader, &self.init, num_values as usize)?;
+        let length = num_values as usize;
+
+        let mut values = Vec::with_capacity(length);
+        decompress_integer(&mut reader, length, &mut values, &mut self.scratch)?;
+        assert_eq!(values.len(), length);
+
+        let mut buffer = reader.into_inner().into_inner();
+        self.iter.swap_buffer(&mut buffer);
+
+        let column: Buffer<i128> = values.into();
+        let column: Buffer<timestamp_tz> = unsafe { std::mem::transmute(column) };
+        let mut col = TimestampTzType::upcast_column(column);
+        if self.data_type.is_nullable() {
+            col = col.wrap_nullable(validity);
+        }
+        Ok((nested, col))
+    }
+}
+
+impl<I> Iterator for TimestampTzNestedIter<I>
+where I: Iterator<Item = Result<(u64, Vec<u8>)>> + PageIterator + Send + Sync
+{
+    type Item = Result<(NestedState, Column)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.iter.next() {
+            Some(Ok((num_values, buffer))) => Some(self.deserialize(num_values, buffer)),
+            Some(Err(err)) => Some(Result::Err(err)),
+            None => None,
+        }
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        match self.iter.nth(n) {
+            Some(Ok((num_values, buffer))) => Some(self.deserialize(num_values, buffer)),
+            Some(Err(err)) => Some(Result::Err(err)),
+            None => None,
+        }
+    }
+}
+
+pub fn read_nested_timestamp_tz<R: NativeReadBuf>(
+    reader: &mut R,
+    data_type: TableDataType,
+    init: Vec<InitNested>,
+    page_metas: Vec<PageMeta>,
+) -> Result<Vec<(NestedState, Column)>> {
+    let mut scratch = vec![];
+    let mut results = Vec::with_capacity(page_metas.len());
+    for page_meta in page_metas {
+        let num_values = page_meta.num_values as usize;
+        let (nested, validity) = read_nested(reader, &init, num_values)?;
+
+        let mut values = Vec::with_capacity(num_values);
+        decompress_integer(reader, num_values, &mut values, &mut scratch)?;
+
+        let column: Buffer<i128> = values.into();
+        let column: Buffer<timestamp_tz> = unsafe { std::mem::transmute(column) };
+        let mut col = TimestampTzType::upcast_column(column);
+        if data_type.is_nullable() {
+            col = col.wrap_nullable(validity);
+        }
+        results.push((nested, col));
+    }
+    Ok(results)
+}