From 09ca1ce8e1b195009051ec95257ca31788e3ae1c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:21:23 -0500 Subject: [PATCH 1/2] fix: annotate `to_timestamp` array results with the execution timezone `handle`/`handle_multiple` built the result array straight from `O::DATA_TYPE`, which for timestamps carries no timezone, while the scalar branch went through `scalar_value(dt, ..)` and did pick up the timezone from the declared return type. With `datafusion.execution.time_zone` set, a string *column* argument therefore produced `Timestamp(u, None)` while `return_type` promised `Timestamp(u, )`, so any plan materializing the column failed, and comparisons against another zoned value reached the Arrow kernel with mismatched types. Re-annotate the array with the declared return type before returning it. Co-Authored-By: Claude Opus 5 --- datafusion/functions/src/datetime/common.rs | 79 +++++++++++++------ .../functions/src/datetime/to_timestamp.rs | 78 ++++++++++++++++-- .../test_files/to_timestamp_timezone.slt | 49 ++++++++++++ 3 files changed, 173 insertions(+), 33 deletions(-) diff --git a/datafusion/functions/src/datetime/common.rs b/datafusion/functions/src/datetime/common.rs index 9a7f94bd5973f..efde66bebad12 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 fe7e1b31fc7e5..02e494275e190 100644 --- a/datafusion/functions/src/datetime/to_timestamp.rs +++ b/datafusion/functions/src/datetime/to_timestamp.rs @@ -836,7 +836,7 @@ mod tests { Array, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, }; - use arrow::array::{ArrayRef, Int64Array, StringBuilder}; + use arrow::array::{ArrayRef, Int64Array, StringArray, StringBuilder}; use arrow::datatypes::{Field, TimeUnit}; use chrono::{DateTime, FixedOffset, Utc}; use datafusion_common::{DataFusionError, assert_contains}; @@ -955,7 +955,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 +1025,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 +1156,55 @@ 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()); + + 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()))); + + // both the single argument and the explicit format overloads + for args in [ + vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![ + "2020-09-08T13:42:29", + ])) as ArrayRef)], + vec![ + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + "2020-09-08 13:42:29", + ])) as ArrayRef), + 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 +1839,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 +1891,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 d48e41d1204de..89c9356c2cd7d 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 From b07fa35d8f6e233914d6e1c77f19ea51ccf746e5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:14:15 -0500 Subject: [PATCH 2/2] test: cover LargeUtf8 and Utf8View array inputs in the timezone regression test `handle`/`handle_multiple` dispatch on each string array flavor separately, so run the timezone assertion over all three. Co-Authored-By: Claude Opus 5 --- .../functions/src/datetime/to_timestamp.rs | 72 +++++++++++-------- 1 file changed, 41 insertions(+), 31 deletions(-) diff --git a/datafusion/functions/src/datetime/to_timestamp.rs b/datafusion/functions/src/datetime/to_timestamp.rs index 02e494275e190..dd855511379aa 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, StringArray, 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}; @@ -1164,41 +1167,48 @@ mod tests { 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()))); - // both the single argument and the explicit format overloads - for args in [ - vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![ - "2020-09-08T13:42:29", - ])) as ArrayRef)], - vec![ - ColumnarValue::Array(Arc::new(StringArray::from(vec![ - "2020-09-08 13:42:29", - ])) as ArrayRef), - 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); + 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); + } } }