[SPARK-56663][SQL] Restore fast path for date_trunc MINUTE/HOUR/DAY

Author: Licht-T

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

@@ -485,6 +485,43 @@ object DateTimeUtils extends SparkDateTimeUtils {
     instantToMicros(truncated.toInstant)
   }
 
+  /**
+   * Fast path for truncating to MINUTE/HOUR/DAY using offset arithmetic instead of
+   * allocating a `ZonedDateTime` per row. The offset is resolved once for `micros`; the
+   * truncation then runs as `floorMod` in local time. We fall back to [[truncToUnit]] when
+   * the offset at the candidate truncated instant differs from the offset at `micros`,
+   * which means the truncation crosses a DST/historical transition and the local-time
+   * alignment we computed is no longer valid (see SPARK-30766/30857). The check is
+   * skipped for fixed-offset zones. Sub-minute offsets (e.g. America/Los_Angeles LMT
+   * -07:52:58 — SPARK-33404) and 30/45-minute offsets (Asia/Kolkata +05:30, Asia/Kathmandu
+   * +05:45) are handled correctly by this path because the offset is applied as part of
+   * the arithmetic; no offset-alignment guard is needed.
+   */
+  private def truncToUnitFast(
+      micros: Long, zoneId: ZoneId, unitMicros: Long, fallbackUnit: ChronoUnit): Long = {
+    val rules = zoneId.getRules
+    val originalSec = Math.floorDiv(micros, MICROS_PER_SECOND)
+    val originalOffsetSec =
+      rules.getOffset(Instant.ofEpochSecond(originalSec)).getTotalSeconds.toLong
+    val offsetMicros = originalOffsetSec * MICROS_PER_SECOND
+    try {
+      val local = Math.addExact(micros, offsetMicros)
+      val truncatedLocal = local - Math.floorMod(local, unitMicros)
+      val candidate = Math.subtractExact(truncatedLocal, offsetMicros)
+      if (!rules.isFixedOffset) {
+        val candidateSec = Math.floorDiv(candidate, MICROS_PER_SECOND)
+        val candidateOffsetSec =
+          rules.getOffset(Instant.ofEpochSecond(candidateSec)).getTotalSeconds.toLong
+        if (candidateOffsetSec != originalOffsetSec) {
+          return truncToUnit(micros, zoneId, fallbackUnit)
+        }
+      }
+      candidate
+    } catch {
+      case _: ArithmeticException => truncToUnit(micros, zoneId, fallbackUnit)
+    }
+  }
+
   /**
    * Returns the trunc date time from original date time and trunc level.
    * Trunc level should be generated using `parseTruncLevel()`, should be between 0 and 9.
@@ -499,9 +536,12 @@ object DateTimeUtils extends SparkDateTimeUtils {
         micros - Math.floorMod(micros, MICROS_PER_MILLIS)
       case TRUNC_TO_SECOND =>
         micros - Math.floorMod(micros, MICROS_PER_SECOND)
-      case TRUNC_TO_MINUTE => truncToUnit(micros, zoneId, ChronoUnit.MINUTES)
-      case TRUNC_TO_HOUR => truncToUnit(micros, zoneId, ChronoUnit.HOURS)
-      case TRUNC_TO_DAY => truncToUnit(micros, zoneId, ChronoUnit.DAYS)
+      case TRUNC_TO_MINUTE =>
+        truncToUnitFast(micros, zoneId, MICROS_PER_MINUTE, ChronoUnit.MINUTES)
+      case TRUNC_TO_HOUR =>
+        truncToUnitFast(micros, zoneId, MICROS_PER_HOUR, ChronoUnit.HOURS)
+      case TRUNC_TO_DAY =>
+        truncToUnitFast(micros, zoneId, MICROS_PER_DAY, ChronoUnit.DAYS)
       case _ => // Try to truncate date levels
         val dDays = microsToDays(micros, zoneId)
         daysToMicros(truncDate(dDays, level), zoneId)

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

@@ -766,6 +766,71 @@ class DateTimeUtilsSuite extends SparkFunSuite with Matchers with SQLHelper {
     }
   }
 
+  test("truncTimestamp with sub-hour zone offsets") {
+    // Asia/Kolkata (+05:30) and Asia/Kathmandu (+05:45) are not aligned to HOUR in UTC.
+    // The fast path applies the offset as part of its arithmetic, so HOUR/DAY truncation
+    // produces the correct local-aligned result without needing the slow path.
+    val kolkata = getZoneId("Asia/Kolkata")
+    val ts = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-01-15T09:42:17.123456+05:30"), kolkata).get
+    val expectedHour = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-01-15T09:00:00+05:30"), kolkata).get
+    assert(DateTimeUtils.truncTimestamp(ts, DateTimeUtils.TRUNC_TO_HOUR, kolkata) === expectedHour)
+    val expectedDay = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-01-15T00:00:00+05:30"), kolkata).get
+    assert(DateTimeUtils.truncTimestamp(ts, DateTimeUtils.TRUNC_TO_DAY, kolkata) === expectedDay)
+
+    val kathmandu = getZoneId("Asia/Kathmandu")
+    val ts2 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-01-15T09:42:17.123456+05:45"), kathmandu).get
+    val expectedHour2 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-01-15T09:00:00+05:45"), kathmandu).get
+    assert(DateTimeUtils.truncTimestamp(
+      ts2, DateTimeUtils.TRUNC_TO_HOUR, kathmandu) === expectedHour2)
+  }
+
+  test("truncTimestamp across DST transitions") {
+    val la = getZoneId("America/Los_Angeles")
+    // Spring-forward in LA: 2024-03-10 02:00 PDT does not exist; 02:30 local maps to
+    // 2024-03-10 03:30 PDT in wall-clock terms. Use an instant just after the transition
+    // so HOUR/DAY truncation candidate falls into the pre-transition offset window.
+    val postSpring = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-03-10T03:30:00-07:00"), la).get
+    val expectedHour = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-03-10T03:00:00-07:00"), la).get
+    assert(DateTimeUtils.truncTimestamp(postSpring, DateTimeUtils.TRUNC_TO_HOUR, la)
+      === expectedHour)
+    val expectedDay = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-03-10T00:00:00-08:00"), la).get
+    assert(DateTimeUtils.truncTimestamp(postSpring, DateTimeUtils.TRUNC_TO_DAY, la)
+      === expectedDay)
+
+    // Fall-back in LA: 2024-11-03 01:30 occurs twice. Truncation to HOUR/DAY should
+    // produce the same wall-clock boundary as the slow path regardless.
+    val postFall = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-11-03T01:30:00-08:00"), la).get
+    val expectedHour2 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-11-03T01:00:00-08:00"), la).get
+    assert(DateTimeUtils.truncTimestamp(postFall, DateTimeUtils.TRUNC_TO_HOUR, la)
+      === expectedHour2)
+    val expectedDay2 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("2024-11-03T00:00:00-07:00"), la).get
+    assert(DateTimeUtils.truncTimestamp(postFall, DateTimeUtils.TRUNC_TO_DAY, la)
+      === expectedDay2)
+  }
+
+  test("SPARK-30766/30857: truncTimestamp before the epoch in HOUR/DAY") {
+    val la = getZoneId("America/Los_Angeles")
+    val ts1 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("1960-02-11T00:01:02.123"), la).get
+    val expectedHour1 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("1960-02-11T00:00:00"), la).get
+    assert(DateTimeUtils.truncTimestamp(ts1, DateTimeUtils.TRUNC_TO_HOUR, la) === expectedHour1)
+    val expectedDay1 = DateTimeUtils.stringToTimestamp(
+      UTF8String.fromString("1960-02-11T00:00:00"), la).get
+    assert(DateTimeUtils.truncTimestamp(ts1, DateTimeUtils.TRUNC_TO_DAY, la) === expectedDay1)
+  }
+
   test("SPARK-51554: time truncation using timeTrunc") {
     // 01:02:03.400500600
     val input = localTimeToNanos(LocalTime.of(1, 2, 3, 400500600))