Re: [I] Make RollupMapper calendar windows DST-aware [airflow]

["DavidAmsellem (via GitHub)" <gi...@apache.org>]

DavidAmsellem commented on issue #68004:
URL: https://github.com/apache/airflow/issues/68004#issuecomment-6314903061

   Before working on this I measured what `DayWindow` does today on DST days, and the spring-forward case behaves differently from what the docstring, `assets.rst` and the strict-xfail test say.
   
   I simulated an hourly producer: one key per real local hour of the day, encoded with the upstream mapper's `input_format`. Then I compared those keys with `RollupMapper(upstream_mapper=StartOfDayMapper(timezone=..., input_format=...), window=DayWindow()).to_upstream(day)` on `main`:
   
   | zone, day | real hours | `%Y-%m-%dT%H` (no offset) | `%Y-%m-%dT%H%z` (with offset) |
   |---|---|---|---|
   | America/New_York 2024-03-10 (spring) | 23 | 23 expected = 23 produced, satisfied | 23 = 23, satisfied |
   | America/New_York 2024-11-03 (fall) | 25 | 24 distinct produced, 24 expected, satisfied | 25 distinct produced, 24 expected; `2024-11-03T01-0400` not expected |
   | Europe/Madrid 2024-03-31 / 2024-10-27 | 23 / 25 | satisfied / satisfied | satisfied / `2024-10-27T02+0200` not expected |
   | Australia/Lord_Howe (30 min shift) 2024-10-06 / 2024-04-07 | 24 / 25 | satisfied / satisfied | satisfied / `2024-04-07T01+1030` not expected |
   | America/Santiago (midnight gap) 2024-09-08 / 2024-04-06 | 23 / 25 | satisfied / satisfied | satisfied / `2024-04-06T23-0300` not expected |
   | America/Havana (midnight gap) 2024-03-10 / 2024-11-03 | 23 / 25 | satisfied / satisfied | satisfied / `2024-11-03T00-0400` not expected |
   
   So:
   
   1. **Spring-forward is not held.** The nonexistent hour's naive step is localised forward onto the next real hour, the two steps collapse into one key, and `to_upstream` returns exactly the 23 keys an hourly producer emits. `test_day_window_rollup_under_yields_on_spring_forward_with_local_tz` asserts `len(upstream_keys) == 24`, but 23 is the right answer for a 23-hour day, so the strict xfail encodes the wrong expectation. The same holds where the gap is at midnight (Santiago, Havana).
   2. **Fall-back without an offset in the key is fine.** The repeated hour produces the same key twice, so 24 distinct keys arrive and 24 are expected.
   3. **Fall-back with an offset in the key is the real gap.** The repeated hour produces two distinct keys, and the window only yields one of them. With `WaitForAll` the run still fires (every expected key arrives, the last one at the end of the day), but one real hour isn't part of the expected set, so it doesn't count towards `MinimumCount` and the UI's required keys show 24 for a 25-hour day.
   
   If that matches your understanding, I'd split it in two:
   
   - **A small PR now:** replace the strict xfail with tests that check the rollup is satisfied on spring-forward and fall-back days for a local-timezone mapper without an offset, add one that pins the fall-back gap with an offset, and correct the `DayWindow` docstring and the `assets.rst` paragraph (the "use UTC for any rollup that crosses a DST boundary" advice is only needed for offset-bearing keys).
   - **This issue then tracks only case 3.** One way: have `RollupMapper.to_upstream` localise the decoded period start with the upstream mapper's timezone and let the temporal windows step through the aware range to the next local period start, so the repeated hour yields both instants and `encode_upstream` (which already handles aware datetimes) renders both keys. That changes the "windows don't touch timezones" contract, so I'd like your view before writing it.
   
   The script I used is below; it runs against `main` with `uv run --project airflow-core python dst_check.py`.
   
   <details><summary>dst_check.py</summary>
   
   ```python
   import pendulum
   from airflow.partition_mappers.base import RollupMapper
   from airflow.partition_mappers.temporal import StartOfDayMapper
   from airflow.partition_mappers.window import DayWindow
   
   
   def check(tz_name, ymd, fmt):
       tz = pendulum.timezone(tz_name)
       t = pendulum.datetime(*ymd, tz="UTC").subtract(days=2)
       produced = []
       for _ in range(24 * 5):
           local = t.in_tz(tz)
           if (local.year, local.month, local.day) == ymd:
               produced.append(local.strftime(fmt))
           t = t.add(hours=1)
       mapper = RollupMapper(
           upstream_mapper=StartOfDayMapper(timezone=tz_name, input_format=fmt), window=DayWindow()
       )
       (down,) = {mapper.to_downstream(k) for k in produced}
       expected = mapper.to_upstream(down)
       print(
           f"{tz_name} {ymd} {fmt}: real hours={len(produced)} distinct produced={len(set(produced))} "
           f"expected={len(expected)} never arrive={sorted(expected - set(produced))} "
           f"not expected={sorted(set(produced) - expected)} satisfied={expected <= set(produced)}"
       )
   
   
   for fmt in ("%Y-%m-%dT%H", "%Y-%m-%dT%H%z"):
       for tz_name, ymd in [
           ("America/New_York", (2024, 3, 10)),
           ("America/New_York", (2024, 11, 3)),
           ("Europe/Madrid", (2024, 3, 31)),
           ("Europe/Madrid", (2024, 10, 27)),
           ("America/Santiago", (2024, 9, 8)),
           ("America/Santiago", (2024, 4, 6)),
       ]:
           check(tz_name, ymd, fmt)
   ```
   
   </details>
   
   ---
   
   Drafted-by: Claude Code (Opus 5.5)
   


-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org