Appendix: Zarr fetch lessons learned
The forecast-vs-retrospective comparison in 6 Forecast vs Retrospective Magnitudes required pulling roughly 20,000 forecast slices — 667 issue dates × 2 reaches × 15 lead days, each from its own per-date Zarr store on s3://geoglows-v2-forecasts/. The first naive implementation took ~26 seconds per issue date with xarray; that projected to ~5 hours sequential or ~40 minutes with 8 threads. Both felt long for what is fundamentally a “read 15 small slices” job. After several iterations the per-date cost dropped to ~7 seconds, and the full archive fetched in ~51 minutes on 8 threads.
Notes on what mattered, for the next person doing this.
What worked
Bypass xarray for the hot path
xr.open_zarr(path).sel(river_id=...).compute() is ergonomic but pays for metadata parsing, dtype decoding, dimension labelling, and a Dask graph on every call. For the same store a thin zarr.open_group(store) + numpy slicing was ~4× faster:
grp = zarr.open_group(store, mode="r")
arr = grp["Qout"][:, :, river_positions] # (members, time, rivers)
time_arr = grp["time"][:] # int32 seconds-since-issuexarray earns its keep when you need labelled dimensions, alignment across files, or lazy compute graphs. For “give me a fixed N-element slice from a known store”, it is overhead.
Slice positionally, not by label
Resolve river_id → array index once (from the metadata table) and then slice with integer positions. Per-call .sel(river_id=...) repeats a coordinate lookup that does not need repeating.
NaN-aware aggregations are mandatory for v2
GEOGloWS v2 ensembles mix two cadences on the same time axis: 51 perturbed members at 3-hourly steps, 1 high-resolution control at 1-hourly steps. The Zarr arrays are padded with NaN to a common grid, so roughly half the cells in Qout are NaN by construction. arr.mean(axis=...) propagates NaN and silently produces all-NaN columns; the correct call is np.nanmean / np.nanmedian / np.nanmax. This was the single most embarrassing debugging session of the project.
time is int32 seconds-since-issue, not nanoseconds
pd.to_datetime(grp["time"][:]) interprets the array as nanoseconds and returns 1970-era timestamps, which then silently align to nothing. The actual encoding is integer seconds offset from the issue date:
lead_day = (time_arr // 86400).astype(int)The Zarr store’s units attribute documents this — worth reading before assuming.
Threading beats async for blocking I/O
asyncio.to_thread(xr.open_zarr, ...) was slower than a plain ThreadPoolExecutor(max_workers=8) despite looking more modern. The underlying boto/HTTP call still blocks; wrapping it in asyncio adds event-loop bookkeeping without unblocking anything. Async helps when the library itself is async (e.g. aiohttp, aioboto3) — it does not help when the inner call is sync.
What did not work
obstore+ zarr v3. The intent was to swap the boto-based store for the faster Rust-basedobstore.zarr.open_group(obstore.S3Store(...))raisedTypeError: Unsupported type for store_like: 'S3Store'. Likely a version-skew issue betweenzarr-python3.x andobstore. Skipped — the numpy + threading approach was already fast enough.- Pre-fetching
Qoutfor all 15 lead days at once. Triedarr = grp["Qout"][:](full read) then slicing in memory. No measurable speedup over slicing per lead day; S3 + zarr already chunks well by time.
Takeaways
- xarray’s ergonomics tax is real on small repeated reads. When the read pattern is “many small slices from many known stores”, drop to
zarr.open_group+ numpy. - Read the store’s coordinate
unitsattribute. Time encodings vary; assuming nanoseconds is a coin flip. - NaN-pad your aggregations. If a Zarr store mixes resolutions on a shared axis,
nan*reductions are the only safe default. - Use threading first, async second. Async only pays off when the underlying library is async too.
The full optimized fetch script lives at analysis/07_forecast_zarr_direct.py. The produced parquet is published to Azure blob at ds-aa-npl-flooding/processed/geoglows/forecast_lead_times_v2.parquet via ocha_stratus and consumed from there by Chapter 6.