From 18d676802a4d5d8fdd4612eca6c8e1c225dabb93 Mon Sep 17 00:00:00 2001 From: Eugene Tolbakov Date: Fri, 29 Mar 2024 18:33:24 +0000 Subject: [PATCH] feat(function): add timestamp epoch integer support for to_timezone (#3620) * feat(function): add timestamp epoch integer support for to_timezone * chore: fmt --- .../src/scalars/timestamp/to_timezone.rs | 111 +++++++++++++----- .../standalone/common/function/time.result | 14 ++- .../cases/standalone/common/function/time.sql | 2 + 3 files changed, 97 insertions(+), 30 deletions(-) diff --git a/src/common/function/src/scalars/timestamp/to_timezone.rs b/src/common/function/src/scalars/timestamp/to_timezone.rs index 1160267dbd5b..1b366ccde3f2 100644 --- a/src/common/function/src/scalars/timestamp/to_timezone.rs +++ b/src/common/function/src/scalars/timestamp/to_timezone.rs @@ -23,7 +23,7 @@ use datatypes::prelude::VectorRef; use datatypes::types::TimestampType; use datatypes::value::Value; use datatypes::vectors::{ - StringVector, TimestampMicrosecondVector, TimestampMillisecondVector, + Int64Vector, StringVector, TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector, TimestampSecondVector, Vector, }; use snafu::{ensure, OptionExt}; @@ -43,6 +43,7 @@ fn convert_to_timezone(arg: &str) -> Option { fn convert_to_timestamp(arg: &Value) -> Option { match arg { Value::Timestamp(ts) => Some(*ts), + Value::Int64(i) => Some(Timestamp::new_millisecond(*i)), _ => None, } } @@ -66,6 +67,8 @@ impl Function for ToTimezoneFunction { fn signature(&self) -> Signature { helper::one_of_sigs2( vec![ + ConcreteDataType::int32_datatype(), + ConcreteDataType::int64_datatype(), ConcreteDataType::timestamp_second_datatype(), ConcreteDataType::timestamp_millisecond_datatype(), ConcreteDataType::timestamp_microsecond_datatype(), @@ -86,39 +89,45 @@ impl Function for ToTimezoneFunction { } ); - // TODO: maybe support epoch timestamp? https://github.com/GreptimeTeam/greptimedb/issues/3477 - let ts = columns[0].data_type().as_timestamp().with_context(|| { - UnsupportedInputDataTypeSnafu { - function: NAME, - datatypes: columns.iter().map(|c| c.data_type()).collect::>(), - } - })?; let array = columns[0].to_arrow_array(); - let times = match ts { - TimestampType::Second(_) => { - let vector = TimestampSecondVector::try_from_arrow_array(array).unwrap(); - (0..vector.len()) - .map(|i| convert_to_timestamp(&vector.get(i))) - .collect::>() - } - TimestampType::Millisecond(_) => { - let vector = TimestampMillisecondVector::try_from_arrow_array(array).unwrap(); - (0..vector.len()) - .map(|i| convert_to_timestamp(&vector.get(i))) - .collect::>() - } - TimestampType::Microsecond(_) => { - let vector = TimestampMicrosecondVector::try_from_arrow_array(array).unwrap(); + let times = match columns[0].data_type() { + ConcreteDataType::Int64(_) | ConcreteDataType::Int32(_) => { + let vector = Int64Vector::try_from_arrow_array(array).unwrap(); (0..vector.len()) .map(|i| convert_to_timestamp(&vector.get(i))) .collect::>() } - TimestampType::Nanosecond(_) => { - let vector = TimestampNanosecondVector::try_from_arrow_array(array).unwrap(); - (0..vector.len()) - .map(|i| convert_to_timestamp(&vector.get(i))) - .collect::>() + ConcreteDataType::Timestamp(ts) => match ts { + TimestampType::Second(_) => { + let vector = TimestampSecondVector::try_from_arrow_array(array).unwrap(); + (0..vector.len()) + .map(|i| convert_to_timestamp(&vector.get(i))) + .collect::>() + } + TimestampType::Millisecond(_) => { + let vector = TimestampMillisecondVector::try_from_arrow_array(array).unwrap(); + (0..vector.len()) + .map(|i| convert_to_timestamp(&vector.get(i))) + .collect::>() + } + TimestampType::Microsecond(_) => { + let vector = TimestampMicrosecondVector::try_from_arrow_array(array).unwrap(); + (0..vector.len()) + .map(|i| convert_to_timestamp(&vector.get(i))) + .collect::>() + } + TimestampType::Nanosecond(_) => { + let vector = TimestampNanosecondVector::try_from_arrow_array(array).unwrap(); + (0..vector.len()) + .map(|i| convert_to_timestamp(&vector.get(i))) + .collect::>() + } + }, + _ => UnsupportedInputDataTypeSnafu { + function: NAME, + datatypes: columns.iter().map(|c| c.data_type()).collect::>(), } + .fail()?, }; let tzs = { @@ -153,7 +162,7 @@ mod tests { use datatypes::timestamp::{ TimestampMicrosecond, TimestampMillisecond, TimestampNanosecond, TimestampSecond, }; - use datatypes::vectors::StringVector; + use datatypes::vectors::{Int64Vector, StringVector}; use super::*; @@ -257,4 +266,48 @@ mod tests { let expect_times: VectorRef = Arc::new(StringVector::from(results)); assert_eq!(expect_times, vector); } + + #[test] + fn test_numerical_to_timezone() { + let f = ToTimezoneFunction; + let results = vec![ + Some("1969-12-31 19:00:00.001"), + None, + Some("1970-01-01 03:00:00.001"), + None, + Some("2024-03-26 23:01:50"), + None, + Some("2024-03-27 06:02:00"), + None, + ]; + let times: Vec> = vec![ + Some(1), + None, + Some(1), + None, + Some(1711508510000), + None, + Some(1711508520000), + None, + ]; + let ts_vector: Int64Vector = Int64Vector::from_owned_iterator(times.into_iter()); + let tzs = vec![ + Some("America/New_York"), + None, + Some("Europe/Moscow"), + None, + Some("America/New_York"), + None, + Some("Europe/Moscow"), + None, + ]; + let args: Vec = vec![ + Arc::new(ts_vector), + Arc::new(StringVector::from(tzs.clone())), + ]; + let vector = f.eval(FunctionContext::default(), &args).unwrap(); + assert_eq!(8, vector.len()); + let expect_times: VectorRef = Arc::new(StringVector::from(results)); + assert_eq!(expect_times, vector); + } } diff --git a/tests/cases/standalone/common/function/time.result b/tests/cases/standalone/common/function/time.result index 7c6815682c1a..25f82f2a48b0 100644 --- a/tests/cases/standalone/common/function/time.result +++ b/tests/cases/standalone/common/function/time.result @@ -54,5 +54,17 @@ select to_timezone('2024-03-29T14:16:43.012345Z'::Timestamp, 'Asia/Shanghai'); select to_timezone(1709992225, 'Asia/Shanghai'); -Error: 3001(EngineExecuteQuery), DataFusion error: Error during planning: Coercion from [Int64, Utf8] to the signature OneOf([Exact([Timestamp(Second, None), Utf8]), Exact([Timestamp(Millisecond, None), Utf8]), Exact([Timestamp(Microsecond, None), Utf8]), Exact([Timestamp(Nanosecond, None), Utf8])]) failed. ++------------------------------------------------------+ +| to_timezone(Int64(1709992225),Utf8("Asia/Shanghai")) | ++------------------------------------------------------+ +| 1970-01-21 02:59:52.225 | ++------------------------------------------------------+ + +select to_timezone(1711508510000::INT64, 'Asia/Shanghai'); + ++---------------------------------------------------------+ +| to_timezone(Int64(1711508510000),Utf8("Asia/Shanghai")) | ++---------------------------------------------------------+ +| 2024-03-27 11:01:50 | ++---------------------------------------------------------+ diff --git a/tests/cases/standalone/common/function/time.sql b/tests/cases/standalone/common/function/time.sql index a8e59132a524..678bb48af329 100644 --- a/tests/cases/standalone/common/function/time.sql +++ b/tests/cases/standalone/common/function/time.sql @@ -11,3 +11,5 @@ select to_timezone('2024-03-29T14:16:43.012345Z', 'Asia/Shanghai'); select to_timezone('2024-03-29T14:16:43.012345Z'::Timestamp, 'Asia/Shanghai'); select to_timezone(1709992225, 'Asia/Shanghai'); + +select to_timezone(1711508510000::INT64, 'Asia/Shanghai');