uros-b commented on code in PR #56848:
URL: https://github.com/apache/spark/pull/56848#discussion_r3949973292


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala:
##########
@@ -1393,3 +1448,235 @@ object DateTimeUtils extends SparkDateTimeUtils {
     c
   }
 }
+
+/**
+ * A zone's transition schedule as primitive arrays: the sorted epoch-second 
transition instants
+ * (`transSec`) and the UTC offset (in seconds) in effect on each window 
`[transSec(i),
+ * transSec(i + 1))` (`offAfter(i)`); `offBefore0` is the offset before the 
first transition. Built
+ * once per zone from the transitions in `[startSec, horizonSec)` -- 
historical plus
+ * rule-generated; outside that range on either side the rules are consulted 
directly rather than
+ * assuming the table's edge windows extend indefinitely. Immutable and shared 
read-only across
+ * tasks via [[ZoneOffsetCache.tableFor]].
+ */
+private[util] final class ZoneTransitionTable(
+    val transSec: Array[Long],
+    val offAfter: Array[Int],
+    val offBefore0: Int,
+    val startSec: Long,
+    val horizonSec: Long) {
+
+  /** Largest `i` with `transSec(i) <= epochSec`, or -1 when before the first 
transition. */
+  def floorIndex(epochSec: Long): Int = {
+    var lo = 0
+    var hi = transSec.length - 1
+    var res = -1
+    while (lo <= hi) {
+      val mid = (lo + hi) >>> 1
+      if (transSec(mid) <= epochSec) {
+        res = mid
+        lo = mid + 1
+      } else {
+        hi = mid - 1
+      }
+    }
+    res
+  }
+}
+
+object ZoneOffsetCache {
+  // The transition table is the same for every task using a given zone, so 
build it once per JVM.
+  private val tables = new java.util.concurrent.ConcurrentHashMap[ZoneId, 
ZoneTransitionTable]()
+  private val tableStart = Instant.parse("1600-01-01T00:00:00Z")
+  private val tableHorizon = Instant.parse("2200-01-01T00:00:00Z")
+
+  private[util] def tableFor(zoneId: ZoneId): ZoneTransitionTable =
+    tables.computeIfAbsent(zoneId, z => buildTable(z))
+
+  private def buildTable(zoneId: ZoneId): ZoneTransitionTable = {
+    val rules = zoneId.getRules
+    val secs = scala.collection.mutable.ArrayBuffer.empty[Long]
+    val offs = scala.collection.mutable.ArrayBuffer.empty[Int]
+    var before0 = 0
+    var seen = false
+    var cur = tableStart
+    var t = rules.nextTransition(cur)
+    while (t != null && t.getInstant.isBefore(tableHorizon)) {
+      if (!seen) {
+        before0 = t.getOffsetBefore.getTotalSeconds
+        seen = true
+      }
+      secs += t.toEpochSecond
+      offs += t.getOffsetAfter.getTotalSeconds
+      // `plusNanos(1)` guarantees progress; transitions are >= 1s apart so 
none is skipped.
+      cur = t.getInstant.plusNanos(1)
+      t = rules.nextTransition(cur)
+    }
+    if (!seen) {
+      before0 = rules.getOffset(tableStart).getTotalSeconds
+    }
+    new ZoneTransitionTable(secs.toArray, offs.toArray, before0, 
tableStart.getEpochSecond,
+      tableHorizon.getEpochSecond)
+  }
+}
+
+/**
+ * Per-task memoization of a zone's UTC offset, used by the 
[[DateTimeUtils.truncTimestamp]] hot
+ * path. The session zone is constant for a query and the offset is 
piecewise-constant between DST
+ * transitions, so a lookup reduces to a range check against a cached 
constant-offset window
+ * `[lo, hi)`.
+ *
+ * The most-recently-used window is held in plain fields, a branch-only fast 
path that temporally
+ * clustered rows -- the common case -- keep hitting. A miss resolves the 
enclosing window
+ * with a single binary search over the shared [[ZoneTransitionTable]] (no 
allocation), which is
+ * itself cheaper than a bare `getOffset`, so even miss-heavy inputs (e.g. 
instants scattered over
+ * many decades in random order) stay close to the uncached path.
+ *
+ * Not thread-safe by design: a fresh instance is created per task (codegen 
mutable state) and used
+ * single-threaded, mirroring how stateful per-row helpers are scoped in 
generated code.
+ */
+class ZoneOffsetCache(val zoneId: ZoneId) {

Review Comment:
   This class, and the cache-taking `truncTimestamp`/`truncTimestampNanos` 
overloads are still `public`, whereas a previous review iteration asked to 
scope this not-thread-safe, per-task-only helper `private[sql]` (still 
module-visible, so the generated Java's FQN reference keeps working).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to