This is an automated email from the ASF dual-hosted git repository.
viirya 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 094e58c48a82 [SPARK-58005][SQL] Add an opt-in lossless Arrow struct
representation for CalendarInterval
094e58c48a82 is described below
commit 094e58c48a82e63240f6433a5f32517ff3d9eb6a
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Wed Jul 8 10:53:40 2026 -0700
[SPARK-58005][SQL] Add an opt-in lossless Arrow struct representation for
CalendarInterval
### What changes were proposed in this pull request?
This extends the opt-in lossless Arrow encoding introduced by SPARK-57975
(#57053) to `CalendarIntervalType`, and hardens the default interval writer's
overflow error:
- **Lossless struct encoding**: with the opt-in flag, a `CalendarInterval`
column maps to an Arrow struct of `(months: int32, days: int32, microseconds:
int64)` -- the type's own field layout, mirroring the default in-memory cache's
`CALENDAR_INTERVAL` `ColumnType`. The components are stored as-is with no unit
conversion, so the full `Long` microsecond domain round-trips. The struct is
tagged through child-field metadata (the geometry/variant pattern) and is
self-describing on read: ` [...]
- **Flag rename**: the parameter is renamed from `losslessTimestampNanos`
to `losslessInternalTypes`, since it now selects the lossless encoding for both
kinds of types whose standard Arrow encoding cannot cover their full Spark
value domain. `ArrowUtils` is `private[sql]`, so the rename has no
compatibility impact; the only intended caller (the Arrow-based Dataset cache,
#56334) wants both types, and the flag expresses one intent: internal storage
wants fidelity.
- **Structured error at the conversion site**: `IntervalMonthDayNanoWriter`
now catches the `Math.multiplyExact(microseconds, 1000L)` overflow exactly at
the conversion and raises the structured `DATETIME_OVERFLOW` (new
`QueryExecutionErrors.calendarIntervalArrowNanosOverflowError`, the same
pattern as `TimestampNTZNanosWriter`'s `timestampNanosEpochNanosOverflowError`)
instead of letting a raw `ArithmeticException: long overflow` escape. Because
the catch is scoped to the single conv [...]
The default `Interval(MONTH_DAY_NANO)` mapping and every existing caller
are unchanged.
### Why are the changes needed?
Spark permits the full `Long` microsecond range in `CalendarInterval`, but
Arrow's `IntervalMonthDayNano` stores the sub-day component as int64
nanoseconds, so any `|microseconds| > Long.MaxValue / 1000` (roughly +/-292
years) is structurally unrepresentable in the standard encoding -- the default
in-memory cache serializer stores the three components raw and has no such
limit. As with the nanosecond timestamps in SPARK-57975, the interchange
mapping must keep the standard encoding fo [...]
### Does this PR introduce _any_ user-facing change?
The lossless encoding itself is opt-in via an internal API parameter and
changes nothing by default. One user-visible improvement on the existing paths:
writing an out-of-range `CalendarInterval` through Arrow (e.g. `toPandas`,
Arrow UDFs) now fails with the structured `DATETIME_OVERFLOW` condition naming
the value and the limit, instead of an opaque `java.lang.ArithmeticException:
long overflow`.
### How was this patch tested?
New tests:
- `ArrowUtilsSuite` "calendar interval lossless struct": schema shape
(struct of int32/int32/int64, non-null children), round-trip, nested
array/struct/map coverage, user-metadata preservation, no misfire on an
untagged struct with the same child names, and the default
`Interval(MONTH_DAY_NANO)` mapping staying unchanged when the flag is off.
- `ArrowWriterSuite` "calendar interval overflow raises DATETIME_OVERFLOW
at the conversion site": the default writer raises the structured condition for
`microseconds = Long.MaxValue / 1000 + 1`.
- `ArrowWriterSuite` "calendar interval lossless struct round-trip covers
the full value domain": write-and-read-back through `ArrowWriter` +
`ArrowColumnVector` for values including `Long.MaxValue` / `Long.MinValue`
microseconds and full-range months/days (all far outside the default mapping's
limit) plus nulls.
- `ArrowWriterSuite` "calendar interval lossless struct round-trip inside
nested types": the same extreme values inside `array<...>`, `struct<...>`, and
`map<int, ...>`.
Existing regression suites pass: `ArrowUtilsSuite`, `ArrowWriterSuite`,
`ArrowConvertersSuite`, `ColumnVectorSuite`, `ColumnarBatchSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
This pull request and its description were written by Claude Code.
Closes #57088 from viirya/interval-arrow-lossless.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
(cherry picked from commit 5ca6b1062887664b16b11f2bfcc015c9e616dc49)
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../org/apache/spark/sql/util/ArrowUtils.scala | 128 +++++++++++++++----
.../spark/sql/vectorized/ArrowColumnVector.java | 32 ++++-
.../catalyst/types/ops/TimestampNanosTypeOps.scala | 2 +-
.../spark/sql/errors/QueryExecutionErrors.scala | 14 ++-
.../spark/sql/execution/arrow/ArrowWriter.scala | 50 +++++++-
.../apache/spark/sql/util/ArrowUtilsSuite.scala | 139 ++++++++++++++++++++-
.../sql/execution/arrow/ArrowWriterSuite.scala | 138 +++++++++++++++++++-
7 files changed, 464 insertions(+), 39 deletions(-)
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 a0a14099d619..e69a0fa7f415 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
@@ -121,11 +121,16 @@ private[sql] object ArrowUtils {
// untouched) and recovered on read in `fromArrowField`.
private val timePrecisionKey = "SPARK::time::precision"
// Marks the epochMicros child of the lossless struct representation of a
nanosecond timestamp
- // (see `toArrowField` with `losslessTimestampNanos = true`). The value is
"ntz" or "ltz" and
+ // (see `toArrowField` with `losslessInternalTypes = true`). The value is
"ntz" or "ltz" and
// distinguishes TimestampNTZNanosType from TimestampLTZNanosType on read;
the precision is
// stored alongside under `timestampNanosPrecisionKey`. The tag lives on a
child field (like the
// geometry/variant struct tags) so it cannot collide with user metadata on
the struct itself.
private val timestampNanosStructKey = "SPARK::timestampNanos::struct"
+ // Marks the months child of the lossless struct representation of a
CalendarInterval (see
+ // `toArrowField` with `losslessInternalTypes = true`). The value is "true";
like
+ // `timestampNanosStructKey`, the tag lives on a child field so it cannot
collide with user
+ // metadata on the struct itself.
+ private val calendarIntervalStructKey = "SPARK::calendarInterval::struct"
private def toArrowMetaData(metadata: Metadata) = {
if (metadata != null && !metadata.isEmpty) {
Map(metadataKey -> metadata.json).asJava
@@ -181,7 +186,7 @@ private[sql] object ArrowUtils {
* - Internal storage (e.g. the Arrow-based Dataset cache) is a closed
write-then-read-back
* loop with no external consumer, where the only requirement is
fidelity to Spark
* semantics, hence this struct.
- * The choice is made per call site via `losslessTimestampNanos` on
`toArrowSchema` /
+ * The choice is made per call site via `losslessInternalTypes` on
`toArrowSchema` /
* `toArrowField` (the same pattern as `largeVarTypes`: one Spark type, two
Arrow encodings,
* selected by the consumer's needs). Only schema construction needs the
flag: the struct is
* self-describing through its child-field tag, so `fromArrowField`,
`ArrowWriter`, and
@@ -213,13 +218,47 @@ private[sql] object ArrowUtils {
new Field("nanosWithinMicro", nanosFieldType,
Seq.empty[Field].asJava)).asJava)
}
+ /**
+ * Builds the lossless Arrow struct representation of a CalendarInterval: a
struct of (months:
+ * int32, days: int32, microseconds: int64) -- the type's own field layout,
mirroring the
+ * default in-memory cache's CALENDAR_INTERVAL ColumnType. The default
Interval(MONTH_DAY_NANO)
+ * mapping multiplies microseconds by 1000 into Arrow's int64 nanosecond
field, so any
+ * |microseconds| > Long.MaxValue / 1000 (roughly +/-292 years) cannot be
represented; this
+ * struct stores the components as-is, so the full Long microsecond domain
round-trips. See
+ * `toTimestampNanosStructField` for why the default interchange mapping
must stay unchanged
+ * and the lossless shape is a per-call-site opt-in for internal storage.
+ */
+ private def toCalendarIntervalStructField(
+ name: String,
+ nullable: Boolean,
+ metadata: Metadata): Field = {
+ val fieldType =
+ new FieldType(nullable, ArrowType.Struct.INSTANCE, null,
toArrowMetaData(metadata))
+ // Tag the months child so `fromArrowField` (and ArrowColumnVector) can
recognize that this
+ // struct represents a CalendarInterval, following the geometry/variant
tag pattern.
+ val monthsFieldType = new FieldType(
+ false,
+ new ArrowType.Int(8 * 4, true),
+ null,
+ Map(calendarIntervalStructKey -> "true").asJava)
+ val daysFieldType = new FieldType(false, new ArrowType.Int(8 * 4, true),
null, null)
+ val microsFieldType = new FieldType(false, new ArrowType.Int(8 * 8, true),
null, null)
+ new Field(
+ name,
+ fieldType,
+ Seq(
+ new Field("months", monthsFieldType, Seq.empty[Field].asJava),
+ new Field("days", daysFieldType, Seq.empty[Field].asJava),
+ new Field("microseconds", microsFieldType,
Seq.empty[Field].asJava)).asJava)
+ }
+
/**
* Maps field from Spark to Arrow. NOTE: timeZoneId required for
TimestampType
*
- * @param losslessTimestampNanos
- * when true, nanosecond timestamps map to the lossless struct
representation covering their
- * full value domain instead of the standard int64 Timestamp(NANOSECOND)
encoding. Only
- * internal-storage callers with no external Arrow consumer (e.g. the
Arrow-based Dataset
+ * @param losslessInternalTypes
+ * when true, types whose standard Arrow encoding cannot cover their full
Spark value domain
+ * (nanosecond timestamps, CalendarInterval) map to lossless struct
representations instead.
+ * Only internal-storage callers with no external Arrow consumer (e.g. the
Arrow-based Dataset
* cache) should pass true; interchange paths must keep the default. See
* `toTimestampNanosStructField` for the full rationale.
*/
@@ -230,7 +269,7 @@ private[sql] object ArrowUtils {
timeZoneId: String,
largeVarTypes: Boolean = false,
metadata: Metadata = Metadata.empty,
- losslessTimestampNanos: Boolean = false): Field = {
+ losslessInternalTypes: Boolean = false): Field = {
dt match {
case ArrayType(elementType, containsNull) =>
val fieldType =
@@ -246,7 +285,7 @@ private[sql] object ArrowUtils {
timeZoneId,
largeVarTypes,
Metadata.empty,
- losslessTimestampNanos)).asJava)
+ losslessInternalTypes)).asJava)
case StructType(fields) =>
val fieldType =
new FieldType(nullable, ArrowType.Struct.INSTANCE, null,
toArrowMetaData(metadata))
@@ -262,7 +301,7 @@ private[sql] object ArrowUtils {
timeZoneId,
largeVarTypes,
field.metadata,
- losslessTimestampNanos)
+ losslessInternalTypes)
}
.toImmutableArraySeq
.asJava)
@@ -283,7 +322,7 @@ private[sql] object ArrowUtils {
timeZoneId,
largeVarTypes,
Metadata.empty,
- losslessTimestampNanos)).asJava)
+ losslessInternalTypes)).asJava)
case udt: UserDefinedType[_] =>
toArrowField(
name,
@@ -292,7 +331,7 @@ private[sql] object ArrowUtils {
timeZoneId,
largeVarTypes,
metadata,
- losslessTimestampNanos)
+ losslessInternalTypes)
case g: GeometryType =>
val fieldType =
new FieldType(nullable, ArrowType.Struct.INSTANCE, null,
toArrowMetaData(metadata))
@@ -346,9 +385,11 @@ private[sql] object ArrowUtils {
Seq(
toArrowField("value", BinaryType, false, timeZoneId,
largeVarTypes),
new Field("metadata", metadataFieldType,
Seq.empty[Field].asJava)).asJava)
- case t: TimestampNTZNanosType if losslessTimestampNanos =>
+ case CalendarIntervalType if losslessInternalTypes =>
+ toCalendarIntervalStructField(name, nullable, metadata)
+ case t: TimestampNTZNanosType if losslessInternalTypes =>
toTimestampNanosStructField(name, isNtz = true, t.precision, nullable,
metadata)
- case t: TimestampLTZNanosType if losslessTimestampNanos =>
+ case t: TimestampLTZNanosType if losslessInternalTypes =>
toTimestampNanosStructField(name, isNtz = false, t.precision,
nullable, metadata)
case t: TimestampNTZNanosType =>
toPrecisionTaggedArrowField(
@@ -420,20 +461,51 @@ private[sql] object ArrowUtils {
}
}
+ // Both lossless-struct recognizers below accept only the exact canonical
shape built by
+ // `toArrowField` (child count, order, types, and nullability), not merely
the presence of the
+ // tag and child names. The struct writers fill children positionally while
ArrowColumnVector's
+ // accessors read them by name, so a permissive match on, say, a tagged but
reordered schema
+ // would silently swap component values. Anything non-canonical falls back
to the generic
+ // struct handling, which is order-faithful.
+ private def isCanonicalStructChild(
+ child: Field,
+ name: String,
+ arrowType: ArrowType): Boolean = {
+ child.getName == name && child.getType == arrowType && !child.isNullable
+ }
+
/**
* Whether the Arrow struct field is the lossless representation of a
nanosecond timestamp built
- * by `toArrowField` with `losslessTimestampNanos = true`. Also callable
from Java
+ * by `toArrowField` with `losslessInternalTypes = true`. Also callable from
Java
* (ArrowColumnVector) to select the timestamp accessor for such structs.
*/
def isTimestampNanosStructField(field: Field): Boolean = {
- field.getType.isInstanceOf[ArrowType.Struct] &&
- field.getChildren.asScala
- .map(_.getName)
- .asJava
- .containsAll(Seq("epochMicros", "nanosWithinMicro").asJava) &&
- field.getChildren.asScala.exists { child =>
- child.getName == "epochMicros" &&
- Set("ntz",
"ltz").contains(child.getMetadata.getOrDefault(timestampNanosStructKey, ""))
+ field.getType.isInstanceOf[ArrowType.Struct] && {
+ val children = field.getChildren
+ children.size == 2 &&
+ isCanonicalStructChild(children.get(0), "epochMicros", new
ArrowType.Int(8 * 8, true)) &&
+ isCanonicalStructChild(
+ children.get(1),
+ "nanosWithinMicro",
+ new ArrowType.Int(8 * 2, true)) &&
+ Set("ntz", "ltz").contains(
+ children.get(0).getMetadata.getOrDefault(timestampNanosStructKey, ""))
+ }
+ }
+
+ /**
+ * Whether the Arrow struct field is the lossless representation of a
CalendarInterval built by
+ * `toArrowField` with `losslessInternalTypes = true`. Also callable from
Java
+ * (ArrowColumnVector) to select the interval accessor for such structs.
+ */
+ def isCalendarIntervalStructField(field: Field): Boolean = {
+ field.getType.isInstanceOf[ArrowType.Struct] && {
+ val children = field.getChildren
+ children.size == 3 &&
+ isCanonicalStructChild(children.get(0), "months", new ArrowType.Int(8 *
4, true)) &&
+ isCanonicalStructChild(children.get(1), "days", new ArrowType.Int(8 * 4,
true)) &&
+ isCanonicalStructChild(children.get(2), "microseconds", new
ArrowType.Int(8 * 8, true)) &&
+ children.get(0).getMetadata.getOrDefault(calendarIntervalStructKey,
"false") == "true"
}
}
@@ -448,6 +520,8 @@ private[sql] object ArrowUtils {
val elementField = field.getChildren().get(0)
val elementType = fromArrowField(elementField)
ArrayType(elementType, containsNull = elementField.isNullable)
+ case ArrowType.Struct.INSTANCE if isCalendarIntervalStructField(field) =>
+ CalendarIntervalType
case ArrowType.Struct.INSTANCE if isTimestampNanosStructField(field) =>
val microsChild = field.getChildren.asScala.find(_.getName ==
"epochMicros").get
val isNtz = microsChild.getMetadata.get(timestampNanosStructKey) ==
"ntz"
@@ -519,16 +593,16 @@ private[sql] object ArrowUtils {
/**
* Maps schema from Spark to Arrow. NOTE: timeZoneId required for
TimestampType in StructType
*
- * @param losslessTimestampNanos
- * see `toArrowField`: opt-in full-domain struct encoding of nanosecond
timestamps for
- * internal storage; interchange paths must keep the default.
+ * @param losslessInternalTypes
+ * see `toArrowField`: opt-in full-domain struct encoding of nanosecond
timestamps and
+ * CalendarInterval for internal storage; interchange paths must keep the
default.
*/
def toArrowSchema(
schema: StructType,
timeZoneId: String,
errorOnDuplicatedFieldNames: Boolean,
largeVarTypes: Boolean,
- losslessTimestampNanos: Boolean = false): Schema = {
+ losslessInternalTypes: Boolean = false): Schema = {
new Schema(schema.map { field =>
toArrowField(
field.name,
@@ -537,7 +611,7 @@ private[sql] object ArrowUtils {
timeZoneId,
largeVarTypes,
field.metadata,
- losslessTimestampNanos)
+ losslessInternalTypes)
}.asJava)
}
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 1ae653d6b725..f44de5ffa9df 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
@@ -229,8 +229,11 @@ public class ArrowColumnVector extends ColumnVector {
} else if (vector instanceof StructVector structVector) {
if (ArrowUtils.isTimestampNanosStructField(structVector.getField())) {
// Lossless struct representation of a nanosecond timestamp
(ArrowUtils.toArrowField with
- // losslessTimestampNanos = true): logically a scalar, so no child
columns are exposed.
+ // losslessInternalTypes = true): logically a scalar, so no child
columns are exposed.
accessor = new TimestampNanosStructAccessor(structVector);
+ } else if
(ArrowUtils.isCalendarIntervalStructField(structVector.getField())) {
+ // Lossless struct representation of a CalendarInterval: also
logically a scalar.
+ accessor = new CalendarIntervalStructAccessor(structVector);
} else {
accessor = new StructAccessor(structVector);
@@ -627,7 +630,7 @@ public class ArrowColumnVector extends ColumnVector {
/**
* Reads the lossless struct representation of a nanosecond timestamp
(epochMicros: int64,
- * nanosWithinMicro: int16), built by ArrowUtils.toArrowField with
losslessTimestampNanos = true.
+ * nanosWithinMicro: int16), built by ArrowUtils.toArrowField with
losslessInternalTypes = true.
* The components are stored as-is (TimestampNanosVal's own layout), so
unlike the int64
* epoch-nanoseconds accessors above there is no decoding arithmetic and no
reduced value domain.
*/
@@ -651,6 +654,31 @@ public class ArrowColumnVector extends ColumnVector {
}
}
+ /**
+ * Reads the lossless struct representation of a CalendarInterval (months:
int32, days: int32,
+ * microseconds: int64), built by ArrowUtils.toArrowField with
losslessInternalTypes = true.
+ * The components are stored as-is, so unlike IntervalMonthDayNanoAccessor
there is no unit
+ * conversion and no reduced value domain.
+ */
+ static class CalendarIntervalStructAccessor extends ArrowVectorAccessor {
+
+ private final IntVector months;
+ private final IntVector days;
+ private final BigIntVector microseconds;
+
+ CalendarIntervalStructAccessor(StructVector vector) {
+ super(vector);
+ this.months = (IntVector) vector.getChild("months");
+ this.days = (IntVector) vector.getChild("days");
+ this.microseconds = (BigIntVector) vector.getChild("microseconds");
+ }
+
+ @Override
+ final CalendarInterval getInterval(int rowId) {
+ return new CalendarInterval(months.get(rowId), days.get(rowId),
microseconds.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 4697c85c660f..8c14a4d6600e 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
@@ -129,7 +129,7 @@ case class TimestampNTZNanosTypeOps(override val t:
TimestampNTZNanosType)
vector match {
case v: TimeStampNanoVector => Some(new TimestampNTZNanosWriter(v))
// The lossless struct representation (ArrowUtils.toArrowField with
- // losslessTimestampNanos = true) is backed by a StructVector; its
writer needs the
+ // losslessInternalTypes = true) is backed by a StructVector; its writer
needs the
// recursively-built child writers, so defer to ArrowWriter's default
matching.
case _ => None
}
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 0c5fc0cf1046..676f3aff4408 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
@@ -54,7 +54,7 @@ import
org.apache.spark.sql.internal.StaticSQLConf.GLOBAL_TEMP_DATABASE
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.array.ByteArrayMethods
-import org.apache.spark.unsafe.types.{TimestampNanosVal, UTF8String}
+import org.apache.spark.unsafe.types.{CalendarInterval, TimestampNanosVal,
UTF8String}
import org.apache.spark.util.{CircularBuffer, Utils}
/**
@@ -2611,6 +2611,18 @@ private[sql] object QueryExecutionErrors extends
QueryErrorsBase with ExecutionE
summary = "")
}
+ def calendarIntervalArrowNanosOverflowError(
+ interval: CalendarInterval): SparkArithmeticException = {
+ new SparkArithmeticException(
+ errorClass = "DATETIME_OVERFLOW",
+ messageParameters = Map(
+ "operation" -> (s"write the interval value $interval as Arrow
IntervalMonthDayNano " +
+ "nanoseconds (the microseconds component must be in
+/-(Long.MaxValue / 1000), " +
+ "roughly +/-292 years)")),
+ context = Array.empty,
+ summary = "")
+ }
+
def timestampNanosEpochNanosOverflowError(
value: TimestampNanosVal, isNtz: Boolean, sink: String):
SparkArithmeticException = {
// Render TIMESTAMP_NTZ values without a zone (LocalDateTime, no trailing
`Z`); TIMESTAMP_LTZ
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 47e5d10926c5..b96e57ce49af 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
@@ -96,8 +96,15 @@ object ArrowWriter {
case (_: DayTimeIntervalType, vector: DurationVector) => new
DurationWriter(vector)
case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) =>
new IntervalMonthDayNanoWriter(vector)
+ // Lossless struct representation of CalendarInterval
(ArrowUtils.toArrowField with
+ // losslessInternalTypes = true).
+ case (CalendarIntervalType, vector: StructVector) =>
+ val children = (0 until vector.size()).map { ordinal =>
+ createFieldWriter(vector.getChildByOrdinal(ordinal))
+ }
+ new CalendarIntervalStructWriter(vector, children.toArray)
// Lossless struct representation of nanosecond timestamps
(ArrowUtils.toArrowField with
- // losslessTimestampNanos = true). The native TimeStampNano(TZ)Vector
writers are created by
+ // losslessInternalTypes = true). The native TimeStampNano(TZ)Vector
writers are created by
// the TypeOps hook; only the struct-backed shape reaches this default
matching.
case (_: TimestampNTZNanosType, vector: StructVector) =>
val children = (0 until vector.size()).map { ordinal =>
@@ -557,7 +564,7 @@ private[arrow] class GeometryWriter(
/**
* Writes a nanosecond timestamp into its lossless Arrow struct representation
* (epochMicros: int64, nanosWithinMicro: int16), built by
`ArrowUtils.toArrowField` with
- * `losslessTimestampNanos = true`. The two components of TimestampNanosVal
are stored as-is with
+ * `losslessInternalTypes = true`. The two components of TimestampNanosVal are
stored as-is with
* no unit conversion, so unlike the Timestamp(NANOSECOND) writers there is no
overflow: the full
* domain of the Spark types (years 0001-9999) round-trips.
*/
@@ -596,6 +603,32 @@ private[arrow] class TimestampLTZNanosStructWriter(
input.getTimestampLTZNanos(ordinal)
}
+/**
+ * Writes a CalendarInterval into its lossless Arrow struct representation
+ * (months: int32, days: int32, microseconds: int64), built by
`ArrowUtils.toArrowField` with
+ * `losslessInternalTypes = true`. The three components are stored as-is with
no unit conversion,
+ * so unlike IntervalMonthDayNanoWriter there is no nanosecond multiplication
and no overflow:
+ * the full Long microsecond domain round-trips.
+ */
+private[arrow] class CalendarIntervalStructWriter(
+ valueVector: StructVector,
+ children: Array[ArrowFieldWriter]) extends StructWriter(valueVector,
children) {
+
+ // Reused across rows; this writer is single-threaded like the vector it
wraps.
+ private val row = new GenericInternalRow(3)
+
+ override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
+ valueVector.setIndexDefined(count)
+ val ci = input.getInterval(ordinal)
+ row.update(0, ci.months)
+ row.update(1, ci.days)
+ row.update(2, ci.microseconds)
+ children(0).write(row, 0)
+ children(1).write(row, 1)
+ children(2).write(row, 2)
+ }
+}
+
private[arrow] class MapWriter(
val valueVector: MapVector,
val structVector: StructVector,
@@ -672,6 +705,17 @@ private[arrow] class IntervalMonthDayNanoWriter(val
valueVector: IntervalMonthDa
override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
val ci = input.getInterval(ordinal)
- valueVector.setSafe(count, ci.months, ci.days,
Math.multiplyExact(ci.microseconds, 1000L))
+ // Arrow's IntervalMonthDayNano stores the sub-day component as int64
nanoseconds, so the
+ // conversion overflows for |microseconds| > Long.MaxValue / 1000.
Translate that into the
+ // structured DATETIME_OVERFLOW at the conversion site (like the
nanosecond timestamp writers
+ // above) so the raw ArithmeticException never escapes -- catching it any
wider risks
+ // re-labeling unrelated arithmetic failures raised by lazily-evaluated
upstream input.
+ val nanos = try {
+ Math.multiplyExact(ci.microseconds, 1000L)
+ } catch {
+ case _: ArithmeticException =>
+ throw QueryExecutionErrors.calendarIntervalArrowNanosOverflowError(ci)
+ }
+ valueVector.setSafe(count, ci.months, ci.days, nanos)
}
}
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 22611341cd37..d60e32896d5c 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,9 @@ package org.apache.spark.sql.util
import java.time.ZoneId
-import org.apache.arrow.vector.types.TimeUnit
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.types.{IntervalUnit, TimeUnit}
import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
import org.apache.spark.{SparkException, SparkFunSuite,
SparkUnsupportedOperationException}
@@ -157,7 +159,7 @@ class ArrowUtilsSuite extends SparkFunSuite {
test("timestamp nanos lossless struct") {
def losslessRoundtrip(schema: StructType, timeZoneId: String = null): Unit
= {
val arrowSchema =
- ArrowUtils.toArrowSchema(schema, timeZoneId, true, false,
losslessTimestampNanos = true)
+ ArrowUtils.toArrowSchema(schema, timeZoneId, true, false,
losslessInternalTypes = true)
assert(ArrowUtils.fromArrowSchema(arrowSchema) === schema)
}
@@ -167,7 +169,7 @@ class ArrowUtilsSuite extends SparkFunSuite {
Seq[DataType](TimestampNTZNanosType(p),
TimestampLTZNanosType(p)).foreach { dt =>
val schema = new StructType().add("value", dt)
val arrowSchema =
- ArrowUtils.toArrowSchema(schema, null, true, false,
losslessTimestampNanos = true)
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
val field = arrowSchema.findField("value")
assert(field.getType === ArrowType.Struct.INSTANCE)
val children = field.getChildren
@@ -219,6 +221,40 @@ class ArrowUtilsSuite extends SparkFunSuite {
assert(ArrowUtils.fromArrowField(taggedStructField(Some("5"))) ===
TimestampNTZNanosType(9))
assert(ArrowUtils.fromArrowField(taggedStructField(None)) ===
TimestampNTZNanosType(9))
+ // Only the exact canonical shape is recognized: the struct writer fills
children
+ // positionally while ArrowColumnVector reads them by name, so a
tagged-but-non-canonical
+ // schema (reordered, wrong width, or extra children) must NOT be treated
as a nanosecond
+ // timestamp -- it falls back to generic (order-faithful) struct handling.
+ def taggedNanosStruct(children: Seq[(String, Int)]): Field = {
+ val fields = children.map { case (name, bitWidth) =>
+ val md = if (name == "epochMicros") {
+ java.util.Collections.singletonMap("SPARK::timestampNanos::struct",
"ntz")
+ } else {
+ null
+ }
+ new Field(
+ name,
+ new FieldType(false, new ArrowType.Int(bitWidth, true), null, md),
+ java.util.Collections.emptyList[Field]())
+ }
+ new Field(
+ "value",
+ new FieldType(true, ArrowType.Struct.INSTANCE, null, null),
+ fields.asJava)
+ }
+ // Reordered children.
+ assert(!ArrowUtils.isTimestampNanosStructField(
+ taggedNanosStruct(Seq("nanosWithinMicro" -> 16, "epochMicros" -> 64))))
+ // Wrong child width.
+ assert(!ArrowUtils.isTimestampNanosStructField(
+ taggedNanosStruct(Seq("epochMicros" -> 64, "nanosWithinMicro" -> 32))))
+ // Extra child.
+ assert(!ArrowUtils.isTimestampNanosStructField(
+ taggedNanosStruct(Seq("epochMicros" -> 64, "nanosWithinMicro" -> 16,
"extra" -> 32))))
+ // The canonical shape built by the same helper is recognized (sanity
check of the helper).
+ assert(ArrowUtils.isTimestampNanosStructField(
+ taggedNanosStruct(Seq("epochMicros" -> 64, "nanosWithinMicro" -> 16))))
+
// A plain struct that merely uses the same child names, but carries no
tag, stays a struct.
val untagged = new StructType().add(
"value",
@@ -236,6 +272,103 @@ class ArrowUtilsSuite extends SparkFunSuite {
assert(defaultSchema.findField("value").getType.isInstanceOf[ArrowType.Timestamp])
}
+ test("calendar interval lossless struct") {
+ def losslessRoundtrip(schema: StructType): Unit = {
+ val arrowSchema =
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
+ assert(ArrowUtils.fromArrowSchema(arrowSchema) === schema)
+ }
+
+ // Top-level: the lossless mapping is a struct of the type's own components
+ // (months: int32, days: int32, microseconds: int64), tagged through child
field metadata.
+ val schema = new StructType().add("value", CalendarIntervalType)
+ val arrowSchema =
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
+ val field = arrowSchema.findField("value")
+ assert(field.getType === ArrowType.Struct.INSTANCE)
+ val children = field.getChildren
+ assert(children.size() === 3)
+ assert(children.get(0).getName === "months")
+ assert(children.get(0).getType === new ArrowType.Int(32, true))
+ assert(!children.get(0).isNullable)
+ assert(children.get(1).getName === "days")
+ assert(children.get(1).getType === new ArrowType.Int(32, true))
+ assert(!children.get(1).isNullable)
+ assert(children.get(2).getName === "microseconds")
+ assert(children.get(2).getType === new ArrowType.Int(64, true))
+ assert(!children.get(2).isNullable)
+ assert(ArrowUtils.isCalendarIntervalStructField(field))
+ assert(ArrowUtils.fromArrowSchema(arrowSchema) === schema)
+
+ // Nested: the flag must reach intervals inside arrays, structs, and maps.
+ losslessRoundtrip(new StructType()
+ .add("arr", ArrayType(CalendarIntervalType))
+ .add("struct", new StructType().add("i", CalendarIntervalType))
+ .add("map", MapType(IntegerType, CalendarIntervalType)))
+
+ // User metadata on the column is preserved alongside the struct tag.
+ val md = new MetadataBuilder().putString("city", "beijing").build()
+ losslessRoundtrip(new StructType().add("value", CalendarIntervalType,
true, md))
+
+ // A plain struct that merely uses the same child names, but carries no
tag, stays a struct.
+ val untagged = new StructType().add(
+ "value",
+ new StructType()
+ .add("months", IntegerType, nullable = false)
+ .add("days", IntegerType, nullable = false)
+ .add("microseconds", LongType, nullable = false))
+ losslessRoundtrip(untagged)
+ assert(
+ ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema(untagged, null,
true, false)) ===
+ untagged)
+
+ // Only the exact canonical shape is recognized:
CalendarIntervalStructWriter fills children
+ // positionally while ArrowColumnVector reads them by name, so a
tagged-but-reordered schema
+ // (which would silently swap months and days) or other non-canonical
shapes must NOT be
+ // treated as a CalendarInterval.
+ def taggedIntervalStruct(children: Seq[(String, Int)]): Field = {
+ val fields = children.map { case (name, bitWidth) =>
+ val md = if (name == "months") {
+
java.util.Collections.singletonMap("SPARK::calendarInterval::struct", "true")
+ } else {
+ null
+ }
+ new Field(
+ name,
+ new FieldType(false, new ArrowType.Int(bitWidth, true), null, md),
+ java.util.Collections.emptyList[Field]())
+ }
+ new Field(
+ "value",
+ new FieldType(true, ArrowType.Struct.INSTANCE, null, null),
+ fields.asJava)
+ }
+ // Reordered children.
+ assert(!ArrowUtils.isCalendarIntervalStructField(
+ taggedIntervalStruct(Seq("days" -> 32, "months" -> 32, "microseconds" ->
64))))
+ // Wrong child width.
+ assert(!ArrowUtils.isCalendarIntervalStructField(
+ taggedIntervalStruct(Seq("months" -> 32, "days" -> 64, "microseconds" ->
64))))
+ // Extra child.
+ assert(!ArrowUtils.isCalendarIntervalStructField(
+ taggedIntervalStruct(
+ Seq("months" -> 32, "days" -> 32, "microseconds" -> 64, "extra" ->
32))))
+ // Missing child.
+ assert(!ArrowUtils.isCalendarIntervalStructField(
+ taggedIntervalStruct(Seq("months" -> 32, "days" -> 32))))
+ // The canonical shape built by the same helper is recognized (sanity
check of the helper).
+ assert(ArrowUtils.isCalendarIntervalStructField(
+ taggedIntervalStruct(Seq("months" -> 32, "days" -> 32, "microseconds" ->
64))))
+
+ // The default mapping is untouched when the flag is off: still
IntervalMonthDayNano.
+ val defaultSchema = ArrowUtils.toArrowSchema(
+ new StructType().add("value", CalendarIntervalType), null, true, false)
+ val defaultType = defaultSchema.findField("value").getType
+ assert(defaultType.isInstanceOf[ArrowType.Interval])
+ assert(
+ defaultType.asInstanceOf[ArrowType.Interval].getUnit ===
IntervalUnit.MONTH_DAY_NANO)
+ }
+
test("time") {
// Arrow's Time type has no precision field, so TIME(p) precision is
preserved via field
// metadata; the Arrow type itself stays Time(NANOSECOND, 64).
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 89caaa00e922..2584193008ff 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
@@ -27,6 +27,7 @@ import org.apache.spark.sql.YearUDT
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.catalyst.encoders.RowEncoder.{encoderFor =>
toRowEncoder}
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
import org.apache.spark.sql.catalyst.util._
import org.apache.spark.sql.catalyst.util.{Geography => InternalGeography,
Geometry => InternalGeometry}
import org.apache.spark.sql.types._
@@ -226,7 +227,7 @@ class ArrowWriterSuite extends SparkFunSuite {
// 0001-9999) must round-trip -- including values that overflow the
default mapping.
def losslessWriter(schema: StructType): ArrowWriter = {
val arrowSchema =
- ArrowUtils.toArrowSchema(schema, null, true, false,
losslessTimestampNanos = true)
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
val root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.rootAllocator)
ArrowWriter.create(root)
}
@@ -273,7 +274,7 @@ class ArrowWriterSuite extends SparkFunSuite {
def losslessWriter(schema: StructType): ArrowWriter = {
val arrowSchema =
- ArrowUtils.toArrowSchema(schema, null, true, false,
losslessTimestampNanos = true)
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
val root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.rootAllocator)
ArrowWriter.create(root)
}
@@ -323,6 +324,139 @@ class ArrowWriterSuite extends SparkFunSuite {
}
}
+ test("calendar interval overflow raises DATETIME_OVERFLOW at the conversion
site") {
+ // The default IntervalMonthDayNano mapping multiplies microseconds by
1000 into Arrow's
+ // int64 nanosecond field. The overflow must surface as the structured
DATETIME_OVERFLOW
+ // (not a raw ArithmeticException), and the translation must be scoped to
the conversion:
+ // an unrelated (Spark)ArithmeticException raised by upstream evaluation
must pass through
+ // unchanged, which the writer guarantees by catching only around
Math.multiplyExact.
+ val schema = new StructType().add("value", CalendarIntervalType, nullable
= true)
+ val writer = ArrowWriter.create(schema, "UTC")
+ // A normal interval still writes fine.
+ writer.write(InternalRow(new CalendarInterval(1, 2, 3000000L)))
+ val tooLarge = new CalendarInterval(0, 0, Long.MaxValue / 1000L + 1L)
+ val e = intercept[SparkArithmeticException] {
+ writer.write(InternalRow(tooLarge))
+ }
+ assert(e.getCondition === "DATETIME_OVERFLOW")
+ assert(e.getMessage.contains("IntervalMonthDayNano"))
+
+ // The other half of the invariant: an arithmetic exception raised by
upstream evaluation
+ // (here, a lazily-throwing getInterval standing in for e.g. an ANSI
DIVIDE_BY_ZERO from a
+ // WholeStageCodegen iterator) must escape unchanged -- the same instance,
not relabeled as
+ // DATETIME_OVERFLOW. SparkArithmeticException extends
ArithmeticException, so this pins the
+ // catch to the Math.multiplyExact expression: a future refactor that
widened the try to
+ // cover input.getInterval would fail here.
+ val upstreamError = new SparkArithmeticException(
+ errorClass = "DIVIDE_BY_ZERO",
+ messageParameters = Map("config" -> "spark.sql.ansi.enabled"),
+ context = Array.empty,
+ summary = "")
+ val throwingRow = new GenericInternalRow(Array[Any](new
CalendarInterval(0, 0, 0L))) {
+ override def getInterval(ordinal: Int): CalendarInterval = throw
upstreamError
+ }
+ val escaped = intercept[SparkArithmeticException] {
+ writer.write(throwingRow)
+ }
+ assert(escaped eq upstreamError,
+ "the upstream exception must escape unchanged, not be wrapped or
relabeled")
+ assert(escaped.getCondition === "DIVIDE_BY_ZERO")
+ writer.root.close()
+ }
+
+ test("calendar interval lossless struct round-trip covers the full value
domain") {
+ // The default IntervalMonthDayNano mapping only covers |microseconds| <=
Long.MaxValue / 1000
+ // (see the DATETIME_OVERFLOW test above). The lossless struct
representation stores the raw
+ // (months, days, microseconds) components, so the full Long microsecond
domain must
+ // round-trip -- including values that overflow the default mapping.
+ def losslessWriter(schema: StructType): ArrowWriter = {
+ val arrowSchema =
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
+ val root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.rootAllocator)
+ ArrowWriter.create(root)
+ }
+
+ val values = Seq(
+ new CalendarInterval(0, 0, 0L),
+ new CalendarInterval(1, 2, 3000000L),
+ new CalendarInterval(-1, -2, -3000000L),
+ // beyond the default mapping's +/-(Long.MaxValue / 1000)
nanosecond-conversion limit
+ new CalendarInterval(0, 0, Long.MaxValue / 1000L + 1L),
+ new CalendarInterval(0, 0, Long.MaxValue),
+ new CalendarInterval(0, 0, Long.MinValue),
+ new CalendarInterval(Int.MaxValue, Int.MaxValue, Long.MaxValue),
+ new CalendarInterval(Int.MinValue, Int.MinValue, Long.MinValue))
+
+ val schema = new StructType().add("value", CalendarIntervalType, nullable
= true)
+ val writer = losslessWriter(schema)
+ (values.map(Option(_)) :+ None).foreach(v =>
writer.write(InternalRow(v.orNull)))
+ writer.finish()
+
+ val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
+ assert(reader.dataType() === CalendarIntervalType)
+ values.zipWithIndex.foreach { case (v, rowId) =>
+ assert(reader.getInterval(rowId) === v)
+ }
+ assert(reader.isNullAt(values.length))
+ writer.root.close()
+ }
+
+ test("calendar interval lossless struct round-trip inside nested types") {
+ val v1 = new CalendarInterval(0, 0, Long.MaxValue)
+ val v2 = new CalendarInterval(Int.MinValue, Int.MinValue, Long.MinValue)
+
+ def losslessWriter(schema: StructType): ArrowWriter = {
+ val arrowSchema =
+ ArrowUtils.toArrowSchema(schema, null, true, false,
losslessInternalTypes = true)
+ val root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.rootAllocator)
+ ArrowWriter.create(root)
+ }
+
+ // array<interval>
+ {
+ val schema = new StructType().add("arr", ArrayType(CalendarIntervalType))
+ val writer = losslessWriter(schema)
+ writer.write(InternalRow(new GenericArrayData(Array[Any](v1, null, v2))))
+ writer.finish()
+ val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
+ val arr = reader.getArray(0)
+ assert(arr.numElements() === 3)
+ assert(arr.getInterval(0) === v1)
+ assert(arr.isNullAt(1))
+ assert(arr.getInterval(2) === v2)
+ writer.root.close()
+ }
+
+ // struct<i: interval>
+ {
+ val schema = new StructType()
+ .add("struct", new StructType().add("i", CalendarIntervalType))
+ val writer = losslessWriter(schema)
+ writer.write(InternalRow(InternalRow(v1)))
+ writer.finish()
+ val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
+ assert(reader.getStruct(0).getInterval(0) === v1)
+ writer.root.close()
+ }
+
+ // map<int, interval>
+ {
+ val schema = new StructType().add("map", MapType(IntegerType,
CalendarIntervalType))
+ val writer = losslessWriter(schema)
+ writer.write(InternalRow(
+ new ArrayBasedMapData(
+ new GenericArrayData(Array[Any](1, 2)),
+ new GenericArrayData(Array[Any](v1, v2)))))
+ writer.finish()
+ val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
+ val map = reader.getMap(0)
+ assert(map.numElements() === 2)
+ assert(map.valueArray().getInterval(0) === v1)
+ assert(map.valueArray().getInterval(1) === v2)
+ writer.root.close()
+ }
+ }
+
test("nested geographies") {
def check(
dt: StructType,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
