Merge branch 'main' into regex-inline-flags

Author: igorpeshansky

integration_tests/src/main/python/orc_cast_test.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# Copyright (c) 2020-2026, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,7 +17,8 @@
 from asserts import assert_gpu_and_cpu_are_equal_collect, assert_gpu_and_cpu_error
 from conftest import is_not_utc
 from data_gen import *
-from marks import allow_non_gpu, datagen_overrides
+from marks import (allow_non_gpu, datagen_overrides, tz_sensitive_test,
+                   validate_execs_in_gpu_plan)
 from pyspark.sql.types import *
 from spark_session import with_cpu_session
 from orc_test import reader_opt_confs
@@ -128,11 +129,61 @@ def test_casting_from_double_to_timestamp(spark_tmp_path, data_gen):
     )
 
 
+@tz_sensitive_test
+@pytest.mark.skipif(not is_not_utc(), reason="non-UTC ORC timestamp regression test")
+@allow_non_gpu(*non_utc_allow_orc_scan)
+@validate_execs_in_gpu_plan('GpuFileSourceScanExec')
+def test_non_utc_timestamp_regressions(spark_tmp_path):
+    integer_path = spark_tmp_path + '/orc_integer_timestamp_regression'
+    double_path = spark_tmp_path + '/orc_double_timestamp_regression'
+    physical_path = spark_tmp_path + '/orc_physical_timestamp_regression'
+
+    with_cpu_session(
+        lambda spark: spark.createDataFrame([(514952012,)], "a long")
+            .write.orc(integer_path)
+    )
+    with_cpu_session(
+        lambda spark: spark.createDataFrame(
+            [(0.0,),
+             (float('nan'),),
+             (float('inf'),),
+             (float('-inf'),),
+             (-8589934591.999999,),
+             (-7953731124.723491,),
+             (-0.0015,),
+             (-0.0005,),
+             (0.0005,),
+             (0.0015,)],
+            "a double")
+            .write.orc(double_path)
+    )
+    with_cpu_session(
+        lambda spark: spark.createDataFrame(
+            [(-2957649381472612,), (-3649379812521628,)], "a long")
+            .selectExpr("timestamp_micros(a) AS a")
+            .write.orc(physical_path)
+    )
+
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: spark.read.schema("a timestamp").orc(integer_path)
+    )
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: spark.read.schema("a timestamp").orc(double_path)
+    )
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: spark.read.orc(physical_path)
+    )
+
+
 @allow_non_gpu(*non_utc_allow_for_test_casting_from_overflow_long)
-def test_casting_from_overflow_double_to_timestamp(spark_tmp_path):
+# Keep both signs beyond the Long microsecond range after any timezone correction.
+@pytest.mark.parametrize("overflow_value", [9223372123255.0, -9223372123255.0])
+def test_casting_from_overflow_double_to_timestamp(spark_tmp_path, overflow_value):
     orc_path = spark_tmp_path + '/orc_casting_from_overflow_double_to_timestamp'
     with_cpu_session(
-        lambda spark: unary_op_df(spark, DoubleGen(min_exp=38)).write.orc(orc_path)
+        lambda spark: spark.createDataFrame(
+            [(overflow_value,)], "a double")
+            .write.orc(orc_path)
     )
     assert_gpu_and_cpu_error(
         df_fun=lambda spark: spark.read.schema("a timestamp").orc(orc_path).collect(),

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcScan.scala

@@ -23,6 +23,7 @@ import java.nio.channels.Channels
 import java.nio.charset.StandardCharsets
 import java.time.ZoneId
 import java.util
+import java.util.TimeZone
 import java.util.regex.Pattern
 
 import scala.annotation.tailrec
@@ -322,8 +323,7 @@ object GpuOrcScan {
       case (DType.BOOL8 | DType.INT8 | DType.INT16 | DType.INT32 | DType.INT64,
       DType.TIMESTAMP_MICROSECONDS) =>
         withResource(OrcCastingShims.castIntegerToTimestamp(col, fromDt)) { timestamp =>
-          GpuTimeZoneDB.fromTimestampToUtcTimestamp(
-            timestamp, ZoneId.systemDefault().normalized())
+          GpuTimeZoneDB.convertOrcFromUtc(timestamp, ZoneId.systemDefault().getId)
         }
 
       // float to bool/integral
@@ -382,16 +382,19 @@ object GpuOrcScan {
         // Math.round half up can be implemented in terms of floor
         // Math.round(x) = n iff x is in [n-0.5, n+0.5) iff x+0.5 is in [n,n+1) iff floor(x+0.5) = n
         //
-        val milliseconds = withResource(Scalar.fromDouble(DateTimeConstants.MILLIS_PER_SECOND)) {
-          thousand =>
-          // ORC assumes value is in seconds
-          withResource(col.mul(thousand, DType.FLOAT64)) { doubleMillis =>
-            withResource(Scalar.fromDouble(0.5)) { half =>
-              withResource(doubleMillis.add(half)) { doubleMillisPlusHalf =>
-                withResource(doubleMillisPlusHalf.floor()) { millis =>
-                  withResource(getOverflowFlags(doubleMillis, millis)) { overflowFlags =>
-                    withResource(Scalar.fromNull(millis.getType)) { nullVal =>
-                      overflowFlags.ifElse(millis, nullVal)
+        val milliseconds = withResource(col.castTo(DType.FLOAT64)) { doubleSeconds =>
+          withResource(convertOrcFloatingPointSeconds(doubleSeconds)) { convertedSeconds =>
+            withResource(Scalar.fromDouble(DateTimeConstants.MILLIS_PER_SECOND)) { thousand =>
+              // ORC applies timezone conversion while the value is still in seconds.
+              withResource(convertedSeconds.mul(thousand, DType.FLOAT64)) { doubleMillis =>
+                withResource(Scalar.fromDouble(0.5)) { half =>
+                  withResource(doubleMillis.add(half)) { doubleMillisPlusHalf =>
+                    withResource(doubleMillisPlusHalf.floor()) { millis =>
+                      withResource(getOverflowFlags(doubleMillis, millis)) { overflowFlags =>
+                        withResource(Scalar.fromNull(millis.getType)) { nullVal =>
+                          overflowFlags.ifElse(millis, nullVal)
+                        }
+                      }
                     }
                   }
                 }
@@ -419,12 +422,11 @@ object GpuOrcScan {
           }
           withResource(Scalar.fromDouble(DateTimeConstants.MICROS_PER_MILLIS)) { thousand =>
             withResource(milliseconds.mul(thousand)) { microseconds =>
-                withResource(microseconds.castTo(DType.INT64)) { longVec =>
-                  withResource(longVec.castTo(DType.TIMESTAMP_MICROSECONDS)) { timestamp =>
-                    GpuTimeZoneDB.fromTimestampToUtcTimestamp(
-                      timestamp, ZoneId.systemDefault().normalized())
-                  }
+              withResource(microseconds.castTo(DType.INT64)) { longVec =>
+                withResource(longVec.castTo(DType.TIMESTAMP_MICROSECONDS)) { timestamp =>
+                  timestamp.incRefCount()
                 }
+              }
             }
           }
         }
@@ -444,6 +446,59 @@ object GpuOrcScan {
     }
   }
 
+  /**
+   * Apply ORC's offset lookup ordering before its millisecond rounding and overflow check.
+   * ORC looks up the offset at `(localMillis - rawOffset)`, while Spark materializes the final
+   * timestamp with java.time rules. Apply only that integral timezone delta to the original
+   * floating-point value so both the DST and historical behavior match Spark's CPU ORC path.
+   */
+  private def convertOrcFloatingPointSeconds(seconds: ColumnVector): ColumnVector = {
+    if (GpuOverrides.isUTCTimezone(ZoneId.systemDefault())) {
+      return seconds.incRefCount()
+    }
+    withResource(Scalar.fromDouble(DateTimeConstants.MICROS_PER_SECOND)) { microsPerSecond =>
+      val localTimestamp = withResource(
+          Scalar.fromDouble(DateTimeConstants.MILLIS_PER_SECOND)) { millisPerSecond =>
+        withResource(seconds.mul(millisPerSecond, DType.FLOAT64)) { doubleMillis =>
+          withResource(doubleMillis.castTo(DType.INT64)) { localMillis =>
+            withResource(localMillis.bitCastTo(DType.TIMESTAMP_MILLISECONDS)) {
+              localMillisTimestamp =>
+              localMillisTimestamp.castTo(DType.TIMESTAMP_MICROSECONDS)
+            }
+          }
+        }
+      }
+      withResource(localTimestamp) { _ =>
+        withResource(localTimestamp.bitCastTo(DType.INT64)) { localMicros =>
+          val rawOffsetMicros = TimeZone.getDefault.getRawOffset.toLong *
+            DateTimeConstants.MICROS_PER_MILLIS
+          withResource(Scalar.fromLong(rawOffsetMicros)) { rawOffset =>
+            withResource(localMicros.sub(rawOffset)) { offsetLookupMicros =>
+              withResource(offsetLookupMicros.castTo(DType.TIMESTAMP_MICROSECONDS)) {
+                offsetLookupTimestamp =>
+                val localAtLookup = GpuTimeZoneDB.fromUtcTimestampToTimestamp(
+                  offsetLookupTimestamp, ZoneId.systemDefault().normalized())
+                withResource(localAtLookup) { _ =>
+                  withResource(localAtLookup.bitCastTo(DType.INT64)) { localAtLookupMicros =>
+                    val offsetSeconds = withResource(
+                        localAtLookupMicros.sub(offsetLookupMicros)) { offsetMicros =>
+                      withResource(offsetMicros.castTo(DType.FLOAT64)) { doubleOffsetMicros =>
+                        doubleOffsetMicros.div(microsPerSecond, DType.FLOAT64)
+                      }
+                    }
+                    withResource(offsetSeconds) { _ =>
+                      seconds.sub(offsetSeconds, DType.FLOAT64)
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   /**
    * Whether the type casting is supported by GPU ORC reading.
    *

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala

@@ -21,7 +21,7 @@ import java.util.Optional
 
 import scala.collection.mutable.ArrayBuffer
 
-import ai.rapids.cudf.{ColumnView, DType, Table}
+import ai.rapids.cudf.{ColumnView, DType, Scalar, Table}
 import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableProducingSeq
 import com.nvidia.spark.rapids.jni.GpuTimeZoneDB
@@ -77,16 +77,17 @@ object GpuOrcTimezoneUtils {
    */
   private def rebaseWithWriterTimezone(
       input: Table, writerTz: String, readerTz: String): Table = {
+    val readerZone = ZoneId.of(readerTz, ZoneId.SHORT_IDS)
     withResource(input) { _ =>
       withResource(GpuTimeZoneDB.buildOrcTimezoneContext(writerTz, readerTz)) { tzCtx =>
         val newColumns = (0 until input.getNumberOfColumns).safeMap { colIdx =>
           val col = input.getColumn(colIdx)
           val dType = col.getType
           if (dType.hasTimeResolution) {
-            GpuTimeZoneDB.convertOrcTimezones(col, tzCtx)
+            convertOrcTimestamp(col, tzCtx, readerZone)
           } else if (dType == DType.LIST || dType == DType.STRUCT) {
             withResource(new ArrayBuffer[ColumnView]) { toClose =>
-              val rebased = rebaseNestedWithWriterTimezone(col, tzCtx, toClose)
+              val rebased = rebaseNestedWithWriterTimezone(col, tzCtx, readerZone, toClose)
               if (rebased eq col) {
                 col.incRefCount()
               } else {
@@ -105,18 +106,64 @@ object GpuOrcTimezoneUtils {
     }
   }
 
+  /**
+   * Match the full Spark ORC timestamp path. Apache ORC uses java.util.TimeZone while decoding,
+   * but Spark materializes the resulting java.sql.Timestamp using java.time rules. Those rule
+   * sets can differ for historical and projected timestamps.
+   */
+  private def convertOrcTimestamp(
+      col: ColumnView,
+      tzCtx: GpuTimeZoneDB.OrcTimezoneContext,
+      readerZone: ZoneId): ai.rapids.cudf.ColumnVector = {
+    withResource(GpuTimeZoneDB.convertOrcTimezones(col, tzCtx)) { orcTimestamp =>
+      val firstTransitionUs = tzCtx.getReaderFirstTransitionUs
+      if (firstTransitionUs == Long.MinValue) {
+        orcTimestamp.incRefCount()
+      } else {
+        val utilMicros = withResource(
+            GpuTimeZoneDB.convertOrcFromUtc(orcTimestamp, tzCtx)) { utilUtc =>
+          utilUtc.castTo(DType.INT64)
+        }
+        withResource(utilMicros) { _ =>
+          val ruleCorrection = withResource(GpuTimeZoneDB.fromTimestampToUtcTimestamp(
+              orcTimestamp, readerZone.normalized())) { zoneUtc =>
+            withResource(zoneUtc.castTo(DType.INT64)) { zoneMicros =>
+              zoneMicros.sub(utilMicros)
+            }
+          }
+          withResource(ruleCorrection) { _ =>
+            val correctedTimestamp = withResource(orcTimestamp.castTo(DType.INT64)) { orcMicros =>
+              withResource(orcMicros.add(ruleCorrection)) { corrected =>
+                corrected.castTo(DType.TIMESTAMP_MICROSECONDS)
+              }
+            }
+            withResource(correctedTimestamp) { _ =>
+              withResource(Scalar.timestampFromLong(
+                  DType.TIMESTAMP_MICROSECONDS, firstTransitionUs)) { firstTransition =>
+                withResource(orcTimestamp.lessThan(firstTransition)) { needsCorrection =>
+                  needsCorrection.ifElse(correctedTimestamp, orcTimestamp)
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   private def rebaseNestedWithWriterTimezone(
       col: ColumnView,
       tzCtx: GpuTimeZoneDB.OrcTimezoneContext,
+      readerZone: ZoneId,
       toClose: ArrayBuffer[ColumnView]): ColumnView = {
     val addToClose = (v: ColumnView) => { toClose += v; v }
     val dType = col.getType
 
     if (dType.hasTimeResolution) {
-      GpuTimeZoneDB.convertOrcTimezones(col, tzCtx)
+      convertOrcTimestamp(col, tzCtx, readerZone)
     } else if (dType == DType.LIST) {
       val child = addToClose(col.getChildColumnView(0))
-      val newChild = rebaseNestedWithWriterTimezone(child, tzCtx, toClose)
+      val newChild = rebaseNestedWithWriterTimezone(child, tzCtx, readerZone, toClose)
       if (newChild ne child) {
         col.replaceListChild(addToClose(newChild))
       } else {
@@ -125,7 +172,7 @@ object GpuOrcTimezoneUtils {
     } else if (dType == DType.STRUCT) {
       val newViews = (0 until col.getNumChildren).map { i =>
         val child = addToClose(col.getChildColumnView(i))
-        val newChild = rebaseNestedWithWriterTimezone(child, tzCtx, toClose)
+        val newChild = rebaseNestedWithWriterTimezone(child, tzCtx, readerZone, toClose)
         if (newChild ne child) addToClose(newChild)
         newChild
       }

sql-plugin/src/main/spark412/scala/com/nvidia/spark/rapids/shims/SparkShims.scala

@@ -17,7 +17,6 @@
 /*** spark-rapids-shim-json-lines
 {"spark": "412"}
 {"spark": "413"}
-{"spark": "420"}
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims
 

sql-plugin/src/main/spark420/scala/com/nvidia/spark/rapids/shims/GroupPartitionsExecMeta.scala

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.
+ *
+ * Licensed 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.
+ */
+
+/*** spark-rapids-shim-json-lines
+{"spark": "420"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+import com.nvidia.spark.rapids._
+
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExec
+
+class GroupPartitionsExecMeta(
+    groupPartitions: GroupPartitionsExec,
+    conf: RapidsConf,
+    parent: Option[RapidsMeta[_, _, _]],
+    rule: DataFromReplacementRule)
+    extends SparkPlanMeta[GroupPartitionsExec](groupPartitions, conf, parent, rule) {
+
+  override def tagPlanForGpu(): Unit = {
+    willNotWorkOnGpu("GroupPartitionsExec is not supported on GPU")
+  }
+
+  override def convertToCpu(): SparkPlan = {
+    // GroupPartitionsExec reads its child's KeyedPartitioning at execution time.
+    // GPU conversions can replace it with UnknownPartitioning, so keep the original
+    // CPU subtree until GroupPartitionsExec has a GPU implementation.
+    groupPartitions
+  }
+
+  override def convertToGpu(): GpuExec = {
+    throw new IllegalStateException("GroupPartitionsExec cannot be converted to GPU")
+  }
+}