feat(datagen): support timestamptz (#13451)

e2e_test/source/basic/datagen.slt

@@ -102,7 +102,7 @@ statement ok
 drop table s1;
 
 statement ok
-create table s1 (v1 struct<v2 int>, t1 timestamp, c1 varchar) with (
+create table s1 (v1 struct<v2 int>, t1 timestamp, z1 timestamptz, c1 varchar) with (
     connector = 'datagen',
     fields.v1.v2.kind = 'random',
     fields.v1.v2.min = '1',
@@ -112,6 +112,10 @@ create table s1 (v1 struct<v2 int>, t1 timestamp, c1 varchar) with (
     fields.t1.max_past = '2h 37min',
     fields.t1.max_past_mode = 'relative',
     fields.t1.seed = '3',
+    fields.z1.kind = 'random',
+    fields.z1.max_past = '2h 37min',
+    fields.z1.max_past_mode = 'relative',
+    fields.z1.seed = '3',
     fields.c1.kind = 'random',
     fields.c1.length = '100',
     fields.c1.seed = '3',

src/common/src/array/data_chunk.rs

@@ -919,6 +919,13 @@ impl DataChunkTestExt for DataChunk {
                                 .generate_datum(offset);
                         array_builder.append(datum);
                     }
+                    DataType::Timestamptz => {
+                        let datum =
+                            FieldGeneratorImpl::with_timestamptz(None, None, None, Self::SEED)
+                                .expect("create timestamptz generator should succeed")
+                                .generate_datum(offset);
+                        array_builder.append(datum);
+                    }
                     _ if data_type.is_numeric() => {
                         let mut data_gen = FieldGeneratorImpl::with_number_random(
                             data_type.clone(),

src/common/src/field_generator/mod.rs

@@ -26,7 +26,7 @@ pub use timestamp::*;
 pub use varchar::*;
 
 use crate::array::{ListValue, StructValue};
-use crate::types::{DataType, Datum, ScalarImpl, Timestamp};
+use crate::types::{DataType, Datum, ScalarImpl, Timestamp, Timestamptz};
 
 pub const DEFAULT_MIN: i16 = i16::MIN;
 pub const DEFAULT_MAX: i16 = i16::MAX;
@@ -96,6 +96,7 @@ pub enum FieldGeneratorImpl {
     VarcharRandomFixedLength(VarcharRandomFixedLengthField),
     VarcharConstant,
     Timestamp(ChronoField<Timestamp>),
+    Timestamptz(ChronoField<Timestamptz>),
     Struct(Vec<(String, FieldGeneratorImpl)>),
     List(Box<FieldGeneratorImpl>, usize),
 }
@@ -189,6 +190,20 @@ impl FieldGeneratorImpl {
         )?))
     }
 
+    pub fn with_timestamptz(
+        base: Option<DateTime<FixedOffset>>,
+        max_past: Option<String>,
+        max_past_mode: Option<String>,
+        seed: u64,
+    ) -> Result<Self> {
+        Ok(FieldGeneratorImpl::Timestamptz(ChronoField::new(
+            base,
+            max_past,
+            max_past_mode,
+            seed,
+        )?))
+    }
+
     pub fn with_varchar(varchar_property: &VarcharProperty, seed: u64) -> Self {
         match varchar_property {
             VarcharProperty::RandomFixedLength(length_option) => {
@@ -235,6 +250,7 @@ impl FieldGeneratorImpl {
             FieldGeneratorImpl::VarcharRandomVariableLength(f) => f.generate(offset),
             FieldGeneratorImpl::VarcharConstant => VarcharConstant::generate_json(),
             FieldGeneratorImpl::Timestamp(f) => f.generate(offset),
+            FieldGeneratorImpl::Timestamptz(f) => f.generate(offset),
             FieldGeneratorImpl::Struct(fields) => {
                 let map = fields
                     .iter_mut()
@@ -267,6 +283,7 @@ impl FieldGeneratorImpl {
             FieldGeneratorImpl::VarcharRandomVariableLength(f) => f.generate_datum(offset),
             FieldGeneratorImpl::VarcharConstant => VarcharConstant::generate_datum(),
             FieldGeneratorImpl::Timestamp(f) => f.generate_datum(offset),
+            FieldGeneratorImpl::Timestamptz(f) => f.generate_datum(offset),
             FieldGeneratorImpl::Struct(fields) => {
                 let data = fields
                     .iter_mut()
@@ -328,6 +345,7 @@ mod tests {
             DataType::Float64,
             DataType::Varchar,
             DataType::Timestamp,
+            DataType::Timestamptz,
         ] {
             let mut generator = match data_type {
                 DataType::Varchar => FieldGeneratorImpl::with_varchar(
@@ -337,6 +355,9 @@ mod tests {
                 DataType::Timestamp => {
                     FieldGeneratorImpl::with_timestamp(None, None, None, seed).unwrap()
                 }
+                DataType::Timestamptz => {
+                    FieldGeneratorImpl::with_timestamptz(None, None, None, seed).unwrap()
+                }
                 _ => FieldGeneratorImpl::with_number_random(data_type, None, None, seed).unwrap(),
             };
 

src/common/src/field_generator/timestamp.rs

@@ -22,7 +22,7 @@ use serde_json::Value;
 use tracing::debug;
 
 use super::DEFAULT_MAX_PAST;
-use crate::types::{Datum, Scalar, Timestamp};
+use crate::types::{Datum, Scalar, Timestamp, Timestamptz};
 
 pub struct ChronoField<T: ChronoFieldInner> {
     max_past: Duration,
@@ -106,3 +106,25 @@ impl ChronoFieldInner for Timestamp {
         Value::String(self.0.to_string())
     }
 }
+
+impl ChronoFieldInner for Timestamptz {
+    fn from_now() -> Self {
+        Timestamptz::from(
+            Utc::now()
+                .duration_round(Duration::microseconds(1))
+                .unwrap(),
+        )
+    }
+
+    fn from_base(base: DateTime<FixedOffset>) -> Self {
+        Timestamptz::from(base)
+    }
+
+    fn minus(&self, duration: Duration) -> Self {
+        Timestamptz::from(self.to_datetime_utc() - duration)
+    }
+
+    fn to_json(&self) -> Value {
+        Value::String(self.to_string())
+    }
+}

src/connector/src/source/datagen/source/reader.rs

@@ -213,7 +213,7 @@ fn generator_from_data_type(
         None => split_index,
     };
     match data_type {
-        DataType::Timestamp => {
+        ty @ (DataType::Timestamp | DataType::Timestamptz) => {
             let max_past_key = format!("fields.{}.max_past", name);
             let max_past_value = fields_option_map.get(&max_past_key).cloned();
             let max_past_mode_key = format!("fields.{}.max_past_mode", name);
@@ -230,12 +230,21 @@ fn generator_from_data_type(
                 None => None,
             };
 
-            FieldGeneratorImpl::with_timestamp(
-                basetime,
-                max_past_value,
-                max_past_mode_value,
-                random_seed,
-            )
+            if ty == DataType::Timestamptz {
+                FieldGeneratorImpl::with_timestamptz(
+                    basetime,
+                    max_past_value,
+                    max_past_mode_value,
+                    random_seed,
+                )
+            } else {
+                FieldGeneratorImpl::with_timestamp(
+                    basetime,
+                    max_past_value,
+                    max_past_mode_value,
+                    random_seed,
+                )
+            }
         }
         DataType::Varchar => {
             let length_key = format!("fields.{}.length", name);