Skip to content

Commit a971f2b

Browse files
alxmrsclaude
alxmrs
and
claude
committed
Lazy registration, parameterized queries, model-concat: address round-3 review
Addresses round-3 PR review (cases 04, 05) and extends the patterns PR-wide: - Demonstrate laziness + pushdown: ERA5 cases (02,03,04,06) now register the whole archive lazily (no .load(), no variable pre-selection) and let projection pushdown read only 2m_temperature and partition pruning read only the queried window. References use lazy .sel(); the up-front "read into memory" steps are gone. - Parameterized queries (no SQL-string interpolation of values): time/region bounds are passed via param_values rather than f-string-formatted into the SQL. Verified pruning still fires with bound parameters. - Forecast skill (05): stack the two models along a `model` dimension into one forecast table and GROUP BY the model column, eliminating the `FROM {table}` identifier interpolation (the injection smell) entirely. The 2m_temperature selection stays because the models store different pressure-level sets (Pangu 13, GraphCast 37) and must be aligned on the common surface field. Reword the "this is the real thing" docstring opening. docs/geospatial.md and the suite README updated to describe lazy registration, parameterized filters, and the model-column forecast query. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9
1 parent e1cbc8b commit a971f2b

7 files changed

Lines changed: 220 additions & 163 deletions

File tree

‎benchmarks/geospatial/02_climatology.py‎

Lines changed: 50 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,17 @@
2525
``da.groupby("time.hour").mean()``. ERA5 is hourly, so grouping by hour of day
2626
gives a clean 24-bin **diurnal cycle**, one sample per day in the window.
2727
28-
Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days,
29-
opened and sliced with a single fluent xarray chain.
28+
The table is the *whole* lazily-opened ARCO-ERA5 archive — nothing is loaded or
29+
pre-selected up front. The query reads only the ``2m_temperature`` column
30+
(projection pushdown) for only the window the parameterized ``WHERE`` asks for
31+
(partition pruning). The bounds are passed as **query parameters**, not
32+
formatted into the SQL string.
3033
"""
3134

3235
from __future__ import annotations
3336

37+
import datetime
38+
3439
import xarray as xr
3540

3641
import xarray_sql as xql
@@ -39,55 +44,73 @@
3944

4045
_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
4146
# A few days over a CONUS-ish box (ERA5 latitude descends; lon is 0–360°E).
42-
_TIME = slice("2020-06-01", "2020-06-03T23")
43-
_LAT = slice(50.0, 25.0)
44-
_LON = slice(235.0, 290.0)
47+
_START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23)
48+
_LAT_N, _LAT_S = 50.0, 25.0
49+
_LON_W, _LON_E = 235.0, 290.0
50+
_PARAMS = {
51+
"start": _START,
52+
"end": _END,
53+
"lat_s": _LAT_S,
54+
"lat_n": _LAT_N,
55+
"lon_w": _LON_W,
56+
"lon_e": _LON_E,
57+
}
4558

4659

4760
def main() -> None:
48-
# Open the full ARCO-ERA5 archive and slice to the window in one chain:
49-
# pick the variable, select the box and days, and read it into memory.
61+
# Open the full ARCO-ERA5 archive lazily — dask off, nothing loaded or
62+
# column-selected here. ERA5 mixes surface (time, lat, lon) and atmospheric
63+
# (… level …) variables, so register it as two tables under an ``era5``
64+
# schema; the query below touches only the surface table's 2m_temperature.
5065
try:
5166
import gcsfs # noqa: F401 — required by the gs:// protocol
5267

53-
with timed("open + slice ERA5 window"):
54-
ds = (
55-
xr.open_zarr(
56-
_URL, chunks=None, storage_options={"token": "anon"}
57-
)[["2m_temperature"]]
58-
.sel(time=_TIME, latitude=_LAT, longitude=_LON)
59-
.load()
60-
)
68+
ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})
6169
except Exception as exc: # noqa: BLE001 — any failure → skip, not crash
6270
raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc
6371

64-
print(
65-
f" ERA5 2m_temperature window: {dict(ds.sizes)} "
66-
f"(diurnal climatology over {ds.sizes['latitude']}×{ds.sizes['longitude']} cells)"
67-
)
68-
6972
ctx = xql.XarrayContext()
70-
ctx.from_dataset("era5", ds, chunks={"time": 24})
73+
with timed("register full ERA5 (lazy)"):
74+
ctx.from_dataset(
75+
"era5",
76+
ds,
77+
chunks={"time": 6},
78+
table_names={
79+
("time", "latitude", "longitude"): "surface",
80+
("time", "level", "latitude", "longitude"): "atmosphere",
81+
},
82+
)
7183

7284
sql = """
7385
SELECT latitude,
7486
longitude,
7587
date_part('hour', time) AS hour,
7688
AVG("2m_temperature") - 273.15 AS clim_c
77-
FROM era5
89+
FROM era5.surface
90+
WHERE time BETWEEN $start AND $end
91+
AND latitude BETWEEN $lat_s AND $lat_n
92+
AND longitude BETWEEN $lon_w AND $lon_e
7893
GROUP BY latitude, longitude, date_part('hour', time)
7994
ORDER BY latitude DESC, longitude, hour
8095
"""
8196
show_sql(sql)
8297

83-
# A climatology is a gridded product, so round-trip the result back to an
98+
# A climatology is a gridded product: round-trip the result back to an
8499
# xarray Dataset keyed by (latitude, longitude, hour) — how it is used.
85-
with timed("SQL diurnal climatology"):
86-
got = ctx.sql(sql).to_dataset(dims=["latitude", "longitude", "hour"])
100+
with timed("SQL diurnal climatology (lazy read, pushdown + pruning)"):
101+
got = ctx.sql(sql, param_values=_PARAMS).to_dataset(
102+
dims=["latitude", "longitude", "hour"]
103+
)
87104

88-
# Array reference: the textbook groupby-over-the-cycle reduction, in °C.
105+
# Array reference: the textbook groupby-over-the-cycle reduction, in °C —
106+
# the same lazy window, materialized only on demand.
89107
with timed("xarray reference"):
90-
ref = ds["2m_temperature"].groupby("time.hour").mean("time") - 273.15
108+
window = ds["2m_temperature"].sel(
109+
time=slice(_START, _END),
110+
latitude=slice(_LAT_N, _LAT_S),
111+
longitude=slice(_LON_W, _LON_E),
112+
)
113+
ref = window.groupby("time.hour").mean("time") - 273.15
91114

92115
assert_grid_close(
93116
"diurnal climatology (°C)", got.clim_c, ref, rtol=1e-4, atol=1e-2

‎benchmarks/geospatial/03_zonal_mean.py‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929

3030
from __future__ import annotations
3131

32+
import datetime
33+
3234
import xarray as xr
3335

3436
import xarray_sql as xql
@@ -38,6 +40,10 @@
3840
_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
3941
# One day of hourly data, global; the WHERE below prunes ERA5 to this window.
4042
_DAY = "2020-06-01"
43+
_START, _END = (
44+
datetime.datetime(2020, 6, 1, 0),
45+
datetime.datetime(2020, 6, 1, 23),
46+
)
4147

4248

4349
def main() -> None:
@@ -70,20 +76,23 @@ def main() -> None:
7076
},
7177
)
7278

73-
sql = f"""
79+
# The window bounds are passed as query parameters, not formatted into the
80+
# SQL string; pruning still kicks in, so only one day is read.
81+
sql = """
7482
SELECT latitude,
7583
AVG("2m_temperature") - 273.15 AS air_mean_c
7684
FROM era5.surface
77-
WHERE time BETWEEN TIMESTAMP '{_DAY}'
78-
AND TIMESTAMP '{_DAY} 23:00:00'
85+
WHERE time BETWEEN $start AND $end
7986
GROUP BY latitude
8087
ORDER BY latitude DESC
8188
"""
8289
show_sql(sql)
8390

8491
# Round-trip the profile back to an xarray Dataset keyed by latitude.
8592
with timed("SQL zonal mean (WHERE-pruned to one day)"):
86-
got = ctx.sql(sql).to_dataset(dims=["latitude"])
93+
got = ctx.sql(
94+
sql, param_values={"start": _START, "end": _END}
95+
).to_dataset(dims=["latitude"])
8796

8897
# Array reference: reduce the same day over the two un-grouped axes.
8998
with timed("xarray reference"):

‎benchmarks/geospatial/04_anomaly.py‎

Lines changed: 48 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -24,73 +24,96 @@
2424
FROM era5 a JOIN clim c
2525
ON (a.latitude, a.longitude, hour(a.time)) = (c.latitude, c.longitude, c.hour)
2626
27-
Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days,
28-
opened and sliced with a single fluent xarray chain (anomalies vs the diurnal
29-
cycle of that window).
27+
The table is the *whole* lazily-opened ARCO-ERA5 archive — nothing loaded or
28+
column-selected up front. Both the climatology CTE and the outer scan read only
29+
``2m_temperature`` (projection pushdown) over only the window the parameterized
30+
``WHERE`` asks for (partition pruning). Bounds are query parameters, not
31+
string-formatted into the SQL.
3032
"""
3133

3234
from __future__ import annotations
3335

36+
import datetime
37+
3438
import xarray as xr
3539

3640
import xarray_sql as xql
3741

3842
from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed
3943

4044
_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
41-
_TIME = slice("2020-06-01", "2020-06-03T23")
42-
_LAT = slice(50.0, 25.0)
43-
_LON = slice(235.0, 290.0)
45+
_START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23)
46+
_LAT_N, _LAT_S = 50.0, 25.0
47+
_LON_W, _LON_E = 235.0, 290.0
48+
_PARAMS = {
49+
"start": _START,
50+
"end": _END,
51+
"lat_s": _LAT_S,
52+
"lat_n": _LAT_N,
53+
"lon_w": _LON_W,
54+
"lon_e": _LON_E,
55+
}
4456

4557

4658
def main() -> None:
4759
try:
4860
import gcsfs # noqa: F401 — required by the gs:// protocol
4961

50-
with timed("open + slice ERA5 window"):
51-
ds = (
52-
xr.open_zarr(
53-
_URL, chunks=None, storage_options={"token": "anon"}
54-
)[["2m_temperature"]]
55-
.sel(time=_TIME, latitude=_LAT, longitude=_LON)
56-
.load()
57-
)
62+
ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})
5863
except Exception as exc: # noqa: BLE001 — any failure → skip, not crash
5964
raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc
6065

61-
print(
62-
f" ERA5 2m_temperature window: {dict(ds.sizes)} (anomaly vs diurnal normals)"
63-
)
64-
6566
ctx = xql.XarrayContext()
66-
ctx.from_dataset("era5", ds, chunks={"time": 24})
67+
with timed("register full ERA5 (lazy)"):
68+
ctx.from_dataset(
69+
"era5",
70+
ds,
71+
chunks={"time": 6},
72+
table_names={
73+
("time", "latitude", "longitude"): "surface",
74+
("time", "level", "latitude", "longitude"): "atmosphere",
75+
},
76+
)
6777

6878
sql = """
6979
WITH clim AS (
7080
SELECT latitude, longitude,
7181
date_part('hour', time) AS hour,
7282
AVG("2m_temperature") AS clim_t
73-
FROM era5
83+
FROM era5.surface
84+
WHERE time BETWEEN $start AND $end
85+
AND latitude BETWEEN $lat_s AND $lat_n
86+
AND longitude BETWEEN $lon_w AND $lon_e
7487
GROUP BY latitude, longitude, date_part('hour', time)
7588
)
7689
SELECT a.time, a.latitude, a.longitude,
7790
a."2m_temperature" - c.clim_t AS anomaly
78-
FROM era5 a
91+
FROM era5.surface a
7992
JOIN clim c
8093
ON a.latitude = c.latitude
8194
AND a.longitude = c.longitude
8295
AND date_part('hour', a.time) = c.hour
96+
WHERE a.time BETWEEN $start AND $end
97+
AND a.latitude BETWEEN $lat_s AND $lat_n
98+
AND a.longitude BETWEEN $lon_w AND $lon_e
8399
ORDER BY a.time, a.latitude DESC, a.longitude
84100
"""
85101
show_sql(sql)
86102

87103
# The anomaly is a gridded field; round-trip it to (time, lat, lon).
88-
with timed("SQL anomaly (climatology CTE self-join)"):
89-
got = ctx.sql(sql).to_dataset(dims=["time", "latitude", "longitude"])
104+
with timed("SQL anomaly (climatology CTE self-join, lazy read)"):
105+
got = ctx.sql(sql, param_values=_PARAMS).to_dataset(
106+
dims=["time", "latitude", "longitude"]
107+
)
90108

91-
# Array reference: grouped broadcast-subtract, in pure xarray.
109+
# Array reference: grouped broadcast-subtract, in pure xarray (lazy window).
92110
with timed("xarray reference"):
93-
grouped = ds["2m_temperature"].groupby("time.hour")
111+
window = ds["2m_temperature"].sel(
112+
time=slice(_START, _END),
113+
latitude=slice(_LAT_N, _LAT_S),
114+
longitude=slice(_LON_W, _LON_E),
115+
)
116+
grouped = window.groupby("time.hour")
94117
ref = grouped - grouped.mean("time")
95118

96119
assert_grid_close(

0 commit comments

Comments
 (0)