feat: route date and timestamp interval arithmetic through codegen dispatch

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

@@ -415,8 +415,8 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
 | --- | --- | --- | --- |
 | `%` | ✅ | Native |  |
 | `*` | ✅ | Native | DayTime interval multiplication routes through the JVM codegen dispatcher; YearMonth and Calendar interval multiplication fall back |
-| `+` | ✅ | Native |  |
-| `-` | ✅ | Native |  |
+| `+` | ✅ | Native | Adding a calendar, year-month or day-time interval to a date or timestamp routes through the JVM codegen dispatcher |
+| `-` | ✅ | Native | `date - date`, `timestamp - timestamp` and subtracting an interval from a date or timestamp route through the JVM codegen dispatcher; `timestamp - timestamp` falls back to Spark in legacy interval mode (`spark.sql.legacy.interval.enabled=true`) because its calendar-interval result can exceed what the dispatcher output can carry |
 | `/` | ✅ | Native |  |
 | `abs` | ✅ | Hybrid | Interval types route through the JVM codegen dispatcher; numeric types run natively |
 | `acos` | ✅ | Native |  |

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

@@ -307,6 +307,11 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
       classOf[MultiplyDTInterval] -> CometMultiplyDTInterval,
       classOf[TimestampAdd] -> CometTimestampAdd,
       classOf[TimestampDiff] -> CometTimestampDiff,
+      classOf[DateAddInterval] -> CometDateAddInterval,
+      classOf[DateAddYMInterval] -> CometDateAddYMInterval,
+      classOf[TimestampAddYMInterval] -> CometTimestampAddYMInterval,
+      classOf[SubtractDates] -> CometSubtractDates,
+      classOf[SubtractTimestamps] -> CometSubtractTimestamps,
       classOf[MicrosToTimestamp] -> CometMicrosToTimestamp,
       classOf[MillisToTimestamp] -> CometMillisToTimestamp,
       classOf[MonthsBetween] -> CometMonthsBetween,

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, Cast, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year}
+import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, Cast, ConvertTimezone, DateAdd, DateAddInterval, DateAddYMInterval, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, SubtractDates, SubtractTimestamps, TimestampAdd, TimestampAddYMInterval, TimestampDiff, 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.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType}
 import org.apache.spark.unsafe.types.UTF8String
@@ -997,6 +997,37 @@ object CometTimestampAdd extends CometCodegenDispatch[TimestampAdd]
 
 object CometTimestampDiff extends CometCodegenDispatch[TimestampDiff]
 
+// Date and timestamp interval arithmetic. `timestamp + day-time or calendar interval` resolves
+// to `TimeAdd` on Spark 3.4 through 4.0 and to `TimestampAddInterval` on 4.1+, so that serde
+// lives in the version shims.
+object CometDateAddInterval extends CometCodegenDispatch[DateAddInterval]
+
+object CometDateAddYMInterval extends CometCodegenDispatch[DateAddYMInterval]
+
+object CometTimestampAddYMInterval extends CometCodegenDispatch[TimestampAddYMInterval]
+
+object CometSubtractDates extends CometCodegenDispatch[SubtractDates]
+
+object CometSubtractTimestamps extends CometCodegenDispatch[SubtractTimestamps] {
+  private val legacyIntervalReason =
+    "In legacy interval mode (`spark.sql.legacy.interval.enabled=true`) the result is a" +
+      " `CalendarIntervalType`, and the JVM codegen dispatcher's calendar-interval output" +
+      " cannot carry a span past about 292 years (see" +
+      " https://github.com/apache/datafusion-comet/issues/5279), so the expression falls" +
+      " back to Spark"
+
+  override def getUnsupportedReasons(): Seq[String] = Seq(legacyIntervalReason)
+
+  // Same `Math.multiplyExact(microseconds, 1000L)` limit `CometMakeInterval` documents as a
+  // compatible note. That one only overflows on extreme arguments; `ts - ts` produces an
+  // arbitrary span from ordinary data, and legacy mode is off by default, so decline it.
+  // Remove this branch once #5279 carries CalendarInterval across the boundary losslessly.
+  override def getSupportLevel(expr: SubtractTimestamps): SupportLevel = expr.dataType match {
+    case CalendarIntervalType => Unsupported(Some(legacyIntervalReason))
+    case _ => Compatible()
+  }
+}
+
 /**
  * Spark's internal `PreciseTimestampConversion` reinterprets a value between the timestamp types
  * (`TimestampType` / `TimestampNTZType`) and `LongType` without losing microsecond precision. It

spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala

@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
 
 import org.apache.comet.expressions.CometEvalMode
-import org.apache.comet.serde.{CometEncode, CometExpressionSerde, CometStringDecode}
+import org.apache.comet.serde.{CometCodegenDispatch, CometEncode, CometExpressionSerde, CometStringDecode}
 import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
 
 /**
@@ -41,7 +41,7 @@ trait CometExprShim {
   def sparkVersionSpecificMathExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map.empty
   def sparkVersionSpecificMiscExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
-    Map.empty
+    Map(classOf[TimeAdd] -> new CometCodegenDispatch[TimeAdd])
   def sparkVersionSpecificMapExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map.empty
 

spark/src/main/spark-3.5/org/apache/comet/shims/CometExprShim.scala

@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
 
 import org.apache.comet.expressions.CometEvalMode
-import org.apache.comet.serde.{CometEncode, CometExpressionSerde, CometStringDecode, CometToPrettyString, CometWidthBucket}
+import org.apache.comet.serde.{CometCodegenDispatch, CometEncode, CometExpressionSerde, CometStringDecode, CometToPrettyString, CometWidthBucket}
 import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
 
 /**
@@ -41,7 +41,9 @@ trait CometExprShim {
   def sparkVersionSpecificMathExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map(classOf[WidthBucket] -> CometWidthBucket)
   def sparkVersionSpecificMiscExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
-    Map(classOf[ToPrettyString] -> CometToPrettyString)
+    Map(
+      classOf[ToPrettyString] -> CometToPrettyString,
+      classOf[TimeAdd] -> new CometCodegenDispatch[TimeAdd])
   def sparkVersionSpecificMapExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map.empty
 

spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala

@@ -19,11 +19,12 @@
 
 package org.apache.comet.shims
 
-import org.apache.spark.sql.catalyst.expressions.EvalMode
+import org.apache.spark.sql.catalyst.expressions.{EvalMode, Expression, TimeAdd}
 import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
 import org.apache.spark.sql.internal.SQLConf
 
 import org.apache.comet.expressions.CometEvalMode
+import org.apache.comet.serde.{CometCodegenDispatch, CometExpressionSerde}
 import org.apache.comet.serde.ExprOuterClass.BinaryOutputStyle
 
 /**
@@ -42,6 +43,12 @@ trait CometExprShim extends Spark4xCometExprShim {
       case _ => BinaryOutputStyle.HEX_DISCRETE
     }
   }
+
+  // Spark 4.0 still spells `timestamp + interval` as `TimeAdd`; 4.1 renames it.
+  override def sparkVersionSpecificMiscExpressions
+      : Map[Class[_ <: Expression], CometExpressionSerde[_]] =
+    super.sparkVersionSpecificMiscExpressions + (classOf[TimeAdd] -> new CometCodegenDispatch[
+      TimeAdd])
 }
 
 object CometEvalModeUtil {

spark/src/main/spark-4.1+/org/apache/comet/serde/CometTimestampAddInterval.scala

@@ -0,0 +1,28 @@
+/*
+ * 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.
+ */
+
+package org.apache.comet.serde
+
+import org.apache.spark.sql.catalyst.expressions.TimestampAddInterval
+
+/**
+ * `timestamp + day-time or calendar interval` resolves to `TimestampAddInterval` on Spark 4.1+
+ * (`TimeAdd` on earlier versions) and runs through the codegen dispatcher.
+ */
+object CometTimestampAddInterval extends CometCodegenDispatch[TimestampAddInterval]

spark/src/main/spark-4.1+/org/apache/comet/shims/CometExprShim.scala

@@ -27,6 +27,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.TimeType
 
 import org.apache.comet.expressions.CometEvalMode
+import org.apache.comet.serde.{CometExpressionSerde, CometTimestampAddInterval}
 import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType}
 
@@ -46,6 +47,12 @@ trait CometExprShim extends Spark4xCometExprShim {
     }
   }
 
+  // Spark 4.1 renames `TimeAdd` to `TimestampAddInterval`.
+  override def sparkVersionSpecificMiscExpressions
+      : Map[Class[_ <: Expression], CometExpressionSerde[_]] =
+    super.sparkVersionSpecificMiscExpressions +
+      (classOf[TimestampAddInterval] -> CometTimestampAddInterval)
+
   // Spark 4.1 introduced TimeType and the make_time / to_time / try_to_time functions.
   // Their planner forms differ from the shared 4.x patterns (DateTimeUtils.makeTime
   // StaticInvoke and ToTimeParser Invoke / TryEval(Invoke)), so they live here rather

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

@@ -0,0 +1,84 @@
+-- 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.
+
+-- date + calendar interval resolves to DateAddInterval and runs through the codegen dispatcher
+-- so results match Spark exactly. With ANSI off, an interval carrying a time part is applied
+-- on the timestamp and the result truncated back to a date; the ANSI error case lives in
+-- date_add_interval_ansi.sql. America/Los_Angeles is pinned so the 25-hour row crosses DST.
+-- Config: spark.sql.session.timeZone=America/Los_Angeles
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true
+-- Config: spark.comet.shuffle.mode=native
+
+statement
+CREATE TABLE test_date_add_interval(d date, y int, m int, dd int, h int, k int) USING parquet
+
+statement
+INSERT INTO test_date_add_interval VALUES
+  (date'2024-01-31', 0, 1, 0, 0, 1),
+  (date'2024-01-31', 1, 1, 1, 0, 1),
+  (date'2024-02-29', 1, 0, 0, 0, 2),
+  (date'2024-03-31', 0, -1, 0, 0, 2),
+  (date'2024-12-31', 0, 0, 1, 0, 3),
+  (date'2024-03-09', 0, 0, 1, 0, 3),
+  (date'2024-03-09', 0, 0, 0, 25, 4),
+  (date'1970-01-01', -1, -1, -1, -1, 4),
+  (date'2024-06-15', NULL, 1, 1, 0, 5),
+  (date'2024-06-15', 0, NULL, 1, 0, 5),
+  (date'2024-06-15', 0, 0, NULL, 0, 6),
+  (date'2024-06-15', 0, 0, 0, NULL, 6),
+  (NULL, 1, 1, 1, 0, 7)
+
+-- column date plus a calendar interval built from columns. Month arithmetic clamps to the end
+-- of the shorter month before the day part is added.
+query
+SELECT d, y, m, dd, d + make_interval(y, m, 0, dd) FROM test_date_add_interval
+
+-- interval on the left
+query
+SELECT make_interval(y, m, 0, dd) + d FROM test_date_add_interval
+
+-- with ANSI off the hour part is applied on the timestamp and truncated away again
+query
+SELECT d, h, d + make_interval(y, m, 0, dd, h) FROM test_date_add_interval
+
+-- The parser rejects interval literals that mix year-month and day-time units unless
+-- spark.sql.legacy.interval.enabled is set, so literal calendar intervals come from
+-- make_interval. Subtraction rewrites to an addition of the negated interval.
+query
+SELECT
+  d + make_interval(1, 0, 0, 1),
+  d + make_interval(-1, 0, 0, -1),
+  d + make_interval(0, 1, 0, 1),
+  d + make_interval(0, 1, 0, 1, 12),
+  d - make_interval(1, 0, 0, 1),
+  d - make_interval(0, 1, 0, 1)
+FROM test_date_add_interval
+
+-- all-literal operands (constant folding is disabled by the test suite)
+query
+SELECT
+  date'2024-01-31' + make_interval(0, 1, 0, 1),
+  date'2024-02-29' + make_interval(1, 0, 0, 1),
+  date'2024-01-31' - make_interval(0, 1, 0, 1),
+  CAST(NULL AS DATE) + make_interval(0, 1, 0, 1),
+  date'2024-01-31' + CAST(NULL AS INTERVAL)
+
+-- date output through native shuffle
+query
+SELECT k, d + make_interval(y, m, 0, dd) AS r
+FROM test_date_add_interval
+DISTRIBUTE BY k

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

@@ -0,0 +1,47 @@
+-- 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.
+
+-- With ANSI on, DateAddInterval keeps day-granular intervals on the date path and rejects an
+-- interval that carries hours, minutes, seconds or fractions of a second. The parser rejects
+-- interval literals that mix year-month and day-time units, so make_interval builds them.
+-- Config: spark.sql.ansi.enabled=true
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true
+-- MinSparkVersion: 4.0
+
+statement
+CREATE TABLE test_date_add_interval_ansi(d date, m int, dd int) USING parquet
+
+statement
+INSERT INTO test_date_add_interval_ansi VALUES
+  (date'2024-01-31', 1, 1),
+  (date'2024-02-29', 12, 0),
+  (date'2024-03-31', -1, -1),
+  (date'2024-06-15', NULL, 1),
+  (NULL, 1, 1)
+
+-- sentinel: a day-granular interval succeeds and asserts native execution
+query
+SELECT d, m, dd, d + make_interval(0, m, 0, dd), d - make_interval(0, 1, 0, 1)
+FROM test_date_add_interval_ansi
+
+-- a NULL interval yields NULL rather than an error
+query
+SELECT d + CAST(NULL AS INTERVAL) FROM test_date_add_interval_ansi
+
+-- an interval with a time part is rejected
+query expect_error(INVALID_INTERVAL_WITH_MICROSECONDS_ADDITION)
+SELECT d + make_interval(0, 1, 0, 1, 12) FROM test_date_add_interval_ansi

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

@@ -0,0 +1,47 @@
+-- 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.
+
+-- With ANSI on, DateAddInterval keeps day-granular intervals on the date path and rejects an
+-- interval that carries hours, minutes, seconds or fractions of a second. The parser rejects
+-- interval literals that mix year-month and day-time units, so make_interval builds them.
+-- Config: spark.sql.ansi.enabled=true
+-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true
+-- MaxSparkVersion: 3.5
+
+statement
+CREATE TABLE test_date_add_interval_ansi(d date, m int, dd int) USING parquet
+
+statement
+INSERT INTO test_date_add_interval_ansi VALUES
+  (date'2024-01-31', 1, 1),
+  (date'2024-02-29', 12, 0),
+  (date'2024-03-31', -1, -1),
+  (date'2024-06-15', NULL, 1),
+  (NULL, 1, 1)
+
+-- sentinel: a day-granular interval succeeds and asserts native execution
+query
+SELECT d, m, dd, d + make_interval(0, m, 0, dd), d - make_interval(0, 1, 0, 1)
+FROM test_date_add_interval_ansi
+
+-- a NULL interval yields NULL rather than an error
+query
+SELECT d + CAST(NULL AS INTERVAL) FROM test_date_add_interval_ansi
+
+-- an interval with a time part is rejected
+query expect_error(Cannot add hours, minutes or seconds, milliseconds, microseconds to a date)
+SELECT d + make_interval(0, 1, 0, 1, 12) FROM test_date_add_interval_ansi