[SPARK-57825][SQL] Add +/- ANSI year-month interval support for TIMESTAMP_NTZ/LTZ(p)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/BinaryArithmeticWithDatetimeResolver.scala

@@ -47,6 +47,7 @@ import org.apache.spark.sql.catalyst.expressions.{
 }
 import org.apache.spark.sql.types.{
   AnsiIntervalType,
+  AnyTimestampNanoType,
   AnyTimestampTypeExpression,
   CalendarIntervalType,
   DatetimeType,
@@ -75,6 +76,10 @@ object BinaryArithmeticWithDatetimeResolver {
           TimestampAddYMInterval(l, r)
         case (_: YearMonthIntervalType, TimestampType | TimestampNTZType) =>
           TimestampAddYMInterval(r, l)
+        case (_: AnyTimestampNanoType, _: YearMonthIntervalType) =>
+          TimestampAddYMInterval(l, r)
+        case (_: YearMonthIntervalType, _: AnyTimestampNanoType) =>
+          TimestampAddYMInterval(r, l)
         case (CalendarIntervalType, CalendarIntervalType) |
              (_: DayTimeIntervalType, _: DayTimeIntervalType) =>
           a
@@ -113,6 +118,9 @@ object BinaryArithmeticWithDatetimeResolver {
         case (TimestampType | TimestampNTZType, _: YearMonthIntervalType) =>
           DatetimeSub(l, r, TimestampAddYMInterval(l,
             UnaryMinus(r, context.evalMode == EvalMode.ANSI)))
+        case (_: AnyTimestampNanoType, _: YearMonthIntervalType) =>
+          DatetimeSub(l, r, TimestampAddYMInterval(l,
+            UnaryMinus(r, context.evalMode == EvalMode.ANSI)))
         case (CalendarIntervalType, CalendarIntervalType) |
              (_: DayTimeIntervalType, _: DayTimeIntervalType) =>
           s

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala

@@ -2749,7 +2749,8 @@ case class TimestampAddYMInterval(
 
   override def toString: String = s"$left + $right"
   override def sql: String = s"${left.sql} + ${right.sql}"
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyTimestampType, YearMonthIntervalType)
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(AnyTimestampType, AnyTimestampNanoType), YearMonthIntervalType)
 
   override def dataType: DataType = timestamp.dataType
 
@@ -2758,16 +2759,25 @@ case class TimestampAddYMInterval(
 
   @transient private lazy val zoneIdInEval: ZoneId = zoneIdForType(left.dataType)
 
-  override def nullSafeEval(micros: Any, months: Any): Any = {
-    timestampAddMonths(micros.asInstanceOf[Long], months.asInstanceOf[Int], zoneIdInEval)
+  override def nullSafeEval(start: Any, months: Any): Any = left.dataType match {
+    case _: AnyTimestampNanoType =>
+      timestampNanosAddMonths(
+        start.asInstanceOf[TimestampNanosVal], months.asInstanceOf[Int], zoneIdInEval)
+    case _ =>
+      timestampAddMonths(start.asInstanceOf[Long], months.asInstanceOf[Int], zoneIdInEval)
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val zid = ctx.addReferenceObj("zoneId", zoneIdInEval, classOf[ZoneId].getName)
     val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (micros, months) => {
-      s"""$dtu.timestampAddMonths($micros, $months, $zid)"""
-    })
+    left.dataType match {
+      case _: AnyTimestampNanoType =>
+        defineCodeGen(ctx, ev, (sd, months) =>
+          s"""$dtu.timestampNanosAddMonths($sd, $months, $zid)""")
+      case _ =>
+        defineCodeGen(ctx, ev, (micros, months) =>
+          s"""$dtu.timestampAddMonths($micros, $months, $zid)""")
+    }
   }
 
   override protected def withNewChildrenInternal(

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala

@@ -283,6 +283,18 @@ object DateTimeUtils extends SparkDateTimeUtils {
     instantToMicros(microsToInstant(micros).atZone(zoneId).plusMonths(months).toInstant)
   }
 
+  /**
+   * Adds a year-month interval expressed in months to a nanosecond-precision timestamp value while
+   * preserving the `nanosWithinMicro` remainder.
+   */
+  def timestampNanosAddMonths(
+      start: TimestampNanosVal,
+      months: Int,
+      zoneId: ZoneId): TimestampNanosVal = {
+    val epochMicros = timestampAddMonths(start.epochMicros, months, zoneId)
+    TimestampNanosVal.fromParts(epochMicros, start.nanosWithinMicro)
+  }
+
   /**
    * Adds a day-time interval expressed in microseconds to a timestamp at the given time zone.
    * It converts the input timestamp to a local timestamp, and adds the interval by:

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala

@@ -2341,6 +2341,72 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     assert(ltzMismatch.errorSubClass == "UNEXPECTED_INPUT_TYPE")
   }
 
+  test("SPARK-57825: add/subtract ANSI year-month interval on nanos timestamps") {
+    val interval = Period.ofYears(1).plusMonths(2)
+    val minusInterval = Period.ofMonths(-1)
+
+    // A year-month shift moves only the month field: the whole fraction (including the
+    // sub-microsecond `789`) and the time of day are carried through unchanged.
+    val ntzType = TimestampNTZNanosType(9)
+    val ntzStart = DateTimeUtils.localDateTimeToTimestampNanos(
+      LocalDateTime.parse("2020-01-02T03:04:05.123456789"), precision = 9)
+    val ntzExpectedAdd = DateTimeUtils.localDateTimeToTimestampNanos(
+      LocalDateTime.parse("2021-03-02T03:04:05.123456789"), precision = 9)
+    val ntzExpectedSub = DateTimeUtils.localDateTimeToTimestampNanos(
+      LocalDateTime.parse("2019-12-02T03:04:05.123456789"), precision = 9)
+
+    checkEvaluation(
+      TimestampAddYMInterval(Literal.create(ntzStart, ntzType), Literal(interval), Some("UTC")),
+      ntzExpectedAdd)
+    checkEvaluation(
+      TimestampAddYMInterval(
+        Literal.create(ntzStart, ntzType),
+        UnaryMinus(Literal(interval)),
+        Some("UTC")),
+      DateTimeUtils.localDateTimeToTimestampNanos(
+        LocalDateTime.parse("2018-11-02T03:04:05.123456789"), precision = 9))
+    checkEvaluation(
+      TimestampAddYMInterval(
+        Literal.create(ntzStart, ntzType), Literal(minusInterval), Some("UTC")),
+      ntzExpectedSub)
+    assert(ntzExpectedAdd.nanosWithinMicro == ntzStart.nanosWithinMicro)
+    assert(ntzExpectedSub.nanosWithinMicro == ntzStart.nanosWithinMicro)
+
+    val ltzType = TimestampLTZNanosType(9)
+    val ltzStart = DateTimeUtils.instantToTimestampNanos(
+      Instant.parse("2020-01-02T03:04:05.123456789Z"), precision = 9)
+    val ltzExpectedAdd = DateTimeUtils.instantToTimestampNanos(
+      Instant.parse("2021-03-02T03:04:05.123456789Z"), precision = 9)
+    val ltzExpectedSub = DateTimeUtils.instantToTimestampNanos(
+      Instant.parse("2019-12-02T03:04:05.123456789Z"), precision = 9)
+
+    checkEvaluation(
+      TimestampAddYMInterval(Literal.create(ltzStart, ltzType), Literal(interval), Some("UTC")),
+      ltzExpectedAdd)
+    checkEvaluation(
+      TimestampAddYMInterval(
+        Literal.create(ltzStart, ltzType),
+        UnaryMinus(Literal(interval)),
+        Some("UTC")),
+      DateTimeUtils.instantToTimestampNanos(
+        Instant.parse("2018-11-02T03:04:05.123456789Z"), precision = 9))
+    checkEvaluation(
+      TimestampAddYMInterval(
+        Literal.create(ltzStart, ltzType), Literal(minusInterval), Some("UTC")),
+      ltzExpectedSub)
+    assert(ltzExpectedAdd.nanosWithinMicro == ltzStart.nanosWithinMicro)
+    assert(ltzExpectedSub.nanosWithinMicro == ltzStart.nanosWithinMicro)
+
+    yearMonthIntervalTypes.foreach { it =>
+      checkConsistencyBetweenInterpretedAndCodegen(
+        (ts: Expression, ym: Expression) => TimestampAddYMInterval(ts, ym, Some("UTC")),
+        ntzType, it)
+      checkConsistencyBetweenInterpretedAndCodegen(
+        (ts: Expression, ym: Expression) => TimestampAddYMInterval(ts, ym, Some("UTC")),
+        ltzType, it)
+    }
+  }
+
   test("SPARK-37552: convert a timestamp_ntz to another time zone") {
     checkEvaluation(
       ConvertTimezone(

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala

@@ -1199,6 +1199,49 @@ class DateTimeUtilsSuite extends SparkFunSuite with Matchers with SQLHelper {
     }
   }
 
+  test("SPARK-57825: timestamp nanos add year-month interval preserves nanosWithinMicro") {
+    def nanos(epochMicros: Long, nanosWithinMicro: Int): TimestampNanosVal =
+      TimestampNanosVal.fromParts(epochMicros, nanosWithinMicro.toShort)
+
+    // The epoch-micros part follows the micro `timestampAddMonths` (including the Jan-31 -> Feb-29
+    // day clamp in a leap year) while the sub-microsecond remainder is carried through unchanged.
+    assert(timestampNanosAddMonths(
+      nanos(date(2020, 1, 31, 12, 0, 0, 123000, LA), 789), 1, LA) ===
+      nanos(date(2020, 2, 29, 12, 0, 0, 123000, LA), 789))
+
+    outstandingZoneIds.foreach { zid =>
+      // The sub-microsecond remainder is preserved for the boundary values 0, 1 and 999.
+      Seq(0, 1, 999).foreach { rem =>
+        // Zero interval is a no-op on both the epoch-micros and the remainder.
+        assert(timestampNanosAddMonths(
+          nanos(date(2021, 3, 18, 19, 44, 1, 123456, zid), rem), 0, zid) ===
+          nanos(date(2021, 3, 18, 19, 44, 1, 123456, zid), rem))
+        // Adding whole years/months shifts only the month field, never the fraction.
+        assert(timestampNanosAddMonths(
+          nanos(date(2020, 1, 2, 3, 4, 5, 123456, zid), rem), 14, zid) ===
+          nanos(date(2021, 3, 2, 3, 4, 5, 123456, zid), rem))
+        // Subtracting months is symmetric.
+        assert(timestampNanosAddMonths(
+          nanos(date(2020, 1, 2, 3, 4, 5, 123456, zid), rem), -1, zid) ===
+          nanos(date(2019, 12, 2, 3, 4, 5, 123456, zid), rem))
+        // Pre-epoch (negative epochMicros) value.
+        assert(timestampNanosAddMonths(
+          nanos(date(1960, 1, 2, 3, 4, 5, 123456, zid), rem), 1, zid) ===
+          nanos(date(1960, 2, 2, 3, 4, 5, 123456, zid), rem))
+      }
+    }
+
+    // Consistency with the micro helper: epochMicros matches `timestampAddMonths` exactly and the
+    // remainder is independent of the interval amount.
+    outstandingZoneIds.foreach { zid =>
+      val start = nanos(date(2020, 1, 2, 3, 4, 5, 123456, zid), 789)
+      val months = 15
+      val result = timestampNanosAddMonths(start, months, zid)
+      assert(result.epochMicros === timestampAddMonths(start.epochMicros, months, zid))
+      assert(result.nanosWithinMicro === start.nanosWithinMicro)
+    }
+  }
+
   test("SPARK-57159: timestampNanosToEpochNanos packs into int64 epoch-nanoseconds") {
     def nanos(epochMicros: Long, nanosWithinMicro: Int): TimestampNanosVal =
       TimestampNanosVal.fromParts(epochMicros, nanosWithinMicro.toShort)

sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ltz-nanos.sql.out

@@ -603,6 +603,41 @@ Project [cast(1960-01-01 19:04:05.123456789 + INTERVAL '0 00:00:00.000001' DAY T
 +- OneRowRelation
 
 
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' YEAR
+-- !query analysis
+Project [2020-01-01 19:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' YEAR#x]
++- OneRowRelation
+
+
+-- !query
+SELECT INTERVAL '1' YEAR + TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC'
+-- !query analysis
+Project [2020-01-01 19:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' YEAR#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH
+-- !query analysis
+Project [2020-01-01 19:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' MONTH#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - INTERVAL '1-2' YEAR TO MONTH
+-- !query analysis
+Project [2020-01-01 19:04:05.123456789 - INTERVAL '1-2' YEAR TO MONTH AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '1960-01-31 03:04:05.123456789 UTC' + INTERVAL '1' MONTH
+-- !query analysis
+Project [1960-01-30 19:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_LTZ '1960-01-30 19:04:05.123456789' + INTERVAL '1' MONTH#x]
++- OneRowRelation
+
+
 -- !query
 SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0)
 -- !query analysis
@@ -627,28 +662,6 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException
 }
 
 
--- !query
-SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH
--- !query analysis
-org.apache.spark.sql.catalyst.ExtendedAnalysisException
-{
-  "errorClass" : "DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES",
-  "sqlState" : "42K09",
-  "messageParameters" : {
-    "left" : "\"TIMESTAMP_LTZ(9)\"",
-    "right" : "\"INTERVAL MONTH\"",
-    "sqlExpr" : "\"(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' MONTH)\""
-  },
-  "queryContext" : [ {
-    "objectType" : "",
-    "objectName" : "",
-    "startIndex" : 8,
-    "stopIndex" : 77,
-    "fragment" : "TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH"
-  } ]
-}
-
-
 -- !query
 SELECT max(c), min(c) FROM VALUES
   (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'),

sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ntz-nanos.sql.out

@@ -530,6 +530,41 @@ Project [cast(1960-01-02 03:04:05.123456789 + INTERVAL '0 00:00:00.000001' DAY T
 +- OneRowRelation
 
 
+-- !query
+SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR
+-- !query analysis
+Project [2020-01-02 03:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR#x]
++- OneRowRelation
+
+
+-- !query
+SELECT INTERVAL '1' YEAR + TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789'
+-- !query analysis
+Project [2020-01-02 03:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH
+-- !query analysis
+Project [2020-01-02 03:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH
+-- !query analysis
+Project [2020-01-02 03:04:05.123456789 - INTERVAL '1-2' YEAR TO MONTH AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH
+-- !query analysis
+Project [1960-01-31 03:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH#x]
++- OneRowRelation
+
+
 -- !query
 SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + make_interval(0, 1, 0, 2, 0, 0, 0)
 -- !query analysis
@@ -554,28 +589,6 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException
 }
 
 
--- !query
-SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH
--- !query analysis
-org.apache.spark.sql.catalyst.ExtendedAnalysisException
-{
-  "errorClass" : "DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES",
-  "sqlState" : "42K09",
-  "messageParameters" : {
-    "left" : "\"TIMESTAMP_NTZ(9)\"",
-    "right" : "\"INTERVAL MONTH\"",
-    "sqlExpr" : "\"(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH)\""
-  },
-  "queryContext" : [ {
-    "objectType" : "",
-    "objectName" : "",
-    "startIndex" : 8,
-    "stopIndex" : 73,
-    "fragment" : "TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH"
-  } ]
-}
-
-
 -- !query
 SELECT max(c), min(c) FROM VALUES
   (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'),

sql/core/src/test/resources/sql-tests/inputs/timestamp-ltz-nanos.sql

@@ -168,11 +168,19 @@ SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' -
   INTERVAL '1 00:04:00.000321' DAY TO SECOND;
 SELECT TIMESTAMP_LTZ '1960-01-02 03:04:05.123456789 UTC' +
   INTERVAL '0 00:00:00.000001' DAY TO SECOND;
--- SPARK-57501: nanos timestamps support only ANSI day-time intervals. A (legacy) calendar interval
--- is rejected by TimestampAddInterval's type check, and a year-month interval has no supported
--- operator overload.
-SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0);
+-- SPARK-57825: TIMESTAMP_LTZ(p) +/- ANSI year-month interval keeps the nanos type/precision and
+-- carries the whole fraction (including the sub-microsecond digits) through unchanged; the month
+-- shift is applied on the session-local wall clock.
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' YEAR;
+-- The interval-first operand order resolves to the same addition.
+SELECT INTERVAL '1' YEAR + TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC';
 SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH;
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - INTERVAL '1-2' YEAR TO MONTH;
+-- Jan-31 -> Feb-29 day clamp on a pre-epoch (leap-year) value.
+SELECT TIMESTAMP_LTZ '1960-01-31 03:04:05.123456789 UTC' + INTERVAL '1' MONTH;
+-- SPARK-57501, SPARK-57825: nanos timestamps support ANSI day-time and year-month intervals; the
+-- legacy calendar interval is still rejected by TimestampAddInterval's type check.
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0);
 
 -- SPARK-57103: MAX / MIN over nanosecond-precision TIMESTAMP_LTZ. The aggregate preserves the
 -- nanosecond type and orders by the sub-microsecond remainder; NULLs are ignored. Values are