From a2a70b4821ea054a5328ce2018842b2a60056a24 Mon Sep 17 00:00:00 2001 From: Martin Hilton Date: Tue, 3 Oct 2023 09:53:54 +0100 Subject: [PATCH 1/6] preserve array type in date_bin and date_trunc functions The result type of date_bin and date_trunc never includes any timezone information. Change this such that the timezone of the resulting array from these functions is copied from the input array. --- datafusion/expr/src/built_in_function.rs | 17 +- .../physical-expr/src/datetime_expressions.rs | 224 +++++++++++++++++- 2 files changed, 224 insertions(+), 17 deletions(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index 58d84545dbb6..d4c7736f40b1 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -618,13 +618,20 @@ impl BuiltinScalarFunction { BuiltinScalarFunction::ConcatWithSeparator => Ok(Utf8), BuiltinScalarFunction::DatePart => Ok(Float64), BuiltinScalarFunction::DateBin | BuiltinScalarFunction::DateTrunc => { - match input_expr_types[1] { - Timestamp(Nanosecond, _) | Utf8 | Null => { + match &input_expr_types[1] { + Timestamp(Nanosecond, None) | Utf8 | Null => { Ok(Timestamp(Nanosecond, None)) } - Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)), - Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)), - Timestamp(Second, _) => Ok(Timestamp(Second, None)), + Timestamp(Nanosecond, Some(tz)) => { + Ok(Timestamp(Nanosecond, Some(Arc::clone(tz)))) + } + Timestamp(Microsecond, tz_opt) => { + Ok(Timestamp(Microsecond, tz_opt.clone())) + } + Timestamp(Millisecond, tz_opt) => { + Ok(Timestamp(Millisecond, tz_opt.clone())) + } + Timestamp(Second, tz_opt) => Ok(Timestamp(Second, tz_opt.clone())), _ => plan_err!( "The {self} function can only accept timestamp as the second arg." ), diff --git a/datafusion/physical-expr/src/datetime_expressions.rs b/datafusion/physical-expr/src/datetime_expressions.rs index 5ce71f4584bb..8abf0ce2f464 100644 --- a/datafusion/physical-expr/src/datetime_expressions.rs +++ b/datafusion/physical-expr/src/datetime_expressions.rs @@ -433,7 +433,8 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result { granularity.as_str(), ) }) - .collect::>()?; + .collect::>()? + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } DataType::Timestamp(TimeUnit::Millisecond, tz_opt) => { @@ -449,7 +450,8 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result { granularity.as_str(), ) }) - .collect::>()?; + .collect::>()? + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } DataType::Timestamp(TimeUnit::Microsecond, tz_opt) => { @@ -465,7 +467,25 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result { granularity.as_str(), ) }) - .collect::>()?; + .collect::>()? + .with_timezone_opt(tz_opt.clone()); + ColumnarValue::Array(Arc::new(array)) + } + DataType::Timestamp(TimeUnit::Nanosecond, tz_opt) => { + let parsed_tz = parse_tz(tz_opt)?; + let array = as_timestamp_nanosecond_array(array)?; + let array = array + .iter() + .map(|x| { + _date_trunc( + TimeUnit::Nanosecond, + &x, + parsed_tz, + granularity.as_str(), + ) + }) + .collect::>()? + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } _ => { @@ -713,35 +733,39 @@ fn date_bin_impl( )) } ColumnarValue::Array(array) => match array.data_type() { - DataType::Timestamp(TimeUnit::Nanosecond, _) => { + DataType::Timestamp(TimeUnit::Nanosecond, tz_opt) => { let array = as_timestamp_nanosecond_array(array)? .iter() .map(f_nanos) - .collect::(); + .collect::() + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } - DataType::Timestamp(TimeUnit::Microsecond, _) => { + DataType::Timestamp(TimeUnit::Microsecond, tz_opt) => { let array = as_timestamp_microsecond_array(array)? .iter() .map(f_micros) - .collect::(); + .collect::() + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } - DataType::Timestamp(TimeUnit::Millisecond, _) => { + DataType::Timestamp(TimeUnit::Millisecond, tz_opt) => { let array = as_timestamp_millisecond_array(array)? .iter() .map(f_millis) - .collect::(); + .collect::() + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } - DataType::Timestamp(TimeUnit::Second, _) => { + DataType::Timestamp(TimeUnit::Second, tz_opt) => { let array = as_timestamp_second_array(array)? .iter() .map(f_secs) - .collect::(); + .collect::() + .with_timezone_opt(tz_opt.clone()); ColumnarValue::Array(Arc::new(array)) } @@ -925,7 +949,9 @@ where mod tests { use std::sync::Arc; - use arrow::array::{ArrayRef, Int64Array, IntervalDayTimeArray, StringBuilder}; + use arrow::array::{ + as_primitive_array, ArrayRef, Int64Array, IntervalDayTimeArray, StringBuilder, + }; use super::*; @@ -1051,6 +1077,91 @@ mod tests { }); } + #[test] + fn test_date_trunc_timezones() { + let cases = vec![ + ( + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T01:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T03:00:00Z", + "2020-09-08T04:00:00Z", + ], + Some("+00".into()), + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + ], + ), + ( + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T01:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T03:00:00Z", + "2020-09-08T04:00:00Z", + ], + None, + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + ], + ), + ( + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T01:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T03:00:00Z", + "2020-09-08T04:00:00Z", + ], + Some("-02".into()), + vec![ + "2020-09-07T02:00:00Z", + "2020-09-07T02:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T02:00:00Z", + ], + ), + ]; + + cases.iter().for_each(|(original, tz_opt, expected)| { + let input = original + .iter() + .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) + .collect::() + .with_timezone_opt(tz_opt.clone()); + let right = expected + .iter() + .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) + .collect::() + .with_timezone_opt(tz_opt.clone()); + let result = date_trunc(&[ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("day".to_string()))), + ColumnarValue::Array(Arc::new(input)), + ]) + .unwrap(); + if let ColumnarValue::Array(result) = result { + assert_eq!( + result.data_type(), + &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()) + ); + let left = as_primitive_array::(&result); + assert_eq!(left, &right); + } else { + panic!("unexpected column type"); + } + }); + } + #[test] fn test_date_bin_single() { use chrono::Duration; @@ -1252,6 +1363,95 @@ mod tests { ); } + #[test] + fn test_date_bin_timezones() { + let cases = vec![ + ( + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T01:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T03:00:00Z", + "2020-09-08T04:00:00Z", + ], + Some("+00".into()), + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + ], + ), + ( + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T01:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T03:00:00Z", + "2020-09-08T04:00:00Z", + ], + None, + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + ], + ), + ( + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T01:00:00Z", + "2020-09-08T02:00:00Z", + "2020-09-08T03:00:00Z", + "2020-09-08T04:00:00Z", + ], + Some("-02".into()), + vec![ + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + "2020-09-08T00:00:00Z", + ], + ), + ]; + + cases.iter().for_each(|(original, tz_opt, expected)| { + let input = original + .iter() + .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) + .collect::() + .with_timezone_opt(tz_opt.clone()); + let right = expected + .iter() + .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) + .collect::() + .with_timezone_opt(tz_opt.clone()); + let result = date_bin(&[ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Array(Arc::new(input)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( + Some(0), + Some("+00".into()), + )), + ]) + .unwrap(); + if let ColumnarValue::Array(result) = result { + assert_eq!( + result.data_type(), + &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()) + ); + let left = as_primitive_array::(&result); + assert_eq!(left, &right); + } else { + panic!("unexpected column type"); + } + }); + } + #[test] fn to_timestamp_invalid_input_type() -> Result<()> { // pass the wrong type of input array to to_timestamp and test From 1a447f4c95f1f1f42c89c2d062371f70c247af83 Mon Sep 17 00:00:00 2001 From: Martin Hilton Date: Tue, 3 Oct 2023 12:45:41 +0100 Subject: [PATCH 2/6] Update datafusion/expr/src/built_in_function.rs Co-authored-by: Alex Huang --- datafusion/expr/src/built_in_function.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index d4c7736f40b1..2e95ef0a6501 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -622,8 +622,8 @@ impl BuiltinScalarFunction { Timestamp(Nanosecond, None) | Utf8 | Null => { Ok(Timestamp(Nanosecond, None)) } - Timestamp(Nanosecond, Some(tz)) => { - Ok(Timestamp(Nanosecond, Some(Arc::clone(tz)))) + Timestamp(Nanosecond, tz_opt) => { + Ok(Timestamp(Nanosecond, tz_opt.clone()))) } Timestamp(Microsecond, tz_opt) => { Ok(Timestamp(Microsecond, tz_opt.clone())) From 5e17c8377b318e924f3a1f7d9336d63a9fc5b47f Mon Sep 17 00:00:00 2001 From: Martin Hilton Date: Tue, 3 Oct 2023 12:50:27 +0100 Subject: [PATCH 3/6] fix: syntax error --- datafusion/expr/src/built_in_function.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index 2e95ef0a6501..70514f52d5f4 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -623,7 +623,7 @@ impl BuiltinScalarFunction { Ok(Timestamp(Nanosecond, None)) } Timestamp(Nanosecond, tz_opt) => { - Ok(Timestamp(Nanosecond, tz_opt.clone()))) + Ok(Timestamp(Nanosecond, tz_opt.clone())) } Timestamp(Microsecond, tz_opt) => { Ok(Timestamp(Microsecond, tz_opt.clone())) From 884be6e106f6a0d3214675830c879d9e0cdc97b7 Mon Sep 17 00:00:00 2001 From: Martin Hilton Date: Tue, 3 Oct 2023 13:11:28 +0100 Subject: [PATCH 4/6] fix: datafusion-cli cargo update --- datafusion-cli/Cargo.lock | 37 +++++++++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock index 1235c0b740bc..d1c8ea266de3 100644 --- a/datafusion-cli/Cargo.lock +++ b/datafusion-cli/Cargo.lock @@ -93,6 +93,7 @@ version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ceb7c683b2f8f40970b70e39ff8be514c95b96fcb9c4af87e1ed2cb2e10801a0" dependencies = [ + "bzip2", "crc32fast", "digest", "lazy_static", @@ -110,6 +111,8 @@ dependencies = [ "thiserror", "typed-builder", "uuid", + "xz2", + "zstd", ] [[package]] @@ -1417,9 +1420,9 @@ checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" [[package]] name = "errno" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "136526188508e25c6fef639d7927dfb3e0e3084488bf202267829cf7fc23dbdd" +checksum = "add4f07d43996f76ef320709726a556a9d4f965d9410d8d0271132d2f8293480" dependencies = [ "errno-dragonfly", "libc", @@ -2122,9 +2125,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.6.3" +version = "2.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" +checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167" [[package]] name = "mimalloc" @@ -2717,9 +2720,9 @@ checksum = "dbb5fb1acd8a1a18b3dd5be62d25485eb770e05afb408a9627d14d451bae12da" [[package]] name = "reqwest" -version = "0.11.20" +version = "0.11.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e9ad3fe7488d7e34558a2033d45a0c90b72d97b4f80705666fea71472e2e6a1" +checksum = "78fdbab6a7e1d7b13cc8ff10197f47986b41c639300cc3c8158cac7847c9bbef" dependencies = [ "base64", "bytes", @@ -2743,6 +2746,7 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", + "system-configuration", "tokio", "tokio-rustls 0.24.1", "tokio-util", @@ -3227,6 +3231,27 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "system-configuration" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba3a3adc5c275d719af8cb4272ea1c4a6d668a777f37e115f6d11ddbc1c8e0e7" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75fb188eb626b924683e3b95e3a48e63551fcfb51949de2f06a9d91dbee93c9" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "tempfile" version = "3.8.0" From f30bb6f0f89d4fff9112052aa945eb79ca3dbdc5 Mon Sep 17 00:00:00 2001 From: Martin Hilton Date: Wed, 4 Oct 2023 09:27:04 +0100 Subject: [PATCH 5/6] review suggestions Add some additional tests suggested in code reviews. --- .../physical-expr/src/datetime_expressions.rs | 79 ++++++++++++++++++- .../sqllogictest/test_files/timestamps.slt | 56 +++++++++++++ 2 files changed, 132 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-expr/src/datetime_expressions.rs b/datafusion/physical-expr/src/datetime_expressions.rs index 8abf0ce2f464..4ab45a9f0c8a 100644 --- a/datafusion/physical-expr/src/datetime_expressions.rs +++ b/datafusion/physical-expr/src/datetime_expressions.rs @@ -1131,6 +1131,40 @@ mod tests { "2020-09-08T02:00:00Z", ], ), + ( + vec![ + "2020-09-08T00:00:00+05", + "2020-09-08T01:00:00+05", + "2020-09-08T02:00:00+05", + "2020-09-08T03:00:00+05", + "2020-09-08T04:00:00+05", + ], + Some("+05".into()), + vec![ + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + ], + ), + ( + vec![ + "2020-09-08T00:00:00+08", + "2020-09-08T01:00:00+08", + "2020-09-08T02:00:00+08", + "2020-09-08T03:00:00+08", + "2020-09-08T04:00:00+08", + ], + Some("+08".into()), + vec![ + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + ], + ), ]; cases.iter().for_each(|(original, tz_opt, expected)| { @@ -1375,6 +1409,7 @@ mod tests { "2020-09-08T04:00:00Z", ], Some("+00".into()), + "1970-01-01T00:00:00Z", vec![ "2020-09-08T00:00:00Z", "2020-09-08T00:00:00Z", @@ -1392,6 +1427,7 @@ mod tests { "2020-09-08T04:00:00Z", ], None, + "1970-01-01T00:00:00Z", vec![ "2020-09-08T00:00:00Z", "2020-09-08T00:00:00Z", @@ -1409,6 +1445,7 @@ mod tests { "2020-09-08T04:00:00Z", ], Some("-02".into()), + "1970-01-01T00:00:00Z", vec![ "2020-09-08T00:00:00Z", "2020-09-08T00:00:00Z", @@ -1417,9 +1454,45 @@ mod tests { "2020-09-08T00:00:00Z", ], ), + ( + vec![ + "2020-09-08T00:00:00+05", + "2020-09-08T01:00:00+05", + "2020-09-08T02:00:00+05", + "2020-09-08T03:00:00+05", + "2020-09-08T04:00:00+05", + ], + Some("+05".into()), + "1970-01-01T00:00:00+05", + vec![ + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + "2020-09-08T00:00:00+05", + ], + ), + ( + vec![ + "2020-09-08T00:00:00+08", + "2020-09-08T01:00:00+08", + "2020-09-08T02:00:00+08", + "2020-09-08T03:00:00+08", + "2020-09-08T04:00:00+08", + ], + Some("+08".into()), + "1970-01-01T00:00:00+08", + vec![ + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + "2020-09-08T00:00:00+08", + ], + ), ]; - cases.iter().for_each(|(original, tz_opt, expected)| { + cases.iter().for_each(|(original, tz_opt, origin, expected)| { let input = original .iter() .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) @@ -1434,8 +1507,8 @@ mod tests { ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), ColumnarValue::Array(Arc::new(input)), ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( - Some(0), - Some("+00".into()), + Some(string_to_timestamp_nanos(origin).unwrap()), + tz_opt.clone(), )), ]) .unwrap(); diff --git a/datafusion/sqllogictest/test_files/timestamps.slt b/datafusion/sqllogictest/test_files/timestamps.slt index bb06c569f081..edafe18caab5 100644 --- a/datafusion/sqllogictest/test_files/timestamps.slt +++ b/datafusion/sqllogictest/test_files/timestamps.slt @@ -1702,3 +1702,59 @@ SELECT TIMESTAMPTZ '2023-03-11 02:00:00 America/Los_Angeles' as ts_geo # postgresql: accepts statement error SELECT TIMESTAMPTZ '2023-03-12 02:00:00 America/Los_Angeles' as ts_geo + + + +########## +## Timezone column tests +########## + +# create a table with a non-UTC time zone. +statement ok +SET TIME ZONE = '+05:00' + +statement ok +CREATE TABLE foo (time TIMESTAMPTZ) AS VALUES + ('2020-01-01T00:00:00+05:00'), + ('2020-01-01T01:00:00+05:00'), + ('2020-01-01T02:00:00+05:00'), + ('2020-01-01T03:00:00+05:00') + +statement ok +SET TIME ZONE = '+00' + +# verify column type +query T +SELECT arrow_typeof(time) FROM foo LIMIT 1 +---- +Timestamp(Nanosecond, Some("+05:00")) + +# check date_trunc +query P +SELECT date_trunc('day', time) FROM foo +---- +2020-01-01T00:00:00+05:00 +2020-01-01T00:00:00+05:00 +2020-01-01T00:00:00+05:00 +2020-01-01T00:00:00+05:00 + +# verify date_trunc column type +query T +SELECT arrow_typeof(date_trunc('day', time)) FROM foo LIMIT 1 +---- +Timestamp(Nanosecond, Some("+05:00")) + +# check date_bin +query P +SELECT date_bin(INTERVAL '1 day', time, '1970-01-01T00:00:00+05:00') FROM foo +---- +2020-01-01T00:00:00+05:00 +2020-01-01T00:00:00+05:00 +2020-01-01T00:00:00+05:00 +2020-01-01T00:00:00+05:00 + +# verify date_trunc column type +query T +SELECT arrow_typeof(date_bin(INTERVAL '1 day', time, '1970-01-01T00:00:00+05:00')) FROM foo LIMIT 1 +---- +Timestamp(Nanosecond, Some("+05:00")) From 3da2677ab3bd19a0e8d0e6e65df28ed9cf0344ab Mon Sep 17 00:00:00 2001 From: Martin Hilton Date: Wed, 4 Oct 2023 09:35:22 +0100 Subject: [PATCH 6/6] fix formatting --- .../physical-expr/src/datetime_expressions.rs | 64 ++++++++++--------- 1 file changed, 33 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-expr/src/datetime_expressions.rs b/datafusion/physical-expr/src/datetime_expressions.rs index 4ab45a9f0c8a..5cf1c21df5c2 100644 --- a/datafusion/physical-expr/src/datetime_expressions.rs +++ b/datafusion/physical-expr/src/datetime_expressions.rs @@ -1492,37 +1492,39 @@ mod tests { ), ]; - cases.iter().for_each(|(original, tz_opt, origin, expected)| { - let input = original - .iter() - .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) - .collect::() - .with_timezone_opt(tz_opt.clone()); - let right = expected - .iter() - .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) - .collect::() - .with_timezone_opt(tz_opt.clone()); - let result = date_bin(&[ - ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), - ColumnarValue::Array(Arc::new(input)), - ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( - Some(string_to_timestamp_nanos(origin).unwrap()), - tz_opt.clone(), - )), - ]) - .unwrap(); - if let ColumnarValue::Array(result) = result { - assert_eq!( - result.data_type(), - &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()) - ); - let left = as_primitive_array::(&result); - assert_eq!(left, &right); - } else { - panic!("unexpected column type"); - } - }); + cases + .iter() + .for_each(|(original, tz_opt, origin, expected)| { + let input = original + .iter() + .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) + .collect::() + .with_timezone_opt(tz_opt.clone()); + let right = expected + .iter() + .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) + .collect::() + .with_timezone_opt(tz_opt.clone()); + let result = date_bin(&[ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Array(Arc::new(input)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( + Some(string_to_timestamp_nanos(origin).unwrap()), + tz_opt.clone(), + )), + ]) + .unwrap(); + if let ColumnarValue::Array(result) = result { + assert_eq!( + result.data_type(), + &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()) + ); + let left = as_primitive_array::(&result); + assert_eq!(left, &right); + } else { + panic!("unexpected column type"); + } + }); } #[test]