feat: add CalendarIntervalType support

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

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;
 

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() {}
+  }
 }

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 =>

spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala

@@ -74,6 +74,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
     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")
   }
 
@@ -88,7 +89,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
     case _: StringType | _: BinaryType => true
     case DateType | TimestampType | TimestampNTZType => true
     case dt if isTimeType(dt) => 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)

spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenInput.scala

@@ -65,7 +65,8 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
     classOf[DurationVector],
     classOf[TimeNanoVector],
     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. */
@@ -142,6 +143,10 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
             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);"
@@ -201,6 +206,10 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
       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)",
@@ -598,6 +607,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
     case IntegerType | DateType => s"getInt($idx)"
     case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType =>
       s"getLong($idx)"
+    case CalendarIntervalType => s"getInterval($idx)"
     case dt if isTimeType(dt) => s"getLong($idx)"
     case FloatType => s"getFloat($idx)"
     case DoubleType => s"getDouble($idx)"
@@ -704,6 +714,11 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
            |      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 dt if isTimeType(dt) =>
         s"""      @Override
            |      public long getLong(int i) {
@@ -859,6 +874,10 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
           s"        case $fi: return ${path}_f$fi.getInt(this.rowIdx);"
         case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType =>
           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 dt if isTimeType(dt) =>
           s"        case $fi: return ${path}_f$fi.getLong(this.rowIdx);"
         case FloatType =>
@@ -916,6 +935,10 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
             isTimeType(f.sparkType) =>
         fieldReadScalar(fi, f.sparkType, 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 =>
@@ -960,6 +983,10 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
       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)",

spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala

@@ -175,6 +175,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
     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
@@ -218,6 +219,14 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
       // 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 if isTimeType(dt) =>
       val set = if (nested) "setSafe" else "set"
       OutputEmit("", s"$targetVec.$set($idx, $source);")
@@ -403,6 +412,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
       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 dt if isTimeType(dt) => s"$target.getLong($idx)"
       case FloatType => s"$target.getFloat($idx)"
       case DoubleType => s"$target.getDouble($idx)"

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)

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -507,7 +507,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
@@ -548,6 +548,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

spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala

@@ -220,7 +220,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging {
     case _: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector | _: BigIntVector |
         _: Float4Vector | _: Float8Vector | _: DecimalVector | _: VarCharVector |
         _: VarBinaryVector | _: DateDayVector | _: DurationVector | _: TimeStampMicroVector |
-        _: TimeStampMicroTZVector =>
+        _: TimeStampMicroTZVector | _: IntervalMonthDayNanoVector =>
       ScalarColumnSpec(v.getClass.asInstanceOf[Class[_ <: ValueVector]], nullable = true)
     case other =>
       throw new UnsupportedOperationException(

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)
     }

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala

@@ -41,7 +41,7 @@ import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
 import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.MutablePair
 import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator}
@@ -412,7 +412,8 @@ object CometShuffleExchangeExec
     def supportedSerializableDataType(dt: DataType): Boolean = dt match {
       case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType |
           _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType |
-          _: TimestampNTZType | _: DecimalType | _: DateType | _: NullType =>
+          _: TimestampNTZType | _: DecimalType | _: DateType | _: NullType |
+          CalendarIntervalType =>
         true
       case dt if isTimeType(dt) =>
         true

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}")

spark/src/test/resources/sql-tests/expressions/datetime/calendar_interval.sql

@@ -0,0 +1,39 @@
+-- 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
+-- Config: spark.comet.exec.shuffle.mode=native
+
+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)
+
+-- Native shuffle materializes ARRAY<INTERVAL>; transform consumes its child interval vector
+-- through the JVM codegen dispatcher.
+query
+SELECT transform(a, x -> x)
+FROM (
+  SELECT * FROM VALUES
+    (1, array(make_interval(1, 2), CAST(NULL AS INTERVAL))),
+    (2, array(make_interval(0, -1, 0, 1, 0, 0, -1)))
+  AS t(id, a)
+  DISTRIBUTE BY id
+)