diff --git a/datafusion/functions/src/datetime/common.rs b/datafusion/functions/src/datetime/common.rs index 9a7f94bd5973..efde66bebad1 100644 --- a/datafusion/functions/src/datetime/common.rs +++ b/datafusion/functions/src/datetime/common.rs @@ -32,7 +32,7 @@ use chrono::{DateTime, TimeZone, Utc}; use datafusion_common::cast::as_generic_string_array; use datafusion_common::{ DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, - internal_datafusion_err, + internal_datafusion_err, internal_err, }; use datafusion_expr::ColumnarValue; @@ -258,27 +258,30 @@ where F: Fn(&str) -> Result, { match &args[0] { - ColumnarValue::Array(a) => match a.data_type() { - DataType::Utf8View => Ok(ColumnarValue::Array(Arc::new( - unary_string_to_primitive_function::<&StringViewArray, O, _>( - &a.as_string_view(), - op, - )?, - ))), - DataType::LargeUtf8 => Ok(ColumnarValue::Array(Arc::new( - unary_string_to_primitive_function::<&GenericStringArray, O, _>( - &a.as_string::(), - op, - )?, - ))), - DataType::Utf8 => Ok(ColumnarValue::Array(Arc::new( - unary_string_to_primitive_function::<&GenericStringArray, O, _>( - &a.as_string::(), - op, - )?, - ))), - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, + ColumnarValue::Array(a) => { + let result: PrimitiveArray = match a.data_type() { + DataType::Utf8View => { + let strings = a.as_string_view(); + unary_string_to_primitive_function(&strings, op)? + } + DataType::LargeUtf8 => { + let strings = a.as_string::(); + unary_string_to_primitive_function(&strings, op)? + } + DataType::Utf8 => { + let strings = a.as_string::(); + unary_string_to_primitive_function(&strings, op)? + } + other => { + return exec_err!( + "Unsupported data type {other:?} for function {name}" + ); + } + }; + Ok(ColumnarValue::Array(Arc::new(with_return_type( + result, dt, + )?))) + } ColumnarValue::Scalar(scalar) => match scalar.try_as_str() { Some(a) => { let result = a @@ -342,9 +345,11 @@ where } } - Ok(ColumnarValue::Array(Arc::new( - strings_to_primitive_function::(args, op, op2, name)?, - ))) + let result = + strings_to_primitive_function::(args, op, op2, name)?; + Ok(ColumnarValue::Array(Arc::new(with_return_type( + result, dt, + )?))) } other => { exec_err!("Unsupported data type {other:?} for function {name}") @@ -543,6 +548,30 @@ where array.iter().map(|x| x.map(&op).transpose()).collect() } +/// Re-annotates `array` with the return type `dt` declared by the calling +/// function. +/// +/// The parsing kernels build the array from `O::DATA_TYPE`, which for timestamps +/// carries no timezone. `dt` is what the function advertised in `return_type`, +/// so it may additionally carry the execution timezone; without this the array +/// and the schema disagree and execution fails when the column is materialized. +fn with_return_type( + array: PrimitiveArray, + dt: &DataType, +) -> Result> { + let from = array.data_type().clone(); + match (&from, dt) { + // Only the timezone may differ, which `with_data_type` adjusts in place. + (DataType::Timestamp(from_unit, _), DataType::Timestamp(to_unit, _)) + if from_unit == to_unit => + { + Ok(array.with_data_type(dt.clone())) + } + _ if &from == dt => Ok(array), + _ => internal_err!("Cannot return {from} array as {dt}"), + } +} + fn scalar_value(dt: &DataType, r: Option) -> Result { match dt { DataType::Date32 => Ok(ScalarValue::Date32(r.and_then(|v| v.to_i32()))), diff --git a/datafusion/functions/src/datetime/to_timestamp.rs b/datafusion/functions/src/datetime/to_timestamp.rs index fe7e1b31fc7e..dd855511379a 100644 --- a/datafusion/functions/src/datetime/to_timestamp.rs +++ b/datafusion/functions/src/datetime/to_timestamp.rs @@ -836,7 +836,10 @@ mod tests { Array, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, }; - use arrow::array::{ArrayRef, Int64Array, StringBuilder}; + use arrow::array::{ + ArrayRef, Int64Array, LargeStringArray, StringArray, StringBuilder, + StringViewArray, + }; use arrow::datatypes::{Field, TimeUnit}; use chrono::{DateTime, FixedOffset, Utc}; use datafusion_common::{DataFusionError, assert_contains}; @@ -955,7 +958,9 @@ mod tests { string_builder.append_null(); ts_builder.append_null(); - let expected_timestamps = &ts_builder.finish() as &dyn Array; + // the test helpers above configure "UTC" as the execution timezone + let expected = ts_builder.finish().with_timezone("UTC"); + let expected_timestamps = &expected as &dyn Array; let string_array = ColumnarValue::Array(Arc::new(string_builder.finish()) as ArrayRef); @@ -1023,7 +1028,9 @@ mod tests { format3_builder.append_value("%+"); ts_builder.append_value(1599572549190850000); - let expected_timestamps = &ts_builder.finish() as &dyn Array; + // the test helpers above configure "UTC" as the execution timezone + let expected = ts_builder.finish().with_timezone("UTC"); + let expected_timestamps = &expected as &dyn Array; let string_array = [ ColumnarValue::Array(Arc::new(date_string_builder.finish()) as ArrayRef), @@ -1152,6 +1159,62 @@ mod tests { Ok(()) } + /// An array result must be annotated with the same timezone that + /// `return_type` advertises, otherwise execution fails when the column is + /// materialized. See . + #[test] + fn to_timestamp_array_respects_execution_timezone() -> Result<()> { + let mut options = ConfigOptions::default(); + options.execution.time_zone = Some("-05:00".to_string()); + + // every string array flavor `handle`/`handle_multiple` dispatch on + let string_arrays: [fn(Vec<&str>) -> ArrayRef; 3] = [ + |values| Arc::new(StringArray::from(values)) as ArrayRef, + |values| Arc::new(LargeStringArray::from(values)) as ArrayRef, + |values| Arc::new(StringViewArray::from(values)) as ArrayRef, + ]; + + for (udf, time_unit) in udfs_and_timeunit() { + let udf = udf.with_updated_config(&options).unwrap(); + let expected = udf.return_type(&[Utf8])?; + assert_eq!(expected, Timestamp(time_unit, Some("-05:00".into()))); + + for string_array in string_arrays { + // both the single argument and the explicit format overloads + for args in [ + vec![ColumnarValue::Array(string_array(vec![ + "2020-09-08T13:42:29", + ]))], + vec![ + ColumnarValue::Array(string_array(vec!["2020-09-08 13:42:29"])), + ColumnarValue::Scalar(ScalarValue::Utf8(Some( + "%Y-%m-%d %H:%M:%S".to_string(), + ))), + ], + ] { + let arg_fields = args + .iter() + .map(|arg| Field::new("arg", arg.data_type(), true).into()) + .collect(); + let result = udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: 1, + return_field: Field::new("f", expected.clone(), true).into(), + config_options: Arc::new(options.clone()), + })?; + + let ColumnarValue::Array(array) = result else { + panic!("expected an array result"); + }; + assert_eq!(array.data_type(), &expected); + } + } + } + + Ok(()) + } + #[test] fn to_timestamp_formats_respects_execution_timezone() -> Result<()> { let udfs = udfs_and_timeunit(); @@ -1786,10 +1849,16 @@ mod tests { micros_builder.append_value(1599572549190850); sec_builder.append_value(1599572549); - let nanos_expected_timestamps = &nanos_builder.finish() as &dyn Array; - let millis_expected_timestamps = &millis_builder.finish() as &dyn Array; - let micros_expected_timestamps = µs_builder.finish() as &dyn Array; - let sec_expected_timestamps = &sec_builder.finish() as &dyn Array; + // the test helpers above configure "UTC" as the execution timezone + let nanos_expected = nanos_builder.finish().with_timezone("UTC"); + let millis_expected = millis_builder.finish().with_timezone("UTC"); + let micros_expected = micros_builder.finish().with_timezone("UTC"); + let sec_expected = sec_builder.finish().with_timezone("UTC"); + + let nanos_expected_timestamps = &nanos_expected as &dyn Array; + let millis_expected_timestamps = &millis_expected as &dyn Array; + let micros_expected_timestamps = µs_expected as &dyn Array; + let sec_expected_timestamps = &sec_expected as &dyn Array; for (func, time_unit) in funcs { // test UTF8 @@ -1832,7 +1901,10 @@ mod tests { .expect("that to_timestamp with format args parsed values without error"); if let ColumnarValue::Array(parsed_array) = parsed_timestamps { assert_eq!(parsed_array.len(), 1); - assert!(matches!(parsed_array.data_type(), Timestamp(_, None))); + assert!(matches!( + parsed_array.data_type(), + Timestamp(_, Some(tz)) if tz.as_ref() == "UTC" + )); match time_unit { Nanosecond => { diff --git a/datafusion/sqllogictest/test_files/to_timestamp_timezone.slt b/datafusion/sqllogictest/test_files/to_timestamp_timezone.slt index d48e41d1204d..89c9356c2cd7 100644 --- a/datafusion/sqllogictest/test_files/to_timestamp_timezone.slt +++ b/datafusion/sqllogictest/test_files/to_timestamp_timezone.slt @@ -199,6 +199,55 @@ select arrow_typeof(to_timestamp(61)) ---- Timestamp(ns, "America/New_York") +## Test 18: Array (non-constant-folded) input must carry the configured timezone +## https://github.com/apache/datafusion/issues/24632 +statement ok +SET datafusion.execution.time_zone = 'UTC'; + +query TP +SELECT arrow_typeof(to_timestamp(s)), to_timestamp(s) FROM (VALUES ('2020-09-08T13:42:29'), ('2020-09-08T13:42:29Z')) AS t(s); +---- +Timestamp(ns, "UTC") 2020-09-08T13:42:29Z +Timestamp(ns, "UTC") 2020-09-08T13:42:29Z + +## The declared type is trusted by the planner, so a comparison against another +## timezone-aware value must reach the Arrow kernel with matching types +query B +SELECT to_timestamp(s) <= to_timestamp('2262-01-01T00:00:00') FROM (VALUES ('2020-09-08T13:42:29')) AS t(s); +---- +true + +statement ok +SET datafusion.execution.time_zone = 'America/New_York'; + +query TP +SELECT arrow_typeof(to_timestamp(s)), to_timestamp(s) FROM (VALUES ('2020-09-08T13:42:29')) AS t(s); +---- +Timestamp(ns, "America/New_York") 2020-09-08T13:42:29-04:00 + +## with an explicit format string +query TP +SELECT arrow_typeof(to_timestamp(s, '%Y-%m-%d %H:%M:%S')), to_timestamp(s, '%Y-%m-%d %H:%M:%S') FROM (VALUES ('2020-09-08 13:42:29')) AS t(s); +---- +Timestamp(ns, "America/New_York") 2020-09-08T13:42:29-04:00 + +## all precision variants +query TTTT +SELECT + arrow_typeof(to_timestamp_seconds(s)), + arrow_typeof(to_timestamp_millis(s)), + arrow_typeof(to_timestamp_micros(s)), + arrow_typeof(to_timestamp_nanos(s)) +FROM (VALUES ('2020-09-08T13:42:29')) AS t(s); +---- +Timestamp(s, "America/New_York") Timestamp(ms, "America/New_York") Timestamp(µs, "America/New_York") Timestamp(ns, "America/New_York") + +## to_date is unaffected by the execution timezone +query TD +SELECT arrow_typeof(to_date(s)), to_date(s) FROM (VALUES ('2020-09-08')) AS t(s); +---- +Date32 2020-09-08 + ## Reset timezone for other tests statement ok RESET datafusion.execution.time_zone