Skip to content

Commit

Permalink
feat(function): add timestamp epoch integer support for to_timezone (#…
Browse files Browse the repository at this point in the history 
…3620)

* feat(function): add timestamp epoch integer support for to_timezone

* chore: fmt
  • Loading branch information
@etolbakov
etolbakov authored Mar 29, 2024
1 parent 93da45f commit 18d6768
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 30 deletions.
111 changes: 82 additions & 29 deletions src/common/function/src/scalars/timestamp/to_timezone.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -43,6 +43,7 @@ fn convert_to_timezone(arg: &str) -> Option<Timezone> {
fn convert_to_timestamp(arg: &Value) -> Option<Timestamp> {
match arg {
Value::Timestamp(ts) => Some(*ts),
Value::Int64(i) => Some(Timestamp::new_millisecond(*i)),
_ => None,
}
}
Expand All @@ -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(),
Expand All @@ -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::<Vec<_>>(),
}
})?;
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::<Vec<_>>()
}
TimestampType::Millisecond(_) => {
let vector = TimestampMillisecondVector::try_from_arrow_array(array).unwrap();
(0..vector.len())
.map(|i| convert_to_timestamp(&vector.get(i)))
.collect::<Vec<_>>()
}
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::<Vec<_>>()
}
TimestampType::Nanosecond(_) => {
let vector = TimestampNanosecondVector::try_from_arrow_array(array).unwrap();
(0..vector.len())
.map(|i| convert_to_timestamp(&vector.get(i)))
.collect::<Vec<_>>()
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::<Vec<_>>()
}
TimestampType::Millisecond(_) => {
let vector = TimestampMillisecondVector::try_from_arrow_array(array).unwrap();
(0..vector.len())
.map(|i| convert_to_timestamp(&vector.get(i)))
.collect::<Vec<_>>()
}
TimestampType::Microsecond(_) => {
let vector = TimestampMicrosecondVector::try_from_arrow_array(array).unwrap();
(0..vector.len())
.map(|i| convert_to_timestamp(&vector.get(i)))
.collect::<Vec<_>>()
}
TimestampType::Nanosecond(_) => {
let vector = TimestampNanosecondVector::try_from_arrow_array(array).unwrap();
(0..vector.len())
.map(|i| convert_to_timestamp(&vector.get(i)))
.collect::<Vec<_>>()
}
},
_ => UnsupportedInputDataTypeSnafu {
function: NAME,
datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
}
.fail()?,
};

let tzs = {
Expand Down Expand Up @@ -153,7 +162,7 @@ mod tests {
use datatypes::timestamp::{
TimestampMicrosecond, TimestampMillisecond, TimestampNanosecond, TimestampSecond,
};
use datatypes::vectors::StringVector;
use datatypes::vectors::{Int64Vector, StringVector};

use super::*;

Expand Down Expand Up @@ -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<Option<i64>> = 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<VectorRef> = 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);
}
}
14 changes: 13 additions & 1 deletion tests/cases/standalone/common/function/time.result
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
+---------------------------------------------------------+

2 changes: 2 additions & 0 deletions tests/cases/standalone/common/function/time.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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');

0 comments on commit 18d6768

Please sign in to comment.