reset

Author: MichaelLeeHZ

Cargo.lock

@@ -14,6 +14,7 @@ workspace = true
 # In alphabetical order
 arrow = { workspace = true }
 async-trait = { workspace = true }
+chrono = { workspace = true }
 common_types = { workspace = true }
 common_util = { workspace = true }
 datafusion = { workspace = true }

query_engine/Cargo.toml

@@ -2,7 +2,8 @@
 
 use std::{mem, sync::Arc};
 
-use arrow::{compute, compute::kernels::cast_utils::string_to_timestamp_nanos};
+use arrow::{compute, compute::kernels::cast_utils::string_to_timestamp_nanos, error::ArrowError};
+use chrono::{prelude::*, LocalResult};
 use datafusion::{
     arrow::datatypes::DataType,
     common::DFSchemaRef,
@@ -281,16 +282,59 @@ impl<'a> ExprRewriter for TypeRewriter<'a> {
 }
 
 fn string_to_timestamp_ms(string: &str) -> Result<ScalarValue> {
+    // TODO(lee): remove following codes after PR(https://github.com/apache/arrow-rs/pull/3787) merged
+    // Because function `string_to_timestamp_nanos` returns a NaiveDateTime's
+    // nanoseconds from a string without a specify time zone, We need to convert
+    // it to local timestamp.
+
+    // without a timezone specifier as a local time, using 'T' as a separator
+    // Example: 2020-09-08T13:42:29.190855
+    if let Ok(ts) = NaiveDateTime::parse_from_str(string, "%Y-%m-%dT%H:%M:%S%.f") {
+        let mills = naive_datetime_to_timestamp(string, ts).map_err(DataFusionError::from)?;
+        return Ok(ScalarValue::TimestampMillisecond(Some(mills), None));
+    }
+
+    // without a timezone specifier as a local time, using ' ' as a separator
+    // Example: 2020-09-08 13:42:29.190855
+    if let Ok(ts) = NaiveDateTime::parse_from_str(string, "%Y-%m-%d %H:%M:%S%.f") {
+        let mills = naive_datetime_to_timestamp(string, ts).map_err(DataFusionError::from)?;
+        return Ok(ScalarValue::TimestampMillisecond(Some(mills), None));
+    }
+
+    let result = string_to_timestamp_nanos(string);
     Ok(ScalarValue::TimestampMillisecond(
         Some(
-            string_to_timestamp_nanos(string)
+            result
                 .map(|t| t / 1_000_000)
                 .map_err(DataFusionError::from)?,
         ),
         None,
     ))
 }
 
+/// Converts the naive datetime (which has no specific timezone) to a
+/// nanosecond epoch timestamp relative to UTC.
+fn naive_datetime_to_timestamp(s: &str, datetime: NaiveDateTime) -> Result<i64, ArrowError> {
+    let l = Local {};
+
+    match l.from_local_datetime(&datetime) {
+        LocalResult::None => Err(ArrowError::CastError(format!(
+            "Error parsing '{s}' as timestamp: local time representation is invalid"
+        ))),
+        LocalResult::Single(local_datetime) => {
+            Ok(local_datetime.with_timezone(&Utc).timestamp_nanos() / 1_000_000)
+        }
+        // Ambiguous times can happen if the timestamp is exactly when
+        // a daylight savings time transition occurs, for example, and
+        // so the datetime could validly be said to be in two
+        // potential offsets. However, since we are about to convert
+        // to UTC anyways, we can pick one arbitrarily
+        LocalResult::Ambiguous(local_datetime, _) => {
+            Ok(local_datetime.with_timezone(&Utc).timestamp_nanos() / 1_000_000)
+        }
+    }
+}
+
 enum TimestampType {
     Second,
     #[allow(dead_code)]
@@ -321,6 +365,7 @@ mod tests {
     };
 
     use super::*;
+    use crate::logical_optimizer::type_conversion;
 
     fn expr_test_schema() -> DFSchemaRef {
         Arc::new(
@@ -445,7 +490,7 @@ mod tests {
 
     #[test]
     fn test_type_conversion_timestamp() {
-        let date_string = "2021-09-07 16:00:00".to_string();
+        let date_string = "2021-09-07T16:00:00Z".to_string();
         let schema = expr_test_schema();
         let mut rewriter = TypeRewriter {
             schemas: vec![&schema],
@@ -498,7 +543,7 @@ mod tests {
         );
 
         // Timestamp c6 between "2021-09-07 16:00:00" and "2021-09-07 17:00:00"
-        let date_string2 = "2021-09-07 17:00:00".to_string();
+        let date_string2 = "2021-09-07T17:00:00Z".to_string();
         let exp = Expr::Between(Between {
             expr: Box::new(col("c6")),
             negated: false,
@@ -530,4 +575,32 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn test_string_to_timestamp_ms() {
+        let date_string = [
+            "2021-09-07T16:00:00+08:00",
+            "2021-09-07 16:00:00+08:00",
+            "2021-09-07T16:00:00Z",
+            "2021-09-07 16:00:00Z",
+        ];
+        let expects: [i64; 4] = [1631001600000, 1631001600000, 1631030400000, 1631030400000];
+        for (index, &string) in date_string.iter().enumerate() {
+            let result = type_conversion::string_to_timestamp_ms(string);
+            if let Ok(ScalarValue::TimestampMillisecond(Some(mills), _)) = result {
+                let expect = *expects.get(index).unwrap();
+                assert_eq!(mills, expect)
+            }
+        }
+
+        let date_string = "2021-09-07 16:00:00".to_string();
+        let d = NaiveDate::from_ymd_opt(2021, 9, 7).unwrap();
+        let t = NaiveTime::from_hms_milli_opt(16, 0, 0, 0).unwrap();
+        let dt = NaiveDateTime::new(d, t);
+        let expect = naive_datetime_to_timestamp(&date_string, dt).unwrap();
+        let result = type_conversion::string_to_timestamp_ms(&date_string);
+        if let Ok(ScalarValue::TimestampMillisecond(Some(mills), _)) = result {
+            assert_eq!(mills, expect)
+        }
+    }
 }