feat: support interval codegen dispatch for nested values and native shuffle

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

@@ -66,6 +66,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
     classOf[TimeNanoVector],
     classOf[TimeStampMicroVector],
     classOf[TimeStampMicroTZVector],
+    classOf[IntervalYearVector],
     classOf[IntervalMonthDayNanoVector])
   private val cometPlainVectorName: String = classOf[CometPlainVector].getName
 
@@ -131,7 +132,8 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
     }
     val intCases = withOrd.collect {
       case (ArrowColumnSpec(cls, _), ord)
-          if cls == classOf[IntVector] || cls == classOf[DateDayVector] =>
+          if cls == classOf[IntVector] || cls == classOf[DateDayVector] ||
+            cls == classOf[IntervalYearVector] =>
         s"      case $ord: return this.col$ord.getInt(this.rowIdx);"
     }
     val longCases = withOrd.collect {
@@ -604,7 +606,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
     case BooleanType => s"getBoolean($idx)"
     case ByteType => s"getByte($idx)"
     case ShortType => s"getShort($idx)"
-    case IntegerType | DateType => s"getInt($idx)"
+    case IntegerType | DateType | _: YearMonthIntervalType => s"getInt($idx)"
     case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType =>
       s"getLong($idx)"
     case CalendarIntervalType => s"getInterval($idx)"
@@ -704,7 +706,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
            |      public short getShort(int i) {
            |        return $childField.getShort(startIndex + i);
            |      }""".stripMargin
-      case IntegerType | DateType =>
+      case IntegerType | DateType | _: YearMonthIntervalType =>
         s"""      @Override
            |      public int getInt(int i) {
            |        return $childField.getInt(startIndex + i);
@@ -870,7 +872,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
           s"        case $fi: return ${path}_f$fi.getByte(this.rowIdx);"
         case ShortType =>
           s"        case $fi: return ${path}_f$fi.getShort(this.rowIdx);"
-        case IntegerType | DateType =>
+        case IntegerType | DateType | _: YearMonthIntervalType =>
           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);"
@@ -924,7 +926,9 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
           fieldReadScalar(fi, ShortType, f.nullable)
       }
     val intCases = scalarOrd.collect {
-      case (f, fi) if f.sparkType == IntegerType || f.sparkType == DateType =>
+      case (f, fi)
+          if f.sparkType == IntegerType || f.sparkType == DateType ||
+            f.sparkType.isInstanceOf[YearMonthIntervalType] =>
         fieldReadScalar(fi, IntegerType, f.nullable)
     }
     val longCases = scalarOrd.collect {

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

@@ -410,8 +410,9 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
       case BooleanType => s"$target.getBoolean($idx)"
       case ByteType => s"$target.getByte($idx)"
       case ShortType => s"$target.getShort($idx)"
-      case IntegerType | DateType => s"$target.getInt($idx)"
-      case LongType | TimestampType | TimestampNTZType => s"$target.getLong($idx)"
+      case IntegerType | DateType | _: YearMonthIntervalType => s"$target.getInt($idx)"
+      case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType =>
+        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)"

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

@@ -40,8 +40,9 @@ private[codegen] object CometSpecializedGettersDispatch {
       case BooleanType => java.lang.Boolean.valueOf(g.getBoolean(ordinal))
       case ByteType => java.lang.Byte.valueOf(g.getByte(ordinal))
       case ShortType => java.lang.Short.valueOf(g.getShort(ordinal))
-      case IntegerType | DateType => java.lang.Integer.valueOf(g.getInt(ordinal))
-      case LongType | TimestampType | TimestampNTZType =>
+      case IntegerType | DateType | _: YearMonthIntervalType =>
+        java.lang.Integer.valueOf(g.getInt(ordinal))
+      case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType =>
         java.lang.Long.valueOf(g.getLong(ordinal))
       case CalendarIntervalType => g.getInterval(ordinal)
       case FloatType => java.lang.Float.valueOf(g.getFloat(ordinal))

spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala

@@ -26,7 +26,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
-import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.types.{ArrayType, DataType, DayTimeIntervalType, MapType, StructType, YearMonthIntervalType}
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
@@ -43,6 +43,16 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
 
   override def enabledConfig: Option[ConfigEntry[Boolean]] = None
 
+  protected final def supportedSinkDataType(dt: DataType): Boolean = dt match {
+    case _: YearMonthIntervalType | _: DayTimeIntervalType => true
+    case StructType(fields) =>
+      fields.nonEmpty && fields.forall(f => supportedSinkDataType(f.dataType))
+    case ArrayType(elementType, _) => supportedSinkDataType(elementType)
+    case MapType(keyType, valueType, _) =>
+      supportedSinkDataType(keyType) && supportedSinkDataType(valueType)
+    case _ => supportedDataType(dt)
+  }
+
   /**
    * The data type to declare for a scan output field. Overridden by sinks whose source carries
    * non-null nested child fields that must be widened to match the planned kernel output types
@@ -54,8 +64,7 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
       op: T,
       builder: Operator.Builder,
       childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
-    val supportedTypes =
-      op.output.forall(a => supportedDataType(a.dataType, allowComplex = true))
+    val supportedTypes = op.output.forall(a => supportedSinkDataType(a.dataType))
 
     if (!supportedTypes) {
       withFallbackReason(op, "Unsupported data type")
@@ -123,8 +132,7 @@ object CometExchangeSink extends CometSink[SparkPlan] {
   private def convertToShuffleScan(
       op: SparkPlan,
       builder: Operator.Builder): Option[OperatorOuterClass.Operator] = {
-    val supportedTypes =
-      op.output.forall(a => supportedDataType(a.dataType, allowComplex = true))
+    val supportedTypes = op.output.forall(a => supportedSinkDataType(a.dataType))
 
     if (!supportedTypes) {
       withFallbackReason(op, "Unsupported data type for shuffle direct read")

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 | _: IntervalMonthDayNanoVector =>
+        _: TimeStampMicroTZVector | _: IntervalYearVector | _: IntervalMonthDayNanoVector =>
       ScalarColumnSpec(v.getClass.asInstanceOf[Class[_ <: ValueVector]], nullable = true)
     case other =>
       throw new UnsupportedOperationException(

spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala

@@ -31,7 +31,7 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNative
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution.{LeafExecNode, LocalTableScanExec}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
-import org.apache.spark.sql.types.{DataType, NullType, StructType}
+import org.apache.spark.sql.types.{DataType, DayTimeIntervalType, NullType, StructType, YearMonthIntervalType}
 
 import com.google.common.base.Objects
 
@@ -138,7 +138,7 @@ object CometLocalTableScanExec extends CometSink[LocalTableScanExec] with DataTy
       dt: DataType,
       name: String,
       fallbackReasons: ListBuffer[String]): Boolean = dt match {
-    case _: NullType => true
+    case _: NullType | _: YearMonthIntervalType | _: DayTimeIntervalType => true
     case _ => super.isTypeSupported(dt, name, fallbackReasons)
   }
 

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, CalendarIntervalType, 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, DayTimeIntervalType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.MutablePair
 import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator}
@@ -413,7 +413,7 @@ object CometShuffleExchangeExec
       case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType |
           _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType |
           _: TimestampNTZType | _: DecimalType | _: DateType | _: NullType |
-          CalendarIntervalType =>
+          _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType =>
         true
       case dt if isTimeType(dt) =>
         true

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

@@ -15,6 +15,9 @@
 -- specific language governing permissions and limitations
 -- under the License.
 
+-- Config: spark.comet.exec.localTableScan.enabled=true
+-- Config: spark.comet.exec.shuffle.mode=native
+
 -- Routes make_dt_interval through the codegen dispatcher; produces DayTimeIntervalType.
 
 statement
@@ -34,6 +37,33 @@ SELECT make_dt_interval(1, 2, 3, 4.5), make_dt_interval(0, 0, 0, 0)
 query
 SELECT make_dt_interval(1), make_dt_interval(1, 2), make_dt_interval()
 
+-- nested interval output through LocalTableScan and codegen dispatch
+query
+SELECT transform(a, x -> x) AS result
+FROM VALUES
+  (array(make_dt_interval(1, 2, 3, 4.5), CAST(NULL AS INTERVAL DAY TO SECOND)))
+AS t(a)
+
+-- top-level, struct, and map interval output through native shuffle
+query
+SELECT d, h, mi, s,
+  make_dt_interval(d, h, mi, s) AS i,
+  named_struct('i', make_dt_interval(d, h, mi, s)) AS st,
+  map('i', make_dt_interval(d, h, mi, s)) AS m
+FROM test_mdi
+DISTRIBUTE BY d
+
+-- null top-level, struct, and map interval output through native shuffle
+query
+SELECT i, st, m
+FROM VALUES
+  (1,
+   CAST(NULL AS INTERVAL DAY TO SECOND),
+   named_struct('i', CAST(NULL AS INTERVAL DAY TO SECOND)),
+   map('i', CAST(NULL AS INTERVAL DAY TO SECOND)))
+AS t(k, i, st, m)
+DISTRIBUTE BY k
+
 -- overflow: days * MICROS_PER_DAY exceeds the int64 microsecond range. makeDayTimeInterval throws
 -- unconditionally (not ANSI-gated); this confirms the dispatched codegen path propagates Spark's
 -- exception. The pattern is the lowercase word so it matches every version: Spark 4.x raises

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

@@ -15,6 +15,9 @@
 -- specific language governing permissions and limitations
 -- under the License.
 
+-- Config: spark.comet.exec.localTableScan.enabled=true
+-- Config: spark.comet.exec.shuffle.mode=native
+
 -- Routes make_ym_interval through the codegen dispatcher; produces YearMonthIntervalType.
 
 statement
@@ -34,6 +37,33 @@ SELECT make_ym_interval(1, 2), make_ym_interval(0, 0), make_ym_interval(-5, 11)
 query
 SELECT make_ym_interval(3), make_ym_interval()
 
+-- nested interval output through LocalTableScan and codegen dispatch
+query
+SELECT transform(a, x -> x) AS result
+FROM VALUES
+  (array(make_ym_interval(1, 2), CAST(NULL AS INTERVAL YEAR TO MONTH)))
+AS t(a)
+
+-- top-level, struct, and map interval output through native shuffle
+query
+SELECT y, m,
+  make_ym_interval(y, m) AS i,
+  named_struct('i', make_ym_interval(y, m)) AS s,
+  map('i', make_ym_interval(y, m)) AS m
+FROM test_myi
+DISTRIBUTE BY y
+
+-- null top-level, struct, and map interval output through native shuffle
+query
+SELECT i, s, m
+FROM VALUES
+  (1,
+   CAST(NULL AS INTERVAL YEAR TO MONTH),
+   named_struct('i', CAST(NULL AS INTERVAL YEAR TO MONTH)),
+   map('i', CAST(NULL AS INTERVAL YEAR TO MONTH)))
+AS t(k, i, s, m)
+DISTRIBUTE BY k
+
 -- overflow: years * 12 exceeds Int range. makeYearMonthInterval throws unconditionally (not
 -- ANSI-gated); this confirms the dispatched codegen path propagates Spark's exception. The
 -- pattern is the lowercase word so it matches every version: Spark 4.x raises