From ac5925febe0df7f31e5e198ddbe10ba0c86fc18e Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 12 Jul 2026 20:15:05 +0800 Subject: [PATCH] Add CalendarIntervalType Arrow support --- native/core/src/execution/serde.rs | 3 + native/proto/src/proto/types.proto | 1 + .../apache/comet/vector/CometPlainVector.java | 75 +++++++++++++++++++ .../org/apache/comet/DataTypeSupport.scala | 4 +- .../codegen/CometBatchKernelCodegen.scala | 3 +- .../CometBatchKernelCodegenInput.scala | 29 ++++++- .../CometBatchKernelCodegenOutput.scala | 10 +++ .../CometSpecializedGettersDispatch.scala | 1 + .../apache/comet/serde/QueryPlanSerde.scala | 3 +- .../comet/execution/arrow/ArrowWriters.scala | 4 +- .../apache/spark/sql/comet/util/Utils.scala | 3 + .../datetime/calendar_interval.sql | 26 +++++++ .../org/apache/comet/CometCodegenSuite.scala | 35 ++++++++- .../arrow/CometArrowStreamSuite.scala | 33 +++++++- 14 files changed, 221 insertions(+), 9 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql diff --git a/native/core/src/execution/serde.rs b/native/core/src/execution/serde.rs index 9957ed8155..f31cb9cd35 100644 --- a/native/core/src/execution/serde.rs +++ b/native/core/src/execution/serde.rs @@ -102,6 +102,9 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType { // Spark's DayTimeIntervalType stores microseconds in an int64, which matches Arrow // Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout. DataTypeId::DayTimeInterval => ArrowDataType::Duration(TimeUnit::Microsecond), + // Spark's CalendarIntervalType stores months, days, and microseconds. Arrow stores the + // same components with nanosecond precision. + DataTypeId::CalendarInterval => ArrowDataType::Interval(IntervalUnit::MonthDayNano), DataTypeId::Null => ArrowDataType::Null, DataTypeId::List => match dt_value .type_info diff --git a/native/proto/src/proto/types.proto b/native/proto/src/proto/types.proto index 6892af8b42..114253ee77 100644 --- a/native/proto/src/proto/types.proto +++ b/native/proto/src/proto/types.proto @@ -62,6 +62,7 @@ message DataType { TIME = 17; YEAR_MONTH_INTERVAL = 18; DAY_TIME_INTERVAL = 19; + CALENDAR_INTERVAL = 20; } DataTypeId type_id = 1; diff --git a/spark/src/main/java/org/apache/comet/vector/CometPlainVector.java b/spark/src/main/java/org/apache/comet/vector/CometPlainVector.java index 8fbef7a490..2e81d5fa5b 100644 --- a/spark/src/main/java/org/apache/comet/vector/CometPlainVector.java +++ b/spark/src/main/java/org/apache/comet/vector/CometPlainVector.java @@ -27,6 +27,8 @@ import org.apache.arrow.vector.*; import org.apache.arrow.vector.util.TransferPair; import org.apache.parquet.Preconditions; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.vectorized.ColumnVector; import org.apache.spark.unsafe.Platform; import org.apache.spark.unsafe.types.UTF8String; @@ -41,6 +43,7 @@ public class CometPlainVector extends CometDecodedVector { private final long valueBufferAddress; private final long offsetBufferAddress; private final boolean isBaseFixedWidthVector; + private final ColumnVector[] intervalChildren; // True when the variable-width offsets are 64-bit (LargeVarChar / LargeVarBinary) rather than // the usual 32-bit. PyArrow UDFs can hand back large_string / large_binary columns. private final boolean isLargeVarWidth; @@ -73,6 +76,16 @@ public CometPlainVector(ValueVector vector, boolean isUuid) { this.offsetBufferAddress = -1; this.isLargeVarWidth = false; } + if (vector instanceof IntervalMonthDayNanoVector) { + this.intervalChildren = + new ColumnVector[] { + new IntervalChildVector(this, 0), + new IntervalChildVector(this, 1), + new IntervalChildVector(this, 2) + }; + } else { + this.intervalChildren = null; + } } @Override @@ -178,6 +191,14 @@ public byte[] getBinary(int rowId) { return result; } + @Override + public ColumnVector getChild(int ordinal) { + if (intervalChildren == null || ordinal < 0 || ordinal >= intervalChildren.length) { + return super.getChild(ordinal); + } + return intervalChildren[ordinal]; + } + @Override public CDataDictionaryProvider getDictionaryProvider() { return null; @@ -204,4 +225,58 @@ private static UUID convertToUuid(byte[] buf) { long leastSigBits = bb.getLong(); return new UUID(mostSigBits, leastSigBits); } + + private static final class IntervalChildVector extends CometVector { + private final CometPlainVector parent; + private final int ordinal; + + private IntervalChildVector(CometPlainVector parent, int ordinal) { + super(ordinal == 2 ? DataTypes.LongType : DataTypes.IntegerType); + this.parent = parent; + this.ordinal = ordinal; + } + + @Override + public int getInt(int rowId) { + return Platform.getInt(null, parent.valueBufferAddress + rowId * 16L + ordinal * 4L); + } + + @Override + public long getLong(int rowId) { + return Platform.getLong(null, parent.valueBufferAddress + rowId * 16L + 8L) / 1000L; + } + + @Override + public boolean isNullAt(int rowId) { + return parent.isNullAt(rowId); + } + + @Override + public boolean hasNull() { + return parent.hasNull(); + } + + @Override + public int numNulls() { + return parent.numNulls(); + } + + @Override + public int numValues() { + return parent.numValues(); + } + + @Override + public ValueVector getValueVector() { + return parent.getValueVector(); + } + + @Override + public CometVector slice(int offset, int length) { + throw new UnsupportedOperationException("Interval child vectors cannot be sliced"); + } + + @Override + public void close() {} + } } diff --git a/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala b/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala index 9f8fc77eba..019b4d629b 100644 --- a/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala +++ b/spark/src/main/scala/org/apache/comet/DataTypeSupport.scala @@ -50,8 +50,8 @@ trait DataTypeSupport { dt match { case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | - BinaryType | StringType | _: DecimalType | DateType | TimestampType | - TimestampNTZType => + BinaryType | StringType | _: DecimalType | DateType | TimestampType | TimestampNTZType | + CalendarIntervalType => true case StructType(fields) => fields.nonEmpty && fields.forall(f => diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 107fb5e7f8..1df1fe5f85 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -73,6 +73,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim { case "VarBinaryVector" => classOf[VarBinaryVector] case "IntervalYearVector" => classOf[IntervalYearVector] case "DurationVector" => classOf[DurationVector] + case "IntervalMonthDayNanoVector" => classOf[IntervalMonthDayNanoVector] case other => throw new IllegalArgumentException(s"unknown Arrow vector class: $other") } @@ -86,7 +87,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim { case _: DecimalType => true case _: StringType | _: BinaryType => true case DateType | TimestampType | TimestampNTZType => true - case _: YearMonthIntervalType | _: DayTimeIntervalType => true + case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true case ArrayType(inner, _) => isSupportedDataType(inner) case st: StructType => st.fields.forall(f => isSupportedDataType(f.dataType)) case mt: MapType => isSupportedDataType(mt.keyType) && isSupportedDataType(mt.valueType) diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala index 09bfc52bd4..64ecb8bf3c 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala @@ -62,7 +62,8 @@ private[codegen] object CometBatchKernelCodegenInput { classOf[Float8Vector], classOf[DateDayVector], classOf[TimeStampMicroVector], - classOf[TimeStampMicroTZVector]) + classOf[TimeStampMicroTZVector], + classOf[IntervalMonthDayNanoVector]) private val cometPlainVectorName: String = classOf[CometPlainVector].getName /** Emit kernel typed-vector field declarations for every level of every input column. */ @@ -137,6 +138,10 @@ private[codegen] object CometBatchKernelCodegenInput { cls == classOf[TimeStampMicroTZVector] => s" case $ord: return this.col$ord.getLong(this.rowIdx);" } + val intervalCases = withOrd.collect { + case (ArrowColumnSpec(cls, _), ord) if cls == classOf[IntervalMonthDayNanoVector] => + s" case $ord: return this.col$ord.getInterval(this.rowIdx);" + } val floatCases = withOrd.collect { case (ArrowColumnSpec(cls, _), ord) if cls == classOf[Float4Vector] => s" case $ord: return this.col$ord.getFloat(this.rowIdx);" @@ -196,6 +201,10 @@ private[codegen] object CometBatchKernelCodegenInput { emitOrdinalSwitch("public long getLong(int ordinal)", "getLong", longCases), emitOrdinalSwitch("public float getFloat(int ordinal)", "getFloat", floatCases), emitOrdinalSwitch("public double getDouble(int ordinal)", "getDouble", doubleCases), + emitOrdinalSwitch( + "public org.apache.spark.unsafe.types.CalendarInterval getInterval(int ordinal)", + "getInterval", + intervalCases), emitOrdinalSwitch( "public org.apache.spark.sql.types.Decimal getDecimal(" + "int ordinal, int precision, int scale)", @@ -592,6 +601,7 @@ private[codegen] object CometBatchKernelCodegenInput { case ShortType => s"getShort($idx)" case IntegerType | DateType => s"getInt($idx)" case LongType | TimestampType | TimestampNTZType => s"getLong($idx)" + case CalendarIntervalType => s"getInterval($idx)" case FloatType => s"getFloat($idx)" case DoubleType => s"getDouble($idx)" case d: DecimalType => s"getDecimal($idx, ${d.precision}, ${d.scale})" @@ -697,6 +707,11 @@ private[codegen] object CometBatchKernelCodegenInput { | public long getLong(int i) { | return $childField.getLong(startIndex + i); | }""".stripMargin + case CalendarIntervalType => + s""" @Override + | public org.apache.spark.unsafe.types.CalendarInterval getInterval(int i) { + |$nullGuard return $childField.getInterval(startIndex + i); + | }""".stripMargin case FloatType => s""" @Override | public float getFloat(int i) { @@ -847,6 +862,10 @@ private[codegen] object CometBatchKernelCodegenInput { s" case $fi: return ${path}_f$fi.getInt(this.rowIdx);" case LongType | TimestampType | TimestampNTZType => s" case $fi: return ${path}_f$fi.getLong(this.rowIdx);" + case CalendarIntervalType => + s""" case $fi: { + |$guard return ${path}_f$fi.getInterval(this.rowIdx); + | }""".stripMargin case FloatType => s" case $fi: return ${path}_f$fi.getFloat(this.rowIdx);" case DoubleType => @@ -900,6 +919,10 @@ private[codegen] object CometBatchKernelCodegenInput { f.sparkType == TimestampNTZType => fieldReadScalar(fi, LongType, f.nullable) } + val intervalCases = scalarOrd.collect { + case (f, fi) if f.sparkType == CalendarIntervalType => + fieldReadScalar(fi, CalendarIntervalType, f.nullable) + } val floatCases = scalarOrd.collect { case (f, fi) if f.sparkType == FloatType => @@ -944,6 +967,10 @@ private[codegen] object CometBatchKernelCodegenInput { structSwitch("public long getLong(int ordinal)", "getLong", longCases), structSwitch("public float getFloat(int ordinal)", "getFloat", floatCases), structSwitch("public double getDouble(int ordinal)", "getDouble", doubleCases), + structSwitch( + "public org.apache.spark.unsafe.types.CalendarInterval getInterval(int ordinal)", + "getInterval", + intervalCases), structSwitch( "public org.apache.spark.sql.types.Decimal getDecimal(" + "int ordinal, int precision, int scale)", diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index fe282cdc65..4fe820166c 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -173,6 +173,7 @@ private[codegen] object CometBatchKernelCodegenOutput { case TimestampNTZType => classOf[TimeStampMicroVector].getName case _: YearMonthIntervalType => classOf[IntervalYearVector].getName case _: DayTimeIntervalType => classOf[DurationVector].getName + case CalendarIntervalType => classOf[IntervalMonthDayNanoVector].getName case _: ArrayType => classOf[ListVector].getName case _: StructType => classOf[StructVector].getName case _: MapType => classOf[MapVector].getName @@ -216,6 +217,14 @@ private[codegen] object CometBatchKernelCodegenOutput { // DayTimeIntervalType -> DurationVector.set(int, long micros). val set = if (nested) "setSafe" else "set" OutputEmit("", s"$targetVec.$set($idx, $source);") + case CalendarIntervalType => + val set = if (nested) "setSafe" else "set" + val interval = ctx.freshName("interval") + OutputEmit( + "", + s"""org.apache.spark.unsafe.types.CalendarInterval $interval = $source; + |$targetVec.$set($idx, $interval.months, $interval.days, + | java.lang.Math.multiplyExact($interval.microseconds, 1000L));""".stripMargin) case dt: DecimalType => // DecimalOutputShortFastPath: precision <= 18 fits in a signed long, so pass the unscaled // value to `setSafe(int, long)` and skip the BigDecimal allocation. @@ -398,6 +407,7 @@ private[codegen] object CometBatchKernelCodegenOutput { case ShortType => s"$target.getShort($idx)" case IntegerType | DateType => s"$target.getInt($idx)" case LongType | TimestampType | TimestampNTZType => s"$target.getLong($idx)" + case CalendarIntervalType => s"$target.getInterval($idx)" case FloatType => s"$target.getFloat($idx)" case DoubleType => s"$target.getDouble($idx)" case dt: DecimalType => s"$target.getDecimal($idx, ${dt.precision}, ${dt.scale})" diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala b/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala index 2f81c58c06..41353fb3d8 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometSpecializedGettersDispatch.scala @@ -43,6 +43,7 @@ private[codegen] object CometSpecializedGettersDispatch { case IntegerType | DateType => java.lang.Integer.valueOf(g.getInt(ordinal)) case LongType | TimestampType | TimestampNTZType => java.lang.Long.valueOf(g.getLong(ordinal)) + case CalendarIntervalType => g.getInterval(ordinal) case FloatType => java.lang.Float.valueOf(g.getFloat(ordinal)) case DoubleType => java.lang.Double.valueOf(g.getDouble(ordinal)) case _: StringType => g.getUTF8String(ordinal) diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 5eee0c5cad..6136bad794 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -504,7 +504,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { def supportedDataType(dt: DataType, allowComplex: Boolean = false): Boolean = dt match { case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | - _: DecimalType | _: DateType | _: BooleanType | _: NullType => + _: DecimalType | _: DateType | _: BooleanType | _: NullType | CalendarIntervalType => true case dt if isTimeType(dt) => true @@ -545,6 +545,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { case dt if isTimeType(dt) => 17 case _: YearMonthIntervalType => 18 case _: DayTimeIntervalType => 19 + case CalendarIntervalType => 20 case dt => logWarning(s"Cannot serialize Spark data type: $dt") return None diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 342441ce28..54b0c24338 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -83,8 +83,8 @@ private[arrow] object ArrowWriter { case (_: YearMonthIntervalType, vector: IntervalYearVector) => new IntervalYearWriter(vector) case (_: DayTimeIntervalType, vector: DurationVector) => new DurationWriter(vector) -// case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) => -// new IntervalMonthDayNanoWriter(vector) + case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) => + new IntervalMonthDayNanoWriter(vector) case (dt, _) => throw QueryExecutionErrors.notSupportTypeError(dt) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 1988867c94..164d53fe19 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -114,6 +114,8 @@ object Utils extends CometTypeShim with Logging { case yi: ArrowType.Interval if yi.getUnit == IntervalUnit.YEAR_MONTH => YearMonthIntervalType() case di: ArrowType.Interval if di.getUnit == IntervalUnit.DAY_TIME => DayTimeIntervalType() + case ci: ArrowType.Interval if ci.getUnit == IntervalUnit.MONTH_DAY_NANO => + CalendarIntervalType case d: ArrowType.Duration if d.getUnit == TimeUnit.MICROSECOND => DayTimeIntervalType() case t: ArrowType.Time if t.getUnit == TimeUnit.NANOSECOND && t.getBitWidth == 64 => // scalastyle:off classforname @@ -162,6 +164,7 @@ object Utils extends CometTypeShim with Logging { // Spark stores DayTimeIntervalType as microseconds in an int64, matching Arrow // Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout. case _: DayTimeIntervalType => new ArrowType.Duration(TimeUnit.MICROSECOND) + case CalendarIntervalType => new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO) case _ => throw new UnsupportedOperationException( s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}") diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql new file mode 100644 index 0000000000..69980d1bf8 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql @@ -0,0 +1,26 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Config: spark.comet.exec.localTableScan.enabled=true + +query +SELECT * FROM VALUES + (make_interval(1, 2, 3, 4, 5, 6, 7.008009)), + (make_interval(30, 25, 0, -100, 40, 80, 299.889987299)), + (make_interval(0, -1, 0, 1, 0, 0, -1)), + (CAST(NULL AS INTERVAL)) +AS test_calendar_interval(i) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 76ffbc8c99..50a31075ff 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -25,12 +25,15 @@ import org.apache.arrow.vector._ import org.apache.spark.{SparkConf, TaskContext} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.api.java.UDF1 -import org.apache.spark.sql.catalyst.expressions.{CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} +import org.apache.spark.sql.catalyst.expressions.{BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String +import org.apache.comet.codegen.CometBatchKernelCodegen +import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec import org.apache.comet.udf.codegen.CometScalaUDFCodegen +import org.apache.comet.vector.CometVector /** * End-to-end correctness for the Arrow-direct codegen dispatcher. Covers the scalar and complex @@ -76,6 +79,36 @@ class CometCodegenSuite } } + test("codegen kernel round-trips CalendarIntervalType") { + val input = new IntervalMonthDayNanoVector("in", CometArrowAllocator) + val field = + CometBatchKernelCodegen.toFfiArrowField("out", CalendarIntervalType, nullable = true) + val output = CometBatchKernelCodegen.allocateOutput(field, 2, 0) + try { + input.allocateNew() + input.setSafe(0, 14, -3, 1234567000L) + input.setNull(1) + input.setValueCount(2) + + val expr = BoundReference(0, CalendarIntervalType, nullable = true) + val spec = ArrowColumnSpec(classOf[IntervalMonthDayNanoVector], nullable = true) + val kernel = CometBatchKernelCodegen.compile(expr, IndexedSeq(spec)).newInstance() + kernel.init(0) + kernel.process(Array(input), output, 2) + output.setValueCount(2) + + val comet = CometVector.getVector(output, null) + val actual = comet.getInterval(0) + assert(actual.months === 14) + assert(actual.days === -3) + assert(actual.microseconds === 1234567L) + assert(comet.getInterval(1) == null) + } finally { + output.close() + input.close() + } + } + test("ScalaUDF over concat(c1, c2) suppresses the null short-circuit") { // Concat is not NullIntolerant. The dispatcher's short-circuit guard inspects every node in // the bound tree and must skip the whole-tree null short-circuit because one child is diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index c423a49d2a..50f723d0bf 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -25,9 +25,13 @@ import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.{BigIntVector, IntVector} +import org.apache.arrow.vector.{BigIntVector, IntervalMonthDayNanoVector, IntVector, VectorSchemaRoot} import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.types.CalendarIntervalType import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.unsafe.types.CalendarInterval import org.apache.comet.vector.{CometPlainVector, CometVector} @@ -51,6 +55,33 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { new ColumnarBatch(vectors.toArray, numRows) } + test("CalendarIntervalType round-trips through Arrow writer and Comet vector") { + val allocator = new RootAllocator(Integer.MAX_VALUE) + val field = Utils.toArrowField("interval", CalendarIntervalType, nullable = true, "UTC") + Utils.fromArrowField(field) shouldBe CalendarIntervalType + val root = VectorSchemaRoot.create(new Schema(Seq(field).asJava), allocator) + try { + val expected = new CalendarInterval(14, -3, 1234567L) + val writer = ArrowWriter.create(root) + writer.write(new GenericInternalRow(Array[Any](expected))) + writer.write(new GenericInternalRow(Array[Any](null))) + writer.finish() + + val arrow = root.getVector(0).asInstanceOf[IntervalMonthDayNanoVector] + IntervalMonthDayNanoVector.getMonths(arrow.getDataBuffer, 0) shouldBe expected.months + IntervalMonthDayNanoVector.getDays(arrow.getDataBuffer, 0) shouldBe expected.days + IntervalMonthDayNanoVector.getNanoseconds(arrow.getDataBuffer, 0) shouldBe + expected.microseconds * 1000L + + val comet = new CometPlainVector(arrow, false) + comet.getInterval(0) shouldBe expected + comet.getInterval(1) shouldBe null + } finally { + root.close() + allocator.close() + } + } + test("reconcileStreamSchema returns expected schema unchanged on empty iterator") { val expected = expectedSchema("c0" -> new ArrowType.Int(64, true)) val (returned, iter) =