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 7d83392fe39..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 @@ -122,6 +125,243 @@ 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]) +@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, 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( + read_orc_df(data_path), + exist_classes=gpu_scan, + 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/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 00000000000..af9ef040270 Binary files /dev/null and b/integration_tests/src/test/resources/before_1582_ts_v2_4.snappy.orc differ 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 61a785b5a9b..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) { @@ -124,11 +144,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 +186,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 +195,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 +218,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 +291,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 +301,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 +317,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..f9489a6ba47 --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTimestampRebaseUtils.scala @@ -0,0 +1,170 @@ +/* + * 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 { + 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 { + 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.rebaseLegacy(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..d286af3197b --- /dev/null +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RebaseDateTimeBridge.scala @@ -0,0 +1,70 @@ +/* + * 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 java.util.TimeZone + +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 = { + 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) + } +}