Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -227,6 +229,16 @@ public ParquetVectorUpdater getUpdater(ColumnDescriptor descriptor, DataType spa
int96RebaseTz);
}
}
} else if (sparkType instanceof TimestampNTZNanosType) {
Comment thread
uros-b marked this conversation as resolved.
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 -> {
Expand Down Expand Up @@ -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;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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)) {
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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))
}
}
}
}

/**
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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._

Expand Down Expand Up @@ -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 =>
Expand Down Expand Up @@ -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") {
Expand Down
Loading