This is an automated email from the ASF dual-hosted git repository.
MaxGekk pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new f377e3050f0b [SPARK-57159][SQL] Add Arrow type mapping for
nanosecond-capable timestamp types
f377e3050f0b is described below
commit f377e3050f0b5fbf5b0936da7b23fe305c84f8df
Author: Maxim Gekk <[email protected]>
AuthorDate: Thu Jun 25 11:34:51 2026 +0200
[SPARK-57159][SQL] Add Arrow type mapping for nanosecond-capable timestamp
types
### What changes were proposed in this pull request?
This PR teaches Spark's Arrow conversion about the nanosecond timestamp
types `TimestampNTZNanosType(p)` and `TimestampLTZNanosType(p)` (`p` in `[7,
9]`), so they can be carried over Arrow like the other timestamp types.
- **Type mapping**: NTZ maps to Arrow `Timestamp(NANOSECOND, null)` and LTZ
to `Timestamp(NANOSECOND, sessionTz)`, and back. Arrow timestamps have no
precision field, so the exact precision is preserved in field metadata
(`SPARK::timestampNanos::precision`) and falls back to `9` when it is missing.
- **Writer/reader**: new `ArrowWriter` field writers pack the value into
int64 epoch-nanoseconds (raising `DATETIME_OVERFLOW` when out of range), and
matching `ArrowColumnVector` accessors decode it back into `TimestampNanosVal`.
- **Reuse**: the epoch-nanos packing
(`DateTimeUtils.timestampNanosToEpochNanos`) and the overflow error are now
shared with the Parquet INT64 path (SPARK-57100) instead of duplicated; the
existing Parquet error message is unchanged.
### Why are the changes needed?
This is the shared Arrow prerequisite for Spark Connect (parent:
SPARK-56822) and also benefits the classic Arrow paths (Arrow result transfer,
`createDataFrame` from Arrow, `mapInArrow`). The Spark <-> Arrow mapping
(`ArrowUtils`) and the row-to-vector writers (`ArrowWriter`) had no support for
the nanosecond timestamp types, so any plan whose schema contained them failed
to serialize.
### Does this PR introduce _any_ user-facing change?
No. The types remain gated behind `spark.sql.timestampNanosTypes.enabled`.
### How was this patch tested?
- `ArrowUtilsSuite`: precision round-trip for `p` in `{7, 8, 9}` (NTZ and
LTZ across multiple session zones), null-tz LTZ error, fallback to `9` when the
precision metadata is absent or present-but-invalid (out of `[7, 9]` or
non-numeric), and that the precision key does not leak into the reconstructed
column `Metadata`.
- `ArrowWriterSuite`: value round-trip (sub-micro `0`/`999`, pre-epoch
instants, large boundaries, nulls) for `p = 9` and `p = 7`, and
`DATETIME_OVERFLOW` for out-of-range values, for both NTZ and LTZ.
- `DateTimeUtilsSuite`: direct unit test for the shared
`timestampNanosToEpochNanos` helper (sub-micro boundaries `0`/`999`, positive
and pre-epoch values, the `floorDiv`/`floorMod` inverse, the `Long.MaxValue`
boundary, and the overflow `ArithmeticException`).
- `ArrowConvertersSuite`: end-to-end Arrow IPC batch round-trip
(`toBatchIterator` -> `fromBatchIterator`) for NTZ and LTZ, including
sub-micro/pre-epoch values and a trailing null.
- `ParquetTimestampNanosSuite`: re-run to confirm the shared-helper
refactor preserves the existing Parquet behavior.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Cursor (Claude Opus 4.8)
Closes #56739 from MaxGekk/nanos-arrow.
Authored-by: Maxim Gekk <[email protected]>
Signed-off-by: Max Gekk <[email protected]>
(cherry picked from commit 686d1c4044ad0405dc6ae2c5cc1249aa47f26451)
Signed-off-by: Max Gekk <[email protected]>
---
.../types/ops/TimestampNanosTypeApiOps.scala | 20 ++++++
.../spark/sql/catalyst/types/ops/TypeApiOps.scala | 9 +++
.../org/apache/spark/sql/util/ArrowUtils.scala | 55 +++++++++++++++++
.../spark/sql/vectorized/ArrowColumnVector.java | 60 ++++++++++++++++++
.../catalyst/types/ops/TimestampNanosTypeOps.scala | 9 +++
.../spark/sql/catalyst/util/DateTimeUtils.scala | 14 +++++
.../spark/sql/errors/QueryExecutionErrors.scala | 6 +-
.../spark/sql/execution/arrow/ArrowWriter.scala | 44 +++++++++++++-
.../sql/catalyst/util/DateTimeUtilsSuite.scala | 33 ++++++++++
.../apache/spark/sql/util/ArrowUtilsSuite.scala | 70 ++++++++++++++++++++-
.../datasources/parquet/ParquetWriteSupport.scala | 9 ++-
.../sql/execution/arrow/ArrowConvertersSuite.scala | 45 +++++++++++++-
.../sql/execution/arrow/ArrowWriterSuite.scala | 71 +++++++++++++++++++++-
13 files changed, 430 insertions(+), 15 deletions(-)
diff --git
a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala
b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala
index 650575d006f1..76bbf0c6f4a8 100644
---
a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala
+++
b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala
@@ -19,6 +19,10 @@ package org.apache.spark.sql.catalyst.types.ops
import java.time.{Instant, LocalDateTime, ZoneId, ZoneOffset}
+import org.apache.arrow.vector.types.TimeUnit
+import org.apache.arrow.vector.types.pojo.ArrowType
+
+import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder
import
org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{InstantNanosEncoder,
LocalDateTimeNanosEncoder}
import org.apache.spark.sql.catalyst.util.TimestampFormatter
@@ -116,6 +120,12 @@ class TimestampNTZNanosTypeApiOps(val t:
TimestampNTZNanosType) extends Timestam
// Mirrors RowEncoder.encoderForDataTypeDefault for TimestampNTZNanosType
(SPARK-57033):
// maps to java.time.LocalDateTime with the column precision.
override protected def nanosEncoder: AgnosticEncoder[_] =
LocalDateTimeNanosEncoder(t.precision)
+
+ // NTZ is zone-less: like TimestampNTZType, the Arrow timestamp carries a
null time zone. The
+ // column precision is not expressible in the Arrow type itself and is
carried in the Arrow
+ // field metadata instead (see ArrowUtils).
+ override def toArrowType(timeZoneId: String): Option[ArrowType] =
+ Some(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null))
}
/**
@@ -154,4 +164,14 @@ class TimestampLTZNanosTypeApiOps(val t:
TimestampLTZNanosType, zoneId: => ZoneI
// Mirrors RowEncoder.encoderForDataTypeDefault for TimestampLTZNanosType
(SPARK-57033):
// maps to java.time.Instant with the column precision.
override protected def nanosEncoder: AgnosticEncoder[_] =
InstantNanosEncoder(t.precision)
+
+ // LTZ is zone-aware: like TimestampType, the Arrow timestamp carries the
session time zone, so
+ // a non-null timeZoneId is mandatory (mirrors ArrowUtils.toArrowTypeDefault
for TimestampType).
+ // The column precision is carried in the Arrow field metadata instead (see
ArrowUtils).
+ override def toArrowType(timeZoneId: String): Option[ArrowType] = {
+ if (timeZoneId == null) {
+ throw SparkException.internalError("Missing timezoneId where it is
mandatory.")
+ }
+ Some(new ArrowType.Timestamp(TimeUnit.NANOSECOND, timeZoneId))
+ }
}
diff --git
a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TypeApiOps.scala
b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TypeApiOps.scala
index 728c6ae40cd4..e94e063f3d09 100644
---
a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TypeApiOps.scala
+++
b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TypeApiOps.scala
@@ -202,6 +202,15 @@ object TypeApiOps {
at match {
case t: ArrowType.Time if t.getUnit == TimeUnit.NANOSECOND &&
t.getBitWidth == 8 * 8 =>
Some(TimeType(TimeType.MICROS_PRECISION))
+ // Nanosecond Arrow timestamps map to the nanosecond-capable Spark
timestamp types. The Arrow
+ // type carries no fractional-second precision, so this precision-less
reverse lookup uses the
+ // canonical maximum precision; the exact precision is recovered from
the Arrow field metadata
+ // by ArrowUtils.fromArrowField when present.
+ case ts: ArrowType.Timestamp
+ if ts.getUnit == TimeUnit.NANOSECOND && ts.getTimezone == null =>
+ Some(TimestampNTZNanosType(TimestampNTZNanosType.MAX_PRECISION))
+ case ts: ArrowType.Timestamp if ts.getUnit == TimeUnit.NANOSECOND =>
+ Some(TimestampLTZNanosType(TimestampLTZNanosType.MAX_PRECISION))
// Add new framework types here
case _ => None
}
diff --git a/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
b/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
index b84336cc0f54..15cf5b23e4ac 100644
--- a/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
+++ b/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
@@ -110,6 +110,11 @@ private[sql] object ArrowUtils {
}
private val metadataKey = "SPARK::metadata::json"
+ // Arrow's Timestamp type carries only (unit, timezone) and has no
fractional-second precision
+ // field, so the precision of the nanosecond timestamp types is stored in
the Arrow field
+ // metadata under this dedicated key (namespaced like `metadataKey`,
separate from the user
+ // metadata blob so user metadata is untouched) and recovered on read in
`fromArrowField`.
+ private val timestampNanosPrecisionKey = "SPARK::timestampNanos::precision"
private def toArrowMetaData(metadata: Metadata) = {
if (metadata != null && !metadata.isEmpty) {
Map(metadataKey -> metadata.json).asJava
@@ -125,6 +130,24 @@ private[sql] object ArrowUtils {
}
}
+ /**
+ * Builds an Arrow field for a nanosecond timestamp type, stashing the
column precision in the
+ * field metadata (alongside the user metadata) so it can be recovered in
`fromArrowField`.
+ */
+ private def toTimestampNanosArrowField(
+ name: String,
+ dt: DataType,
+ precision: Int,
+ nullable: Boolean,
+ timeZoneId: String,
+ largeVarTypes: Boolean,
+ metadata: Metadata): Field = {
+ val base =
Option(toArrowMetaData(metadata)).map(_.asScala.toMap).getOrElse(Map.empty)
+ val md = (base + (timestampNanosPrecisionKey -> precision.toString)).asJava
+ val fieldType = new FieldType(nullable, toArrowType(dt, timeZoneId,
largeVarTypes), null, md)
+ new Field(name, fieldType, Seq.empty[Field].asJava)
+ }
+
/** Maps field from Spark to Arrow. NOTE: timeZoneId required for
TimestampType */
def toArrowField(
name: String,
@@ -231,6 +254,24 @@ private[sql] object ArrowUtils {
Seq(
toArrowField("value", BinaryType, false, timeZoneId,
largeVarTypes),
new Field("metadata", metadataFieldType,
Seq.empty[Field].asJava)).asJava)
+ case t: TimestampNTZNanosType =>
+ toTimestampNanosArrowField(
+ name,
+ t,
+ t.precision,
+ nullable,
+ timeZoneId,
+ largeVarTypes,
+ metadata)
+ case t: TimestampLTZNanosType =>
+ toTimestampNanosArrowField(
+ name,
+ t,
+ t.precision,
+ nullable,
+ timeZoneId,
+ largeVarTypes,
+ metadata)
case dataType =>
val fieldType = new FieldType(
nullable,
@@ -310,6 +351,20 @@ private[sql] object ArrowUtils {
StructField(child.getName, dt, child.isNullable,
fromArrowMetaData(child.getMetadata))
}
StructType(fields.toArray)
+ // Recover the exact precision of nanosecond timestamps from the field
metadata written by
+ // `toTimestampNanosArrowField`. Foreign Arrow data (or an out-of-range
value) has no usable
+ // key, so fall back to the canonical maximum precision via
`fromArrowType`.
+ case ts: ArrowType.Timestamp if ts.getUnit == TimeUnit.NANOSECOND =>
+ val precision =
Option(field.getMetadata.get(timestampNanosPrecisionKey))
+ .flatMap(s => scala.util.Try(s.toInt).toOption)
+ .filter { p =>
+ p >= TimestampNTZNanosType.MIN_PRECISION && p <=
TimestampNTZNanosType.MAX_PRECISION
+ }
+ precision match {
+ case Some(p) if ts.getTimezone == null => TimestampNTZNanosType(p)
+ case Some(p) => TimestampLTZNanosType(p)
+ case None => fromArrowType(ts)
+ }
case arrowType => fromArrowType(arrowType)
}
}
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java
index 8a47e93724d9..3267daea6bcc 100644
---
a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java
@@ -25,11 +25,13 @@ import
org.apache.arrow.vector.holders.NullableVarCharHolder;
import org.apache.spark.SparkUnsupportedOperationException;
import org.apache.spark.annotation.DeveloperApi;
+import org.apache.spark.sql.catalyst.util.DateTimeConstants;
import org.apache.spark.sql.catalyst.util.STUtils;
import org.apache.spark.sql.util.ArrowUtils;
import org.apache.spark.sql.types.*;
import org.apache.spark.unsafe.types.BinaryView;
import org.apache.spark.unsafe.types.CalendarInterval;
+import org.apache.spark.unsafe.types.TimestampNanosVal;
import org.apache.spark.unsafe.types.UTF8String;
/**
@@ -123,6 +125,18 @@ public class ArrowColumnVector extends ColumnVector {
return accessor.getInterval(rowId);
}
+ @Override
+ public TimestampNanosVal getTimestampNTZNanos(int rowId) {
+ if (isNullAt(rowId)) return null;
+ return accessor.getTimestampNanos(rowId);
+ }
+
+ @Override
+ public TimestampNanosVal getTimestampLTZNanos(int rowId) {
+ if (isNullAt(rowId)) return null;
+ return accessor.getTimestampNanos(rowId);
+ }
+
@Override
public byte[] getBinary(int rowId) {
if (isNullAt(rowId)) return null;
@@ -204,6 +218,10 @@ public class ArrowColumnVector extends ColumnVector {
accessor = new TimestampNTZAccessor(timeStampMicroVector);
} else if (vector instanceof TimeNanoVector timeNanoVector) {
accessor = new TimeNanoAccessor(timeNanoVector);
+ } else if (vector instanceof TimeStampNanoTZVector timeStampNanoTZVector) {
+ accessor = new TimestampLTZNanosAccessor(timeStampNanoTZVector);
+ } else if (vector instanceof TimeStampNanoVector timeStampNanoVector) {
+ accessor = new TimestampNTZNanosAccessor(timeStampNanoVector);
} else if (vector instanceof MapVector mapVector) {
accessor = new MapAccessor(mapVector);
} else if (vector instanceof ListVector listVector) {
@@ -280,6 +298,10 @@ public class ArrowColumnVector extends ColumnVector {
throw SparkUnsupportedOperationException.apply();
}
+ TimestampNanosVal getTimestampNanos(int rowId) {
+ throw SparkUnsupportedOperationException.apply();
+ }
+
Decimal getDecimal(int rowId, int precision, int scale) {
throw SparkUnsupportedOperationException.apply();
}
@@ -559,6 +581,44 @@ public class ArrowColumnVector extends ColumnVector {
}
}
+ // Decodes a single int64 of epoch-nanoseconds back into the (epochMicros,
nanosWithinMicro)
+ // pair. floorDiv/floorMod keep nanosWithinMicro in [0, 999] for pre-epoch
(negative) values too.
+ private static TimestampNanosVal decodeEpochNanos(long nanos) {
+ return TimestampNanosVal.fromTrustedRowBytes(
+ Math.floorDiv(nanos, DateTimeConstants.NANOS_PER_MICROS),
+ (short) Math.floorMod(nanos, DateTimeConstants.NANOS_PER_MICROS));
+ }
+
+ static class TimestampNTZNanosAccessor extends ArrowVectorAccessor {
+
+ private final TimeStampNanoVector accessor;
+
+ TimestampNTZNanosAccessor(TimeStampNanoVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final TimestampNanosVal getTimestampNanos(int rowId) {
+ return decodeEpochNanos(accessor.get(rowId));
+ }
+ }
+
+ static class TimestampLTZNanosAccessor extends ArrowVectorAccessor {
+
+ private final TimeStampNanoTZVector accessor;
+
+ TimestampLTZNanosAccessor(TimeStampNanoTZVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final TimestampNanosVal getTimestampNanos(int rowId) {
+ return decodeEpochNanos(accessor.get(rowId));
+ }
+ }
+
static class ArrayAccessor extends ArrowVectorAccessor {
private final ListVector accessor;
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeOps.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeOps.scala
index 48628619cdb4..6ecebf9a3fe0 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeOps.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeOps.scala
@@ -19,12 +19,15 @@ package org.apache.spark.sql.catalyst.types.ops
import java.time.{Instant, LocalDateTime}
+import org.apache.arrow.vector.{TimeStampNanoTZVector, TimeStampNanoVector,
ValueVector}
+
import org.apache.spark.SparkIllegalArgumentException
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Expression, Literal,
MutableTimestampNanos, MutableValue}
import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
import org.apache.spark.sql.catalyst.types.{PhysicalDataType,
PhysicalTimestampLTZNanosType, PhysicalTimestampNTZNanosType}
import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.execution.arrow.{ArrowFieldWriter,
TimestampLTZNanosWriter, TimestampNTZNanosWriter}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ObjectType, TimestampLTZNanosType,
TimestampNTZNanosType}
import org.apache.spark.unsafe.types.TimestampNanosVal
@@ -121,6 +124,9 @@ case class TimestampNTZNanosTypeOps(override val t:
TimestampNTZNanosType)
"timestampNanosToLocalDateTime",
path :: Nil,
returnNullable = false))
+
+ override def createArrowFieldWriter(vector: ValueVector):
Option[ArrowFieldWriter] =
+ Some(new TimestampNTZNanosWriter(vector.asInstanceOf[TimeStampNanoVector]))
}
/**
@@ -171,4 +177,7 @@ case class TimestampLTZNanosTypeOps(override val t:
TimestampLTZNanosType)
"timestampNanosToInstant",
path :: Nil,
returnNullable = false))
+
+ override def createArrowFieldWriter(vector: ValueVector):
Option[ArrowFieldWriter] =
+ Some(new
TimestampLTZNanosWriter(vector.asInstanceOf[TimeStampNanoTZVector]))
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index 481f99709858..86cd01a0f3b5 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -317,6 +317,20 @@ object DateTimeUtils extends SparkDateTimeUtils {
TimestampNanosVal.fromParts(epochMicros, start.nanosWithinMicro)
}
+ /**
+ * Packs a [[TimestampNanosVal]] (epoch micros + nanos within the micro)
into a single int64 of
+ * nanoseconds since the epoch, the representation used by the Arrow
nanosecond timestamp vectors
+ * and the Parquet INT64 epoch-nanoseconds physical type. Throws
[[ArithmeticException]] when the
+ * value falls outside the int64 epoch-nanosecond range; callers translate
that into a
+ * `DATETIME_OVERFLOW` error naming their sink (see
+ *
[[org.apache.spark.sql.errors.QueryExecutionErrors.timestampNanosEpochNanosOverflowError]]).
+ */
+ def timestampNanosToEpochNanos(value: TimestampNanosVal): Long = {
+ Math.addExact(
+ Math.multiplyExact(value.epochMicros, NANOS_PER_MICROS),
+ value.nanosWithinMicro.toLong)
+ }
+
/**
* Adds a full interval (months, days, microseconds) to a timestamp
represented as the number of
* microseconds since 1970-01-01 00:00:00Z.
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
index c3873f0b2881..b4b37c05f8c2 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
@@ -2615,8 +2615,8 @@ private[sql] object QueryExecutionErrors extends
QueryErrorsBase with ExecutionE
summary = "")
}
- def parquetTimestampNanosOverflowError(
- value: TimestampNanosVal, isNtz: Boolean): SparkArithmeticException = {
+ def timestampNanosEpochNanosOverflowError(
+ value: TimestampNanosVal, isNtz: Boolean, sink: String):
SparkArithmeticException = {
// Render TIMESTAMP_NTZ values without a zone (LocalDateTime, no trailing
`Z`); TIMESTAMP_LTZ
// values are absolute instants and render as UTC with a trailing `Z`.
val rendered =
@@ -2625,7 +2625,7 @@ private[sql] object QueryExecutionErrors extends
QueryErrorsBase with ExecutionE
new SparkArithmeticException(
errorClass = "DATETIME_OVERFLOW",
messageParameters = Map(
- "operation" -> (s"write the timestamp value $rendered as Parquet INT64
" +
+ "operation" -> (s"write the timestamp value $rendered as $sink " +
"epoch-nanoseconds " +
"(supported range: 1677-09-21T00:12:43.145224192Z to
2262-04-11T23:47:16.854775807Z)")),
context = Array.empty,
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala
index 1e2ae058ac63..6030eee94a6c 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala
@@ -25,8 +25,8 @@ import org.apache.arrow.vector.complex._
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.SpecializedGetters
import org.apache.spark.sql.catalyst.types.ops.TypeOps
-import org.apache.spark.sql.catalyst.util.STUtils
-import org.apache.spark.sql.errors.ExecutionErrors
+import org.apache.spark.sql.catalyst.util.{DateTimeUtils, STUtils}
+import org.apache.spark.sql.errors.{ExecutionErrors, QueryExecutionErrors}
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.ArrowUtils
@@ -390,6 +390,46 @@ private[sql] class TimeWriter(
}
}
+private[sql] class TimestampNTZNanosWriter(
+ val valueVector: TimeStampNanoVector) extends ArrowFieldWriter {
+
+ override def setNull(): Unit = {
+ valueVector.setNull(count)
+ }
+
+ override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
+ val v = input.getTimestampNTZNanos(ordinal)
+ val nanos = try {
+ DateTimeUtils.timestampNanosToEpochNanos(v)
+ } catch {
+ case _: ArithmeticException =>
+ throw QueryExecutionErrors.timestampNanosEpochNanosOverflowError(
+ v, isNtz = true, sink = "Arrow INT64")
+ }
+ valueVector.setSafe(count, nanos)
+ }
+}
+
+private[sql] class TimestampLTZNanosWriter(
+ val valueVector: TimeStampNanoTZVector) extends ArrowFieldWriter {
+
+ override def setNull(): Unit = {
+ valueVector.setNull(count)
+ }
+
+ override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
+ val v = input.getTimestampLTZNanos(ordinal)
+ val nanos = try {
+ DateTimeUtils.timestampNanosToEpochNanos(v)
+ } catch {
+ case _: ArithmeticException =>
+ throw QueryExecutionErrors.timestampNanosEpochNanosOverflowError(
+ v, isNtz = false, sink = "Arrow INT64")
+ }
+ valueVector.setSafe(count, nanos)
+ }
+}
+
private[arrow] class ArrayWriter(
val valueVector: ListVector,
val elementWriter: ArrowFieldWriter) extends ArrowFieldWriter {
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
index 6db96d8d366a..34fc5286722c 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
@@ -1199,6 +1199,39 @@ class DateTimeUtilsSuite extends SparkFunSuite with
Matchers with SQLHelper {
}
}
+ test("SPARK-57159: timestampNanosToEpochNanos packs into int64
epoch-nanoseconds") {
+ def nanos(epochMicros: Long, nanosWithinMicro: Int): TimestampNanosVal =
+ TimestampNanosVal.fromParts(epochMicros, nanosWithinMicro.toShort)
+
+ // Packs (epochMicros, nanosWithinMicro) as epochMicros * 1000 +
nanosWithinMicro, including the
+ // sub-microsecond remainder boundaries 0 and 999.
+ assert(timestampNanosToEpochNanos(nanos(0L, 0)) === 0L)
+ assert(timestampNanosToEpochNanos(nanos(0L, 999)) === 999L)
+ assert(timestampNanosToEpochNanos(nanos(1L, 0)) === NANOS_PER_MICROS)
+ assert(timestampNanosToEpochNanos(nanos(1234567L, 7)) === 1234567L *
NANOS_PER_MICROS + 7L)
+ // Pre-epoch (negative epochMicros) values pack linearly too.
+ assert(timestampNanosToEpochNanos(nanos(-1L, 0)) === -NANOS_PER_MICROS)
+ assert(timestampNanosToEpochNanos(nanos(-1234567L, 13)) === -1234567L *
NANOS_PER_MICROS + 13L)
+
+ // Round-trips with the documented inverse (floorDiv/floorMod) that the
Arrow reader uses to
+ // reconstruct the pair, for positive and pre-epoch values alike.
+ Seq(nanos(0L, 0), nanos(0L, 999), nanos(1234567L, 7), nanos(-1234567L,
13)).foreach { v =>
+ val packed = timestampNanosToEpochNanos(v)
+ assert(Math.floorDiv(packed, NANOS_PER_MICROS) === v.epochMicros)
+ assert(Math.floorMod(packed, NANOS_PER_MICROS) ===
v.nanosWithinMicro.toLong)
+ }
+
+ // The maximum representable instant packs exactly to Long.MaxValue.
+ assert(timestampNanosToEpochNanos(
+ nanos(Long.MaxValue / NANOS_PER_MICROS, (Long.MaxValue %
NANOS_PER_MICROS).toInt)) ===
+ Long.MaxValue)
+
+ // Out-of-range values raise ArithmeticException (callers translate it
into DATETIME_OVERFLOW).
+ intercept[ArithmeticException] {
+ timestampNanosToEpochNanos(nanos(Long.MaxValue / NANOS_PER_MICROS + 1,
0))
+ }
+ }
+
test("SPARK-34903: subtract timestamps") {
DateTimeTestUtils.outstandingZoneIds.foreach { zid =>
Seq(
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
index 8610a8018b27..16682f981633 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
@@ -19,7 +19,8 @@ package org.apache.spark.sql.util
import java.time.ZoneId
-import org.apache.arrow.vector.types.pojo.ArrowType
+import org.apache.arrow.vector.types.TimeUnit
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
import org.apache.spark.{SparkException, SparkFunSuite,
SparkUnsupportedOperationException}
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.LA
@@ -86,6 +87,73 @@ class ArrowUtilsSuite extends SparkFunSuite {
roundtripWithTz(LA.getId)
}
+ test("timestamp nanos") {
+ // NTZ is zone-independent (null Arrow timezone); precision preserved via
field metadata.
+ Seq(7, 8, 9).foreach { p =>
+ val schema = new StructType().add("value", TimestampNTZNanosType(p))
+ val arrowSchema = ArrowUtils.toArrowSchema(schema, null, true, false)
+ val fieldType =
arrowSchema.findField("value").getType.asInstanceOf[ArrowType.Timestamp]
+ assert(fieldType.getUnit === TimeUnit.NANOSECOND)
+ assert(fieldType.getTimezone === null)
+ assert(ArrowUtils.fromArrowSchema(arrowSchema) === schema)
+ }
+
+ // LTZ is zone-aware: it requires a non-null session time zone; precision
preserved.
+ def roundtripLtz(timeZoneId: String): Unit = {
+ Seq(7, 8, 9).foreach { p =>
+ val schema = new StructType().add("value", TimestampLTZNanosType(p))
+ val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId, true,
false)
+ val fieldType =
arrowSchema.findField("value").getType.asInstanceOf[ArrowType.Timestamp]
+ assert(fieldType.getUnit === TimeUnit.NANOSECOND)
+ assert(fieldType.getTimezone === timeZoneId)
+ assert(ArrowUtils.fromArrowSchema(arrowSchema) === schema)
+ }
+ }
+ roundtripLtz(ZoneId.systemDefault().getId)
+ roundtripLtz("Asia/Tokyo")
+ roundtripLtz("UTC")
+ roundtripLtz(LA.getId)
+
+ // LTZ without a time zone is an error, mirroring TimestampType.
+ checkError(
+ exception = intercept[SparkException] {
+ ArrowUtils.toArrowSchema(
+ new StructType().add("value", TimestampLTZNanosType(9)), null, true,
false)
+ },
+ condition = "INTERNAL_ERROR",
+ parameters = Map("message" -> "Missing timezoneId where it is
mandatory."))
+
+ // Fallback: a nanosecond Arrow timestamp without precision metadata maps
to canonical p=9.
+ def nanosField(timeZoneId: String): Field = new Field(
+ "value",
+ new FieldType(true, new ArrowType.Timestamp(TimeUnit.NANOSECOND,
timeZoneId), null, null),
+ java.util.Collections.emptyList[Field]())
+ assert(ArrowUtils.fromArrowField(nanosField(null)) ===
TimestampNTZNanosType(9))
+ assert(ArrowUtils.fromArrowField(nanosField("UTC")) ===
TimestampLTZNanosType(9))
+
+ // Fallback also covers a present-but-invalid precision key (out of [7, 9]
or non-numeric):
+ // the value is unusable, so the type maps to the canonical p=9 just like
the no-metadata case.
+ def nanosFieldWithPrecision(timeZoneId: String, precision: String): Field
= new Field(
+ "value",
+ new FieldType(
+ true,
+ new ArrowType.Timestamp(TimeUnit.NANOSECOND, timeZoneId),
+ null,
+ java.util.Collections.singletonMap("SPARK::timestampNanos::precision",
precision)),
+ java.util.Collections.emptyList[Field]())
+ assert(
+ ArrowUtils.fromArrowField(nanosFieldWithPrecision(null, "5")) ===
TimestampNTZNanosType(9))
+ assert(
+ ArrowUtils.fromArrowField(nanosFieldWithPrecision("UTC", "x")) ===
TimestampLTZNanosType(9))
+
+ // The precision metadata key does not leak into the reconstructed column
Metadata.
+ val md = new MetadataBuilder().putString("city", "beijing").build()
+ val schemaWithMeta =
+ new StructType().add("value", TimestampNTZNanosType(7), nullable = true,
md)
+ assert(ArrowUtils.fromArrowSchema(
+ ArrowUtils.toArrowSchema(schemaWithMeta, null, true, false)) ===
schemaWithMeta)
+ }
+
test("array") {
roundtrip(ArrayType(IntegerType, containsNull = true))
roundtrip(ArrayType(IntegerType, containsNull = false))
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetWriteSupport.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetWriteSupport.scala
index 48e57f1d6bd3..641a563cd7c1 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetWriteSupport.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetWriteSupport.scala
@@ -34,7 +34,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.{SPARK_LEGACY_DATETIME_METADATA_KEY,
SPARK_LEGACY_INT96_METADATA_KEY, SPARK_TIMEZONE_METADATA_KEY,
SPARK_VERSION_METADATA_KEY}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.SpecializedGetters
-import org.apache.spark.sql.catalyst.util.{DateTimeConstants, DateTimeUtils,
STUtils}
+import org.apache.spark.sql.catalyst.util.{DateTimeUtils, STUtils}
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.datasources.DataSourceUtils
import
org.apache.spark.sql.execution.datasources.parquet.types.ops.ParquetTypeOps
@@ -193,12 +193,11 @@ class ParquetWriteSupport extends
WriteSupport[InternalRow] with Logging {
private def timestampNanosToEpochNanos(value: TimestampNanosVal, isNtz:
Boolean): Long = {
try {
- Math.addExact(
- Math.multiplyExact(value.epochMicros,
DateTimeConstants.NANOS_PER_MICROS),
- value.nanosWithinMicro.toLong)
+ DateTimeUtils.timestampNanosToEpochNanos(value)
} catch {
case _: ArithmeticException =>
- throw QueryExecutionErrors.parquetTimestampNanosOverflowError(value,
isNtz)
+ throw QueryExecutionErrors.timestampNanosEpochNanosOverflowError(
+ value, isNtz, sink = "Parquet INT64")
}
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala
index 0d7d15b8fcd0..e3d5e2e4c5a2 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala
@@ -36,9 +36,9 @@ import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.classic.DataFrame
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.sql.types.{ArrayType, BinaryType, Decimal,
IntegerType, NullType, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, Decimal,
IntegerType, NullType, StringType, StructField, StructType,
TimestampLTZNanosType, TimestampNTZNanosType}
import org.apache.spark.sql.util.ArrowUtils
-import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.unsafe.types.{TimestampNanosVal, UTF8String}
import org.apache.spark.util.Utils
@@ -1431,6 +1431,47 @@ class ArrowConvertersSuite extends SharedSparkSession {
assert(count == inputRows.length)
}
+ test("SPARK-57159: roundtrip arrow batches with nanosecond timestamps") {
+ withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") {
+ Seq[(DataType, String)](
+ (TimestampNTZNanosType(9), null),
+ (TimestampLTZNanosType(9), "UTC")).foreach { case (dt, timeZoneId) =>
+ val values = Seq(
+ TimestampNanosVal.fromParts(0L, 0.toShort),
+ TimestampNanosVal.fromParts(0L, 999.toShort),
+ TimestampNanosVal.fromParts(1234567L, 7.toShort),
+ // pre-epoch instant with a sub-microsecond remainder
+ TimestampNanosVal.fromParts(-1234567L, 13.toShort))
+ // A trailing null exercises the null path.
+ val inputRows = values.map(v => InternalRow(v)) :+ InternalRow(null)
+ val schema = StructType(Seq(StructField("value", dt, nullable = true)))
+
+ val ctx = TaskContext.empty()
+ val batchIter = ArrowConverters.toBatchIterator(
+ inputRows.iterator, schema, 5, timeZoneId, true, false, ctx)
+ // The output iterator reuses a mutable row, so read each row before
advancing.
+ val outputRowIter = ArrowConverters.fromBatchIterator(
+ batchIter, schema, timeZoneId, true, false, ctx)
+
+ var count = 0
+ outputRowIter.zipWithIndex.foreach { case (row, i) =>
+ if (i < values.length) {
+ val got = dt match {
+ case _: TimestampNTZNanosType => row.getTimestampNTZNanos(0)
+ case _: TimestampLTZNanosType => row.getTimestampLTZNanos(0)
+ }
+ assert(got.epochMicros === values(i).epochMicros)
+ assert(got.nanosWithinMicro === values(i).nanosWithinMicro)
+ } else {
+ assert(row.isNullAt(0))
+ }
+ count += 1
+ }
+ assert(count === inputRows.length)
+ }
+ }
+ }
+
test("ArrowBatchStreamWriter roundtrip") {
val inputRows = (0 until 9).map(InternalRow(_)) :+ InternalRow(null)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala
index eb826a9e2357..e4a22ff18846 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala
@@ -21,7 +21,7 @@ import scala.jdk.CollectionConverters._
import org.apache.arrow.vector.VectorSchemaRoot
-import org.apache.spark.SparkFunSuite
+import org.apache.spark.{SparkArithmeticException, SparkFunSuite}
import org.apache.spark.sql.Row
import org.apache.spark.sql.YearUDT
import org.apache.spark.sql.catalyst.InternalRow
@@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.util.{Geography =>
InternalGeography, Geome
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.ArrowUtils
import org.apache.spark.sql.vectorized._
-import org.apache.spark.unsafe.types.{BinaryView, CalendarInterval, UTF8String}
+import org.apache.spark.unsafe.types.{BinaryView, CalendarInterval,
TimestampNanosVal, UTF8String}
import org.apache.spark.util.MaybeNull
class ArrowWriterSuite extends SparkFunSuite {
@@ -152,6 +152,73 @@ class ArrowWriterSuite extends SparkFunSuite {
check(new YearUDT, Seq(2020, 2021, null, 2022))
}
+ test("timestamp nanos round-trip") {
+ // Decompose an int64 epoch-nanoseconds value into the (epochMicros,
nanosWithinMicro) pair,
+ // matching how the Arrow reader reconstructs it.
+ def fromEpochNanos(nanos: Long): TimestampNanosVal =
+ TimestampNanosVal.fromParts(Math.floorDiv(nanos, 1000L),
Math.floorMod(nanos, 1000L).toShort)
+
+ val values = Seq(
+ TimestampNanosVal.fromParts(0L, 0.toShort),
+ TimestampNanosVal.fromParts(0L, 999.toShort),
+ TimestampNanosVal.fromParts(1234567L, 7.toShort),
+ // pre-epoch instant with a sub-microsecond remainder
+ TimestampNanosVal.fromParts(-1234567L, 13.toShort),
+ // large positive/negative epoch-nanoseconds within the representable
range
+ fromEpochNanos(9000000000000000000L),
+ fromEpochNanos(-9000000000000000000L))
+
+ def check(dt: DataType, timeZoneId: String): Unit = {
+ val schema = new StructType().add("value", dt, nullable = true)
+ val writer = ArrowWriter.create(schema, timeZoneId)
+ assert(writer.schema === schema)
+ // Append a trailing null to exercise the null path.
+ (values.map(Option(_)) :+ None).foreach { v =>
+ writer.write(InternalRow(v.orNull))
+ }
+ writer.finish()
+
+ val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
+ values.zipWithIndex.foreach { case (v, rowId) =>
+ val got = dt match {
+ case _: TimestampNTZNanosType => reader.getTimestampNTZNanos(rowId)
+ case _: TimestampLTZNanosType => reader.getTimestampLTZNanos(rowId)
+ }
+ assert(got.epochMicros === v.epochMicros)
+ assert(got.nanosWithinMicro === v.nanosWithinMicro)
+ }
+ assert(reader.isNullAt(values.length))
+ writer.root.close()
+ }
+
+ check(TimestampNTZNanosType(9), null)
+ check(TimestampLTZNanosType(9), "UTC")
+ // The value path packs the full nanosecond value regardless of the column
precision (precision
+ // is carried in the Arrow field metadata, not the value), so p=7
round-trips identically to
+ // p=9; exercising it guards the value path against a future
precision-enforcing change.
+ check(TimestampNTZNanosType(7), null)
+ check(TimestampLTZNanosType(7), "UTC")
+ }
+
+ test("timestamp nanos out of range raises DATETIME_OVERFLOW") {
+ def check(dt: DataType, timeZoneId: String): Unit = {
+ val schema = new StructType().add("value", dt, nullable = true)
+ val writer = ArrowWriter.create(schema, timeZoneId)
+ // epochMicros past the int64 epoch-nanosecond range overflows when
packed, but still
+ // renders as a valid Instant/LocalDateTime in the error message.
+ val tooLarge = TimestampNanosVal.fromParts(Long.MaxValue / 1000L + 1L,
0.toShort)
+ val e = intercept[SparkArithmeticException] {
+ writer.write(InternalRow(tooLarge))
+ }
+ assert(e.getCondition === "DATETIME_OVERFLOW")
+ assert(e.getMessage.contains("Arrow INT64"))
+ writer.root.close()
+ }
+
+ check(TimestampNTZNanosType(9), null)
+ check(TimestampLTZNanosType(9), "UTC")
+ }
+
test("nested geographies") {
def check(
dt: StructType,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
