fix: annotate to_timestamp array results with the execution timezone

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<O::Native>,
 {
     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<i64>, O, _>(
-                    &a.as_string::<i64>(),
-                    op,
-                )?,
-            ))),
-            DataType::Utf8 => Ok(ColumnarValue::Array(Arc::new(
-                unary_string_to_primitive_function::<&GenericStringArray<i32>, O, _>(
-                    &a.as_string::<i32>(),
-                    op,
-                )?,
-            ))),
-            other => exec_err!("Unsupported data type {other:?} for function {name}"),
-        },
+        ColumnarValue::Array(a) => {
+            let result: PrimitiveArray<O> = 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::<i64>();
+                    unary_string_to_primitive_function(&strings, op)?
+                }
+                DataType::Utf8 => {
+                    let strings = a.as_string::<i32>();
+                    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::<O, _, _>(args, op, op2, name)?,
-                )))
+                let result =
+                    strings_to_primitive_function::<O, _, _>(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<O: ArrowPrimitiveType>(
+    array: PrimitiveArray<O>,
+    dt: &DataType,
+) -> Result<PrimitiveArray<O>> {
+    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<i64>) -> Result<ScalarValue> {
     match dt {
         DataType::Date32 => Ok(ScalarValue::Date32(r.and_then(|v| v.to_i32()))),

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 <https://github.com/apache/datafusion/issues/24632>.
+    #[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 = &micros_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 = &micros_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 => {

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