feat: support multiply_ym_interval with YearMonth interval codegen dispatch

docs/source/user-guide/latest/datatypes.md

@@ -80,14 +80,14 @@ the tables below and may be reconsidered based on demand:
 
 ## Interval
 
-Interval types fall back to Spark today. Native acceleration is tracked by
+Calendar interval types fall back to Spark today. Native acceleration is tracked by
 [#4540](https://github.com/apache/datafusion-comet/issues/4540).
 
-| Type                    | Status | Notes             |
-| ----------------------- | ------ | ----------------- |
-| `YearMonthIntervalType` | 🔜     | Tracked by #4540. |
-| `DayTimeIntervalType`   | 🔜     | Tracked by #4540. |
-| `CalendarIntervalType`  | 🔜     | Tracked by #4540. |
+| Type                    | Status | Notes |
+| ----------------------- | ------ | ----- |
+| `YearMonthIntervalType` | ✅     |       |
+| `DayTimeIntervalType`   | ✅     |       |
+| `CalendarIntervalType`  | ✅     |       |
 
 ## Complex
 

docs/source/user-guide/latest/expressions.md

@@ -411,7 +411,7 @@ expression-level). The `outer` variants are wired but marked `Incompatible`; the
 | Function | Status | Implementation | Notes |
 | --- | --- | --- | --- |
 | `%` | ✅ | Native |  |
-| `*` | ✅ | Native | Interval multiplication falls back |
+| `*` | ✅ | Hybrid | YearMonth and DayTime interval multiplication routes through the JVM codegen dispatcher; Calendar interval multiplication falls back |
 | `+` | ✅ | Native |  |
 | `-` | ✅ | Native |  |
 | `/` | ✅ | Native |  |

native/core/src/execution/planner.rs

@@ -38,7 +38,9 @@ use crate::execution::{
 };
 use crate::jvm_bridge::{jni_call, JVMClasses};
 use arrow::compute::CastOptions;
-use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit, DECIMAL128_MAX_PRECISION};
+use arrow::datatypes::{
+    DataType, Field, FieldRef, IntervalUnit, Schema, TimeUnit, DECIMAL128_MAX_PRECISION,
+};
 use arrow::ffi_stream::FFI_ArrowArrayStream;
 use datafusion::functions_aggregate::bit_and_or_xor::{bit_and_udaf, bit_or_udaf, bit_xor_udaf};
 use datafusion::functions_aggregate::count::count_udaf;
@@ -104,8 +106,8 @@ use datafusion::physical_expr::LexOrdering;
 use crate::parquet::parquet_exec::init_datasource_exec;
 use arrow::array::{
     new_empty_array, Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array,
-    Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray,
-    NullArray, StringBuilder, TimestampMicrosecondArray,
+    Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
+    IntervalYearMonthArray, ListArray, NullArray, StringBuilder, TimestampMicrosecondArray,
 };
 use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
 use arrow::row::{OwnedRow, RowConverter, SortField};
@@ -412,6 +414,9 @@ impl PhysicalPlanner {
                         DataType::Time64(TimeUnit::Nanosecond) => {
                             ScalarValue::Time64Nanosecond(None)
                         }
+                        DataType::Interval(IntervalUnit::YearMonth) => {
+                            ScalarValue::IntervalYearMonth(None)
+                        }
                         DataType::Duration(TimeUnit::Microsecond) => {
                             ScalarValue::DurationMicrosecond(None)
                         }
@@ -427,9 +432,12 @@ impl PhysicalPlanner {
                         Value::IntVal(value) => match data_type {
                             DataType::Int32 => ScalarValue::Int32(Some(*value)),
                             DataType::Date32 => ScalarValue::Date32(Some(*value)),
+                            DataType::Interval(IntervalUnit::YearMonth) => {
+                                ScalarValue::IntervalYearMonth(Some(*value))
+                            }
                             dt => {
                                 return Err(GeneralError(format!(
-                                    "Expected either 'Int32' or 'Date32' for IntVal, but found {dt:?}"
+                                    "Expected either 'Int32', 'Date32', or 'Interval(YearMonth)' for IntVal, but found {dt:?}"
                                 )))
                             }
                         },
@@ -4061,6 +4069,10 @@ fn literal_to_array_ref(
             list_literal.int_values.into(),
             Some(nulls.clone().into()),
         ))),
+        DataType::Interval(IntervalUnit::YearMonth) => Ok(Arc::new(IntervalYearMonthArray::new(
+            list_literal.int_values.into(),
+            Some(nulls.clone().into()),
+        ))),
         DataType::Timestamp(TimeUnit::Microsecond, None) => {
             Ok(Arc::new(TimestampMicrosecondArray::new(
                 list_literal.long_values.into(),

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

@@ -58,6 +58,7 @@ private[codegen] object CometBatchKernelCodegenInput extends CometTypeShim {
     classOf[TinyIntVector],
     classOf[SmallIntVector],
     classOf[IntVector],
+    classOf[IntervalYearVector],
     classOf[BigIntVector],
     classOf[Float4Vector],
     classOf[Float8Vector],
@@ -131,7 +132,9 @@ 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 +607,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 +707,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 +873,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,8 +927,11 @@ 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 =>
-        fieldReadScalar(fi, IntegerType, f.nullable)
+      case (f, fi)
+          if f.sparkType == IntegerType ||
+            f.sparkType == DateType ||
+            f.sparkType.isInstanceOf[YearMonthIntervalType] =>
+        fieldReadScalar(fi, f.sparkType, f.nullable)
     }
     val longCases = scalarOrd.collect {
       case (f, fi)

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

@@ -305,6 +305,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
       classOf[MicrosToTimestamp] -> CometMicrosToTimestamp,
       classOf[MillisToTimestamp] -> CometMillisToTimestamp,
       classOf[MonthsBetween] -> CometMonthsBetween,
+      classOf[MultiplyYMInterval] -> CometMultiplyYMInterval,
       classOf[Minute] -> CometMinute,
       classOf[NextDay] -> CometNextDay,
       classOf[PreciseTimestampConversion] -> CometPreciseTimestampConversion,
@@ -527,7 +528,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 | CalendarIntervalType =>
+        _: DecimalType | _: DateType | _: BooleanType | _: NullType =>
       true
     case dt if isTimeType(dt) =>
       true

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

@@ -21,7 +21,7 @@ package org.apache.comet.serde
 
 import java.util.Locale
 
-import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year}
+import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType}
 import org.apache.spark.unsafe.types.UTF8String
@@ -952,6 +952,8 @@ object CometGetTimestamp extends CometCodegenDispatch[GetTimestamp]
 
 object CometMakeYMInterval extends CometCodegenDispatch[MakeYMInterval]
 
+object CometMultiplyYMInterval extends CometCodegenDispatch[MultiplyYMInterval]
+
 object CometMakeDTInterval extends CometCodegenDispatch[MakeDTInterval]
 
 object CometMultiplyDTInterval extends CometCodegenDispatch[MultiplyDTInterval]

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

@@ -24,7 +24,7 @@ import java.lang
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Literal}
 import org.apache.spark.sql.catalyst.util.ArrayData
-import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, NullType, ShortType, StringType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, NullType, ShortType, StringType, TimestampNTZType, TimestampType, YearMonthIntervalType}
 import org.apache.spark.unsafe.types.UTF8String
 
 import com.google.protobuf.ByteString
@@ -41,24 +41,27 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging {
     "Not all data types are supported for literal values")
 
   override def getSupportLevel(expr: Literal): SupportLevel = {
-
+    val dataType = expr.dataType
     if (supportedDataType(
-        expr.dataType,
+        dataType,
         allowComplex = expr.value == null ||
 
           // Nested literal support for native reader
           // can be tracked https://github.com/apache/datafusion-comet/issues/1937
-          (expr.dataType
+          (dataType
             .isInstanceOf[ArrayType] && (!isComplexType(
-            expr.dataType.asInstanceOf[ArrayType].elementType) || expr.dataType
+            dataType.asInstanceOf[ArrayType].elementType) || dataType
             .asInstanceOf[ArrayType]
             .elementType
             .isInstanceOf[ArrayType])))) {
       Compatible(None)
     } else {
-      expr.dataType match {
-        case _: DayTimeIntervalType => Compatible(None)
-        case _ => Unsupported(Some(s"Unsupported data type ${expr.dataType}"))
+      dataType match {
+        // Keep interval types out of QueryPlanSerde.supportedDataType, which gates broader native
+        // paths.
+        case _: DayTimeIntervalType | _: YearMonthIntervalType | CalendarIntervalType =>
+          Compatible(None)
+        case _ => Unsupported(Some(s"Unsupported data type $dataType"))
       }
     }
   }
@@ -80,7 +83,8 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging {
         case _: BooleanType => exprBuilder.setBoolVal(value.asInstanceOf[Boolean])
         case _: ByteType => exprBuilder.setByteVal(value.asInstanceOf[Byte])
         case _: ShortType => exprBuilder.setShortVal(value.asInstanceOf[Short])
-        case _: IntegerType | _: DateType => exprBuilder.setIntVal(value.asInstanceOf[Int])
+        case _: IntegerType | _: DateType | _: YearMonthIntervalType =>
+          exprBuilder.setIntVal(value.asInstanceOf[Int])
         case _: LongType | _: TimestampType | _: TimestampNTZType | _: DayTimeIntervalType =>
           exprBuilder.setLongVal(value.asInstanceOf[Long])
         case dt if isTimeType(dt) =>
@@ -153,7 +157,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with Logging {
             else null.asInstanceOf[Integer])
           listLiteralBuilder.addNullMask(casted != null)
         })
-      case IntegerType | DateType =>
+      case IntegerType | DateType | _: YearMonthIntervalType =>
         array.foreach(v => {
           val casted = v.asInstanceOf[Integer]
           listLiteralBuilder.addIntValues(casted)

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

@@ -19,6 +19,7 @@
 
 package org.apache.comet.serde.operator
 
+import scala.collection.mutable.ListBuffer
 import scala.jdk.CollectionConverters._
 
 import org.apache.spark.sql.comet.{CometNativeExec, CometSinkPlaceHolder}
@@ -28,7 +29,7 @@ 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.comet.CometConf
+import org.apache.comet.{CometConf, DataTypeSupport}
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
 import org.apache.comet.ConfigEntry
 import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass}
@@ -39,7 +40,7 @@ import org.apache.comet.serde.QueryPlanSerde.{serializeDataType, supportedDataTy
  * CometSink is the base class for transformations from a Spark operator to a Comet operator where
  * the native plan is a ScanExec that will read data from the Comet operator running the JVM.
  */
-abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
+abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] with DataTypeSupport {
 
   override def enabledConfig: Option[ConfigEntry[Boolean]] = None
 
@@ -54,8 +55,9 @@ 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 =>
+      supportedDataType(a.dataType, allowComplex = true) ||
+        isTypeSupported(a.dataType, a.name, ListBuffer.empty))
 
     if (!supportedTypes) {
       withFallbackReason(op, "Unsupported data type")

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

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

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

@@ -0,0 +1,64 @@
+-- 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.
+
+-- Routes multiply_ym_interval through the codegen dispatcher; produces YearMonthIntervalType.
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true
+
+statement
+CREATE TABLE test_multiply_ym_interval(y int, m int, i int, l long, f float, d double, dec decimal(10,2)) USING parquet
+
+statement
+INSERT INTO test_multiply_ym_interval VALUES
+  (1, 2, 2, CAST(3 AS BIGINT), CAST(1.5 AS FLOAT), CAST(2.5 AS DOUBLE), CAST(2.50 AS DECIMAL(10, 2))),
+  (-1, 1, -2, CAST(-3 AS BIGINT), CAST(-1.5 AS FLOAT), CAST(-2.5 AS DOUBLE), CAST(-2.50 AS DECIMAL(10, 2))),
+  (0, 6, 0, CAST(0 AS BIGINT), CAST(0.5 AS FLOAT), CAST(0.5 AS DOUBLE), CAST(0.50 AS DECIMAL(10, 2))),
+  (2, -6, NULL, NULL, NULL, NULL, NULL)
+
+query
+SELECT
+  make_ym_interval(y, m) * i,
+  make_ym_interval(y, m) * l,
+  make_ym_interval(y, m) * f,
+  make_ym_interval(y, m) * d,
+  make_ym_interval(y, m) * dec
+FROM test_multiply_ym_interval
+
+-- literal interval input
+query
+SELECT INTERVAL '1-2' YEAR TO MONTH * i FROM test_multiply_ym_interval
+
+-- numeric on the left is normalized by Spark to multiply_ym_interval.
+query
+SELECT i * make_ym_interval(y, m), 2 * INTERVAL '1-2' YEAR TO MONTH
+FROM test_multiply_ym_interval
+
+-- literal multipliers, including half-up rounding for fractional months.
+query
+SELECT
+  make_ym_interval(1, 2) * 2,
+  make_ym_interval(1, 2) * 1.5D,
+  make_ym_interval(1, 2) * CAST(1.50 AS DECIMAL(10, 2)),
+  make_ym_interval(-1, 1) * 1.5D
+
+-- null interval input
+query
+SELECT make_ym_interval(NULL, m) * 2 FROM test_multiply_ym_interval
+
+-- 178956970 years and 7 months is Int.MaxValue months. Multiplication overflows regardless of
+-- ANSI mode; the lowercase pattern matches Spark 3.x and 4.x error messages.
+query expect_error(overflow)
+SELECT make_ym_interval(178956970, 7) * 2