refactor: Optimize timestamp/timestamptz function by DayLUT (#19031)

Author: TCeason

src/query/expression/src/utils/date_helper.rs

@@ -14,6 +14,7 @@
 
 use std::sync::LazyLock;
 
+use databend_common_column::types::timestamp_tz;
 use databend_common_exception::Result;
 use databend_common_timezone::fast_components_from_timestamp;
 use databend_common_timezone::fast_utc_from_local;
@@ -260,6 +261,18 @@ fn date_from_components(c: &DateTimeComponents) -> Option<Date> {
     Date::new(c.year as i16, c.month as i8, c.day as i8).ok()
 }
 
+#[inline]
+pub fn timestamp_tz_local_micros(value: timestamp_tz) -> Option<i64> {
+    let offset = value.micros_offset()?;
+    value.timestamp().checked_add(offset)
+}
+
+#[inline]
+pub fn timestamp_tz_components_via_lut(value: timestamp_tz) -> Option<DateTimeComponents> {
+    let local = timestamp_tz_local_micros(value)?;
+    fast_components_from_timestamp(local, &TimeZone::UTC)
+}
+
 fn datetime_from_components(c: &DateTimeComponents) -> Option<DateTime> {
     let date = date_from_components(c)?;
     Some(date.at(
@@ -1025,55 +1038,39 @@ pub fn today_date(now: &Zoned, tz: &TimeZone) -> i32 {
 // The working hours of all departments of The State Council are from 8 a.m. to 12 p.m. and from 1:30 p.m. to 5:30 p.m. The winter working hours will be implemented after September 17th.
 pub fn calc_date_to_timestamp(val: i32, tz: &TimeZone) -> std::result::Result<i64, String> {
     let ts = (val as i64) * 24 * 3600 * MICROS_PER_SEC;
-    let z = ts.to_timestamp(tz);
-
-    let tomorrow = z.date().tomorrow();
-    let yesterday = z.date().yesterday();
-
-    // If there were no yesterday or tomorrow, it might be the limit value.
-    // e.g. 9999-12-31
-    if tomorrow.is_err() || yesterday.is_err() {
-        let tz_offset_micros = tz
-            .to_timestamp(date(1970, 1, 1).at(0, 0, 0, 0))
-            .unwrap()
-            .as_microsecond();
-        return Ok(ts + tz_offset_micros);
-    }
-
-    // tomorrow midnight
-    let tomorrow_date = tomorrow.map_err(|e| format!("Calc tomorrow midnight with error {}", e))?;
-
-    let tomorrow_zoned = tomorrow_date.to_zoned(tz.clone()).unwrap_or(z.clone());
-    let tomorrow_is_dst = tz.to_offset_info(tomorrow_zoned.timestamp()).dst().is_dst();
-
-    // yesterday midnight
-    let yesterday_date =
-        yesterday.map_err(|e| format!("Calc yesterday midnight with error {}", e))?;
-    let yesterday_zoned = yesterday_date.to_zoned(tz.clone()).unwrap_or(z.clone());
-    let yesterday_is_std = tz
-        .to_offset_info(yesterday_zoned.timestamp())
-        .dst()
-        .is_std();
-
-    // today midnight
-    let today_datetime_midnight = z.date().to_datetime(Time::midnight());
-    let today_zoned = today_datetime_midnight
-        .to_zoned(tz.clone())
-        .map_err(|e| format!("Calc today midnight with error {}", e))?;
-    let today_is_dst = tz.to_offset_info(today_zoned.timestamp()).dst().is_dst();
-
-    let tz_offset_micros = tz
-        .to_timestamp(date(1970, 1, 1).at(0, 0, 0, 0))
-        .unwrap()
-        .as_microsecond();
-
-    let base_res = ts + tz_offset_micros;
+    let local_date = val.to_date(tz);
+    let year = i32::from(local_date.year());
+    let month = local_date.month() as u8;
+    let day = local_date.day() as u8;
+
+    if let Some(micros) = fast_utc_from_local(tz, year, month, day, 0, 0, 0, 0) {
+        return Ok(micros);
+    }
+
+    let midnight = local_date.to_datetime(Time::midnight());
+    match midnight.to_zoned(tz.clone()) {
+        Ok(zoned) => Ok(zoned.timestamp().as_microsecond()),
+        Err(_err) => {
+            for minutes in 1..=1440 {
+                let delta = SignedDuration::from_secs((minutes * 60) as i64);
+                if let Ok(adj) = midnight.checked_add(delta) {
+                    if let Ok(zoned) = adj.to_zoned(tz.clone()) {
+                        return Ok(zoned.timestamp().as_microsecond());
+                    }
+                } else {
+                    break;
+                }
+            }
 
-    // Origin：(today_is_dst && tomorrow_is_dst && !yesterday_is_std) || (today_is_dst && !tomorrow_is_dst && yesterday_is_std)
-    if today_is_dst && (tomorrow_is_dst != yesterday_is_std) {
-        Ok(base_res - 3600 * MICROS_PER_SEC)
-    } else {
-        Ok(base_res)
+            // The timezone database might not have explicit rules for extremely
+            // old/new dates, so fall back to the legacy behavior that applies the
+            // canonical offset we use for 1970-01-01.
+            let tz_offset_micros = tz
+                .to_timestamp(date(1970, 1, 1).at(0, 0, 0, 0))
+                .unwrap()
+                .as_microsecond();
+            Ok(ts + tz_offset_micros)
+        }
     }
 }
 

src/query/functions/benches/bench.rs

@@ -213,6 +213,26 @@ mod datetime_fast_path {
         });
     }
 
+    #[divan::bench]
+    fn convert_timezone(bencher: divan::Bencher) {
+        let expr = build_expr("convert_timezone('America/Los_Angeles', ts)", &[(
+            "ts",
+            DataType::Timestamp,
+        )]);
+        let data = &*SAMPLES;
+        let block = DataBlock::new(vec![data.timestamp_entry()], data.rows());
+        let func_ctx = FunctionContext {
+            tz: TimeZone::get("Asia/Shanghai").unwrap(),
+            ..Default::default()
+        };
+        let evaluator = Evaluator::new(&block, &func_ctx, &BUILTIN_FUNCTIONS);
+
+        bencher.bench(|| {
+            let value = evaluator.run(&expr).unwrap();
+            divan::black_box(value);
+        });
+    }
+
     fn build_expr(sql: &str, columns: &[(&str, DataType)]) -> Expr {
         let raw_expr = parser::parse_raw_expr(sql, columns);
         type_check::check(&raw_expr, &BUILTIN_FUNCTIONS).unwrap()

src/query/functions/src/scalars/timestamp/src/datetime.rs

@@ -235,10 +235,7 @@ fn register_convert_timezone(registry: &mut FunctionRegistry) {
                         return;
                     }
                 }
-                // Convert source timestamp from source timezone to target timezone
-                let p_src_timestamp = src_timestamp.to_timestamp(&ctx.func_ctx.tz);
-                let src_dst_from_utc = p_src_timestamp.offset().seconds();
-
+                let source_tz = &ctx.func_ctx.tz;
                 let t_tz = match TimeZone::get(target_tz) {
                     Ok(tz) => tz,
                     Err(e) => {
@@ -251,17 +248,33 @@ fn register_convert_timezone(registry: &mut FunctionRegistry) {
                     }
                 };
 
-                let result_timestamp = p_src_timestamp
-                    .with_time_zone(t_tz.clone())
-                    .timestamp()
-                    .as_microsecond();
-                let target_dst_from_utc = p_src_timestamp
-                    .with_time_zone(t_tz.clone())
-                    .offset()
-                    .seconds();
+                let source_components = fast_components_from_timestamp(src_timestamp, source_tz);
+                let target_components = fast_components_from_timestamp(src_timestamp, &t_tz);
+
+                let (instant_micros, src_dst_from_utc, target_dst_from_utc) =
+                    if let (Some(src_comp), Some(target_comp)) =
+                        (source_components, target_components)
+                    {
+                        (
+                            src_timestamp,
+                            src_comp.offset_seconds,
+                            target_comp.offset_seconds,
+                        )
+                    } else {
+                        // Fall back to the slower Jiff conversion for timestamps
+                        // outside the LUT coverage (e.g. <1900 or >2299).
+                        let src_zoned = src_timestamp.to_timestamp(source_tz);
+                        let target_zoned = src_zoned.with_time_zone(t_tz.clone());
+                        (
+                            target_zoned.timestamp().as_microsecond(),
+                            src_zoned.offset().seconds(),
+                            target_zoned.offset().seconds(),
+                        )
+                    };
+
                 let offset_as_micros_sec = (target_dst_from_utc - src_dst_from_utc) as i64;
                 match offset_as_micros_sec.checked_mul(MICROS_PER_SEC) {
-                    Some(offset) => match result_timestamp.checked_add(offset) {
+                    Some(offset) => match instant_micros.checked_add(offset) {
                         Some(res) => output.push(res),
                         None => {
                             ctx.set_error(output.len(), "calc final time error".to_string());
@@ -759,6 +772,13 @@ fn register_timestamp_to_timestamp_tz(registry: &mut FunctionRegistry) {
         ctx: &mut EvalContext,
     ) -> Value<TimestampTzType> {
         vectorize_with_builder_1_arg::<TimestampType, TimestampTzType>(|val, output, ctx| {
+            if let Some(components) = fast_components_from_timestamp(val, &ctx.func_ctx.tz) {
+                let offset = components.offset_seconds;
+                let ts_tz = timestamp_tz::new(val - (offset as i64 * MICROS_PER_SEC), offset);
+                output.push(ts_tz);
+                return;
+            }
+
             let ts = match Timestamp::from_microsecond(val) {
                 Ok(ts) => ts,
                 Err(err) => {
@@ -1053,15 +1073,22 @@ fn register_timestamp_tz_to_date(registry: &mut FunctionRegistry) {
     }
 
     fn calc_timestamp_tz_to_date(val: timestamp_tz) -> Result<i32, String> {
-        let offset = Offset::from_seconds(val.seconds_offset()).map_err(|err| err.to_string())?;
-
-        Ok(val
-            .timestamp()
-            .to_timestamp(&TimeZone::fixed(offset))
-            .date()
-            .since((Unit::Day, Date::new(1970, 1, 1).unwrap()))
-            .unwrap()
-            .get_days())
+        if let Some(days) = timestamp_tz_components_via_lut(val)
+            .and_then(|c| days_from_components(c.year, c.month, c.day))
+        {
+            Ok(days)
+        } else {
+            let offset =
+                Offset::from_seconds(val.seconds_offset()).map_err(|err| err.to_string())?;
+
+            Ok(val
+                .timestamp()
+                .to_timestamp(&TimeZone::fixed(offset))
+                .date()
+                .since((Unit::Day, Date::new(1970, 1, 1).unwrap()))
+                .unwrap()
+                .get_days())
+        }
     }
 }