Fall back to CPU for non-UTC ORC timestamp writes

cuDF's ORC writer always stamps writerTimezone="UTC" in the stripe footer
and cannot record the JVM writer timezone (rapidsai/cudf#23422). Because
ORC's timestamp type is timezone-agnostic, a GPU-written file in a non-UTC
JVM is read back shifted by the zone offset by a CPU ORC reader. Restore the
write-side guard so non-UTC timestamp ORC writes fall back to CPU, and update
the test to assert the fallback instead of xfail.

Signed-off-by: Chong Gao <chongg@nvidia.com>

Author: Chong Gao

integration_tests/src/main/python/orc_test.py

@@ -1118,24 +1118,25 @@ def test_orc_not_support_timestamp_ltz(std_input_path):
                              conf={},
                              error_message="ParseException")
 
-# test ORC reader and writer with the same timezone
-# the `tz_sensitive_test` mark guarantees the write and read are in the same timezone
-# The `spark.sql.session.timeZone` here does not impact reader and writer timezone, but any way, we test it.
-# For the tests that reader and writer timezones are different, refer to `OrcTimezoneSuite`
+# Timestamp writes: in UTC the GPU writes on the GPU; in a non-UTC JVM the GPU write must fall
+# back to CPU. cuDF's ORC writer always stamps writerTimezone="UTC" in the stripe footer and
+# cannot record the JVM writer timezone (https://github.com/rapidsai/cudf/issues/23422), so a
+# GPU-written non-UTC file would be read back shifted by the zone offset by a CPU ORC reader.
+# The `tz_sensitive_test` mark runs this in both UTC and non-UTC JVM timezones.
 @tz_sensitive_test
 @ignore_order(local=True)
-@pytest.mark.xfail(
-    is_not_utc(),
-    reason="https://github.com/rapidsai/cudf/issues/23422")
 def test_orc_gpu_write_cpu_read_timestamp_in_non_utc_timezone(spark_tmp_path):
     data_path = spark_tmp_path + "/ORC_GPU_WRITE_TZ"
-    assert_gpu_and_cpu_writes_are_equal_collect(
-        lambda spark, path: (
-            spark.range(3)
-                .selectExpr("CAST(1593604800 + id AS TIMESTAMP) AS ts")
-                .write.orc(path)),
-        lambda spark, path: spark.read.orc(path),
-        data_path)
+    write_func = lambda spark, path: (
+        spark.range(3)
+            .selectExpr("CAST(1593604800 + id AS TIMESTAMP) AS ts")
+            .write.orc(path))
+    read_func = lambda spark, path: spark.read.orc(path)
+    if is_not_utc():
+        # Non-UTC: the timestamp write must fall back to CPU (DataWritingCommandExec).
+        assert_gpu_fallback_write(write_func, read_func, data_path, 'DataWritingCommandExec')
+    else:
+        assert_gpu_and_cpu_writes_are_equal_collect(write_func, read_func, data_path)
 
 
 @pytest.mark.parametrize("reader_confs", reader_opt_confs, ids=idfn)

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuOrcFileFormat.scala

@@ -16,6 +16,8 @@
 
 package org.apache.spark.sql.rapids
 
+import java.time.ZoneId
+
 import ai.rapids.cudf._
 import com.nvidia.spark.rapids._
 import com.nvidia.spark.rapids.jni.fileio.RapidsFileIO
@@ -31,6 +33,7 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.execution.datasources.FileFormat
 import org.apache.spark.sql.execution.datasources.orc.{OrcFileFormat, OrcOptions, OrcUtils}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.rapids.execution.TrampolineUtil
 import org.apache.spark.sql.types._
 
@@ -82,8 +85,19 @@ object GpuOrcFileFormat extends Logging {
         t.isInstanceOf[BooleanType])
     }
 
-    // ORC writing always uses UTC internally (cuDF writes writerTimezone="UTC").
-    // The reader side handles timezone conversion. No write-side timezone restriction needed.
+    // cuDF's ORC writer always stamps writerTimezone="UTC" in the stripe footer and cannot
+    // record the actual JVM writer timezone (https://github.com/rapidsai/cudf/issues/23422).
+    // Because ORC's `timestamp` type is timezone-agnostic, a file written on the GPU in a
+    // non-UTC JVM is read back shifted by the zone offset by a CPU ORC reader. Fall back to CPU
+    // for non-UTC timestamp writes so the output stays interoperable. Reads use the JVM default
+    // (systemDefault) zone, so the write-side gate matches the reader on the same check.
+    val types = schema.map(_.dataType).toSet
+    if (types.exists(GpuOverrides.isOrContainsTimestamp) &&
+        !GpuOverrides.isUTCTimezone(ZoneId.systemDefault())) {
+      meta.willNotWorkOnGpu("Writing ORC timestamps is only supported in the UTC timezone " +
+        s"(JVM: ${ZoneId.systemDefault()}, session: ${SQLConf.get.sessionLocalTimeZone}). " +
+        "See https://github.com/rapidsai/cudf/issues/23422")
+    }
 
     if (hasBools && !meta.conf.isOrcBoolTypeEnabled) {
       meta.willNotWorkOnGpu("Nullable Booleans can not work in certain cases with ORC writer." +