Licht-T commented on code in PR #56848:
URL: https://github.com/apache/spark/pull/56848#discussion_r3950221445
##########
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:
Done, with one wrinkle: Scala emits no static forwarder for `private[sql]`
members, which broke the generated code's `DateTimeUtils.truncTimestamp(...)`
call (caught by the codegen tests). The `TruncInstant` codegen now calls
through `DateTimeUtils$.MODULE$` instead.
--
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]
