improve Spark from_utc_timestamp compatibility

The timezone spellings accepted by Spark are different than those
accepted by Arrow's parser. As a result, `from_utc_timestamp` in comet
isn't safe to use, since it would fail on things like Java short IDs,
prefixed offsets with second precision, and SystemV-style IDs.
Additionally, Arrow will accept offsets beyond the Spark limit of 18
hours.

This change improves compatibility by parsing timezone IDs with the
rules implemented by Spark. As part of doing this, I added
SparkFromUtcTimestampExpr so that computed timezones evaluated for rows
with non-null timestamps, avoiding failing when a timezone is invalid
but wouldn't be required anyway. (Full compatibility with spark
behaviour seems quite complex and I haven't attempted it here.)

It also adds more documentation recording remaining differences:

- timezone validation and evaluation order can differ from Spark
- chrono accepts a narrower calendar range than Spark implemented range
- chrono-tz doesn't currently perform DST transitions past 2099

This does not achieve full compatibility with Spark due to those issues,
but it makes the region of incompatibility a lot smaller.

Author: nat-openai

datafusion/functions/src/datetime/to_local_time.rs

@@ -15,7 +15,6 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::ops::Add;
 use std::sync::Arc;
 
 use arrow::array::timezone::Tz;
@@ -30,8 +29,8 @@ use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
 
 use datafusion_common::cast::as_primitive_array;
 use datafusion_common::{
-    Result, ScalarValue, exec_err, internal_datafusion_err, internal_err,
-    utils::take_function_args,
+    Result, ScalarValue, exec_datafusion_err, exec_err, internal_datafusion_err,
+    internal_err, utils::take_function_args,
 };
 use datafusion_expr::{
     Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
@@ -307,7 +306,10 @@ fn to_local_time(time_value: &ColumnarValue) -> Result<ColumnarValue> {
 /// ```
 ///
 /// See `test_adjust_to_local_time()` for example
-pub fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+pub fn adjust_to_local_time<T: ArrowTimestampType>(
+    ts: i64,
+    tz: impl TimeZone + Copy,
+) -> Result<i64> {
     fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
     where
         F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
@@ -337,12 +339,18 @@ pub fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i6
         .fix()
         .local_minus_utc() as i64;
 
-    let adjusted_date_time = date_time.add(
-        // This should not fail under normal circumstances as the
-        // maximum possible offset is 26 hours (93,600 seconds)
-        TimeDelta::try_seconds(offset_seconds)
-            .ok_or_else(|| internal_datafusion_err!("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000"))?,
-    );
+    let adjusted_date_time = date_time
+        .checked_add_signed(
+            // This should not fail under normal circumstances as the
+            // maximum possible offset is 26 hours (93,600 seconds)
+            TimeDelta::try_seconds(offset_seconds)
+                .ok_or_else(|| internal_datafusion_err!("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000"))?,
+        )
+        .ok_or_else(|| {
+            exec_datafusion_err!(
+                "Timezone adjustment of timestamp {ts} by {offset_seconds} seconds exceeds Chrono's calendar range"
+            )
+        })?;
 
     // convert the naive datetime back to i64
     match T::UNIT {

datafusion/spark/Cargo.toml

@@ -106,3 +106,7 @@ name = "sha2"
 [[bench]]
 harness = false
 name = "floor"
+
+[[bench]]
+harness = false
+name = "from_utc_timestamp"

datafusion/spark/benches/from_utc_timestamp.rs

@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measure the physical evaluator with a timezone column. The narrow and wide batches use
+//! identical timestamp and timezone columns; only the number of unused Int64 columns differs.
+//! Mixed-null timestamps expose any cost from filtering those unrelated columns.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, StringArray, TimestampMicrosecondArray};
+use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use arrow::record_batch::RecordBatch;
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use datafusion::physical_expr::expressions::Column;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_spark::function::datetime::from_utc_timestamp::SparkFromUtcTimestampExpr;
+
+fn batch(rows: usize, columns: usize, half_null: bool) -> RecordBatch {
+    let mut fields = vec![
+        Field::new(
+            "timestamp",
+            DataType::Timestamp(TimeUnit::Microsecond, None),
+            true,
+        ),
+        Field::new("timezone", DataType::Utf8, false),
+    ];
+    let mut arrays: Vec<ArrayRef> = vec![
+        Arc::new(TimestampMicrosecondArray::from_iter((0..rows).map(|row| {
+            (!half_null || row % 2 != 0).then_some(row as i64 * 86_400_000_000)
+        }))),
+        Arc::new(StringArray::from_iter_values(
+            (0..rows).map(|row| if (row / 2) % 2 == 0 { "PST" } else { "EST" }),
+        )),
+    ];
+
+    for column in 2..columns {
+        fields.push(Field::new(
+            format!("unused_{column}"),
+            DataType::Int64,
+            false,
+        ));
+        arrays.push(Arc::new(Int64Array::from_iter_values(
+            (0..rows).map(|row| (column * rows + row) as i64),
+        )));
+    }
+
+    RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let rows = 8192;
+    let expr = SparkFromUtcTimestampExpr::new(
+        Arc::new(Column::new("timestamp", 0)),
+        Arc::new(Column::new("timezone", 1)),
+    );
+    let mut group = c.benchmark_group("from_utc_timestamp_column");
+    group.throughput(Throughput::Elements(rows as u64));
+
+    for columns in [2, 64, 256] {
+        for (half_null, name) in [(false, "no_nulls"), (true, "half_nulls")] {
+            let batch = batch(rows, columns, half_null);
+            let result = expr.evaluate(&batch).unwrap().into_array(rows).unwrap();
+            assert_eq!(result.len(), rows);
+            assert_eq!(result.null_count(), if half_null { rows / 2 } else { 0 });
+
+            group.bench_with_input(
+                BenchmarkId::new(format!("{columns}_columns"), name),
+                &batch,
+                |b, batch| {
+                    b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap()));
+                },
+            );
+        }
+    }
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);