diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java index 75af98a2affd1..8efb8eb5c6c82 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java @@ -30,6 +30,7 @@ import org.apache.spark.SparkUnsupportedOperationException; import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.catalyst.util.DateTimeUtils$; import org.apache.spark.sql.catalyst.util.RebaseDateTime; import org.apache.spark.sql.execution.datasources.DataSourceUtils; import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException; @@ -40,6 +41,7 @@ import java.math.BigDecimal; import java.math.BigInteger; import java.math.RoundingMode; +import java.nio.ByteOrder; import java.time.ZoneId; import java.time.ZoneOffset; import java.util.Arrays; @@ -227,6 +229,16 @@ public ParquetVectorUpdater getUpdater(ColumnDescriptor descriptor, DataType spa int96RebaseTz); } } + } else if (sparkType instanceof TimestampNTZNanosType) { + return new Int96AsTimestampNanosUpdater( + false, false, null, null, ((TimestampNTZNanosType) sparkType).precision()); + } else if (sparkType instanceof TimestampLTZNanosType) { + final boolean failIfRebase = "EXCEPTION".equals(int96RebaseMode); + final boolean rebase = !"CORRECTED".equals(int96RebaseMode); + final ZoneId tz = shouldConvertTimestamps() ? convertTz : null; + return new Int96AsTimestampNanosUpdater( + rebase, failIfRebase, int96RebaseTz, tz, + ((TimestampLTZNanosType) sparkType).precision()); } } case BINARY -> { @@ -1536,6 +1548,83 @@ public void decodeSingleDictionaryId( } } + // Reads a legacy INT96 timestamp column as a nanosecond timestamp, into the two-child + // (epochMicros, nanosWithinMicro) vector -- the vectorized read side of widening a legacy INT96 + // timestamp to nanosecond precision. INT96 stores nanoseconds-of-day, so a foreign file (e.g. + // Impala/Hive) can carry true sub-microsecond digits; binaryToSQLTimestamp floors to micros, so + // the sub-micro remainder is recovered straight from the raw INT96 and truncated to the read + // precision (kept in lock-step with the row-based int96AsNanosConverter). The LTZ family applies + // the same INT96 Julian rebase and timezone conversion as the TimestampType path; the NTZ family + // passes rebase=false and convertTz=null (mirrors the TimestampNTZType path). + private static class Int96AsTimestampNanosUpdater implements ParquetVectorUpdater { + private final boolean rebase; + private final boolean failIfRebase; + private final String timeZone; + private final ZoneId convertTz; + private final int precision; + + Int96AsTimestampNanosUpdater( + boolean rebase, boolean failIfRebase, String timeZone, ZoneId convertTz, int precision) { + this.rebase = rebase; + this.failIfRebase = failIfRebase; + this.timeZone = timeZone; + this.convertTz = convertTz; + this.precision = precision; + } + + @Override + public void readValues( + int total, + int offset, + WritableColumnVector values, + VectorizedValuesReader valuesReader) { + for (int i = 0; i < total; i++) { + readValue(offset + i, values, valuesReader); + } + } + + @Override + public void skipValues(int total, VectorizedValuesReader valuesReader) { + valuesReader.skipFixedLenByteArray(total, 12); + } + + @Override + public void readValue( + int offset, + WritableColumnVector values, + VectorizedValuesReader valuesReader) { + putInt96AsNanos(offset, values, valuesReader.readBinary(12)); + } + + @Override + public void decodeSingleDictionaryId( + int offset, + WritableColumnVector values, + WritableColumnVector dictionaryIds, + Dictionary dictionary) { + putInt96AsNanos(offset, values, dictionary.decodeToBinary(dictionaryIds.getDictId(offset))); + } + + private void putInt96AsNanos(int offset, WritableColumnVector values, Binary binary) { + long micros = ParquetRowConverter.binaryToSQLTimestamp(binary); + if (rebase) { + micros = rebaseInt96(micros, failIfRebase, timeZone); + } + if (convertTz != null) { + micros = DateTimeUtils.convertTz(micros, convertTz, UTC); + } + // INT96 stores nanoseconds-of-day (first 8 bytes, little-endian). Recover the sub-micro + // remainder (1000 ns per micro), truncate to the read precision (matching + // int96AsNanosConverter), so a foreign nanosecond INT96 is not silently floored to micros. + long timeOfDayNanos = binary.toByteBuffer().order(ByteOrder.LITTLE_ENDIAN).getLong(); + int rawNanosWithinMicro = (int) (timeOfDayNanos % 1000L); + short nanosWithinMicro = (short) DateTimeUtils$.MODULE$ + .truncateNanosWithinMicroToPrecision(rawNanosWithinMicro, precision); + values.getChild(0).putLong(offset, micros); + values.getChild(1).putShort(offset, nanosWithinMicro); + } + } + private static class FixedLenByteArrayUpdater implements ParquetVectorUpdater { private final int arrayLen; diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala index 746d8b46fb462..701751e959509 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala @@ -18,21 +18,22 @@ package org.apache.spark.sql.execution.datasources.parquet.types.ops import java.lang.{Long => JLong} -import java.time.{Instant, LocalDateTime, ZoneId} +import java.nio.ByteOrder +import java.time.{Instant, LocalDateTime, ZoneId, ZoneOffset} import org.apache.parquet.column.{ColumnDescriptor, Dictionary} -import org.apache.parquet.io.api.{Converter, RecordConsumer} +import org.apache.parquet.io.api.{Binary, Converter, RecordConsumer} import org.apache.parquet.schema.{LogicalTypeAnnotation, Type, Types} import org.apache.parquet.schema.LogicalTypeAnnotation.{TimestampLogicalTypeAnnotation, TimeUnit} -import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64 +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.{INT64, INT96} import org.apache.parquet.schema.Type.Repetition import org.apache.spark.sql.catalyst.expressions.SpecializedGetters -import org.apache.spark.sql.catalyst.util.DateTimeUtils +import org.apache.spark.sql.catalyst.util.{DateTimeConstants, DateTimeUtils} import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.datasources.DataSourceUtils -import org.apache.spark.sql.execution.datasources.parquet.{HasParentContainerUpdater, ParentContainerUpdater, ParquetPrimitiveConverter, ParquetToSparkSchemaConverter, ParquetVectorUpdater, VectorizedValuesReader} +import org.apache.spark.sql.execution.datasources.parquet.{HasParentContainerUpdater, ParentContainerUpdater, ParquetPrimitiveConverter, ParquetRowConverter, ParquetToSparkSchemaConverter, ParquetVectorUpdater, VectorizedValuesReader} import org.apache.spark.sql.execution.vectorized.WritableColumnVector import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, TimestampLTZNanosType, TimestampNTZNanosType} @@ -90,8 +91,8 @@ private[parquet] trait TimestampNanosParquetOps extends ParquetTypeOps { override def isBatchReadSupported(sqlConf: SQLConf): Boolean = true // Vectorized-decode only the canonical INT64 TIMESTAMP(NANOS) encoding here; return None for - // anything else. The factory then handles an INT64 TIMESTAMP(MICROS) column (promoting it to - // nanos) and raises SchemaColumnConvertNotSupportedException for the rest. + // anything else. The factory then promotes an INT64 TIMESTAMP(MICROS) or a legacy INT96 timestamp + // column to nanos, and raises SchemaColumnConvertNotSupportedException for the rest. override def getVectorUpdater(descriptor: ColumnDescriptor): Option[ParquetVectorUpdater] = { val parquetType = descriptor.getPrimitiveType if (TimestampNanosParquetOps.isNanosTimestamp(parquetType)) { @@ -154,6 +155,8 @@ private[parquet] trait TimestampNanosParquetOps extends ParquetTypeOps { nanosConverter(updater) } else if (TimestampNanosParquetOps.isMicrosTimestamp(parquetType, isAdjustedToUTC)) { microsAsNanosConverter(updater, datetimeRebaseSpec) + } else if (TimestampNanosParquetOps.isInt96Timestamp(parquetType)) { + int96AsNanosConverter(updater, convertTz, int96RebaseSpec) } else { throw QueryExecutionErrors.cannotCreateParquetConverterForDataTypeError( sparkType, parquetType.toString) @@ -183,6 +186,39 @@ private[parquet] trait TimestampNanosParquetOps extends ParquetTypeOps { this.updater.set(TimestampNanosVal.fromParts(rebase(value), 0.toShort)) } } + + private def int96AsNanosConverter( + updater: ParentContainerUpdater, + convertTz: Option[ZoneId], + int96RebaseSpec: RebaseSpec): Converter with HasParentContainerUpdater = { + // INT96 carries no logical annotation, so the requested type's family decides handling (mirrors + // the INT96 arms of ParquetRowConverter): LTZ applies the INT96 Julian rebase and any timezone + // conversion; NTZ applies neither. INT96 stores nanoseconds-of-day, so a foreign file (e.g. + // Impala/Hive) can carry true sub-microsecond digits; binaryToSQLTimestamp floors to micros, so + // the sub-micro remainder is recovered straight from the raw INT96 as nanosWithinMicro -- a + // whole-microsecond rebase/timezone shift never perturbs it -- and truncated to the read + // precision. The read side of widening a legacy INT96 timestamp to nanosecond precision. + val int96Rebase: Long => Long = + if (isNtz) identity + else DataSourceUtils.createTimestampRebaseFuncInRead(int96RebaseSpec, "Parquet INT96") + new ParquetPrimitiveConverter(updater) { + override def addBinary(value: Binary): Unit = { + val julianMicros = ParquetRowConverter.binaryToSQLTimestamp(value) + val micros = if (isNtz) { + julianMicros + } else { + val gregorianMicros = int96Rebase(julianMicros) + convertTz.map(DateTimeUtils.convertTz(gregorianMicros, _, ZoneOffset.UTC)) + .getOrElse(gregorianMicros) + } + val timeOfDayNanos = value.toByteBuffer.order(ByteOrder.LITTLE_ENDIAN).getLong + val rawNanosWithinMicro = (timeOfDayNanos % DateTimeConstants.NANOS_PER_MICROS).toInt + val nanosWithinMicro = + DateTimeUtils.truncateNanosWithinMicroToPrecision(rawNanosWithinMicro, precision).toShort + this.updater.set(TimestampNanosVal.fromParts(micros, nanosWithinMicro)) + } + } + } } /** @@ -250,6 +286,20 @@ private[ops] object TimestampNanosParquetOps { case _ => false }) + /** + * Whether the Parquet field is a legacy INT96 timestamp column. INT96 carries no logical + * annotation (and thus no time-zone family), so only the physical type is checked; the requested + * nanos type's family (LTZ / NTZ) decides the rebase / timezone handling, exactly as the INT96 + * arms of [[ParquetRowConverter]] do for the microsecond timestamp types. + * + * Unlike the annotated micros path ([[isMicrosTimestamp]]), there is no time-zone family to match + * against, so an INT96 file can be requested as either LTZ or NTZ nanos -- mirroring Spark's + * existing INT96 -> TimestampType / TimestampNTZType reads. The same-family guard therefore + * applies only to the annotated micros path, not here. + */ + private[ops] def isInt96Timestamp(parquetType: Type): Boolean = + parquetType.isPrimitive && parquetType.asPrimitiveType.getPrimitiveTypeName == INT96 + // Repacks an externalized nanos filter value into the signed INT64 epoch-nanoseconds the write // path produces. Conversion is at precision 9 (a lossless repack): the literal has already been // floored to the column precision upstream, so no sub-microsecond digits are dropped here. The diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala index 2e2a7e18316cd..f35fc578f6529 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala @@ -18,18 +18,20 @@ package org.apache.spark.sql.execution.datasources.parquet import java.io.File +import java.nio.{ByteBuffer, ByteOrder} import org.apache.hadoop.fs.Path import org.apache.parquet.example.data.simple.SimpleGroupFactory import org.apache.parquet.hadoop.ParquetFileWriter.Mode import org.apache.parquet.hadoop.example.ExampleParquetWriter +import org.apache.parquet.io.api.Binary import org.apache.parquet.schema.{LogicalTypeAnnotation, MessageType, Types} import org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit -import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64 +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.{INT64, INT96} -import org.apache.spark.{SparkArithmeticException, SparkException} +import org.apache.spark.{SparkArithmeticException, SparkException, SparkUpgradeException} import org.apache.spark.sql.{AnalysisException, QueryTest, Row} -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ @@ -67,6 +69,46 @@ class ParquetTimestampNanosSuite extends QueryTest with ParquetTest with SharedS } } + // Builds a 12-byte INT96 value: nanoseconds-of-day (little-endian long) followed by the Julian + // day (little-endian int), the on-disk layout ParquetRowConverter.binaryToSQLTimestamp reads. + private def int96Binary(julianDay: Int, timeOfDayNanos: Long): Binary = { + val buf = ByteBuffer.allocate(12).order(ByteOrder.LITTLE_ENDIAN) + buf.putLong(timeOfDayNanos) + buf.putInt(julianDay) + Binary.fromConstantByteArray(buf.array()) + } + + // Writes a foreign INT96 timestamp column (a raw INT96 with no logical annotation, the layout + // Impala/Hive emit). INT96 carries no time-zone family in the schema, so the requested read type + // decides LTZ vs NTZ. dictionaryEnabled toggles dictionary vs plain encoding so both the + // vectorized updater's dictionary-decode and plain-read branches can be exercised. + private def writeForeignInt96Parquet( + file: File, + values: Seq[Binary], + dictionaryEnabled: Boolean): Unit = { + val schema: MessageType = Types.buildMessage() + .optional(INT96) + .named("ts") + .named("spark_schema") + val conf = spark.sessionState.newHadoopConf() + val writer = ExampleParquetWriter.builder(new Path(file.toURI)) + .withType(schema) + .withConf(conf) + .withDictionaryEncoding(dictionaryEnabled) + .withWriteMode(Mode.OVERWRITE) + .build() + try { + val factory = new SimpleGroupFactory(schema) + values.foreach { v => + val group = factory.newGroup() + group.add("ts", v) + writer.write(group) + } + } finally { + writer.close() + } + } + test("SPARK-57102: Spark write/read round-trips nanos value and precision") { withNanosEnabled { Seq("true", "false").foreach { vectorized => @@ -136,6 +178,77 @@ class ParquetTimestampNanosSuite extends QueryTest with ParquetTest with SharedS } } + test("INT96 vectorized read preserves sub-microsecond nanos from a foreign file") { + // Spark only writes micro-aligned INT96, so neither a Spark round-trip nor the widening suite + // (which writes through Spark) ever carries sub-microsecond digits. A foreign file (Impala/ + // Hive) can: INT96 stores nanoseconds-of-day, and the default vectorized reader + // (Int96AsTimestampNanosUpdater#putInt96AsNanos) must recover the sub-micro remainder rather + // than floor to micros. Julian day 2440588 is 1970-01-01; 45296123456789 ns-of-day is + // 12:34:56.123456789, whose .789 remainder a micros floor would drop. Runs both families, + // both readers (withAllParquetReaders) and dictionary on/off (the updater's dictionary-decode + // and plain-read branches); the row-based decode is pinned in TimestampNanosParquetOpsSuite. + withNanosEnabled { + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val binary = int96Binary(julianDay = 2440588, timeOfDayNanos = 45296123456789L) + Seq(true, false).foreach { dictionaryEnabled => + withAllParquetReaders { + withTempPath { dir => + val file = new File(dir, "int96.parquet") + writeForeignInt96Parquet(file, Seq.fill(4)(binary), dictionaryEnabled) + Seq( + TimestampNTZNanosType(9) -> "TIMESTAMP_NTZ '1970-01-01 12:34:56.123456789'", + TimestampLTZNanosType(9) -> "TIMESTAMP_LTZ '1970-01-01 12:34:56.123456789'" + ).foreach { case (readType, literal) => + withClue(s"readType=$readType dictionary=$dictionaryEnabled") { + val read = spark.read + .schema(StructType(Seq(StructField("ts", readType)))) + .parquet(file.getCanonicalPath) + assert(read.schema("ts").dataType === readType) + val expected = spark.sql(s"SELECT $literal AS ts").collect().head + checkAnswer(read, Seq.fill(4)(expected)) + } + } + } + } + } + } + } + } + + test("INT96 read of an ancient value fails under EXCEPTION rebase for LTZ") { + // A pre-1582 (Julian) INT96 read as the LTZ family under int96RebaseModeInRead=EXCEPTION must + // refuse the ambiguous value instead of silently rebasing -- exercising the vectorized + // updater's failIfRebase arm (Int96AsTimestampNanosUpdater rebase=true, failIfRebase=true) end + // to end; the row-path rebase is pinned in TimestampNanosParquetOpsSuite. Julian day 2200000 + // is ~year 1311. Both readers are covered by withAllParquetReaders. + withNanosEnabled { + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + SQLConf.PARQUET_INT96_REBASE_MODE_IN_READ.key -> LegacyBehaviorPolicy.EXCEPTION.toString) { + val ancient = int96Binary(julianDay = 2200000, timeOfDayNanos = 0L) + withAllParquetReaders { + withTempPath { dir => + val file = new File(dir, "int96_ancient.parquet") + writeForeignInt96Parquet(file, Seq(ancient), dictionaryEnabled = false) + // The vectorized reader throws the SparkUpgradeException directly, while the row-based + // reader wraps it in a SparkException (FAILED_READ_FILE); catch their common ancestor + // and look for the upgrade exception anywhere in the cause chain. + val e = intercept[Exception] { + spark.read + .schema(StructType(Seq(StructField("ts", TimestampLTZNanosType(9))))) + .parquet(file.getCanonicalPath) + .collect() + } + assert( + Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null) + .exists(_.isInstanceOf[SparkUpgradeException]), + s"expected a SparkUpgradeException in the cause chain of: $e") + } + } + } + } + } + test("SPARK-57102: explicit lower-precision read schema truncates sub-precision nanos") { withNanosEnabled { withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala index 7a72d485b429b..e8dc1d1fdadd2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala @@ -202,7 +202,8 @@ class ParquetTypeWideningSuite // Widening TIMESTAMP(6) to nanosecond precision: INT64 TIMESTAMP(MICROS) files are read as nanos // by promoting each micros value to (epochMicros, 0). Source must be TIMESTAMP(MICROS) (what - // Delta writes), hence the explicit output type; INT96/MILLIS aren't supported (see below). + // Delta writes), hence the explicit output type. INT96 has its own loop below; MILLIS stays + // unsupported. // Values stay on the micros grid and include a pre-1582 date (LTZ Julian rebase, LEGACY mode) and // a far-future date past the int64 epoch-nanos range (~2262). Requires the nanos preview flag. for { @@ -225,9 +226,36 @@ class ParquetTypeWideningSuite } } + // The same widening from a legacy INT96 timestamp column. INT96 has no logical unit; each value + // decodes to micros (binaryToSQLTimestamp) and is promoted to (epochMicros, nanosWithinMicro). + // Spark writes micro-aligned INT96, so a round-trip carries no sub-microsecond digits; foreign + // nanosecond INT96 (whose sub-micro remainder is preserved) is covered by + // TimestampNanosParquetOpsSuite. INT96 is only produced for the LTZ family. Both a CORRECTED and + // a LEGACY (Julian) INT96 rebase are exercised: LEGACY over the pre-1582 value actually runs the + // INT96 read rebase, which the earlier CORRECTED-only pinning skipped. for { - outputTimestampType <- - Seq(ParquetOutputTimestampType.INT96, ParquetOutputTimestampType.TIMESTAMP_MILLIS) + toType: DataType <- Seq( + TimestampLTZNanosType(TimestampLTZNanosType.NANOS_PRECISION), + TimestampLTZNanosType(7)) + int96RebaseMode <- Seq(LegacyBehaviorPolicy.CORRECTED, LegacyBehaviorPolicy.LEGACY) + } + test(s"parquet widening conversion TimestampType (int96, $int96RebaseMode) -> $toType") { + withSQLConf( + SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> ParquetOutputTimestampType.INT96.toString, + SQLConf.PARQUET_INT96_REBASE_MODE_IN_WRITE.key -> int96RebaseMode.toString, + SQLConf.PARQUET_INT96_REBASE_MODE_IN_READ.key -> int96RebaseMode.toString) { + checkAllParquetReaders( + values = Seq( + "2020-01-01 12:34:56.123456", "1312-02-27 01:02:03.654321", "5138-11-16 09:46:40"), + fromType = TimestampType, + toType = toType, + expectError = false) + } + } + + for { + outputTimestampType <- Seq(ParquetOutputTimestampType.TIMESTAMP_MILLIS) } test(s"unsupported parquet conversion TimestampType ($outputTimestampType) -> nanos") { withSQLConf( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala index 598a84f3dbe35..b1c934332b68e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala @@ -17,21 +17,22 @@ package org.apache.spark.sql.execution.datasources.parquet.types.ops -import java.time.{Instant, LocalDateTime, ZoneOffset} +import java.nio.{ByteBuffer, ByteOrder} +import java.time.{Instant, LocalDateTime, ZoneId, ZoneOffset} import org.apache.parquet.column.ColumnDescriptor import org.apache.parquet.filter2.predicate.FilterApi import org.apache.parquet.filter2.predicate.SparkFilterApi.longColumn -import org.apache.parquet.io.api.PrimitiveConverter +import org.apache.parquet.io.api.{Binary, PrimitiveConverter} import org.apache.parquet.schema.{LogicalTypeAnnotation, Type, Types} import org.apache.parquet.schema.LogicalTypeAnnotation.{TimestampLogicalTypeAnnotation, TimeUnit} -import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.{INT32, INT64} +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.{INT32, INT64, INT96} import org.apache.parquet.schema.Type.Repetition.REQUIRED -import org.apache.spark.{SparkArithmeticException, SparkFunSuite, SparkRuntimeException} +import org.apache.spark.{SparkArithmeticException, SparkFunSuite, SparkRuntimeException, SparkUpgradeException} import org.apache.spark.sql.catalyst.util.{DateTimeConstants, DateTimeUtils} import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec -import org.apache.spark.sql.execution.datasources.parquet.ParentContainerUpdater +import org.apache.spark.sql.execution.datasources.parquet.{ParentContainerUpdater, ParquetRowConverter} import org.apache.spark.sql.internal.LegacyBehaviorPolicy import org.apache.spark.sql.types.{TimestampLTZNanosType, TimestampNTZNanosType} import org.apache.spark.unsafe.types.TimestampNanosVal @@ -167,6 +168,61 @@ class TimestampNanosParquetOpsSuite extends SparkFunSuite { TimestampNanosVal.fromParts(farFutureMicros, 0.toShort)) } + // ---------- INT96 read (widening a legacy INT96 timestamp to nanos) ---------- + + test("isInt96Timestamp matches only the INT96 physical type") { + assert(TimestampNanosParquetOps.isInt96Timestamp(Types.primitive(INT96, REQUIRED).named("c"))) + // INT64 TIMESTAMP(NANOS) and a raw INT64 are not INT96. + assert(!TimestampNanosParquetOps.isInt96Timestamp(Types.primitive(INT64, REQUIRED) + .as(LogicalTypeAnnotation.timestampType(true, TimeUnit.NANOS)).named("c"))) + assert(!TimestampNanosParquetOps.isInt96Timestamp(Types.primitive(INT64, REQUIRED).named("c"))) + } + + test("INT96 read preserves sub-microsecond nanos (both families, precision 9)") { + // INT96 stores nanoseconds-of-day, so a foreign file (e.g. Impala/Hive) can carry true + // sub-microsecond digits. binaryToSQLTimestamp floors to micros; the converter recovers the + // remainder (here ...789) from the raw INT96 instead of hardcoding 0. + val binary = int96Binary(julianDay = 2451545, timeOfDayNanos = 45296123456789L) + val expectedMicros = ParquetRowConverter.binaryToSQLTimestamp(binary) + assert(decodeInt96(ltz, binary) === TimestampNanosVal.fromParts(expectedMicros, 789.toShort)) + assert(decodeInt96(ntz, binary) === TimestampNanosVal.fromParts(expectedMicros, 789.toShort)) + } + + test("INT96 read truncates sub-precision nanos at an explicit lower read precision") { + // nanosWithinMicro 789 -> truncated to 700 at precision 7 (matching the nanos read path). + val ltz7 = TimestampLTZNanosParquetOps(TimestampLTZNanosType(7)) + val binary = int96Binary(julianDay = 2451545, timeOfDayNanos = 45296123456789L) + val expectedMicros = ParquetRowConverter.binaryToSQLTimestamp(binary) + assert(decodeInt96(ltz7, binary) === TimestampNanosVal.fromParts(expectedMicros, 700.toShort)) + } + + test("INT96 read applies the INT96 rebase for LTZ but never for NTZ") { + // A pre-1582 Julian day: for LTZ, LEGACY rebase (Julian -> proleptic Gregorian) shifts the + // value, CORRECTED does not, and EXCEPTION refuses the ancient value. NTZ never rebases, so its + // decode is identical across modes. + val ancient = int96Binary(julianDay = 2200000, timeOfDayNanos = 0L) + val ltzCorrected = decodeInt96(ltz, ancient, RebaseSpec(LegacyBehaviorPolicy.CORRECTED)) + .asInstanceOf[TimestampNanosVal] + val ltzLegacy = decodeInt96(ltz, ancient, RebaseSpec(LegacyBehaviorPolicy.LEGACY)) + .asInstanceOf[TimestampNanosVal] + assert(ltzLegacy.epochMicros != ltzCorrected.epochMicros) + intercept[SparkUpgradeException] { + decodeInt96(ltz, ancient, RebaseSpec(LegacyBehaviorPolicy.EXCEPTION)) + } + assert(decodeInt96(ntz, ancient, RebaseSpec(LegacyBehaviorPolicy.LEGACY)) === + decodeInt96(ntz, ancient, RebaseSpec(LegacyBehaviorPolicy.CORRECTED))) + } + + test("INT96 read applies the timezone conversion for LTZ but not for NTZ") { + // convertTz shifts the LTZ instant by the zone offset; NTZ ignores it (wall-clock semantics). + val binary = int96Binary(julianDay = 2451545, timeOfDayNanos = 45296000000000L) + val zone = ZoneId.of("America/Los_Angeles") + val ltzNoTz = decodeInt96(ltz, binary).asInstanceOf[TimestampNanosVal] + val ltzTz = decodeInt96(ltz, binary, convertTz = Some(zone)).asInstanceOf[TimestampNanosVal] + assert(ltzTz.epochMicros != ltzNoTz.epochMicros) + assert(decodeInt96(ntz, binary, convertTz = Some(zone)) === decodeInt96(ntz, binary)) + } + // ---------- (epochMicros, nanosWithinMicro) -> INT64 epoch-nanos packing ---------- test("timestampNanosToEpochNanos combines micros and sub-micro nanos") { @@ -350,6 +406,34 @@ class TimestampNanosParquetOpsSuite extends SparkFunSuite { captured } + // Builds a 12-byte INT96 value: nanoseconds-of-day (little-endian long) followed by the Julian + // day (little-endian int), the on-disk layout ParquetRowConverter.binaryToSQLTimestamp reads. + private def int96Binary(julianDay: Int, timeOfDayNanos: Long): Binary = { + val buf = ByteBuffer.allocate(12).order(ByteOrder.LITTLE_ENDIAN) + buf.putLong(timeOfDayNanos) + buf.putInt(julianDay) + Binary.fromConstantByteArray(buf.array()) + } + + private val int96Field: Type = Types.primitive(INT96, REQUIRED).named("c") + + // Builds the extended converter over an INT96 field, feeds one crafted INT96 binary through + // addBinary, and returns the decoded TimestampNanosVal. `rebaseSpec` drives the INT96 read rebase + // and `convertTz` the optional timezone conversion (both LTZ-only; NTZ ignores them). + private def decodeInt96( + ops: TimestampNanosParquetOps, + binary: Binary, + rebaseSpec: RebaseSpec = RebaseSpec(LegacyBehaviorPolicy.CORRECTED), + convertTz: Option[ZoneId] = None): Any = { + var captured: Any = null + val updater = new ParentContainerUpdater { + override def set(value: Any): Unit = captured = value + } + val converter = ops.newConverter(int96Field, updater, null, convertTz, rebaseSpec, rebaseSpec) + converter.asInstanceOf[PrimitiveConverter].addBinary(binary) + captured + } + // ---------- vectorized read updater (getVectorUpdater / getVectorUpdaterOrNull) ---------- private def nanosTimestampColumn(isAdjustedToUTC: Boolean): ColumnDescriptor = {