From cc1d1018d4dca1fd86849dfbd9f452e439480625 Mon Sep 17 00:00:00 2001 From: Allen Xu Date: Wed, 9 Sep 2026 15:54:52 +0800 Subject: [PATCH 1/5] Fix legacy ORC timestamp rebasing Signed-off-by: Allen Xu --- .../spark/rapids/GpuOrcTimezoneUtils.scala | 118 +++++++++---- .../rapids/GpuTimestampRebaseUtils.scala | 164 +++++++++++++++++ .../sql/rapids/RebaseDateTimeBridge.scala | 62 +++++++ .../rapids/GpuTimestampRebaseSuite.scala | 167 ++++++++++++++++++ .../spark/rapids/OrcCalendarSuite.scala | 50 +++++- .../rapids/timezone/OrcTimezoneSuite.scala | 4 + .../sql/rapids/utils/RapidsTestSettings.scala | 6 - 7 files changed, 522 insertions(+), 49 deletions(-) create mode 100644 sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala create mode 100644 sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala create mode 100644 tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala index 61a785b5a9b..d969df6a3c5 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala @@ -124,11 +124,19 @@ object GpuOrcTimezoneUtils { val readerZone = ZoneId.of(readerTz, ZoneId.SHORT_IDS) withResource(input) { _ => if (containsOrcTimestamp(input)) { - withResource(GpuTimeZoneDB.buildOrcTimezoneContext(writerTz, readerTz)) { tzCtx => - rebaseColumns(input, Some(tzCtx), readerZone, writerUsedProlepticGregorian) + val legacyTimestampRebase = if (writerUsedProlepticGregorian) { + None + } else { + Some(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(readerZone.getId)) + } + withResource(legacyTimestampRebase) { legacyRebase => + withResource(GpuTimeZoneDB.buildOrcTimezoneContext(writerTz, readerTz)) { tzCtx => + rebaseColumns(input, Some(tzCtx), readerZone, legacyRebase, + writerUsedProlepticGregorian) + } } } else { - rebaseColumns(input, None, readerZone, writerUsedProlepticGregorian) + rebaseColumns(input, None, readerZone, None, writerUsedProlepticGregorian) } } } @@ -158,6 +166,8 @@ object GpuOrcTimezoneUtils { input: Table, tzCtx: Option[GpuTimeZoneDB.OrcTimezoneContext], readerZone: ZoneId, + legacyTimestampRebase: Option[ + GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext], writerUsedProlepticGregorian: Boolean): Table = { val newColumns = (0 until input.getNumberOfColumns).safeMap { colIdx => val col = input.getColumn(colIdx) @@ -165,11 +175,12 @@ object GpuOrcTimezoneUtils { if (dType == DType.TIMESTAMP_DAYS && !writerUsedProlepticGregorian) { DateTimeRebase.rebaseJulianToGregorian(col) } else if (dType.hasTimeResolution) { - convertOrcTimestamp(col, tzCtx.get, readerZone) + convertOrcTimestamp(col, tzCtx.get, readerZone, legacyTimestampRebase) } else if (dType == DType.LIST || dType == DType.STRUCT) { withResource(new ArrayBuffer[ColumnView]) { toClose => val rebased = rebaseNestedWithWriterTimezone( - col, tzCtx, readerZone, writerUsedProlepticGregorian, toClose) + col, tzCtx, readerZone, legacyTimestampRebase, + writerUsedProlepticGregorian, toClose) if (rebased eq col) { col.incRefCount() } else { @@ -187,45 +198,70 @@ 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. + * Match the full Spark ORC timestamp path after reconstructing the writer-specific ORC epoch. + * Legacy-calendar files use Spark's timezone-specific Julian-to-Gregorian rebase map. Files + * already written with the proleptic calendar only need the java.util.TimeZone versus java.time + * rule correction from the ORC materialization path. */ private def convertOrcTimestamp( col: ColumnView, tzCtx: GpuTimeZoneDB.OrcTimezoneContext, - readerZone: ZoneId): ai.rapids.cudf.ColumnVector = { + readerZone: ZoneId, + legacyTimestampRebase: Option[ + GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext]): 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) - } - val ruleCorrection = withResource(utilMicros) { utilMicros => - withResource(GpuTimeZoneDB.fromTimestampToUtcTimestamp( - orcTimestamp, readerZone.normalized())) { zoneUtc => - withResource(zoneUtc.castTo(DType.INT64)) { zoneMicros => - zoneMicros.sub(utilMicros) - } + legacyTimestampRebase match { + case Some(rebase) => rebase.rebase(orcTimestamp) + case None => correctOrcTimestampRules(orcTimestamp, tzCtx, readerZone) + } + } + } + + /** Correct only the historical java.util.TimeZone/java.time rule difference. */ + private def correctOrcTimestampRules( + orcTimestamp: ColumnVector, + tzCtx: GpuTimeZoneDB.OrcTimezoneContext, + readerZone: ZoneId): ColumnVector = { + val firstTransitionUs = tzCtx.getReaderFirstTransitionUs + if (firstTransitionUs == Long.MinValue) { + orcTimestamp.incRefCount() + } else { + withResource(correctOrcTimestampRulesBeforeTransition( + orcTimestamp, tzCtx, readerZone)) { correctedTimestamp => + withResource(Scalar.timestampFromLong( + DType.TIMESTAMP_MICROSECONDS, firstTransitionUs)) { firstTransition => + withResource(orcTimestamp.lessThan(firstTransition)) { needsCorrection => + needsCorrection.ifElse(correctedTimestamp, orcTimestamp) } } - val correctedTimestamp = withResource(ruleCorrection) { ruleCorrection => - withResource(orcTimestamp.castTo(DType.INT64)) { orcMicros => - withResource(orcMicros.add(ruleCorrection)) { corrected => - corrected.castTo(DType.TIMESTAMP_MICROSECONDS) - } - } + } + } + } + + private def correctOrcTimestampRulesBeforeTransition( + orcTimestamp: ColumnVector, + tzCtx: GpuTimeZoneDB.OrcTimezoneContext, + readerZone: ZoneId): ColumnVector = { + withResource(computeOrcTimestampRuleCorrection( + orcTimestamp, tzCtx, readerZone)) { ruleCorrection => + 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 computeOrcTimestampRuleCorrection( + orcTimestamp: ColumnVector, + tzCtx: GpuTimeZoneDB.OrcTimezoneContext, + readerZone: ZoneId): ColumnVector = { + withResource(GpuTimeZoneDB.convertOrcFromUtc(orcTimestamp, tzCtx)) { utilUtc => + withResource(GpuTimeZoneDB.fromTimestampToUtcTimestamp( + orcTimestamp, readerZone.normalized())) { zoneUtc => + withResource(Seq(utilUtc, zoneUtc).safeMap(_.castTo(DType.INT64))) { + case Seq(utilMicros, zoneMicros) => + zoneMicros.sub(utilMicros) } } } @@ -235,6 +271,8 @@ object GpuOrcTimezoneUtils { col: ColumnView, tzCtx: Option[GpuTimeZoneDB.OrcTimezoneContext], readerZone: ZoneId, + legacyTimestampRebase: Option[ + GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext], writerUsedProlepticGregorian: Boolean, toClose: ArrayBuffer[ColumnView]): ColumnView = { val addToClose = (v: ColumnView) => { toClose += v; v } @@ -243,11 +281,12 @@ object GpuOrcTimezoneUtils { if (dType == DType.TIMESTAMP_DAYS && !writerUsedProlepticGregorian) { DateTimeRebase.rebaseJulianToGregorian(col) } else if (dType.hasTimeResolution) { - convertOrcTimestamp(col, tzCtx.get, readerZone) + convertOrcTimestamp(col, tzCtx.get, readerZone, legacyTimestampRebase) } else if (dType == DType.LIST) { val child = addToClose(col.getChildColumnView(0)) val newChild = rebaseNestedWithWriterTimezone( - child, tzCtx, readerZone, writerUsedProlepticGregorian, toClose) + child, tzCtx, readerZone, legacyTimestampRebase, + writerUsedProlepticGregorian, toClose) if (newChild ne child) { col.replaceListChild(addToClose(newChild)) } else { @@ -258,7 +297,8 @@ object GpuOrcTimezoneUtils { val newViews = (0 until col.getNumChildren).map { i => val child = addToClose(col.getChildColumnView(i)) val newChild = rebaseNestedWithWriterTimezone( - child, tzCtx, readerZone, writerUsedProlepticGregorian, toClose) + child, tzCtx, readerZone, legacyTimestampRebase, + writerUsedProlepticGregorian, toClose) if (newChild ne child) { childChanged = true addToClose(newChild) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala new file mode 100644 index 00000000000..d981a4d5b5a --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala @@ -0,0 +1,164 @@ +/* + * 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. + */ + +package com.nvidia.spark.rapids + +import ai.rapids.cudf.{BinaryOp, ColumnVector, ColumnView, DType} +import ai.rapids.cudf.{HostColumnVector, OrderByArg, Table} +import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} + +import org.apache.spark.sql.rapids.RebaseDateTimeBridge + +private[rapids] object GpuTimestampRebaseUtils { + + private def retainOrCopy(input: ColumnView): ColumnVector = input match { + case columnVector: ColumnVector => columnVector.incRefCount() + case _ => input.copyToColumnVector() + } + + private def isModernOrAllNull(input: ColumnView): Boolean = { + withResource(input.min()) { minValue => + !minValue.isValid || minValue.getLong >= RebaseDateTimeBridge.lastSwitchJulianTs + } + } + + final class JulianToGregorianMicrosContext( + timeZoneId: String, + switches: Option[Table], + paddedDiffs: Option[Table]) extends AutoCloseable { + + private def rebaseOnHost(input: ColumnView): ColumnVector = { + val rowCount = input.getRowCount.toInt + withResource(input.copyToHost()) { hostInput => + withResource(HostColumnVector.builder(DType.TIMESTAMP_MICROSECONDS, rowCount)) { builder => + var row = 0 + while (row < rowCount) { + if (hostInput.isNull(row)) { + builder.appendNull() + } else { + builder.append(RebaseDateTimeBridge.rebaseJulianToGregorianMicros( + timeZoneId, hostInput.getLong(row))) + } + row += 1 + } + withResource(builder.build()) { hostOutput => + hostOutput.copyToDevice() + } + } + } + } + + private def rebaseWithSearchColumn( + input: ColumnView, + searchColumn: ColumnVector): ColumnVector = { + withResource(new Table(searchColumn)) { searchTable => + withResource(switches.get.upperBound(searchTable, OrderByArg.asc(0, false))) { indices => + val hasBeforeFirstSwitch = withResource(indices.min()) { minIndex => + minIndex.isValid && minIndex.getInt == 0 + } + if (hasBeforeFirstSwitch) { + // Spark's precomputed maps intentionally stop at the Common Era boundary. Preserve + // exact Spark semantics for rarer BCE data by using its Calendar-based slow path. + rebaseOnHost(input) + } else { + withResource(paddedDiffs.get.gather(indices)) { gatheredDiffs => + input.binaryOp( + BinaryOp.ADD, gatheredDiffs.getColumn(0), DType.TIMESTAMP_MICROSECONDS) + } + } + } + } + } + + def rebase(input: ColumnView): ColumnVector = { + require(input.getType == DType.TIMESTAMP_MICROSECONDS, + s"expected TIMESTAMP_MICROSECONDS but found ${input.getType}") + if (input.getRowCount == 0 || isModernOrAllNull(input)) { + retainOrCopy(input) + } else if (switches.isEmpty) { + // Spark's bundled map can lag valid IDs added by newer JDK timezone databases. + rebaseOnHost(input) + } else { + input match { + case columnVector: ColumnVector => + rebaseWithSearchColumn(input, columnVector) + case _ => + withResource(input.copyToColumnVector()) { searchColumn => + rebaseWithSearchColumn(input, searchColumn) + } + } + } + } + + override def close(): Unit = { + try { + switches.foreach(_.close()) + } finally { + paddedDiffs.foreach(_.close()) + } + } + } + + final class LazyJulianToGregorianMicrosContext(timeZoneId: String) extends AutoCloseable { + private var delegate: JulianToGregorianMicrosContext = _ + + def rebase(input: ColumnView): ColumnVector = { + require(input.getType == DType.TIMESTAMP_MICROSECONDS, + s"expected TIMESTAMP_MICROSECONDS but found ${input.getType}") + if (input.getRowCount == 0 || isModernOrAllNull(input)) { + retainOrCopy(input) + } else { + if (delegate == null) { + delegate = createJulianToGregorianMicrosContext(timeZoneId) + } + delegate.rebase(input) + } + } + + override def close(): Unit = if (delegate != null) { + delegate.close() + delegate = null + } + } + + def createJulianToGregorianMicrosContext( + timeZoneId: String): JulianToGregorianMicrosContext = { + RebaseDateTimeBridge.getJulianToGregorianMicros(timeZoneId).map { info => + require(info.switches.nonEmpty, s"empty Spark timestamp rebase map for '$timeZoneId'") + require(info.switches.length == info.diffs.length, + s"invalid Spark timestamp rebase map for '$timeZoneId'") + + val switchTable = withResource( + ColumnVector.timestampMicroSecondsFromLongs(info.switches: _*)) { switchColumn => + new Table(switchColumn) + } + closeOnExcept(switchTable) { _ => + // upperBound returns 0 before the first switch and k + 1 at switch k. The leading + // sentinel aligns every valid upper-bound index directly with its Spark rebase diff. + val padded = new Array[Long](info.diffs.length + 1) + System.arraycopy(info.diffs, 0, padded, 1, info.diffs.length) + val diffTable = withResource( + ColumnVector.durationMicroSecondsFromLongs(padded: _*)) { diffColumn => + new Table(diffColumn) + } + new JulianToGregorianMicrosContext( + timeZoneId, Some(switchTable), Some(diffTable)) + } + }.getOrElse { + new JulianToGregorianMicrosContext(timeZoneId, None, None) + } + } +} diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala new file mode 100644 index 00000000000..e0e7215962a --- /dev/null +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala @@ -0,0 +1,62 @@ +/* + * 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. + */ + +package org.apache.spark.sql.rapids + +import java.time.{ZoneId, ZoneOffset} + +import org.apache.spark.sql.catalyst.util.RebaseDateTime + +/** Access Spark's runtime-specific timestamp rebase records from the sql package. */ +object RebaseDateTimeBridge { + final case class RebaseInfo(switches: Array[Long], diffs: Array[Long]) + + private val MicrosPerSecond = 1000000L + + private lazy val julianToGregorianMicros = + RebaseDateTime.loadRebaseRecords("julian-gregorian-rebase-micros.json") + + val lastSwitchJulianTs: Long = RebaseDateTime.lastSwitchJulianTs + + private def copyInfo(timeZoneId: String): Option[RebaseInfo] = { + julianToGregorianMicros.get(timeZoneId).map { info => + RebaseInfo(info.switches.clone(), info.diffs.clone()) + } + } + + def getJulianToGregorianMicros(timeZoneId: String): Option[RebaseInfo] = { + copyInfo(timeZoneId).orElse { + val zoneId = ZoneId.of(timeZoneId, ZoneId.SHORT_IDS) + copyInfo(zoneId.getId).orElse { + zoneId.normalized() match { + case offset: ZoneOffset => + copyInfo("UTC").map { utcInfo => + // A fixed-offset local midnight is shifted by the inverse offset from UTC. + val switchShift = Math.multiplyExact( + -offset.getTotalSeconds.toLong, MicrosPerSecond) + RebaseInfo( + utcInfo.switches.map(switch => Math.addExact(switch, switchShift)), + utcInfo.diffs) + } + case _ => None + } + } + } + } + + def rebaseJulianToGregorianMicros(timeZoneId: String, micros: Long): Long = + RebaseDateTime.rebaseJulianToGregorianMicros(timeZoneId, micros) +} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala new file mode 100644 index 00000000000..a0d353348a1 --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala @@ -0,0 +1,167 @@ +/* + * 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. + */ + +package com.nvidia.spark.rapids + +import java.time.ZoneId + +import ai.rapids.cudf.ColumnVector +import com.nvidia.spark.rapids.Arm.withResource +import com.nvidia.spark.rapids.CudfTestHelper.assertColumnsAreEqual +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.catalyst.util.RebaseDateTime + +class GpuTimestampRebaseSuite extends AnyFunSuite { + private val LosAngeles = "America/Los_Angeles" + + private val losAngelesInputs: Array[java.lang.Long] = Array( + null, + -30578137077000001L, + -30578137077000000L, + -30578137076999999L, + -30578137076876544L, + -12219264000000001L, + -12219264000000000L, + -12219263999999999L, + -2717640000000001L, + -2717640000000000L, + -2717639999999999L, + -2208988800000000L) + + private val losAngelesExpected: Array[java.lang.Long] = Array( + null, + -30578655899000001L, + -30578655899000000L, + -30578655898999999L, + -30578655898876544L, + -12220128422000001L, + -12219264422000000L, + -12219264421999999L, + -2717640422000001L, + -2717640000000000L, + -2717639999999999L, + -2208988800000000L) + + private def assertRebase( + timeZoneId: String, + input: Array[java.lang.Long], + expected: Array[java.lang.Long]): Unit = { + withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(timeZoneId)) { + context => + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(input: _*)) { inputColumn => + withResource(context.rebase(inputColumn)) { actual => + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(expected: _*)) { + expectedColumn => + assertColumnsAreEqual(expectedColumn, actual) + } + } + } + } + } + + test("timezone-specific timestamp rebase matches Spark at every critical boundary") { + losAngelesInputs.zip(losAngelesExpected).foreach { case (input, expected) => + if (input != null) { + assert(RebaseDateTime.rebaseJulianToGregorianMicros(LosAngeles, input) === expected) + } + } + assertRebase(LosAngeles, losAngelesInputs, losAngelesExpected) + } + + test("UTC timestamp rebase remains distinct from Los Angeles historical rules") { + val input = Array[java.lang.Long](-30578137076876544L) + val expected = Array[java.lang.Long](-30578655476876544L) + assert(RebaseDateTime.rebaseJulianToGregorianMicros("UTC", input.head) === expected.head) + assertRebase("UTC", input, expected) + } + + test("fixed-offset and short timezone IDs match Spark") { + val fixedOffset = "GMT+05:30" + val fixedOffsetInputs: Array[java.lang.Long] = Array( + -30578137076876544L, + -12219312600000001L, + -12219312600000000L, + -12219312599999999L) + val fixedOffsetExpected: Array[java.lang.Long] = Array( + -30578655476876544L, + -12220176600000001L, + -12219312600000000L, + -12219312599999999L) + fixedOffsetInputs.zip(fixedOffsetExpected).foreach { case (input, expected) => + assert(RebaseDateTime.rebaseJulianToGregorianMicros(fixedOffset, input) === expected) + } + assertRebase(fixedOffset, fixedOffsetInputs, fixedOffsetExpected) + + assert(RebaseDateTime.rebaseJulianToGregorianMicros( + "PST", losAngelesInputs(4)) === losAngelesExpected(4)) + assertRebase("PST", Array(losAngelesInputs(4)), Array(losAngelesExpected(4))) + } + + test("valid JDK timezone missing from Spark's bundled map uses the exact fallback") { + val timeZoneId = "Europe/Kyiv" + assume(ZoneId.getAvailableZoneIds.contains(timeZoneId), + s"$timeZoneId is not available in this JDK timezone database") + val input: Array[java.lang.Long] = Array( + -30578137076876544L, + -12219264000000001L, + -12219264000000000L, + 0L) + val expected: Array[java.lang.Long] = Array( + -30578655600876544L, + -12219264124000001L, + -12219264124000000L, + 0L) + input.zip(expected).foreach { case (value, rebased) => + assert(RebaseDateTime.rebaseJulianToGregorianMicros(timeZoneId, value) === rebased) + } + assertRebase(timeZoneId, input, expected) + assertRebase(timeZoneId, Array[java.lang.Long](null, 0L), + Array[java.lang.Long](null, 0L)) + } + + test("timestamp rebase preserves empty and sliced columns") { + assertRebase(LosAngeles, Array.empty[java.lang.Long], Array.empty[java.lang.Long]) + + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs( + 0L, losAngelesInputs(4), losAngelesInputs(5), 0L)) { input => + withResource(input.subVector(1, 3)) { slicedInput => + withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(LosAngeles)) { + context => + withResource(context.rebase(slicedInput)) { actual => + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs( + losAngelesExpected(4), losAngelesExpected(5))) { expected => + assertColumnsAreEqual(expected, actual) + } + } + } + } + } + } + + test("timestamp rebase uses Spark's exact fallback before the first map boundary") { + val input: Array[java.lang.Long] = Array( + -62135740800000001L, + -100000000000000000L) + val expected: Array[java.lang.Long] = Array( + -62135568422000001L, + -99999050022000000L) + input.zip(expected).foreach { case (value, rebased) => + assert(RebaseDateTime.rebaseJulianToGregorianMicros(LosAngeles, value) === rebased) + } + assertRebase(LosAngeles, input, expected) + } +} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala index 92c320e4ebe..a69846ebf97 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala @@ -19,6 +19,7 @@ package com.nvidia.spark.rapids import java.io.File import java.nio.file.{Files, StandardCopyOption} import java.time.{LocalDate, ZoneId} +import java.util.TimeZone import ai.rapids.cudf.{ColumnVector, Table} import com.nvidia.spark.rapids.Arm.{withResource, withResourceIfAllowed} @@ -35,6 +36,8 @@ import org.apache.spark.sql.rapids.shims.TrampolineConnectShims.SparkSession class OrcCalendarSuite extends SparkQueryCompareTestSuite { private val legacyDateResource = "test-data/before_1582_date_v2_4.snappy.orc" + private val legacyTimestampResource = "test-data/before_1582_ts_v2_4.snappy.orc" + private val legacyTimestampValue = "1001-01-01 01:02:03.123456" private val dateValue = LocalDate.of(1200, 1, 1).toEpochDay private val modernDateValue = LocalDate.of(2000, 1, 1).toEpochDay @@ -51,12 +54,12 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { .set("spark.sql.files.maxPartitionBytes", (1L << 30).toString) } - private def readLegacyDateResource(spark: SparkSession) = { + private def readLegacyResource(spark: SparkSession, resourceName: String) = { val resource = Option(Thread.currentThread().getContextClassLoader - .getResource(legacyDateResource)).getOrElse { - throw new IllegalStateException(s"Missing Spark test resource: $legacyDateResource") + .getResource(resourceName)).getOrElse { + throw new IllegalStateException(s"Missing Spark test resource: $resourceName") } - val file = File.createTempFile("spark-24-date", ".orc") + val file = File.createTempFile("spark-24-calendar", ".orc") file.deleteOnExit() val input = resource.openStream() try { @@ -67,6 +70,9 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { spark.read.orc(file.getCanonicalPath) } + private def readLegacyDateResource(spark: SparkSession) = + readLegacyResource(spark, legacyDateResource) + private def setDate(vector: DateColumnVector): Unit = { vector.setUsingProlepticCalendar(true) vector.vector(0) = dateValue @@ -155,6 +161,42 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { } } + for { + v1SourceList <- Seq("orc", "") + useChunkedReader <- Seq(false, true) + } { + test(s"read Spark 2.4 legacy ORC timestamp, source list is ($v1SourceList), " + + s"chunked=$useChunkedReader") { + val originalTimeZone = TimeZone.getDefault + val conf = calendarConf(RapidsReaderType.PERFILE, useChunkedReader, v1SourceList) + val expectedScan = if (v1SourceList == "orc") "GpuFileSourceScanExec" else "GpuBatchScan" + try { + val (fromCpu, fromGpu) = runOnCpuAndGpu( + spark => { + TimeZone.setDefault(TimeZone.getTimeZone("America/Los_Angeles")) + spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + readLegacyResource(spark, legacyTimestampResource) + }, + identity, + conf = conf, + repart = 0, + skipCanonicalizationCheck = true, + existClasses = expectedScan) + compareResults( + sort = false, + floatEpsilon = 0.0, + fromCpu = fromCpu, + fromGpu = fromGpu) + Seq(fromCpu, fromGpu).foreach { rows => + assert(rows.length === 1) + assert(rows.head.getTimestamp(0).toString === legacyTimestampValue) + } + } finally { + TimeZone.setDefault(originalTimeZone) + } + } + } + for { readerType <- Seq(RapidsReaderType.COALESCING, RapidsReaderType.MULTITHREADED) useChunkedReader <- Seq(false, true) diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala index c562b0402fa..18405b10b6f 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala @@ -88,6 +88,8 @@ class OrcTimezoneSuite extends SparkQueryCompareTestSuite { // Exact Asia/Shanghai writer=reader reproducer for the ORC epoch borrow correction. private val ShanghaiEpochBorrowTsUs = -7713116127L + // Exact SPARK-31284 value after reader-timezone Julian-to-Gregorian rebasing. + private val Spark31284TimestampUs = -30578655898876544L // Exact pre-first-transition values from non-UTC schema-evolution failures. private val newYorkHistoricalTsUs = -2957649381472612L private val shanghaiHistoricalTsUs = -3649379812521628L @@ -141,6 +143,8 @@ class OrcTimezoneSuite extends SparkQueryCompareTestSuite { ParisFirstTransitionLocalUs, ParisFirstTransitionLocalUs + 1L) Seq( + Spark31284TimestampUs, + Spark31284TimestampUs, newYorkHistoricalTsUs, shanghaiHistoricalTsUs, ShanghaiEpochBorrowTsUs, diff --git a/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala b/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala index bf239bc6008..ab1f0b3e506 100644 --- a/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala +++ b/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala @@ -253,9 +253,6 @@ class RapidsTestSettings extends BackendTestSettings { .exclude("Propagate Hadoop configs from orc options to underlying file system", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/11602. " + "Recovery trigger: GPU ORC writes propagate data-source options; P1.")) - .exclude("SPARK-31284: compatibility with Spark 2.4 in reading timestamps", - KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15471. " + - "Recovery trigger: GPU ORC legacy timestamp reads match Spark CPU; P0.")) .exclude("SPARK-31284, SPARK-31423: rebasing timestamps in write", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15473. " + "Recovery trigger: GPU ORC legacy timestamp round trips match Spark CPU; P0.")) @@ -271,9 +268,6 @@ class RapidsTestSettings extends BackendTestSettings { .exclude("Propagate Hadoop configs from orc options to underlying file system", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/11602. " + "Recovery trigger: GPU ORC writes propagate data-source options; P1.")) - .exclude("SPARK-31284: compatibility with Spark 2.4 in reading timestamps", - KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15471. " + - "Recovery trigger: GPU ORC legacy timestamp reads match Spark CPU; P0.")) .exclude("SPARK-31284, SPARK-31423: rebasing timestamps in write", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15473. " + "Recovery trigger: GPU ORC legacy timestamp round trips match Spark CPU; P0.")) From 58833a3e5111e7f4ad89f13e8328ccc113342cf7 Mon Sep 17 00:00:00 2001 From: Allen Xu Date: Wed, 9 Sep 2026 16:29:36 +0800 Subject: [PATCH 2/5] Move legacy ORC regression coverage to integration tests Signed-off-by: Allen Xu --- integration_tests/src/main/python/orc_test.py | 19 ++ .../resources/before_1582_ts_v2_4.snappy.orc | Bin 0 -> 251 bytes .../rapids/GpuTimestampRebaseSuite.scala | 167 ------------------ .../spark/rapids/OrcCalendarSuite.scala | 50 +----- .../rapids/timezone/OrcTimezoneSuite.scala | 4 - .../sql/rapids/utils/RapidsTestSettings.scala | 6 + 6 files changed, 29 insertions(+), 217 deletions(-) create mode 100644 integration_tests/src/test/resources/before_1582_ts_v2_4.snappy.orc delete mode 100644 tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala diff --git a/integration_tests/src/main/python/orc_test.py b/integration_tests/src/main/python/orc_test.py index 7d83392fe39..d3842b0f81e 100644 --- a/integration_tests/src/main/python/orc_test.py +++ b/integration_tests/src/main/python/orc_test.py @@ -122,6 +122,25 @@ def test_basic_read(std_input_path, name, read_func, v1_enabled_list, orc_impl, read_func(std_input_path + '/' + name), conf=all_confs) + +@pytest.mark.parametrize('v1_enabled_list', ['', 'orc']) +@pytest.mark.parametrize('vectorized_reader', [False, True]) +@tz_sensitive_test +def test_orc_read_spark_2_4_legacy_timestamp( + std_input_path, v1_enabled_list, vectorized_reader): + data_path = std_input_path + '/before_1582_ts_v2_4.snappy.orc' + all_confs = { + 'spark.sql.sources.useV1SourceList': v1_enabled_list, + 'spark.sql.orc.impl': 'native', + 'spark.sql.orc.enableVectorizedReader': vectorized_reader, + } + gpu_scan = 'GpuFileSourceScanExec' if v1_enabled_list == 'orc' else 'GpuBatchScanExec' + assert_cpu_and_gpu_are_equal_collect_with_capture( + read_orc_df(data_path), + exist_classes=gpu_scan, + conf=all_confs, + require_non_empty=True) + # ORC does not support negative scale for decimal. So here is "decimal_gens_no_neg". # Otherwise it will get the below exception. # ... diff --git a/integration_tests/src/test/resources/before_1582_ts_v2_4.snappy.orc b/integration_tests/src/test/resources/before_1582_ts_v2_4.snappy.orc new file mode 100644 index 0000000000000000000000000000000000000000..af9ef040270ac3f79570e275d33940972934b35d GIT binary patch literal 251 zcmeYdau#G@;9?VE;b074a0N0IxY!uLKuC;((Mv#L=As?lA3yUdKV36Pj(~J_TK$=m2QGyFdvjS-jAtnw6AZC(c4hRl`SOd0^ zO9Hzk5kkET3>?e?rBYl#(m+B;!9asiiIamzh>=l>siatmiG$OC(F$ya6;3l8Cir - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(input: _*)) { inputColumn => - withResource(context.rebase(inputColumn)) { actual => - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(expected: _*)) { - expectedColumn => - assertColumnsAreEqual(expectedColumn, actual) - } - } - } - } - } - - test("timezone-specific timestamp rebase matches Spark at every critical boundary") { - losAngelesInputs.zip(losAngelesExpected).foreach { case (input, expected) => - if (input != null) { - assert(RebaseDateTime.rebaseJulianToGregorianMicros(LosAngeles, input) === expected) - } - } - assertRebase(LosAngeles, losAngelesInputs, losAngelesExpected) - } - - test("UTC timestamp rebase remains distinct from Los Angeles historical rules") { - val input = Array[java.lang.Long](-30578137076876544L) - val expected = Array[java.lang.Long](-30578655476876544L) - assert(RebaseDateTime.rebaseJulianToGregorianMicros("UTC", input.head) === expected.head) - assertRebase("UTC", input, expected) - } - - test("fixed-offset and short timezone IDs match Spark") { - val fixedOffset = "GMT+05:30" - val fixedOffsetInputs: Array[java.lang.Long] = Array( - -30578137076876544L, - -12219312600000001L, - -12219312600000000L, - -12219312599999999L) - val fixedOffsetExpected: Array[java.lang.Long] = Array( - -30578655476876544L, - -12220176600000001L, - -12219312600000000L, - -12219312599999999L) - fixedOffsetInputs.zip(fixedOffsetExpected).foreach { case (input, expected) => - assert(RebaseDateTime.rebaseJulianToGregorianMicros(fixedOffset, input) === expected) - } - assertRebase(fixedOffset, fixedOffsetInputs, fixedOffsetExpected) - - assert(RebaseDateTime.rebaseJulianToGregorianMicros( - "PST", losAngelesInputs(4)) === losAngelesExpected(4)) - assertRebase("PST", Array(losAngelesInputs(4)), Array(losAngelesExpected(4))) - } - - test("valid JDK timezone missing from Spark's bundled map uses the exact fallback") { - val timeZoneId = "Europe/Kyiv" - assume(ZoneId.getAvailableZoneIds.contains(timeZoneId), - s"$timeZoneId is not available in this JDK timezone database") - val input: Array[java.lang.Long] = Array( - -30578137076876544L, - -12219264000000001L, - -12219264000000000L, - 0L) - val expected: Array[java.lang.Long] = Array( - -30578655600876544L, - -12219264124000001L, - -12219264124000000L, - 0L) - input.zip(expected).foreach { case (value, rebased) => - assert(RebaseDateTime.rebaseJulianToGregorianMicros(timeZoneId, value) === rebased) - } - assertRebase(timeZoneId, input, expected) - assertRebase(timeZoneId, Array[java.lang.Long](null, 0L), - Array[java.lang.Long](null, 0L)) - } - - test("timestamp rebase preserves empty and sliced columns") { - assertRebase(LosAngeles, Array.empty[java.lang.Long], Array.empty[java.lang.Long]) - - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs( - 0L, losAngelesInputs(4), losAngelesInputs(5), 0L)) { input => - withResource(input.subVector(1, 3)) { slicedInput => - withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(LosAngeles)) { - context => - withResource(context.rebase(slicedInput)) { actual => - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs( - losAngelesExpected(4), losAngelesExpected(5))) { expected => - assertColumnsAreEqual(expected, actual) - } - } - } - } - } - } - - test("timestamp rebase uses Spark's exact fallback before the first map boundary") { - val input: Array[java.lang.Long] = Array( - -62135740800000001L, - -100000000000000000L) - val expected: Array[java.lang.Long] = Array( - -62135568422000001L, - -99999050022000000L) - input.zip(expected).foreach { case (value, rebased) => - assert(RebaseDateTime.rebaseJulianToGregorianMicros(LosAngeles, value) === rebased) - } - assertRebase(LosAngeles, input, expected) - } -} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala index a69846ebf97..92c320e4ebe 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala @@ -19,7 +19,6 @@ package com.nvidia.spark.rapids import java.io.File import java.nio.file.{Files, StandardCopyOption} import java.time.{LocalDate, ZoneId} -import java.util.TimeZone import ai.rapids.cudf.{ColumnVector, Table} import com.nvidia.spark.rapids.Arm.{withResource, withResourceIfAllowed} @@ -36,8 +35,6 @@ import org.apache.spark.sql.rapids.shims.TrampolineConnectShims.SparkSession class OrcCalendarSuite extends SparkQueryCompareTestSuite { private val legacyDateResource = "test-data/before_1582_date_v2_4.snappy.orc" - private val legacyTimestampResource = "test-data/before_1582_ts_v2_4.snappy.orc" - private val legacyTimestampValue = "1001-01-01 01:02:03.123456" private val dateValue = LocalDate.of(1200, 1, 1).toEpochDay private val modernDateValue = LocalDate.of(2000, 1, 1).toEpochDay @@ -54,12 +51,12 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { .set("spark.sql.files.maxPartitionBytes", (1L << 30).toString) } - private def readLegacyResource(spark: SparkSession, resourceName: String) = { + private def readLegacyDateResource(spark: SparkSession) = { val resource = Option(Thread.currentThread().getContextClassLoader - .getResource(resourceName)).getOrElse { - throw new IllegalStateException(s"Missing Spark test resource: $resourceName") + .getResource(legacyDateResource)).getOrElse { + throw new IllegalStateException(s"Missing Spark test resource: $legacyDateResource") } - val file = File.createTempFile("spark-24-calendar", ".orc") + val file = File.createTempFile("spark-24-date", ".orc") file.deleteOnExit() val input = resource.openStream() try { @@ -70,9 +67,6 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { spark.read.orc(file.getCanonicalPath) } - private def readLegacyDateResource(spark: SparkSession) = - readLegacyResource(spark, legacyDateResource) - private def setDate(vector: DateColumnVector): Unit = { vector.setUsingProlepticCalendar(true) vector.vector(0) = dateValue @@ -161,42 +155,6 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { } } - for { - v1SourceList <- Seq("orc", "") - useChunkedReader <- Seq(false, true) - } { - test(s"read Spark 2.4 legacy ORC timestamp, source list is ($v1SourceList), " + - s"chunked=$useChunkedReader") { - val originalTimeZone = TimeZone.getDefault - val conf = calendarConf(RapidsReaderType.PERFILE, useChunkedReader, v1SourceList) - val expectedScan = if (v1SourceList == "orc") "GpuFileSourceScanExec" else "GpuBatchScan" - try { - val (fromCpu, fromGpu) = runOnCpuAndGpu( - spark => { - TimeZone.setDefault(TimeZone.getTimeZone("America/Los_Angeles")) - spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") - readLegacyResource(spark, legacyTimestampResource) - }, - identity, - conf = conf, - repart = 0, - skipCanonicalizationCheck = true, - existClasses = expectedScan) - compareResults( - sort = false, - floatEpsilon = 0.0, - fromCpu = fromCpu, - fromGpu = fromGpu) - Seq(fromCpu, fromGpu).foreach { rows => - assert(rows.length === 1) - assert(rows.head.getTimestamp(0).toString === legacyTimestampValue) - } - } finally { - TimeZone.setDefault(originalTimeZone) - } - } - } - for { readerType <- Seq(RapidsReaderType.COALESCING, RapidsReaderType.MULTITHREADED) useChunkedReader <- Seq(false, true) diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala index 18405b10b6f..c562b0402fa 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/timezone/OrcTimezoneSuite.scala @@ -88,8 +88,6 @@ class OrcTimezoneSuite extends SparkQueryCompareTestSuite { // Exact Asia/Shanghai writer=reader reproducer for the ORC epoch borrow correction. private val ShanghaiEpochBorrowTsUs = -7713116127L - // Exact SPARK-31284 value after reader-timezone Julian-to-Gregorian rebasing. - private val Spark31284TimestampUs = -30578655898876544L // Exact pre-first-transition values from non-UTC schema-evolution failures. private val newYorkHistoricalTsUs = -2957649381472612L private val shanghaiHistoricalTsUs = -3649379812521628L @@ -143,8 +141,6 @@ class OrcTimezoneSuite extends SparkQueryCompareTestSuite { ParisFirstTransitionLocalUs, ParisFirstTransitionLocalUs + 1L) Seq( - Spark31284TimestampUs, - Spark31284TimestampUs, newYorkHistoricalTsUs, shanghaiHistoricalTsUs, ShanghaiEpochBorrowTsUs, diff --git a/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala b/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala index ab1f0b3e506..bf239bc6008 100644 --- a/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala +++ b/tests/src/test/spark330/scala/org/apache/spark/sql/rapids/utils/RapidsTestSettings.scala @@ -253,6 +253,9 @@ class RapidsTestSettings extends BackendTestSettings { .exclude("Propagate Hadoop configs from orc options to underlying file system", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/11602. " + "Recovery trigger: GPU ORC writes propagate data-source options; P1.")) + .exclude("SPARK-31284: compatibility with Spark 2.4 in reading timestamps", + KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15471. " + + "Recovery trigger: GPU ORC legacy timestamp reads match Spark CPU; P0.")) .exclude("SPARK-31284, SPARK-31423: rebasing timestamps in write", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15473. " + "Recovery trigger: GPU ORC legacy timestamp round trips match Spark CPU; P0.")) @@ -268,6 +271,9 @@ class RapidsTestSettings extends BackendTestSettings { .exclude("Propagate Hadoop configs from orc options to underlying file system", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/11602. " + "Recovery trigger: GPU ORC writes propagate data-source options; P1.")) + .exclude("SPARK-31284: compatibility with Spark 2.4 in reading timestamps", + KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15471. " + + "Recovery trigger: GPU ORC legacy timestamp reads match Spark CPU; P0.")) .exclude("SPARK-31284, SPARK-31423: rebasing timestamps in write", KNOWN_ISSUE("https://github.com/NVIDIA/cudf-spark/issues/15473. " + "Recovery trigger: GPU ORC legacy timestamp round trips match Spark CPU; P0.")) From 8a6c316c13c3611092677cf2da958aa880cc1221 Mon Sep 17 00:00:00 2001 From: Allen Xu Date: Wed, 9 Sep 2026 18:05:01 +0800 Subject: [PATCH 3/5] Address ORC rebase retry review feedback Signed-off-by: Allen Xu --- integration_tests/src/main/python/orc_test.py | 5 +++- .../com/nvidia/spark/rapids/GpuOrcScan.scala | 23 ++++++++++--------- .../spark/rapids/GpuOrcTimezoneUtils.scala | 22 +++++++++++++++++- .../rapids/GpuTimestampRebaseUtils.scala | 10 ++++++-- 4 files changed, 45 insertions(+), 15 deletions(-) diff --git a/integration_tests/src/main/python/orc_test.py b/integration_tests/src/main/python/orc_test.py index d3842b0f81e..cddd45b5ffa 100644 --- a/integration_tests/src/main/python/orc_test.py +++ b/integration_tests/src/main/python/orc_test.py @@ -125,14 +125,17 @@ def test_basic_read(std_input_path, name, read_func, v1_enabled_list, orc_impl, @pytest.mark.parametrize('v1_enabled_list', ['', 'orc']) @pytest.mark.parametrize('vectorized_reader', [False, True]) +@pytest.mark.parametrize('chunked_reader', [False, True]) +@inject_oom @tz_sensitive_test def test_orc_read_spark_2_4_legacy_timestamp( - std_input_path, v1_enabled_list, vectorized_reader): + std_input_path, v1_enabled_list, vectorized_reader, chunked_reader): data_path = std_input_path + '/before_1582_ts_v2_4.snappy.orc' all_confs = { 'spark.sql.sources.useV1SourceList': v1_enabled_list, 'spark.sql.orc.impl': 'native', 'spark.sql.orc.enableVectorizedReader': vectorized_reader, + 'spark.rapids.sql.reader.chunked': chunked_reader, } gpu_scan = 'GpuFileSourceScanExec' if v1_enabled_list == 'orc' else 'GpuBatchScanExec' assert_cpu_and_gpu_are_equal_collect_with_capture( diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcScan.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcScan.scala index 61236ae8691..80fb22b03b7 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcScan.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcScan.scala @@ -3132,14 +3132,23 @@ object MakeOrcTableProducer extends Logging { tableSchema, splits, debugDumpPrefix, debugDumpAlways, writerTimezone, writerUsedProlepticGregorian) } else { - val table = withResource(buffer) { _ => + val rebased = withResource(buffer) { _ => try { RmmRapidsRetryIterator.withRetryNoSplit[Table] { - NvtxIdWithMetrics(NvtxRegistry.ORC_DECODE, metrics(GPU_DECODE_TIME)) { + val table = NvtxIdWithMetrics(NvtxRegistry.ORC_DECODE, metrics(GPU_DECODE_TIME)) { Table.readORC(parseOpts, buffer, offset, bufferSize) } + closeOnExcept(table) { _ => + if (readDataSchema.length < table.getNumberOfColumns) { + throw new QueryExecutionException(s"Expected ${readDataSchema.length} columns " + + s"but read ${table.getNumberOfColumns} from ${splits.mkString("; ")}") + } + } + GpuOrcTimezoneUtils.rebaseOrcDateTime( + table, writerTimezone, writerUsedProlepticGregorian) } } catch { + case e: QueryExecutionException => throw e case e: Exception => val dumpMsg = debugDumpPrefix.map { prefix => if (!debugDumpAlways) { @@ -3152,15 +3161,7 @@ object MakeOrcTableProducer extends Logging { throw new IOException(s"Error when processing ${splits.mkString("; ")}$dumpMsg", e) } } - closeOnExcept(table) { _ => - if (readDataSchema.length < table.getNumberOfColumns) { - throw new QueryExecutionException(s"Expected ${readDataSchema.length} columns " + - s"but read ${table.getNumberOfColumns} from ${splits.mkString("; ")}") - } - } metrics(NUM_OUTPUT_BATCHES) += 1 - val rebased = GpuOrcTimezoneUtils.rebaseOrcDateTime( - table, writerTimezone, writerUsedProlepticGregorian) val evolvedSchemaTable = SchemaUtils.evolveSchemaIfNeededAndClose(rebased, tableSchema, readDataSchema, isSchemaCaseSensitive, Some(GpuOrcScan.castColumnTo)) GpuMetric.recordOutputBatchBytes(evolvedSchemaTable, metrics.get(GPU_OUTPUT_BATCH_BYTES)) @@ -3221,7 +3222,7 @@ case class OrcTableReader( } } metrics(NUM_OUTPUT_BATCHES) += 1 - val rebased = GpuOrcTimezoneUtils.rebaseOrcDateTime( + val rebased = GpuOrcTimezoneUtils.rebaseOrcDateTimeWithRetry( table, writerTimezone, writerUsedProlepticGregorian) val evolvedSchemaTable = SchemaUtils.evolveSchemaIfNeededAndClose(rebased, catalystTableSchema, readDataSchema, isSchemaCaseSensitive, Some(GpuOrcScan.castColumnTo)) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala index d969df6a3c5..64c22e48767 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOrcTimezoneUtils.scala @@ -22,12 +22,32 @@ import java.util.Optional import scala.collection.mutable.ArrayBuffer import ai.rapids.cudf.{ColumnVector, ColumnView, DType, Scalar, Table} -import com.nvidia.spark.rapids.Arm.withResource +import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableProducingSeq import com.nvidia.spark.rapids.jni.{DateTimeRebase, GpuTimeZoneDB} object GpuOrcTimezoneUtils { + /** + * Rebase ORC date/time values inside an idempotent retry scope. + * + * The decoded table is made spillable before entering the retry block. Each attempt + * materializes its own table reference, which [[rebaseOrcDateTime]] consumes, so an OOM + * during any of the rebase allocations can safely retry without rereading the ORC chunk. + */ + private[rapids] def rebaseOrcDateTimeWithRetry( + input: Table, + writerTimezone: ZoneId, + writerUsedProlepticGregorian: Boolean): Table = { + val spillable = closeOnExcept(input) { _ => + SpillableTable(input, SpillPriorities.ACTIVE_BATCHING_PRIORITY) + } + RmmRapidsRetryIterator.withRetryNoSplit(spillable) { attempt => + rebaseOrcDateTime( + attempt.getTable(), writerTimezone, writerUsedProlepticGregorian) + } + } + /** Resolve an ORC stripe footer timezone once at the metadata boundary. */ private[rapids] def resolveWriterTimezone(writerTimezone: String): ZoneId = { if (writerTimezone.isEmpty) { diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala index d981a4d5b5a..f9489a6ba47 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala @@ -88,7 +88,13 @@ private[rapids] object GpuTimestampRebaseUtils { s"expected TIMESTAMP_MICROSECONDS but found ${input.getType}") if (input.getRowCount == 0 || isModernOrAllNull(input)) { retainOrCopy(input) - } else if (switches.isEmpty) { + } else { + rebaseLegacy(input) + } + } + + private[rapids] def rebaseLegacy(input: ColumnView): ColumnVector = { + if (switches.isEmpty) { // Spark's bundled map can lag valid IDs added by newer JDK timezone databases. rebaseOnHost(input) } else { @@ -124,7 +130,7 @@ private[rapids] object GpuTimestampRebaseUtils { if (delegate == null) { delegate = createJulianToGregorianMicrosContext(timeZoneId) } - delegate.rebase(input) + delegate.rebaseLegacy(input) } } From 74251601ac296a596ec349c88dd083bef84df357 Mon Sep 17 00:00:00 2001 From: Allen Xu Date: Thu, 10 Sep 2026 17:19:41 +0800 Subject: [PATCH 4/5] Cover legacy ORC timezone rebasing boundaries Signed-off-by: Allen Xu --- docs/compatibility.md | 15 +- integration_tests/src/main/python/orc_test.py | 4 +- .../sql/rapids/RebaseDateTimeBridge.scala | 12 +- .../rapids/GpuTimestampRebaseSuite.scala | 138 ++++++++++++ .../spark/rapids/OrcCalendarSuite.scala | 197 +++++++++++++++++- 5 files changed, 359 insertions(+), 7 deletions(-) create mode 100644 tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala diff --git a/docs/compatibility.md b/docs/compatibility.md index 610a6f0b489..557b1249173 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -251,9 +251,18 @@ See https://github.com/NVIDIA/cudf-spark/issues/7246 ## ORC The ORC format has fairly complete support for both reads and writes. There are only a few known -issues. The first is for reading timestamps and dates around the transition between Julian and -Gregorian calendars as described [here](https://github.com/NVIDIA/cudf-spark/issues/131). A -similar issue exists for writing dates as described +issues. The reader supports rebasing legacy-calendar ORC timestamps, including Spark 2.4 files, +to Spark's proleptic Gregorian calendar. Rebasing follows the reader JVM's default timezone +after applying the file's writer timezone; `spark.sql.session.timeZone` does not select the +rebase timezone. This also applies to timestamps nested in arrays and structs. Values before +the Common Era and timezones absent from Spark's rebase map use Spark's CPU implementation +within the GPU reader. + +This addresses the legacy timestamp case in [#15471](https://github.com/NVIDIA/cudf-spark/issues/15471). +The broader set of historical date/timestamp reading limitations remains tracked in +[#131](https://github.com/NVIDIA/cudf-spark/issues/131). In particular, proleptic Gregorian ORC +timestamps near the October 1582 cutover can still differ from CPU when the writer and reader +use different timezones. A similar issue exists for writing dates as described [here](https://github.com/NVIDIA/cudf-spark/issues/139). Writing timestamps, however only appears to work for dates after the epoch as described [here](https://github.com/NVIDIA/cudf-spark/issues/140). diff --git a/integration_tests/src/main/python/orc_test.py b/integration_tests/src/main/python/orc_test.py index cddd45b5ffa..6bf38fbc776 100644 --- a/integration_tests/src/main/python/orc_test.py +++ b/integration_tests/src/main/python/orc_test.py @@ -126,16 +126,18 @@ def test_basic_read(std_input_path, name, read_func, v1_enabled_list, orc_impl, @pytest.mark.parametrize('v1_enabled_list', ['', 'orc']) @pytest.mark.parametrize('vectorized_reader', [False, True]) @pytest.mark.parametrize('chunked_reader', [False, True]) +@pytest.mark.parametrize('session_timezone', ['UTC', 'Asia/Shanghai']) @inject_oom @tz_sensitive_test def test_orc_read_spark_2_4_legacy_timestamp( - std_input_path, v1_enabled_list, vectorized_reader, chunked_reader): + std_input_path, v1_enabled_list, vectorized_reader, chunked_reader, session_timezone): data_path = std_input_path + '/before_1582_ts_v2_4.snappy.orc' all_confs = { 'spark.sql.sources.useV1SourceList': v1_enabled_list, 'spark.sql.orc.impl': 'native', 'spark.sql.orc.enableVectorizedReader': vectorized_reader, 'spark.rapids.sql.reader.chunked': chunked_reader, + 'spark.sql.session.timeZone': session_timezone, } gpu_scan = 'GpuFileSourceScanExec' if v1_enabled_list == 'orc' else 'GpuBatchScanExec' assert_cpu_and_gpu_are_equal_collect_with_capture( diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala index e0e7215962a..d286af3197b 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.rapids import java.time.{ZoneId, ZoneOffset} +import java.util.TimeZone import org.apache.spark.sql.catalyst.util.RebaseDateTime @@ -57,6 +58,13 @@ object RebaseDateTimeBridge { } } - def rebaseJulianToGregorianMicros(timeZoneId: String, micros: Long): Long = - RebaseDateTime.rebaseJulianToGregorianMicros(timeZoneId, micros) + def rebaseJulianToGregorianMicros(timeZoneId: String, micros: Long): Long = { + val sparkTimeZoneId = ZoneId.of(timeZoneId, ZoneId.SHORT_IDS) match { + // ZoneId normalizes short IDs such as EST to -05:00. Spark's Calendar fallback + // uses TimeZone.getTimeZone(String), which silently treats that spelling as GMT. + case offset: ZoneOffset => TimeZone.getTimeZone(offset).getID + case _ => timeZoneId + } + RebaseDateTime.rebaseJulianToGregorianMicros(sparkTimeZoneId, micros) + } } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala new file mode 100644 index 00000000000..d7fb91e24ff --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala @@ -0,0 +1,138 @@ +/* + * 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. + */ + +package com.nvidia.spark.rapids + +import java.time.ZoneId +import java.util.TimeZone + +import scala.collection.JavaConverters._ + +import ai.rapids.cudf.{ColumnVector, ColumnView} +import com.nvidia.spark.rapids.Arm.withResource +import com.nvidia.spark.rapids.CudfTestHelper.assertColumnsAreEqual +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.catalyst.util.RebaseDateTime +import org.apache.spark.sql.rapids.RebaseDateTimeBridge + +class GpuTimestampRebaseSuite extends AnyFunSuite { + private val timeZones = Seq( + "UTC", "America/Los_Angeles", "Asia/Shanghai", "PST", "EST", "GMT+05:30", "GMT-03:30") + + private def assertRebase( + timeZone: String, + values: Array[java.lang.Long], + rebase: ColumnView => ColumnVector): Unit = { + // Use Spark's public implementation as the oracle, including its Calendar slow path. + val expected = values.map { value => + if (value == null) null else java.lang.Long.valueOf( + RebaseDateTime.rebaseJulianToGregorianMicros(timeZone, value.longValue())) + } + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(values: _*)) { input => + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(expected: _*)) { cpu => + withResource(rebase(input)) { gpu => + assertColumnsAreEqual(cpu, gpu) + } + } + } + } + + for (timeZone <- timeZones) { + test(s"timestamp rebase matches Spark at every map boundary in $timeZone") { + val info = RebaseDateTimeBridge.getJulianToGregorianMicros(timeZone).get + // Include every calendar and historical timezone transition, not only October 1582. + // Keep this batch free of nulls and BCE values so it exercises the GPU lookup path. + val values = info.switches.flatMap(switch => Seq(switch - 1L, switch, switch + 1L)) + .filter(_ >= info.switches.head).distinct.map(java.lang.Long.valueOf) + withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(timeZone)) { + context => assertRebase(timeZone, values, context.rebase) + } + } + + test(s"timestamp rebase preserves nulls and modern values in a legacy batch in $timeZone") { + val values = Array[java.lang.Long](null, -30578137076876544L, 0L, + RebaseDateTime.lastSwitchJulianTs - 1L, RebaseDateTime.lastSwitchJulianTs, null) + withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(timeZone)) { + context => assertRebase(timeZone, values, context.rebase) + } + } + + test(s"timestamp rebase matches Spark before the first map boundary in $timeZone") { + val first = RebaseDateTimeBridge.getJulianToGregorianMicros(timeZone).get.switches.head + val values = Array[java.lang.Long](first - 1L, first, first + 1L, + -100000000000000000L, 0L, null) + withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(timeZone)) { + context => assertRebase(timeZone, values, context.rebase) + } + } + } + + test("timestamp rebase preserves empty, all-null and modern columns") { + for (values <- Seq(Array.empty[java.lang.Long], Array[java.lang.Long](null, null), + Array[java.lang.Long](0L, RebaseDateTime.lastSwitchJulianTs, null))) { + withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext("UTC")) { + context => assertRebase("UTC", values, context.rebase) + } + } + } + + test("timestamp rebase supports sliced column views") { + val values = Array[java.lang.Long](-30578137076876544L, null, 0L) + val timeZone = "America/Los_Angeles" + withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(timeZone)) { + context => + assertRebase(timeZone, values, _ => { + withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs( + (Array[java.lang.Long](0L) ++ values ++ Array[java.lang.Long](0L)): _*)) { + padded => + withResource(padded.subVector(1, values.length + 1)) { slice => + context.rebase(slice) + } + } + }) + } + } + + test("a missing rebase map uses Spark's host implementation") { + val values = Array[java.lang.Long](null, -30578137076876544L, + -100000000000000000L, 0L) + // Exercise the missing-map branch even if a future Spark tzdb covers all JDK IDs. + withResource(new GpuTimestampRebaseUtils.JulianToGregorianMicrosContext( + "America/Los_Angeles", None, None)) { context => + assertRebase("America/Los_Angeles", values, context.rebase) + } + // Also exercise factory selection for actual JDK IDs absent from the runtime Spark map. + ZoneId.getAvailableZoneIds.asScala.toSeq.sorted + .find(RebaseDateTimeBridge.getJulianToGregorianMicros(_).isEmpty).foreach { timeZone => + withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(timeZone)) { + context => assertRebase(timeZone, values, context.rebase) + } + } + } + + test("a normalized fixed-offset ID retains the JVM timezone in the BCE fallback") { + val timeZone = "EST" + val readerZone = TimeZone.getTimeZone(timeZone).toZoneId.getId + // One microsecond before Julian March 1, 101 BCE at midnight in EST. Resolving + // the normalized "-05:00" through TimeZone.getTimeZone(String) silently uses GMT. + val values = Array[java.lang.Long](-65317950000000001L, -100000000000000000L, + -30578137076876544L, 0L, null) + withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(readerZone)) { + context => assertRebase(timeZone, values, context.rebase) + } + } +} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala index 92c320e4ebe..cf33546f34a 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala @@ -18,18 +18,26 @@ package com.nvidia.spark.rapids import java.io.File import java.nio.file.{Files, StandardCopyOption} +import java.sql.Timestamp import java.time.{LocalDate, ZoneId} +import java.util.TimeZone + +import scala.collection.JavaConverters._ import ai.rapids.cudf.{ColumnVector, Table} import com.nvidia.spark.rapids.Arm.{withResource, withResourceIfAllowed} import com.nvidia.spark.rapids.RapidsReaderType.RapidsReaderType import org.apache.hadoop.fs.Path import org.apache.hadoop.hive.ql.exec.vector.{ - DateColumnVector, ListColumnVector, LongColumnVector, StructColumnVector} + DateColumnVector, ListColumnVector, LongColumnVector, StructColumnVector, TimestampColumnVector} import org.apache.orc.{OrcFile, TypeDescription} +import org.apache.orc.impl.RecordReaderImpl import org.apache.spark.SparkConf +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.SpecializedGetters import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback import org.apache.spark.sql.rapids.shims.TrampolineConnectShims.SparkSession class OrcCalendarSuite extends SparkQueryCompareTestSuite { @@ -125,6 +133,193 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { writeCalendarFile(spark, base, id = 1, writerUsedProlepticGregorian = true) } + private def writeTimestampCalendarFile( + spark: SparkSession, + base: File, + proleptic: Boolean, + includeBce: Boolean, + includeProlepticCutover: Boolean): Int = { + val schema = TypeDescription.fromString( + "struct,timestamps:array>") + val path = new Path(base.getCanonicalPath, s"timestamps-$proleptic.orc") + val writer = OrcFile.createWriter(path, + OrcFile.writerOptions(spark.sparkContext.hadoopConfiguration) + .setSchema(schema).setProlepticGregorian(proleptic)) + val cutover = Timestamp.valueOf("1582-10-15 00:00:00").getTime * 1000L + val bceValues = if (includeBce) Seq[java.lang.Long](-100000000000000000L) else Seq.empty + // The legacy reader fix covers the cutover. Proleptic cross-timezone cutover reads + // still have a separate #131 discrepancy, preserved in the ignored regression below. + val cutoverValues = if (!proleptic || includeProlepticCutover) { + Seq[java.lang.Long](cutover - 1L, cutover, cutover + 1L) + } else { + Seq.empty + } + val values = bceValues ++ Seq[java.lang.Long](null, + Timestamp.valueOf("1001-01-01 01:02:03.123456").getTime * 1000L + 456L) ++ + cutoverValues ++ Seq[java.lang.Long](0L, 946684800123456L) + + def setTimestamp(vector: TimestampColumnVector, row: Int, value: java.lang.Long): Unit = { + // Supply hybrid-calendar input and let ORC convert it when writing a proleptic file. + vector.setUsingProlepticCalendar(false) + if (value == null) { + vector.noNulls = false + vector.isNull(row) = true + } else { + vector.time(row) = Math.floorDiv(value.longValue(), 1000L) + vector.nanos(row) = Math.floorMod(value.longValue(), 1000000L).toInt * 1000 + } + } + + try { + val batch = schema.createRowBatch() + val nested = batch.cols(4).asInstanceOf[StructColumnVector] + val timestamps = batch.cols(5).asInstanceOf[ListColumnVector] + values.zipWithIndex.foreach { case (value, row) => + batch.cols(0).asInstanceOf[LongColumnVector].vector(row) = + row * 2 + (if (proleptic) 1 else 0) + setTimestamp(batch.cols(1).asInstanceOf[TimestampColumnVector], row, value) + setTimestamp(batch.cols(2).asInstanceOf[TimestampColumnVector], row, 946684800123456L) + setTimestamp(batch.cols(3).asInstanceOf[TimestampColumnVector], row, null) + setTimestamp(nested.fields(0).asInstanceOf[TimestampColumnVector], row, value) + nested.noNulls = false + nested.isNull(row) = row == values.size - 1 + timestamps.offsets(row) = row * 2L + timestamps.lengths(row) = 2L + setTimestamp(timestamps.child.asInstanceOf[TimestampColumnVector], row * 2, value) + setTimestamp(timestamps.child.asInstanceOf[TimestampColumnVector], row * 2 + 1, null) + } + timestamps.childCount = values.size * 2 + batch.size = values.size + writer.addRowBatch(batch) + } finally { + writer.close() + } + withResourceIfAllowed(OrcFile.createReader(path, + OrcFile.readerOptions(spark.sparkContext.hadoopConfiguration))) { reader => + assert(reader.writerUsedProlepticGregorian() === proleptic) + val rows = reader.rows().asInstanceOf[RecordReaderImpl] + try { + assert(reader.getStripes.size() > 0) + reader.getStripes.asScala.foreach { stripe => + assert(rows.readStripeFooter(stripe).getWriterTimezone === TimeZone.getDefault.getID) + } + } finally { + rows.close() + } + } + values.size + } + + private val timestampZonePairs = Seq( + ("UTC", "America/Los_Angeles", false), + ("America/Los_Angeles", "UTC", false), + ("PST", "Asia/Shanghai", false), + ("Asia/Shanghai", "PST", false), + ("EST", "America/Los_Angeles", false), + ("UTC", "EST", false), + ("GMT+05:30", "UTC", false), + ("UTC", "GMT-03:30", false), + ("UTC", "EST", true)) + + // Cover every reader family and both scan APIs, CPU reader modes and chunking modes. + // The original Spark 2.4 fixture additionally covers the full V1/vectorized/chunked product. + private val timestampReaderModes = Seq( + (RapidsReaderType.PERFILE, false, "orc", false), + (RapidsReaderType.PERFILE, true, "", true), + (RapidsReaderType.COALESCING, false, "", true), + (RapidsReaderType.MULTITHREADED, true, "orc", false)) + + private def readTimestampMicros( + spark: SparkSession, + base: File, + gpuScan: Option[String]): Array[Row] = { + val frame = spark.read.orc(base.getCanonicalPath) + gpuScan.foreach(ExecutionPlanCaptureCallback.assertContains(frame, _)) + def micros(row: SpecializedGetters, ordinal: Int): java.lang.Long = { + if (row.isNullAt(ordinal)) null else java.lang.Long.valueOf(row.getLong(ordinal)) + } + // Read Catalyst's microseconds directly, before java.sql.Timestamp materialization. + // Preserve parent/child null masks as well as every array element. + frame.queryExecution.executedPlan.executeCollect().map { row => + val nested = if (row.isNullAt(4)) null else Row(micros(row.getStruct(4, 1), 0)) + val array = row.getArray(5) + Row(row.getInt(0), micros(row, 1), micros(row, 2), micros(row, 3), nested, + (0 until array.numElements()).map(micros(array, _))) + } + } + + private def checkTimestampCalendars( + writerZone: String, + readerZone: String, + includeBce: Boolean, + readerType: RapidsReaderType, + chunked: Boolean, + v1SourceList: String, + vectorized: Boolean, + includeProlepticCutover: Boolean = false): Unit = { + val originalTimeZone = TimeZone.getDefault + val scanClass = if (v1SourceList == "orc") "GpuFileSourceScanExec" else "GpuBatchScan" + val conf = calendarConf(readerType, chunked, v1SourceList) + .set("spark.sql.orc.impl", "native") + .set("spark.sql.orc.enableVectorizedReader", vectorized.toString) + try { + withTempPath { base => + val rowCount = withCpuSparkSession(spark => { + TimeZone.setDefault(TimeZone.getTimeZone(writerZone)) + assert(base.mkdirs()) + writeTimestampCalendarFile(spark, base, proleptic = false, + includeBce = includeBce, includeProlepticCutover = includeProlepticCutover) + + writeTimestampCalendarFile(spark, base, proleptic = true, + includeBce = includeBce, includeProlepticCutover = includeProlepticCutover) + }, conf) + val sessionZones = Seq(readerZone, if (readerZone == "UTC") "Asia/Shanghai" else "UTC") + val results = sessionZones.map { sessionZone => + withClue(s"JVM=$readerZone, session=$sessionZone: ") { + def read(spark: SparkSession, gpuScan: Option[String]): Array[Row] = { + // Set this inside the session callback: session initialization resets the JVM TZ. + TimeZone.setDefault(TimeZone.getTimeZone(readerZone)) + spark.conf.set("spark.sql.session.timeZone", sessionZone) + readTimestampMicros(spark, base, gpuScan) + } + val cpu = withCpuSparkSession(read(_, None), conf) + val gpu = withGpuSparkSession(read(_, Some(scanClass)), conf) + assert(cpu.length === rowCount) + compareResults(sort = true, floatEpsilon = 0.0, fromCpu = cpu, fromGpu = gpu) + gpu + } + } + // Changing only the SQL session zone must not change the stored timestamp micros. + compareResults(sort = true, floatEpsilon = 0.0, + fromCpu = results.head, fromGpu = results.last) + } + } finally { + TimeZone.setDefault(originalTimeZone) + } + } + + for { + (writerZone, readerZone, includeBce) <- timestampZonePairs + (readerType, chunked, v1SourceList, vectorized) <- timestampReaderModes + } { + test(s"read mixed ORC timestamp calendars from $writerZone in $readerZone with " + + s"$readerType, chunked=$chunked, source=($v1SourceList), vectorized=$vectorized, " + + s"BCE=$includeBce") { + checkTimestampCalendars(writerZone, readerZone, includeBce, + readerType, chunked, v1SourceList, vectorized) + } + } + + // KNOWN_ISSUE: https://github.com/NVIDIA/cudf-spark/issues/131 (P2). + // Recover when proleptic ORC cutover conversion matches CPU across writer/reader zones. + // This also fails with main's timezone converter; the legacy rows match after rebasing. + ignore("proleptic ORC cutover from America/Los_Angeles to UTC (#131)") { + checkTimestampCalendars("America/Los_Angeles", "UTC", includeBce = false, + readerType = RapidsReaderType.PERFILE, chunked = false, + v1SourceList = "orc", vectorized = false, + includeProlepticCutover = true) + } + test("proleptic nested ORC date rebase reuses the unchanged struct column") { withGpuSparkSession { _ => withResource(ColumnVector.daysFromInts(0)) { dateColumn => From de169e597b1034705dccc7ec370c21ddfe7a0611 Mon Sep 17 00:00:00 2001 From: Allen Xu Date: Fri, 11 Sep 2026 11:56:44 +0800 Subject: [PATCH 5/5] Move legacy ORC timestamp coverage to integration tests Signed-off-by: Allen Xu --- integration_tests/src/main/python/orc_test.py | 216 ++++++++++++++++++ .../rapids/GpuTimestampRebaseSuite.scala | 138 ----------- .../spark/rapids/OrcCalendarSuite.scala | 197 +--------------- 3 files changed, 217 insertions(+), 334 deletions(-) delete mode 100644 tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala diff --git a/integration_tests/src/main/python/orc_test.py b/integration_tests/src/main/python/orc_test.py index 6bf38fbc776..dcde9763b0e 100644 --- a/integration_tests/src/main/python/orc_test.py +++ b/integration_tests/src/main/python/orc_test.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json + import pytest +from py4j.java_gateway import get_field, set_field from asserts import * from conftest import is_not_utc @@ -146,6 +149,219 @@ def test_orc_read_spark_2_4_legacy_timestamp( conf=all_confs, require_non_empty=True) + +def _orc_timestamp_rebase_records(spark): + # Load the resource from the runtime Spark JAR, without calling private Scala methods. + stream = spark._jvm.java.lang.Thread.currentThread().getContextClassLoader() \ + .getResourceAsStream("julian-gregorian-rebase-micros.json") + assert stream is not None + # A public wrapper keeps Py4J from reflecting into JDK-private JAR stream classes. + stream = spark._jvm.java.io.BufferedInputStream(stream) + try: + records = json.loads(spark._jvm.org.apache.commons.io.IOUtils.toString(stream, "UTF-8")) + return {record['tz']: record['switches'] for record in records} + finally: + stream.close() + + +def _orc_timestamp_rebase_switches(spark, timezone): + # Records choose inputs only; the actual ORC CPU reader is the result oracle. + records = _orc_timestamp_rebase_records(spark) + zone = spark._jvm.java.util.TimeZone.getTimeZone(timezone).toZoneId() + for name in (timezone, zone.toString()): + if name in records: + return [value * 1000000 for value in records[name]] + normalized = zone.normalized().toString() + if normalized == 'Z' or normalized.startswith(('+', '-')): + offset = spark._jvm.java.time.ZoneOffset.of(normalized).getTotalSeconds() + return [(value - offset) * 1000000 for value in records['UTC']] + return None + + +def _orc_timestamp_values(spark, timezone, sample, proleptic): + if sample == 'empty': + return [] + if sample == 'all-null': + return [None, None] + if sample == 'modern': + return [0, 946684800123456, None] + if sample == 'boundaries' and not proleptic: + switches = _orc_timestamp_rebase_switches(spark, timezone) + if switches is None: + pytest.skip("The reader timezone has no Spark rebase map boundaries") + # Keep the scalar timestamp column free of null/BCE values to exercise GPU lookup. + return sorted({value + delta for value in switches for delta in (-1, 0, 1) + if value + delta >= switches[0]}) + if sample == 'bce' and not proleptic: + switches = _orc_timestamp_rebase_switches(spark, timezone) + first_boundary = [switches[0] - 1, switches[0], switches[0] + 1] if switches else [] + # The first value catches normalized EST becoming GMT in Spark's Calendar fallback. + return [-65317950000000001, -100000000000000000, + 0, None] + first_boundary + timestamp = spark._jvm.java.sql.Timestamp + values = [None, timestamp.valueOf("1001-01-01 01:02:03.123456").getTime() * 1000 + 456] + if not proleptic or sample == 'proleptic-cutover': + cutover = timestamp.valueOf("1582-10-15 00:00:00").getTime() * 1000 + values += [cutover - 1, cutover, cutover + 1] + if sample in ('mixed-bce', 'missing-map'): + values.insert(0, -100000000000000000) + # The proleptic cross-timezone cutover remains #131, covered by the strict xfail below. + return values + [0, 946684800123456] + + +def _write_orc_timestamp_calendar(spark, path, values, proleptic): + jvm = spark._jvm + orc = jvm.org.apache.orc + hadoop_conf = spark.sparkContext._jsc.hadoopConfiguration() + path = jvm.org.apache.hadoop.fs.Path(path) + schema = orc.TypeDescription.fromString( + "struct,timestamps:array>") + writer = orc.OrcFile.createWriter(path, orc.OrcFile.writerOptions(hadoop_conf) + .setSchema(schema).setProlepticGregorian(proleptic)) + + def set_timestamp(vector, row, micros): + # Supply hybrid-calendar input; ORC converts it when writing a proleptic file. + vector.setUsingProlepticCalendar(False) + if micros is None: + set_field(vector, 'noNulls', False) + get_field(vector, 'isNull')[row] = True + else: + value = jvm.java.sql.Timestamp(micros // 1000) + value.setNanos((micros % 1000000) * 1000) + vector.set(row, value) + + try: + batch = schema.createRowBatch(max(1024, len(values))) + columns = get_field(batch, 'cols') + nested = columns[4] + timestamps = columns[5] + child = get_field(timestamps, 'child') + child.ensureSize(len(values) * 2, False) + for row, value in enumerate(values): + get_field(columns[0], 'vector')[row] = row * 2 + int(proleptic) + set_timestamp(columns[1], row, value) + set_timestamp(columns[2], row, 946684800123456) + set_timestamp(columns[3], row, None) + set_timestamp(get_field(nested, 'fields')[0], row, value) + set_field(nested, 'noNulls', False) + get_field(nested, 'isNull')[row] = row == len(values) - 1 + get_field(timestamps, 'offsets')[row] = row * 2 + get_field(timestamps, 'lengths')[row] = 2 + set_timestamp(child, row * 2, value) + set_timestamp(child, row * 2 + 1, None) + set_field(timestamps, 'childCount', len(values) * 2) + set_field(batch, 'size', len(values)) + writer.addRowBatch(batch) + finally: + writer.close() + + reader = orc.OrcFile.createReader(path, orc.OrcFile.readerOptions(hadoop_conf)) + try: + assert reader.writerUsedProlepticGregorian() == proleptic + rows = reader.rows() + try: + assert reader.getNumberOfRows() == len(values) + for stripe in reader.getStripes(): + assert rows.readStripeFooter(stripe).getWriterTimezone() == \ + jvm.java.util.TimeZone.getDefault().getID() + finally: + rows.close() + finally: + reader.close() + + +def _read_orc_timestamp_micros(spark, path, reader_zone, gpu_scan=None): + # The harness configures the reader JVM timezone for both driver and executors. + assert spark._jvm.java.util.TimeZone.getDefault().getID() == reader_zone + frame = spark.read.orc(path) + # Read Catalyst micros directly: Python datetime cannot represent BCE timestamps. + # Preserve parent/child null masks and microsecond precision, without a SQL cast. + def micros(row, ordinal): + return None if row.isNullAt(ordinal) else row.getLong(ordinal) + + results = [] + for row in frame._jdf.queryExecution().executedPlan().executeCollect(): + nested = None if row.isNullAt(4) else [micros(row.getStruct(4, 1), 0)] + array = row.getArray(5) + results.append((row.getInt(0), micros(row, 1), micros(row, 2), micros(row, 3), + nested, [micros(array, i) for i in range(array.numElements())])) + if gpu_scan: + spark._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.assertContains( + frame._jdf, gpu_scan) + return sorted(results, key=lambda row: row[0]) + + +_legacy_orc_timezones = ['UTC', 'America/Los_Angeles', 'Asia/Shanghai', 'PST', 'EST', + 'GMT+05:30', 'GMT-03:30'] +_legacy_orc_reader_modes = [('PERFILE', False, 'orc', False), ('PERFILE', True, '', True), + ('COALESCING', False, '', True), ('MULTITHREADED', True, 'orc', False)] + + +@pytest.mark.parametrize('writer_zone,sample,reader_mode', + [(zone, 'mixed', mode) for zone in _legacy_orc_timezones for mode in _legacy_orc_reader_modes] + + [('reader', sample, _legacy_orc_reader_modes[0]) + for sample in ['boundaries', 'bce', 'mixed-bce', 'empty', 'all-null', 'modern', 'missing-map']] + + [pytest.param('America/Los_Angeles', 'proleptic-cutover', _legacy_orc_reader_modes[0], + marks=pytest.mark.xfail(strict=True, raises=AssertionError, + reason='https://github.com/NVIDIA/cudf-spark/issues/131: ' + 'proleptic cross-timezone cutover differs from CPU'))], + ids=lambda value: '-'.join(map(str, value)) if isinstance(value, tuple) else None) +@tz_sensitive_test +def test_orc_read_legacy_timestamp_calendars(spark_tmp_path, writer_zone, sample, reader_mode): + jvm = spark_jvm() + original_zone = jvm.java.util.TimeZone.getDefault() + reader_zone = original_zone.getID() + # Read in the harness's JVM timezone, including on distributed executors. The existing + # timezone CI matrix varies it independently of the writer and SQL session timezones. + if writer_zone == 'reader': + writer_zone = reader_zone + if sample == 'proleptic-cutover' and reader_zone != 'UTC': + pytest.skip("The preserved #131 reproducer requires a UTC reader JVM") + if sample == 'missing-map': + if with_cpu_session(lambda spark: _orc_timestamp_rebase_switches(spark, reader_zone)): + pytest.skip("The reader JVM timezone has a Spark rebase map") + print('Missing Spark rebase map:', reader_zone) + reader_type, chunked, v1_sources, vectorized = reader_mode + conf = {'spark.rapids.sql.format.orc.reader.type': reader_type, + 'spark.rapids.sql.reader.chunked': chunked, + 'spark.sql.sources.useV1SourceList': v1_sources, + 'spark.sql.orc.impl': 'native', + 'spark.sql.orc.enableVectorizedReader': vectorized, + 'spark.sql.files.maxPartitionBytes': str(1 << 30)} + data_path = spark_tmp_path + '/legacy_timestamp_calendars' + + def write(spark): + # Only fixture generation changes the driver's timezone; no Spark tasks run here. + jvm.java.util.TimeZone.setDefault(jvm.java.util.TimeZone.getTimeZone(writer_zone)) + try: + count = 0 + for proleptic in (False, True): + values = _orc_timestamp_values(spark, writer_zone, sample, proleptic) + _write_orc_timestamp_calendar(spark, data_path + '/{}.orc'.format(proleptic), + values, proleptic) + count += len(values) + return count + finally: + jvm.java.util.TimeZone.setDefault(original_zone) + + row_count = with_cpu_session(write, conf) + gpu_scan = 'GpuFileSourceScanExec' if v1_sources == 'orc' else 'GpuBatchScanExec' + session_zones = [reader_zone, 'Asia/Shanghai' if reader_zone == 'UTC' else 'UTC'] + results = [] + for session_zone in session_zones: + read_conf = copy_and_update(conf, {'spark.sql.session.timeZone': session_zone}) + cpu = with_cpu_session( + lambda spark: _read_orc_timestamp_micros(spark, data_path, reader_zone), read_conf) + gpu = with_gpu_session( + lambda spark: _read_orc_timestamp_micros(spark, data_path, reader_zone, gpu_scan), + read_conf) + assert len(cpu) == row_count + assert_equal(cpu, gpu) + results.append(gpu) + # Changing only the SQL session timezone must not change the timestamp microseconds. + assert_equal(results[0], results[1]) + # ORC does not support negative scale for decimal. So here is "decimal_gens_no_neg". # Otherwise it will get the below exception. # ... diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala deleted file mode 100644 index d7fb91e24ff..00000000000 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuTimestampRebaseSuite.scala +++ /dev/null @@ -1,138 +0,0 @@ -/* - * 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. - */ - -package com.nvidia.spark.rapids - -import java.time.ZoneId -import java.util.TimeZone - -import scala.collection.JavaConverters._ - -import ai.rapids.cudf.{ColumnVector, ColumnView} -import com.nvidia.spark.rapids.Arm.withResource -import com.nvidia.spark.rapids.CudfTestHelper.assertColumnsAreEqual -import org.scalatest.funsuite.AnyFunSuite - -import org.apache.spark.sql.catalyst.util.RebaseDateTime -import org.apache.spark.sql.rapids.RebaseDateTimeBridge - -class GpuTimestampRebaseSuite extends AnyFunSuite { - private val timeZones = Seq( - "UTC", "America/Los_Angeles", "Asia/Shanghai", "PST", "EST", "GMT+05:30", "GMT-03:30") - - private def assertRebase( - timeZone: String, - values: Array[java.lang.Long], - rebase: ColumnView => ColumnVector): Unit = { - // Use Spark's public implementation as the oracle, including its Calendar slow path. - val expected = values.map { value => - if (value == null) null else java.lang.Long.valueOf( - RebaseDateTime.rebaseJulianToGregorianMicros(timeZone, value.longValue())) - } - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(values: _*)) { input => - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs(expected: _*)) { cpu => - withResource(rebase(input)) { gpu => - assertColumnsAreEqual(cpu, gpu) - } - } - } - } - - for (timeZone <- timeZones) { - test(s"timestamp rebase matches Spark at every map boundary in $timeZone") { - val info = RebaseDateTimeBridge.getJulianToGregorianMicros(timeZone).get - // Include every calendar and historical timezone transition, not only October 1582. - // Keep this batch free of nulls and BCE values so it exercises the GPU lookup path. - val values = info.switches.flatMap(switch => Seq(switch - 1L, switch, switch + 1L)) - .filter(_ >= info.switches.head).distinct.map(java.lang.Long.valueOf) - withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(timeZone)) { - context => assertRebase(timeZone, values, context.rebase) - } - } - - test(s"timestamp rebase preserves nulls and modern values in a legacy batch in $timeZone") { - val values = Array[java.lang.Long](null, -30578137076876544L, 0L, - RebaseDateTime.lastSwitchJulianTs - 1L, RebaseDateTime.lastSwitchJulianTs, null) - withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(timeZone)) { - context => assertRebase(timeZone, values, context.rebase) - } - } - - test(s"timestamp rebase matches Spark before the first map boundary in $timeZone") { - val first = RebaseDateTimeBridge.getJulianToGregorianMicros(timeZone).get.switches.head - val values = Array[java.lang.Long](first - 1L, first, first + 1L, - -100000000000000000L, 0L, null) - withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(timeZone)) { - context => assertRebase(timeZone, values, context.rebase) - } - } - } - - test("timestamp rebase preserves empty, all-null and modern columns") { - for (values <- Seq(Array.empty[java.lang.Long], Array[java.lang.Long](null, null), - Array[java.lang.Long](0L, RebaseDateTime.lastSwitchJulianTs, null))) { - withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext("UTC")) { - context => assertRebase("UTC", values, context.rebase) - } - } - } - - test("timestamp rebase supports sliced column views") { - val values = Array[java.lang.Long](-30578137076876544L, null, 0L) - val timeZone = "America/Los_Angeles" - withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(timeZone)) { - context => - assertRebase(timeZone, values, _ => { - withResource(ColumnVector.timestampMicroSecondsFromBoxedLongs( - (Array[java.lang.Long](0L) ++ values ++ Array[java.lang.Long](0L)): _*)) { - padded => - withResource(padded.subVector(1, values.length + 1)) { slice => - context.rebase(slice) - } - } - }) - } - } - - test("a missing rebase map uses Spark's host implementation") { - val values = Array[java.lang.Long](null, -30578137076876544L, - -100000000000000000L, 0L) - // Exercise the missing-map branch even if a future Spark tzdb covers all JDK IDs. - withResource(new GpuTimestampRebaseUtils.JulianToGregorianMicrosContext( - "America/Los_Angeles", None, None)) { context => - assertRebase("America/Los_Angeles", values, context.rebase) - } - // Also exercise factory selection for actual JDK IDs absent from the runtime Spark map. - ZoneId.getAvailableZoneIds.asScala.toSeq.sorted - .find(RebaseDateTimeBridge.getJulianToGregorianMicros(_).isEmpty).foreach { timeZone => - withResource(GpuTimestampRebaseUtils.createJulianToGregorianMicrosContext(timeZone)) { - context => assertRebase(timeZone, values, context.rebase) - } - } - } - - test("a normalized fixed-offset ID retains the JVM timezone in the BCE fallback") { - val timeZone = "EST" - val readerZone = TimeZone.getTimeZone(timeZone).toZoneId.getId - // One microsecond before Julian March 1, 101 BCE at midnight in EST. Resolving - // the normalized "-05:00" through TimeZone.getTimeZone(String) silently uses GMT. - val values = Array[java.lang.Long](-65317950000000001L, -100000000000000000L, - -30578137076876544L, 0L, null) - withResource(new GpuTimestampRebaseUtils.LazyJulianToGregorianMicrosContext(readerZone)) { - context => assertRebase(timeZone, values, context.rebase) - } - } -} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala index cf33546f34a..92c320e4ebe 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/OrcCalendarSuite.scala @@ -18,26 +18,18 @@ package com.nvidia.spark.rapids import java.io.File import java.nio.file.{Files, StandardCopyOption} -import java.sql.Timestamp import java.time.{LocalDate, ZoneId} -import java.util.TimeZone - -import scala.collection.JavaConverters._ import ai.rapids.cudf.{ColumnVector, Table} import com.nvidia.spark.rapids.Arm.{withResource, withResourceIfAllowed} import com.nvidia.spark.rapids.RapidsReaderType.RapidsReaderType import org.apache.hadoop.fs.Path import org.apache.hadoop.hive.ql.exec.vector.{ - DateColumnVector, ListColumnVector, LongColumnVector, StructColumnVector, TimestampColumnVector} + DateColumnVector, ListColumnVector, LongColumnVector, StructColumnVector} import org.apache.orc.{OrcFile, TypeDescription} -import org.apache.orc.impl.RecordReaderImpl import org.apache.spark.SparkConf -import org.apache.spark.sql.Row -import org.apache.spark.sql.catalyst.expressions.SpecializedGetters import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback import org.apache.spark.sql.rapids.shims.TrampolineConnectShims.SparkSession class OrcCalendarSuite extends SparkQueryCompareTestSuite { @@ -133,193 +125,6 @@ class OrcCalendarSuite extends SparkQueryCompareTestSuite { writeCalendarFile(spark, base, id = 1, writerUsedProlepticGregorian = true) } - private def writeTimestampCalendarFile( - spark: SparkSession, - base: File, - proleptic: Boolean, - includeBce: Boolean, - includeProlepticCutover: Boolean): Int = { - val schema = TypeDescription.fromString( - "struct,timestamps:array>") - val path = new Path(base.getCanonicalPath, s"timestamps-$proleptic.orc") - val writer = OrcFile.createWriter(path, - OrcFile.writerOptions(spark.sparkContext.hadoopConfiguration) - .setSchema(schema).setProlepticGregorian(proleptic)) - val cutover = Timestamp.valueOf("1582-10-15 00:00:00").getTime * 1000L - val bceValues = if (includeBce) Seq[java.lang.Long](-100000000000000000L) else Seq.empty - // The legacy reader fix covers the cutover. Proleptic cross-timezone cutover reads - // still have a separate #131 discrepancy, preserved in the ignored regression below. - val cutoverValues = if (!proleptic || includeProlepticCutover) { - Seq[java.lang.Long](cutover - 1L, cutover, cutover + 1L) - } else { - Seq.empty - } - val values = bceValues ++ Seq[java.lang.Long](null, - Timestamp.valueOf("1001-01-01 01:02:03.123456").getTime * 1000L + 456L) ++ - cutoverValues ++ Seq[java.lang.Long](0L, 946684800123456L) - - def setTimestamp(vector: TimestampColumnVector, row: Int, value: java.lang.Long): Unit = { - // Supply hybrid-calendar input and let ORC convert it when writing a proleptic file. - vector.setUsingProlepticCalendar(false) - if (value == null) { - vector.noNulls = false - vector.isNull(row) = true - } else { - vector.time(row) = Math.floorDiv(value.longValue(), 1000L) - vector.nanos(row) = Math.floorMod(value.longValue(), 1000000L).toInt * 1000 - } - } - - try { - val batch = schema.createRowBatch() - val nested = batch.cols(4).asInstanceOf[StructColumnVector] - val timestamps = batch.cols(5).asInstanceOf[ListColumnVector] - values.zipWithIndex.foreach { case (value, row) => - batch.cols(0).asInstanceOf[LongColumnVector].vector(row) = - row * 2 + (if (proleptic) 1 else 0) - setTimestamp(batch.cols(1).asInstanceOf[TimestampColumnVector], row, value) - setTimestamp(batch.cols(2).asInstanceOf[TimestampColumnVector], row, 946684800123456L) - setTimestamp(batch.cols(3).asInstanceOf[TimestampColumnVector], row, null) - setTimestamp(nested.fields(0).asInstanceOf[TimestampColumnVector], row, value) - nested.noNulls = false - nested.isNull(row) = row == values.size - 1 - timestamps.offsets(row) = row * 2L - timestamps.lengths(row) = 2L - setTimestamp(timestamps.child.asInstanceOf[TimestampColumnVector], row * 2, value) - setTimestamp(timestamps.child.asInstanceOf[TimestampColumnVector], row * 2 + 1, null) - } - timestamps.childCount = values.size * 2 - batch.size = values.size - writer.addRowBatch(batch) - } finally { - writer.close() - } - withResourceIfAllowed(OrcFile.createReader(path, - OrcFile.readerOptions(spark.sparkContext.hadoopConfiguration))) { reader => - assert(reader.writerUsedProlepticGregorian() === proleptic) - val rows = reader.rows().asInstanceOf[RecordReaderImpl] - try { - assert(reader.getStripes.size() > 0) - reader.getStripes.asScala.foreach { stripe => - assert(rows.readStripeFooter(stripe).getWriterTimezone === TimeZone.getDefault.getID) - } - } finally { - rows.close() - } - } - values.size - } - - private val timestampZonePairs = Seq( - ("UTC", "America/Los_Angeles", false), - ("America/Los_Angeles", "UTC", false), - ("PST", "Asia/Shanghai", false), - ("Asia/Shanghai", "PST", false), - ("EST", "America/Los_Angeles", false), - ("UTC", "EST", false), - ("GMT+05:30", "UTC", false), - ("UTC", "GMT-03:30", false), - ("UTC", "EST", true)) - - // Cover every reader family and both scan APIs, CPU reader modes and chunking modes. - // The original Spark 2.4 fixture additionally covers the full V1/vectorized/chunked product. - private val timestampReaderModes = Seq( - (RapidsReaderType.PERFILE, false, "orc", false), - (RapidsReaderType.PERFILE, true, "", true), - (RapidsReaderType.COALESCING, false, "", true), - (RapidsReaderType.MULTITHREADED, true, "orc", false)) - - private def readTimestampMicros( - spark: SparkSession, - base: File, - gpuScan: Option[String]): Array[Row] = { - val frame = spark.read.orc(base.getCanonicalPath) - gpuScan.foreach(ExecutionPlanCaptureCallback.assertContains(frame, _)) - def micros(row: SpecializedGetters, ordinal: Int): java.lang.Long = { - if (row.isNullAt(ordinal)) null else java.lang.Long.valueOf(row.getLong(ordinal)) - } - // Read Catalyst's microseconds directly, before java.sql.Timestamp materialization. - // Preserve parent/child null masks as well as every array element. - frame.queryExecution.executedPlan.executeCollect().map { row => - val nested = if (row.isNullAt(4)) null else Row(micros(row.getStruct(4, 1), 0)) - val array = row.getArray(5) - Row(row.getInt(0), micros(row, 1), micros(row, 2), micros(row, 3), nested, - (0 until array.numElements()).map(micros(array, _))) - } - } - - private def checkTimestampCalendars( - writerZone: String, - readerZone: String, - includeBce: Boolean, - readerType: RapidsReaderType, - chunked: Boolean, - v1SourceList: String, - vectorized: Boolean, - includeProlepticCutover: Boolean = false): Unit = { - val originalTimeZone = TimeZone.getDefault - val scanClass = if (v1SourceList == "orc") "GpuFileSourceScanExec" else "GpuBatchScan" - val conf = calendarConf(readerType, chunked, v1SourceList) - .set("spark.sql.orc.impl", "native") - .set("spark.sql.orc.enableVectorizedReader", vectorized.toString) - try { - withTempPath { base => - val rowCount = withCpuSparkSession(spark => { - TimeZone.setDefault(TimeZone.getTimeZone(writerZone)) - assert(base.mkdirs()) - writeTimestampCalendarFile(spark, base, proleptic = false, - includeBce = includeBce, includeProlepticCutover = includeProlepticCutover) + - writeTimestampCalendarFile(spark, base, proleptic = true, - includeBce = includeBce, includeProlepticCutover = includeProlepticCutover) - }, conf) - val sessionZones = Seq(readerZone, if (readerZone == "UTC") "Asia/Shanghai" else "UTC") - val results = sessionZones.map { sessionZone => - withClue(s"JVM=$readerZone, session=$sessionZone: ") { - def read(spark: SparkSession, gpuScan: Option[String]): Array[Row] = { - // Set this inside the session callback: session initialization resets the JVM TZ. - TimeZone.setDefault(TimeZone.getTimeZone(readerZone)) - spark.conf.set("spark.sql.session.timeZone", sessionZone) - readTimestampMicros(spark, base, gpuScan) - } - val cpu = withCpuSparkSession(read(_, None), conf) - val gpu = withGpuSparkSession(read(_, Some(scanClass)), conf) - assert(cpu.length === rowCount) - compareResults(sort = true, floatEpsilon = 0.0, fromCpu = cpu, fromGpu = gpu) - gpu - } - } - // Changing only the SQL session zone must not change the stored timestamp micros. - compareResults(sort = true, floatEpsilon = 0.0, - fromCpu = results.head, fromGpu = results.last) - } - } finally { - TimeZone.setDefault(originalTimeZone) - } - } - - for { - (writerZone, readerZone, includeBce) <- timestampZonePairs - (readerType, chunked, v1SourceList, vectorized) <- timestampReaderModes - } { - test(s"read mixed ORC timestamp calendars from $writerZone in $readerZone with " + - s"$readerType, chunked=$chunked, source=($v1SourceList), vectorized=$vectorized, " + - s"BCE=$includeBce") { - checkTimestampCalendars(writerZone, readerZone, includeBce, - readerType, chunked, v1SourceList, vectorized) - } - } - - // KNOWN_ISSUE: https://github.com/NVIDIA/cudf-spark/issues/131 (P2). - // Recover when proleptic ORC cutover conversion matches CPU across writer/reader zones. - // This also fails with main's timezone converter; the legacy rows match after rebasing. - ignore("proleptic ORC cutover from America/Los_Angeles to UTC (#131)") { - checkTimestampCalendars("America/Los_Angeles", "UTC", includeBce = false, - readerType = RapidsReaderType.PERFILE, chunked = false, - v1SourceList = "orc", vectorized = false, - includeProlepticCutover = true) - } - test("proleptic nested ORC date rebase reuses the unchanged struct column") { withGpuSparkSession { _ => withResource(ColumnVector.daysFromInts(0)) { dateColumn =>