Extract a fixed offset from `Tz` once and apply it to array's timestamps
- Dominant language
- Rust
- Stars
- 3.6k
- Forks
- 1.3k
- Avg merge
- 2d 18h
- Merged PRs (30d)
- 169
Description
## Problem statement
Computations dealing with arrays of date/time values in a time zone go by obtaining the `Tz` value and then calling its methods per element.
This is inefficient for two reasons:
1. A per-element dispatch on the `TzInner` enum, which may or may not be optimized out of the loop;
2. For named time zones, a lookup into the chrono-tz generated dispatch and method call(s), which are again subject to optimization. Some of the time zones, most importantly UTC, amount to applying a fixed offset invariantly to all values.
## Candidates for optimization
Extracted and summarized with Claude:
### 1. `date_part` extraction (year/month/day/hour/etc.)
[arrow-arith/src/temporal.rs:418-488](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-arith/src/temporal.rs#L418-L488)
Four near-identical impls for `TimestampSecondType` / `MillisecondType` / `MicrosecondType` / `NanosecondType`. `tz` is resolved once via `get_tz(...)`, then each element does:
```rust
timestamp_s_to_datetime(d)
.map(|c| Utc.from_utc_datetime(&c).with_timezone(&tz))
.map(map_func)
```
`with_timezone` invokes `Tz::offset_from_utc_datetime` per row.
### 2. Timestamp ± interval arithmetic
[arrow-arith/src/numeric.rs:437](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-arith/src/numeric.rs#L437) resolves `l_tz` once, then per-element ops (e.g. [Op::Add at numeric.rs:456-465](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-arith/src/numeric.rs#L456-L465)) call `T::add_year_month(l, r, l_tz)`, which forwards to [add_year_months in arrow-array/src/types.rs:412-419](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-array/src/types.rs#L412-L419) → `as_datetime_with_timezone` (candidate 6) → `.with_timezone(&tz)`.
### 3. Timestamp display / pretty-print
[arrow-cast/src/display.rs:755-763](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-cast/src/display.rs#L755-L763), `write_timestamp`, invoked once per row by the [timestamp_display! macro](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-cast/src/display.rs#L783-L812):
```rust
let date = Utc.from_utc_datetime(&naive).with_timezone(&tz);
```
### 4. Timestamp → Time32/Time64 cast
[arrow-cast/src/cast/mod.rs:615-631](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-cast/src/cast/mod.rs#L615-L631) (`as_time_res_with_timezone`), called per element from many `(Timestamp(_, tz), Time32/Time64(_))` match arms, e.g. [cast/mod.rs:1974-1985](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-cast/src/cast/mod.rs#L1974-L1985). `tz` is parsed once per cast call, then `as_datetime_with_timezone` (candidate 6) is called inside `try_unary` for every element.
### 5. Timestamp → Date32 cast
[arrow-cast/src/cast/mod.rs:633-657](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-cast/src/cast/mod.rs#L633-L657) (`timestamp_to_date32`): `tz` parsed once, then `array.try_unary(|x| as_datetime_with_timezone::(x, tz)...)` per element.
### 6. Timestamp → Timestamp with different timezone
[arrow-cast/src/cast/mod.rs:2629-2648](https://github.com/apache/arrow-rs/blob/e1e71ea89b211758e20254296008f96e744ed138/arrow-cast/src/cast/mod.rs#L2629-L2648) (`adjust_timestamp_to_timezone`):
```rust
let adjust = |o| {
let local = as_datetime::(o)?;
let offset = to_tz.offset_from_local_datetime(&local).single()?;
T::from_naive_datetime(local - offset.fix(), None)
};
```
`to_tz` is loop-invariant across the whole `unary_opt`/`try_unary` call. This is the cleanest optimization target: if `to_tz` is a fixed offset, `offset` is also loop-invariant and can be computed once outside the closure.
## Proposed solution
Add a public accessor on `Tz`, e.g. `pub fn fixed_offset(&self) -> Option`, in `arrow-array/src/timezone.rs`:
- Returns `Some` immediately for `TzInner::Offset`.
- Can also return `Some` on recognized named time zone variants, primarily `UTC` and its aliases, `Etc/GMT*` etc..
Each candidate site would then check this once outside its per-element loop and, when `Some(offset)`, apply the offset via plain arithmetic (`NaiveDateTime + Duration` / integer add) instead of dispatching through `chrono::TimeZone` per element.
### Other aliternatives considered
Replace chrono with jiff as proposed in #9183. Jiff has the requisite method on its `TimeZone` type.
That is a much wider change and requires API breaks.
## Additional context
In testing of a similar optimization done outsize of arrow-rs, elimination of the chrono-tz dispatch from offsetting a timestamp for a timestamp-to-string array cast resulted in 13-18% improvement.
Contributor guide
Research direction
Start with arrow-array/src/timezone.rs and the Tz implementation, then inspect the candidate paths in arrow-arith/src/temporal.rs, arrow-arith/src/numeric.rs, and arrow-cast/src/cast/mod.rs and display.rs. Done means fixed-offset time zones are detected once and applied without per-element timezone dispatch, while the listed timestamp operations retain their existing behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data-engineering, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100