Support timestamp and interval arithmetic

datafusion/common/src/scalar.rs

@@ -975,20 +975,20 @@ pub const NANOSECOND_MODE: bool = true;
 ///   interval will have the type [`IntervalDayTimeType`].
 /// - When subtracting timestamps at microseconds/nanoseconds precision, the
 ///   output interval will have the type [`IntervalMonthDayNanoType`].
-fn ts_sub_to_interval<const TIME_MOD: bool>(
+fn ts_sub_to_interval<const TIME_MODE: bool>(
     lhs_ts: i64,
     rhs_ts: i64,
     lhs_tz: &Option<String>,
     rhs_tz: &Option<String>,
 ) -> Result<ScalarValue> {
-    let (parsed_lhs_tz, parsed_rhs_tz) =
-        (parse_timezones(lhs_tz), parse_timezones(rhs_tz));
+    let parsed_lhs_tz = parse_timezones(lhs_tz)?;
+    let parsed_rhs_tz = parse_timezones(rhs_tz)?;
 
-    let lhs_dt = with_timezone_to_naive_datetime::<TIME_MOD>(lhs_ts, &parsed_lhs_tz)?;
-    let rhs_dt = with_timezone_to_naive_datetime::<TIME_MOD>(rhs_ts, &parsed_rhs_tz)?;
-    let delta_secs = lhs_dt.signed_duration_since(rhs_dt);
+    let (naive_lhs, naive_rhs) =
+        calculate_naives::<TIME_MODE>(lhs_ts, parsed_lhs_tz, rhs_ts, parsed_rhs_tz)?;
+    let delta_secs = naive_lhs.signed_duration_since(naive_rhs);
 
-    match TIME_MOD {
+    match TIME_MODE {
         MILLISECOND_MODE => {
             let as_millisecs = delta_secs.num_milliseconds();
             Ok(ScalarValue::new_interval_dt(
@@ -1011,18 +1011,82 @@ fn ts_sub_to_interval<const TIME_MOD: bool>(
     }
 }
 
-// This function parses the timezone from string to Tz.
-// If it cannot parse or timezone field is [`None`], it returns [`None`].
-pub fn parse_timezones(tz: &Option<String>) -> Option<Tz> {
+/// This function parses the timezone from string to Tz.
+/// If it cannot parse or timezone field is [`None`], it returns [`None`].
+pub fn parse_timezones(tz: &Option<String>) -> Result<Option<Tz>> {
     if let Some(tz) = tz {
-        let parsed_tz: Option<Tz> = FromStr::from_str(tz)
-            .map_err(|_| {
-                DataFusionError::Execution("cannot parse given timezone".to_string())
-            })
-            .ok();
-        parsed_tz
+        let parsed_tz: Tz = FromStr::from_str(tz).map_err(|_| {
+            DataFusionError::Execution("cannot parse given timezone".to_string())
+        })?;
+        Ok(Some(parsed_tz))
     } else {
-        None
+        Ok(None)
+    }
+}
+
+/// This function takes two timestamps with an optional timezone,
+/// and returns the duration between them. If one of the timestamps
+/// has a [`None`] timezone, the other one is also treated as having [`None`].
+pub fn calculate_naives<const TIME_MODE: bool>(
+    lhs_ts: i64,
+    parsed_lhs_tz: Option<Tz>,
+    rhs_ts: i64,
+    parsed_rhs_tz: Option<Tz>,
+) -> Result<(NaiveDateTime, NaiveDateTime)> {
+    let err = || {
+        DataFusionError::Execution(String::from(
+            "error while converting Int64 to DateTime in timestamp subtraction",
+        ))
+    };
+    match (parsed_lhs_tz, parsed_rhs_tz, TIME_MODE) {
+        (Some(lhs_tz), Some(rhs_tz), MILLISECOND_MODE) => {
+            let lhs = arrow_array::temporal_conversions::as_datetime_with_timezone::<
+                arrow_array::types::TimestampMillisecondType,
+            >(lhs_ts, rhs_tz)
+            .ok_or_else(err)?
+            .naive_local();
+            let rhs = arrow_array::temporal_conversions::as_datetime_with_timezone::<
+                arrow_array::types::TimestampMillisecondType,
+            >(rhs_ts, lhs_tz)
+            .ok_or_else(err)?
+            .naive_local();
+            Ok((lhs, rhs))
+        }
+        (Some(lhs_tz), Some(rhs_tz), NANOSECOND_MODE) => {
+            let lhs = arrow_array::temporal_conversions::as_datetime_with_timezone::<
+                arrow_array::types::TimestampNanosecondType,
+            >(lhs_ts, rhs_tz)
+            .ok_or_else(err)?
+            .naive_local();
+            let rhs = arrow_array::temporal_conversions::as_datetime_with_timezone::<
+                arrow_array::types::TimestampNanosecondType,
+            >(rhs_ts, lhs_tz)
+            .ok_or_else(err)?
+            .naive_local();
+            Ok((lhs, rhs))
+        }
+        (_, _, MILLISECOND_MODE) => {
+            let lhs = arrow_array::temporal_conversions::as_datetime::<
+                arrow_array::types::TimestampMillisecondType,
+            >(lhs_ts)
+            .ok_or_else(err)?;
+            let rhs = arrow_array::temporal_conversions::as_datetime::<
+                arrow_array::types::TimestampMillisecondType,
+            >(rhs_ts)
+            .ok_or_else(err)?;
+            Ok((lhs, rhs))
+        }
+        (_, _, NANOSECOND_MODE) => {
+            let lhs = arrow_array::temporal_conversions::as_datetime::<
+                arrow_array::types::TimestampNanosecondType,
+            >(lhs_ts)
+            .ok_or_else(err)?;
+            let rhs = arrow_array::temporal_conversions::as_datetime::<
+                arrow_array::types::TimestampNanosecondType,
+            >(rhs_ts)
+            .ok_or_else(err)?;
+            Ok((lhs, rhs))
+        }
     }
 }
 
@@ -1112,9 +1176,13 @@ pub fn milliseconds_add_array<const INTERVAL_MODE: i8>(
     interval: i128,
     sign: i32,
 ) -> Result<i64> {
-    let secs = ts_ms / 1000;
-    let nsecs = ((ts_ms % 1000) * 1_000_000) as u32;
-    do_date_time_math_array::<INTERVAL_MODE>(secs, nsecs, interval, sign)
+    let mut secs = ts_ms / 1000;
+    let mut nsecs = ((ts_ms % 1000) * 1_000_000) as i32;
+    if nsecs < 0 {
+        secs -= 1;
+        nsecs += 1_000_000_000;
+    }
+    do_date_time_math_array::<INTERVAL_MODE>(secs, nsecs as u32, interval, sign)
         .map(|dt| dt.timestamp_millis())
 }
 
@@ -1131,9 +1199,13 @@ pub fn microseconds_add_array<const INTERVAL_MODE: i8>(
     interval: i128,
     sign: i32,
 ) -> Result<i64> {
-    let secs = ts_us / 1_000_000;
-    let nsecs = ((ts_us % 1_000_000) * 1000) as u32;
-    do_date_time_math_array::<INTERVAL_MODE>(secs, nsecs, interval, sign)
+    let mut secs = ts_us / 1_000_000;
+    let mut nsecs = ((ts_us % 1_000_000) * 1000) as i32;
+    if nsecs < 0 {
+        secs -= 1;
+        nsecs += 1_000_000_000;
+    }
+    do_date_time_math_array::<INTERVAL_MODE>(secs, nsecs as u32, interval, sign)
         .map(|dt| dt.timestamp_nanos() / 1000)
 }
 
@@ -1150,9 +1222,13 @@ pub fn nanoseconds_add_array<const INTERVAL_MODE: i8>(
     interval: i128,
     sign: i32,
 ) -> Result<i64> {
-    let secs = ts_ns / 1_000_000_000;
-    let nsecs = (ts_ns % 1_000_000_000) as u32;
-    do_date_time_math_array::<INTERVAL_MODE>(secs, nsecs, interval, sign)
+    let mut secs = ts_ns / 1_000_000_000;
+    let mut nsecs = (ts_ns % 1_000_000_000) as i32;
+    if nsecs < 0 {
+        secs -= 1;
+        nsecs += 1_000_000_000;
+    }
+    do_date_time_math_array::<INTERVAL_MODE>(secs, nsecs as u32, interval, sign)
         .map(|dt| dt.timestamp_nanos())
 }
 

datafusion/core/tests/sql/set_variable.rs

@@ -427,7 +427,10 @@ async fn set_time_zone_bad_time_zone_format() {
             .unwrap();
 
     let err = pretty_format_batches(&result).err().unwrap().to_string();
-    assert_eq!(err, "Parser error: Invalid timezone \"08:00\": only offset based timezones supported without chrono-tz feature");
+    assert_eq!(
+        err,
+        "Parser error: Invalid timezone \"08:00\": '08:00' is not a valid timezone"
+    );
 
     plan_and_collect(&ctx, "SET TIME ZONE = '08'")
         .await
@@ -440,22 +443,16 @@ async fn set_time_zone_bad_time_zone_format() {
             .unwrap();
 
     let err = pretty_format_batches(&result).err().unwrap().to_string();
-    assert_eq!(err, "Parser error: Invalid timezone \"08\": only offset based timezones supported without chrono-tz feature");
+    assert_eq!(
+        err,
+        "Parser error: Invalid timezone \"08\": '08' is not a valid timezone"
+    );
 
     // we dont support named time zone yet
     plan_and_collect(&ctx, "SET TIME ZONE = 'Asia/Taipei'")
         .await
         .unwrap();
 
-    // casting UTF-8 to TimestampTZ isn't supported yet, add Timestamp as the middle layer for now
-    let result =
-        plan_and_collect(&ctx, "SELECT '2000-01-01T00:00:00'::TIMESTAMP::TIMESTAMPTZ")
-            .await
-            .unwrap();
-
-    let err = pretty_format_batches(&result).err().unwrap().to_string();
-    assert_eq!(err, "Parser error: Invalid timezone \"Asia/Taipei\": only offset based timezones supported without chrono-tz feature");
-
     // this is invalid even after we support named time zone
     plan_and_collect(&ctx, "SET TIME ZONE = 'Asia/Taipei2'")
         .await
@@ -467,5 +464,5 @@ async fn set_time_zone_bad_time_zone_format() {
             .await
             .unwrap();
     let err = pretty_format_batches(&result).err().unwrap().to_string();
-    assert_eq!(err, "Parser error: Invalid timezone \"Asia/Taipei2\": only offset based timezones supported without chrono-tz feature");
+    assert_eq!(err, "Parser error: Invalid timezone \"Asia/Taipei2\": 'Asia/Taipei2' is not a valid timezone");
 }

datafusion/core/tests/sqllogictests/test_files/arrow_typeof.slt

@@ -279,7 +279,7 @@ query error Cannot automatically convert Interval\(DayTime\) to Interval\(MonthD
 ---
 select arrow_cast(interval '30 minutes', 'Interval(MonthDayNano)');
 
-query error DataFusion error: Error during planning: Cannot automatically convert Utf8 to Interval\(MonthDayNano\)
+query error DataFusion error: Arrow error: Cast error: Casting from Utf8 to Interval\(MonthDayNano\) not supported
 select arrow_cast('30 minutes', 'Interval(MonthDayNano)');
 
 

datafusion/physical-expr/Cargo.toml

@@ -44,6 +44,7 @@ unicode_expressions = ["unicode-segmentation"]
 [dependencies]
 ahash = { version = "0.8", default-features = false, features = ["runtime-rng"] }
 arrow = { workspace = true }
+arrow-array = { version = "34.0.0", default-features = false, features = ["chrono-tz"] }
 arrow-buffer = { workspace = true }
 arrow-schema = { workspace = true }
 blake2 = { version = "^0.10.2", optional = true }