[WIP][SPARK-57832][SQL] Support timestamp - timestamp subtraction for nanosecond-precision timestamps

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

@@ -145,7 +145,9 @@ object BinaryArithmeticWithDatetimeResolver {
             TimestampAddInterval(l, UnaryMinus(r, context.evalMode == EvalMode.ANSI))), l.dataType)
         case _
           if AnyTimestampTypeExpression.unapply(l) ||
-            AnyTimestampTypeExpression.unapply(r) =>
+            AnyTimestampTypeExpression.unapply(r) ||
+            AnyTimestampNanoType.acceptsType(l.dataType) ||
+            AnyTimestampNanoType.acceptsType(r.dataType) =>
           SubtractTimestamps(l, r)
         case (_, DateType) => SubtractDates(l, r)
         case (DateType, dt) if dt != StringType => DateSub(l, r)

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

@@ -302,6 +302,35 @@ abstract class TypeCoercionHelper {
         }
     }
 
+  /** Whether `dt` is on the LTZ/NTZ timestamp fractional-precision axis (a micro or nanos type). */
+  private def isTimestampFamily(dt: DataType): Boolean =
+    TimestampFamily.fractionalPrecision(dt).isDefined
+
+  /**
+   * Common operand type for [[SubtractTimestamps]] over two differing timestamp-family operands.
+   * The operands are widened to the larger of the two fractional-second precisions (the micro types
+   * count as 6, the nanos types carry their own precision `p` in [7, 9]) and unified in one
+   * time-zone family:
+   *   - a cross-family pair unifies in the no-time-zone (NTZ) family, mirroring the microsecond
+   *     precedent where TIMESTAMP - TIMESTAMP_NTZ coerces both operands to TIMESTAMP_NTZ;
+   *   - a same-family pair keeps that family, so a TIMESTAMP - TIMESTAMP_LTZ(p) style pair still
+   *     subtracts in the session time zone (DST-aware) exactly as a pure LTZ pair does.
+   * The subtraction reads only each operand's epochMicros and always yields a microsecond-grid
+   * DayTimeIntervalType, so widening the precision never changes the numeric result -- it only keeps
+   * the two operands the same concrete type. For a pure-micro cross-family pair this returns
+   * TimestampNTZType, identical to the pre-nanos behavior.
+   */
+  private def subtractTimestampsCommonType(dt1: DataType, dt2: DataType): DataType = {
+    val p = math.max(
+      TimestampFamily.fractionalPrecision(dt1).getOrElse(6),
+      TimestampFamily.fractionalPrecision(dt2).getOrElse(6))
+    if (TimestampFamily.isLtz(dt1) && TimestampFamily.isLtz(dt2)) {
+      if (p <= 6) TimestampType else TimestampLTZNanosType(p)
+    } else {
+      if (p <= 6) TimestampNTZType else TimestampNTZNanosType(p)
+    }
+  }
+
   /**
    * Type coercion helper that matches agaist [[In]] and [[InSubquery]] expressions in order to
    * type coerce LHS and RHS to expected types.
@@ -738,14 +767,18 @@ abstract class TypeCoercionHelper {
         d.copy(startDate = Cast(d.startDate, DateType))
       case d @ DateSub(StringTypeExpression(), _) => d.copy(startDate = Cast(d.startDate, DateType))
 
-      case s @ SubtractTimestamps(DateTypeExpression(), AnyTimestampTypeExpression(), _, _) =>
+      case s @ SubtractTimestamps(DateTypeExpression(), r, _, _)
+          if isTimestampFamily(r.dataType) =>
         s.copy(left = Cast(s.left, s.right.dataType))
-      case s @ SubtractTimestamps(AnyTimestampTypeExpression(), DateTypeExpression(), _, _) =>
+      case s @ SubtractTimestamps(l, DateTypeExpression(), _, _)
+          if isTimestampFamily(l.dataType) =>
         s.copy(right = Cast(s.right, s.left.dataType))
-      case s @ SubtractTimestamps(AnyTimestampTypeExpression(), AnyTimestampTypeExpression(), _, _)
-          if s.left.dataType != s.right.dataType =>
-        val newLeft = castIfNotSameType(s.left, TimestampNTZType)
-        val newRight = castIfNotSameType(s.right, TimestampNTZType)
+      case s @ SubtractTimestamps(l, r, _, _)
+          if isTimestampFamily(l.dataType) && isTimestampFamily(r.dataType) &&
+            l.dataType != r.dataType =>
+        val commonType = subtractTimestampsCommonType(l.dataType, r.dataType)
+        val newLeft = castIfNotSameType(s.left, commonType)
+        val newRight = castIfNotSameType(s.right, commonType)
         s.copy(left = newLeft, right = newRight)
 
       case t @ TimestampAddInterval(StringTypeExpression(), _, _) =>
@@ -766,14 +799,18 @@ abstract class TypeCoercionHelper {
       case d @ DateSub(AnyTimestampTypeExpression(), _) =>
         d.copy(startDate = Cast(d.startDate, DateType))
 
-      case s @ SubtractTimestamps(DateTypeExpression(), AnyTimestampTypeExpression(), _, _) =>
+      case s @ SubtractTimestamps(DateTypeExpression(), r, _, _)
+          if isTimestampFamily(r.dataType) =>
         s.copy(left = Cast(s.left, s.right.dataType))
-      case s @ SubtractTimestamps(AnyTimestampTypeExpression(), DateTypeExpression(), _, _) =>
+      case s @ SubtractTimestamps(l, DateTypeExpression(), _, _)
+          if isTimestampFamily(l.dataType) =>
         s.copy(right = Cast(s.right, s.left.dataType))
-      case s @ SubtractTimestamps(AnyTimestampTypeExpression(), AnyTimestampTypeExpression(), _, _)
-          if s.left.dataType != s.right.dataType =>
-        val newLeft = castIfNotSameType(s.left, TimestampNTZType)
-        val newRight = castIfNotSameType(s.right, TimestampNTZType)
+      case s @ SubtractTimestamps(l, r, _, _)
+          if isTimestampFamily(l.dataType) && isTimestampFamily(r.dataType) &&
+            l.dataType != r.dataType =>
+        val commonType = subtractTimestampsCommonType(l.dataType, r.dataType)
+        val newLeft = castIfNotSameType(s.left, commonType)
+        val newRight = castIfNotSameType(s.right, commonType)
         s.copy(left = newLeft, right = newRight)
 
       case other => other

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

@@ -4523,7 +4523,13 @@ case class SubtractTimestamps(
   def this(endTimestamp: Expression, startTimestamp: Expression) =
     this(endTimestamp, startTimestamp, SQLConf.get.legacyIntervalEnabled)
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyTimestampType, AnyTimestampType)
+  // Nanosecond-precision timestamps are accepted alongside the microsecond types. The difference is
+  // always reported on the microsecond grid (a DayTimeIntervalType has microsecond resolution), so
+  // each operand contributes only its epochMicros; the sub-microsecond remainder is truncated.
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(
+      TypeCollection(AnyTimestampType, AnyTimestampNanoType),
+      TypeCollection(AnyTimestampType, AnyTimestampNanoType))
   override def dataType: DataType =
     if (legacyInterval) CalendarIntervalType else DayTimeIntervalType()
 
@@ -4532,6 +4538,13 @@ case class SubtractTimestamps(
 
   @transient private lazy val zoneIdInEval: ZoneId = zoneIdForType(left.dataType)
 
+  // For the nanosecond carrier the child value is a boxed TimestampNanosVal, so read its
+  // epochMicros; for the microsecond timestamp types it is already a boxed Long.
+  private def toMicros(value: Any): Long = value match {
+    case v: TimestampNanosVal => v.epochMicros
+    case n => n.asInstanceOf[Long]
+  }
+
   @transient
   private lazy val evalFunc: (Long, Long) => Any = if (legacyInterval) {
     (leftMicros, rightMicros) =>
@@ -4541,17 +4554,29 @@ case class SubtractTimestamps(
       subtractTimestamps(leftMicros, rightMicros, zoneIdInEval)
   }
 
-  override def nullSafeEval(leftMicros: Any, rightMicros: Any): Any = {
-    evalFunc(leftMicros.asInstanceOf[Long], rightMicros.asInstanceOf[Long])
+  override def nullSafeEval(leftTs: Any, rightTs: Any): Any = {
+    evalFunc(toMicros(leftTs), toMicros(rightTs))
   }
 
-  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = if (legacyInterval) {
-    defineCodeGen(ctx, ev, (end, start) =>
-      s"new org.apache.spark.unsafe.types.CalendarInterval(0, 0, $end - $start)")
-  } else {
-    val zid = ctx.addReferenceObj("zoneId", zoneIdInEval, classOf[ZoneId].getName)
-    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-    defineCodeGen(ctx, ev, (l, r) => s"""$dtu.subtractTimestamps($l, $r, $zid)""")
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    // The nanosecond carrier exposes epochMicros as a public field; the microsecond types are
+    // already primitive longs. Reduce each operand to microseconds before subtracting.
+    def toMicrosCode(e: Expression): String => String = e.dataType match {
+      case _: AnyTimestampNanoType => c => s"$c.epochMicros"
+      case _ => c => c
+    }
+    val leftMicros = toMicrosCode(left)
+    val rightMicros = toMicrosCode(right)
+    if (legacyInterval) {
+      defineCodeGen(ctx, ev, (end, start) =>
+        s"new org.apache.spark.unsafe.types.CalendarInterval(0, 0, " +
+          s"${leftMicros(end)} - ${rightMicros(start)})")
+    } else {
+      val zid = ctx.addReferenceObj("zoneId", zoneIdInEval, classOf[ZoneId].getName)
+      val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+      defineCodeGen(ctx, ev, (l, r) =>
+        s"""$dtu.subtractTimestamps(${leftMicros(l)}, ${rightMicros(r)}, $zid)""")
+    }
   }
 
   override def toString: String = s"($left - $right)"

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionSuite.scala

@@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.plans.ReferenceAllColumns
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.{Rule, RuleExecutor}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.CalendarInterval
@@ -436,6 +437,48 @@ abstract class TypeCoercionSuiteBase extends AnalysisTest {
     ruleTest(rule,
       SubtractTimestamps(timestampNTZLiteral, timestampLiteral),
       SubtractTimestamps(timestampNTZLiteral, Cast(timestampLiteral, TimestampNTZType)))
+
+    // SPARK-57832: subtraction accepts nanosecond-precision timestamps. A DATE operand takes the
+    // nanos type of the other side; a timestamp pair that differs only in precision or family is
+    // widened to a common type before being handed to SubtractTimestamps.
+    val ntzNanos9 = Literal.create(
+      DateTimeUtils.localDateTimeToTimestampNanos(
+        LocalDateTime.parse("2021-01-01T00:00:00"), precision = 9),
+      TimestampNTZNanosType(9))
+    val ltzNanos9 = Literal.create(
+      DateTimeUtils.instantToTimestampNanos(
+        java.time.Instant.parse("2021-01-01T00:00:00Z"), precision = 9),
+      TimestampLTZNanosType(9))
+    val ntzNanos7 = Literal.create(
+      DateTimeUtils.localDateTimeToTimestampNanos(
+        LocalDateTime.parse("2021-01-01T00:00:00"), precision = 7),
+      TimestampNTZNanosType(7))
+
+    // DATE - nanos and nanos - DATE cast the DATE side to the nanos type.
+    ruleTest(rule,
+      SubtractTimestamps(dateLiteral, ntzNanos9),
+      SubtractTimestamps(Cast(dateLiteral, TimestampNTZNanosType(9)), ntzNanos9))
+    ruleTest(rule,
+      SubtractTimestamps(ntzNanos9, dateLiteral),
+      SubtractTimestamps(ntzNanos9, Cast(dateLiteral, TimestampNTZNanosType(9))))
+    // Same-precision same-family pair is already the same type -> left untouched.
+    ruleTest(rule,
+      SubtractTimestamps(ntzNanos9, ntzNanos9),
+      SubtractTimestamps(ntzNanos9, ntzNanos9))
+    // Cross-family nanos pair (LTZ vs NTZ) unifies in the NTZ family at the max precision.
+    ruleTest(rule,
+      SubtractTimestamps(ltzNanos9, ntzNanos7),
+      SubtractTimestamps(
+        Cast(ltzNanos9, TimestampNTZNanosType(9)), Cast(ntzNanos7, TimestampNTZNanosType(9))))
+    // Micro NTZ vs nanos NTZ: same family, widen precision to the nanos type.
+    ruleTest(rule,
+      SubtractTimestamps(timestampNTZLiteral, ntzNanos9),
+      SubtractTimestamps(Cast(timestampNTZLiteral, TimestampNTZNanosType(9)), ntzNanos9))
+    // Micro LTZ (TIMESTAMP) vs nanos LTZ: same LTZ family, widen precision to the nanos LTZ type
+    // so the subtraction still runs in the session time zone.
+    ruleTest(rule,
+      SubtractTimestamps(timestampLiteral, ltzNanos9),
+      SubtractTimestamps(Cast(timestampLiteral, TimestampLTZNanosType(9)), ltzNanos9))
   }
 
   test("datetime comparison") {

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

@@ -2407,6 +2407,91 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("SPARK-57832: subtract nanosecond-precision timestamps") {
+    // The difference between two nanosecond timestamps is reported on the microsecond grid: only
+    // each operand's epochMicros participates, so the sub-microsecond remainder is truncated. The
+    // two operands below share the same wall clock down to the microsecond and differ only within
+    // it (789 vs 111), so a nanosecond-aware subtraction and a micros-grid one must agree on zero.
+    val ntzType = TimestampNTZNanosType(9)
+    val ntzLeft = DateTimeUtils.localDateTimeToTimestampNanos(
+      LocalDateTime.parse("2020-01-02T03:04:05.123456789"), precision = 9)
+    val ntzRight = DateTimeUtils.localDateTimeToTimestampNanos(
+      LocalDateTime.parse("2020-01-01T03:04:05.000000111"), precision = 9)
+    // 1 day + 0.123456 s; the 789/111 sub-microsecond digits drop out.
+    checkEvaluation(
+      SubtractTimestamps(
+        Literal.create(ntzLeft, ntzType),
+        Literal.create(ntzRight, ntzType),
+        legacyInterval = false,
+        timeZoneId = Some("UTC")),
+      Duration.ofDays(1).plus(123456, ChronoUnit.MICROS))
+    // Two values inside the same microsecond subtract to exactly zero (remainder truncated).
+    checkEvaluation(
+      SubtractTimestamps(
+        Literal.create(ntzLeft, ntzType),
+        Literal.create(
+          DateTimeUtils.localDateTimeToTimestampNanos(
+            LocalDateTime.parse("2020-01-02T03:04:05.123456001"), precision = 9),
+          ntzType),
+        legacyInterval = false,
+        timeZoneId = Some("UTC")),
+      Duration.ZERO)
+    // Legacy calendar-interval result carries the same microsecond difference.
+    checkEvaluation(
+      SubtractTimestamps(
+        Literal.create(ntzLeft, ntzType),
+        Literal.create(ntzRight, ntzType),
+        legacyInterval = true,
+        timeZoneId = Some("UTC")),
+      new CalendarInterval(0, 0, MICROS_PER_DAY + 123456L))
+
+    // LTZ nanos: subtraction reads the local wall clock at the session zone. Evaluated at UTC the
+    // instants and their local date-times coincide, so the difference matches the NTZ case above.
+    val ltzType = TimestampLTZNanosType(9)
+    val ltzLeft = DateTimeUtils.instantToTimestampNanos(
+      Instant.parse("2020-01-02T03:04:05.123456789Z"), precision = 9)
+    val ltzRight = DateTimeUtils.instantToTimestampNanos(
+      Instant.parse("2020-01-01T03:04:05.000000111Z"), precision = 9)
+    checkEvaluation(
+      SubtractTimestamps(
+        Literal.create(ltzLeft, ltzType),
+        Literal.create(ltzRight, ltzType),
+        legacyInterval = false,
+        timeZoneId = Some("UTC")),
+      Duration.ofDays(1).plus(123456, ChronoUnit.MICROS))
+
+    // Pre-epoch operand exercises the negative-epoch path; the result stays on the micros grid.
+    val ntzPreEpoch = DateTimeUtils.localDateTimeToTimestampNanos(
+      LocalDateTime.parse("1960-01-01T00:00:00.000000999"), precision = 9)
+    checkEvaluation(
+      SubtractTimestamps(
+        Literal.create(ntzLeft, ntzType),
+        Literal.create(ntzPreEpoch, ntzType),
+        legacyInterval = false,
+        timeZoneId = Some("UTC")),
+      Duration.between(
+        LocalDateTime.parse("1960-01-01T00:00:00"),
+        LocalDateTime.parse("2020-01-02T03:04:05.123456")))
+
+    // NULL operands propagate.
+    checkEvaluation(
+      SubtractTimestamps(
+        Literal.create(null, ntzType),
+        Literal.create(ntzRight, ntzType),
+        legacyInterval = false,
+        timeZoneId = Some("UTC")),
+      null)
+
+    Seq(false, true).foreach { legacy =>
+      checkConsistencyBetweenInterpretedAndCodegen(
+        (l: Expression, r: Expression) => SubtractTimestamps(l, r, legacy, Some("UTC")),
+        ntzType, ntzType)
+      checkConsistencyBetweenInterpretedAndCodegen(
+        (l: Expression, r: Expression) => SubtractTimestamps(l, r, legacy, Some("UTC")),
+        ltzType, ltzType)
+    }
+  }
+
   test("SPARK-37552: convert a timestamp_ntz to another time zone") {
     checkEvaluation(
       ConvertTimezone(

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

@@ -662,6 +662,54 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException
 }
 
 
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC'
+-- !query analysis
+Project [(2020-01-01 19:04:05.123456789 - 2019-12-31 19:04:05.000000111) AS (TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP_LTZ '2019-12-31 19:04:05.000000111')#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456001 UTC'
+-- !query analysis
+Project [(2020-01-01 19:04:05.123456789 - 2020-01-01 19:04:05.123456001) AS (TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP_LTZ '2020-01-01 19:04:05.123456001')#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC'
+-- !query analysis
+Project [(2019-12-31 19:04:05.000000111 - 2020-01-01 19:04:05.123456789) AS (TIMESTAMP_LTZ '2019-12-31 19:04:05.000000111' - TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789')#x]
++- OneRowRelation
+
+
+-- !query
+SELECT ('2020-01-02 03:04:05.1234567 UTC' :: timestamp_ltz(7)) - ('2020-01-01 03:04:05.000000009 UTC' :: timestamp_ltz(9))
+-- !query analysis
+Project [(cast(cast(2020-01-02 03:04:05.1234567 UTC as timestamp_ltz(7)) as timestamp_ltz(9)) - cast(2020-01-01 03:04:05.000000009 UTC as timestamp_ltz(9))) AS (CAST(2020-01-02 03:04:05.1234567 UTC AS TIMESTAMP_LTZ(7)) - CAST(2020-01-01 03:04:05.000000009 UTC AS TIMESTAMP_LTZ(9)))#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05 UTC'
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC' - TIMESTAMP_LTZ '1960-01-01 00:00:00.000000999 UTC'
+-- !query analysis
+Project [(2019-12-31 16:00:00.123456789 - 1959-12-31 16:00:00.000000999) AS (TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789' - TIMESTAMP_LTZ '1959-12-31 16:00:00.000000999')#x]
++- OneRowRelation
+
+
+-- !query
+SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - CAST(NULL AS timestamp_ltz(9))
+-- !query analysis
+Project [(2020-01-01 19:04:05.123456789 - cast(null as timestamp_ltz(9))) AS (TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - CAST(NULL AS TIMESTAMP_LTZ(9)))#x]
++- OneRowRelation
+
+
 -- !query
 SELECT max(c), min(c) FROM VALUES
   (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'),