From 475b075ec78ef9e3a1318114ed768feebaf3858a Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 23 Jun 2026 18:31:51 +0300 Subject: [PATCH 01/26] Add geospatial SQL benchmark suite (benchmarks/geospatial/) A suite of worked examples arguing that core geospatial/climate operations we associate with the array paradigm are, underneath, relational operations. Each case poses the operation in SQL against xarray-sql and asserts the result matches an xarray/array reference to floating-point tolerance. Cases (all validated): - 01 NDVI: per-pixel column algebra + CASE mask, on a real Sentinel-2 L2A scene in Zarr (ESA EOPF sample service). - 02 climatology: diurnal cycle as GROUP BY lat, lon, hour (ARCO-ERA5). - 03 zonal mean: GROUP BY latitude over the full ARCO-ERA5 archive, WHERE-pruned to one day. - 04 anomaly: climatology CTE self-JOIN (ARCO-ERA5). - 05 forecast skill: Pangu-Weather and GraphCast vs ERA5 (WeatherBench2), a forecast<->truth JOIN on valid_time = init + lead; reproduces GraphCast beating Pangu at every lead. - 06 zonal stats: raster x vector range JOIN of full ERA5 against continental bounding boxes. - 07 reprojection: PROJ-backed scalar UDF (ST_Transform-style), returning a {lon, lat} struct to keep PROJ single-threaded. - 08 regridding: bilinear regrid as a sparse weight-table JOIN, validated against xarray .interp(). Includes a shared harness (_harness.py: timing, peak memory, assert_allclose, SQL echo), ERA5 helpers (_era5.py), a suite README, and a narrative write-up (docs/geospatial.md) wired into the docs nav. Scripts carry PEP 723 inline metadata so they are runnable standalone via `uv run`. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/01_ndvi.py | 142 +++++++++++++ benchmarks/geospatial/02_climatology.py | 101 +++++++++ benchmarks/geospatial/03_zonal_mean.py | 91 +++++++++ benchmarks/geospatial/04_anomaly.py | 109 ++++++++++ benchmarks/geospatial/05_forecast_skill.py | 183 +++++++++++++++++ benchmarks/geospatial/06_zonal_vector.py | 135 ++++++++++++ benchmarks/geospatial/07_reproject_udf.py | 166 +++++++++++++++ benchmarks/geospatial/08_regrid_weights.py | 158 ++++++++++++++ benchmarks/geospatial/README.md | 74 +++++++ benchmarks/geospatial/_era5.py | 70 +++++++ benchmarks/geospatial/_harness.py | 120 +++++++++++ docs/geospatial.md | 227 +++++++++++++++++++++ zensical.toml | 1 + 13 files changed, 1577 insertions(+) create mode 100644 benchmarks/geospatial/01_ndvi.py create mode 100644 benchmarks/geospatial/02_climatology.py create mode 100644 benchmarks/geospatial/03_zonal_mean.py create mode 100644 benchmarks/geospatial/04_anomaly.py create mode 100644 benchmarks/geospatial/05_forecast_skill.py create mode 100644 benchmarks/geospatial/06_zonal_vector.py create mode 100644 benchmarks/geospatial/07_reproject_udf.py create mode 100644 benchmarks/geospatial/08_regrid_weights.py create mode 100644 benchmarks/geospatial/README.md create mode 100644 benchmarks/geospatial/_era5.py create mode 100644 benchmarks/geospatial/_harness.py create mode 100644 docs/geospatial.md diff --git a/benchmarks/geospatial/01_ndvi.py b/benchmarks/geospatial/01_ndvi.py new file mode 100644 index 0000000..9c9e89e --- /dev/null +++ b/benchmarks/geospatial/01_ndvi.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "zarr>=3", +# "numpy", +# ] +# /// +"""NDVI + cloud/nodata masking — "apply_ufunc over a raster" is column algebra. + +The Normalized Difference Vegetation Index is the workhorse of optical remote +sensing: ``NDVI = (NIR - Red) / (NIR + Red)``, computed per pixel, with invalid +pixels (nodata / clouds) masked out. The array paradigm reaches for +``xarray.apply_ufunc`` (the coiled/benchmarks #1545 "vectorized operations" +case) to broadcast this over a whole scene. + +But a per-pixel formula over two bands is just *column arithmetic over two +columns*, and "mask the invalid pixels" is just ``CASE WHEN``:: + + SELECT x, y, + CASE WHEN red = 0 OR nir = 0 THEN NULL + ELSE (nir_refl - red_refl) / (nir_refl + red_refl) + END AS ndvi + FROM scene + +Each pixel is one row; the ufunc is the SELECT expression. No broadcasting +rules, no apply_ufunc signature — the SQL *is* the index definition. + +Dataset: a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample +service (bands B04=red, B08=NIR at 10 m, stored as uint16 with +``reflectance = DN * 0.0001 - 0.1``). We read one cloud-light window so the +case stays bounded. Requires network; skips cleanly if the store is offline. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +import xarray_sql as xql + +from _harness import CaseSkipped, check_close, run_case, show_sql, timed + +# A specific Sentinel-2 L2A product from the EOPF sample service (STAC +# collection ``sentinel-2-l2a``, tile T32TLQ over the Italian Alps). Hard-coded +# so the case needs no STAC client; if the sample service rotates this product +# out, the open below fails and the case skips. +_BASE = ( + "https://objectstore.eodc.eu:2222/e05ab01a9d56408d82ac32d69a5aae2a:" + "202505-s02msil2a/03/products/cpm_v256/" + "S2B_MSIL2A_20250503T102559_N0511_R108_T32TLQ_20250503T132157.zarr" +) +_R10M = _BASE + "/measurements/reflectance/r10m" + +# proj:transform of the 10 m grid (EPSG:32632), from the STAC asset metadata. +_X0, _Y0, _RES = 300_000.0, 5_000_040.0, 10.0 + +# Reflectance encoding (STAC raster:scale / raster:offset), DN -> reflectance. +_SCALE, _OFFSET = 0.0001, -0.1 + +# A 1024×1024 window (~105 km²) offset into the tile to land on vegetated +# terrain (median NDVI ≈ 0.33 in early May) rather than the nodata border. +_ROW0, _COL0, _N = 2_500, 6_500, 1024 + + +def _load_scene() -> xr.Dataset: + """Read the B04/B08 window from the EOPF Zarr store into an xr.Dataset.""" + try: + import zarr + + b04 = zarr.open_array(_R10M + "/b04", mode="r") + b08 = zarr.open_array(_R10M + "/b08", mode="r") + rows = slice(_ROW0, _ROW0 + _N) + cols = slice(_COL0, _COL0 + _N) + red = np.asarray(b04[rows, cols]) + nir = np.asarray(b08[rows, cols]) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"EOPF Sentinel-2 Zarr unavailable ({exc})") from exc + + # Pixel-center coordinates from the affine transform. + x = _X0 + (np.arange(_COL0, _COL0 + _N) + 0.5) * _RES + y = _Y0 - (np.arange(_ROW0, _ROW0 + _N) + 0.5) * _RES + return xr.Dataset( + {"red": (["y", "x"], red), "nir": (["y", "x"], nir)}, + coords={"y": y.astype("float64"), "x": x.astype("float64")}, + ).chunk({"y": 256, "x": 256}) + + +def main() -> None: + ds = _load_scene() + n = ds.sizes["y"] * ds.sizes["x"] + print( + f" Sentinel-2 L2A scene window: {dict(ds.sizes)} " + f"({n:,} pixels, B04=red/B08=NIR, uint16 DN)" + ) + + ctx = xql.XarrayContext() + ctx.from_dataset("scene", ds, chunks={"y": 256, "x": 256}) + + # DN -> reflectance happens inline; CASE masks nodata (DN == 0). + sql = """ + SELECT x, y, + CASE + WHEN red = 0 OR nir = 0 THEN NULL + ELSE ( (CAST(nir AS DOUBLE) * 0.0001 - 0.1) + - (CAST(red AS DOUBLE) * 0.0001 - 0.1) ) + / ( (CAST(nir AS DOUBLE) * 0.0001 - 0.1) + + (CAST(red AS DOUBLE) * 0.0001 - 0.1) ) + END AS ndvi + FROM scene + """ + show_sql(sql) + + with timed("SQL NDVI"): + got = ctx.sql(sql).to_dataset(dims=["y", "x"]).ndvi + + # Array reference: the same formula via xarray broadcasting + .where mask. + with timed("xarray reference"): + red = ds.red.astype("float64") * _SCALE + _OFFSET + nir = ds.nir.astype("float64") * _SCALE + _OFFSET + ref = ((nir - red) / (nir + red)).where((ds.red != 0) & (ds.nir != 0)) + + # Align both to the same (y, x) grid before comparing. + got = got.sortby(["y", "x"]) + ref = ref.sortby(["y", "x"]) + check_close("NDVI (per-pixel)", got, ref, rtol=1e-6, atol=1e-6) + + valid = np.isfinite(got.values) + print( + f"\n NDVI over {valid.sum():,} valid pixels: " + f"min {np.nanmin(got.values):.3f}, " + f"mean {np.nanmean(got.values):.3f}, " + f"max {np.nanmax(got.values):.3f}" + ) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "NDVI: per-pixel column algebra + CASE mask") + ) diff --git a/benchmarks/geospatial/02_climatology.py b/benchmarks/geospatial/02_climatology.py new file mode 100644 index 0000000..e294b74 --- /dev/null +++ b/benchmarks/geospatial/02_climatology.py @@ -0,0 +1,101 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "pandas", +# "gcsfs", +# "zarr>=3", +# ] +# /// +"""Diurnal climatology — the "rechunk + grouped reduction" that is a GROUP BY. + +A *climatology* is the average value for each time-of-cycle, computed +independently at every location: "what is the typical temperature here at 06:00 +local cycle?" In the array paradigm (and in the coiled/benchmarks #1545 +write-up) this is the canonical painful workload — load native Zarr chunks, +*rechunk* to put all of time in one chunk ("pencils"), run a grouped reduction +over the calendar, then rechunk back to "pancakes" for output. + +The rechunking exists only to serve the array layout. The *operation* is:: + + SELECT latitude, longitude, hour_of_day, AVG("2m_temperature") + GROUP BY latitude, longitude, hour_of_day + +Group by location and time-of-cycle, average the rest. Same answer as +``da.groupby("time.hour").mean()`` — here the **diurnal cycle** (mean +temperature by hour of day) over a region, at ERA5's 0.25° resolution. + +Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days, +read once into memory (see ``_era5.load_window``). +""" + +from __future__ import annotations + +import xarray_sql as xql + +from _era5 import load_window, open_era5 +from _harness import check_close, run_case, show_sql, timed + +# A few days over a CONUS-ish box (ERA5 latitude descends; lon is 0–360°E). +_TIME = slice("2020-06-01", "2020-06-03T23") +_LAT = slice(50.0, 25.0) +_LON = slice(235.0, 290.0) + + +def main() -> None: + full = open_era5() + with timed("read ERA5 window into memory"): + ds = load_window( + full, "2m_temperature", time=_TIME, latitude=_LAT, longitude=_LON + ) + print( + f" ERA5 2m_temperature window: {dict(ds.sizes)} " + f"(diurnal climatology over {ds.sizes['latitude']}×{ds.sizes['longitude']} cells)" + ) + + ctx = xql.XarrayContext() + ctx.from_dataset("era5", ds, chunks={"time": 24}) + + sql = """ + SELECT latitude, + longitude, + date_part('hour', time) AS hour, + AVG("2m_temperature") - 273.15 AS clim_c + FROM era5 + GROUP BY latitude, longitude, date_part('hour', time) + """ + show_sql(sql) + + with timed("SQL diurnal climatology"): + got = ctx.sql(sql).to_pandas() + + # Array reference: the textbook groupby-over-the-cycle reduction. + with timed("xarray reference"): + clim = ds["2m_temperature"].groupby("time.hour").mean("time") - 273.15 + ref = clim.to_dataframe(name="clim_c").reset_index() + + merged = got.merge( + ref, on=["latitude", "longitude", "hour"], validate="one_to_one" + ) + assert len(merged) == len(got) == len(ref), ( + f"row-count mismatch: sql={len(got)} ref={len(ref)} merged={len(merged)}" + ) + check_close( + "diurnal climatology (°C)", + merged["clim_c_x"], + merged["clim_c_y"], + rtol=1e-4, + atol=1e-3, + ) + + print( + f"\n {len(got):,} climatology cells (latitude × longitude × hour-of-day)." + ) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Climatology: GROUP BY lat, lon, hour (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/03_zonal_mean.py b/benchmarks/geospatial/03_zonal_mean.py new file mode 100644 index 0000000..06c95b4 --- /dev/null +++ b/benchmarks/geospatial/03_zonal_mean.py @@ -0,0 +1,91 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "gcsfs", +# "zarr>=3", +# ] +# /// +"""Zonal mean — the array reduction that is secretly a GROUP BY. + +A *zonal mean* averages a field around each circle of latitude (over all +longitudes, and here over a day of hours too), collapsing a 3-D field to a 1-D +profile of value-vs-latitude — the classic pole-to-pole temperature curve. In +the array paradigm this is ``da.mean(dim=["longitude", "time"])``, a reduction +over two axes. + +Relationally it is nothing more than:: + + SELECT latitude, AVG("2m_temperature") GROUP BY latitude + +The "axes" we reduce over are just the columns we *don't* group by. Same answer, +and the SQL reads like the plain-English definition of a zonal mean. + +Dataset: the full **ARCO-ERA5** archive (0.25° global, 1.3M hourly timesteps). +The table is the whole reanalysis; ``WHERE time …`` prunes it to one day, and +the GROUP BY produces a 721-point global temperature profile. +""" + +from __future__ import annotations + + +import xarray_sql as xql + +from _era5 import open_era5, register_era5 +from _harness import check_close, run_case, show_sql, timed + +# One day of hourly data, global; the WHERE below prunes ERA5 to this window. +_DAY = "2020-06-01" + + +def main() -> None: + ds = open_era5() + print( + f" ARCO-ERA5: {ds.sizes['time']:,} hourly timesteps, " + f"{ds.sizes['latitude']}×{ds.sizes['longitude']} grid, " + f"{len(ds.data_vars)} variables (no pre-slicing)" + ) + + ctx = xql.XarrayContext() + with timed("register full ERA5"): + register_era5(ctx, ds) + + sql = f""" + SELECT latitude, + AVG("2m_temperature") - 273.15 AS air_mean_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '{_DAY}' + AND TIMESTAMP '{_DAY} 23:00:00' + GROUP BY latitude + ORDER BY latitude DESC + """ + show_sql(sql) + + with timed("SQL zonal mean (WHERE-pruned to one day)"): + got = ctx.sql(sql).to_pandas() + + # Array reference: reduce the same day over the two un-grouped axes. + with timed("xarray reference"): + window = ds["2m_temperature"].sel(time=_DAY) + ref = (window.mean(dim=["longitude", "time"]) - 273.15).sortby( + "latitude", ascending=False + ) + + check_close( + "zonal mean (2m_temp vs latitude, °C)", + got["air_mean_c"], + ref, + rtol=1e-4, + atol=1e-3, + ) + + print("\n Global temperature profile (every 10th parallel):") + print(got.iloc[::72].to_string(index=False)) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Zonal mean: GROUP BY latitude (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/04_anomaly.py b/benchmarks/geospatial/04_anomaly.py new file mode 100644 index 0000000..96df0df --- /dev/null +++ b/benchmarks/geospatial/04_anomaly.py @@ -0,0 +1,109 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "pandas", +# "gcsfs", +# "zarr>=3", +# ] +# /// +"""Temperature anomaly — "broadcast-subtract the climatology" is a self-JOIN. + +An *anomaly* is the departure of each observation from its climatological +normal: ``anomaly(t) = T(t) − climatology(hour-of-day(t))`` at each cell. The +array paradigm computes the climatology, then leans on xarray's grouped +broadcasting to line it back up with every timestep: +``ds.groupby("time.hour") - climatology``. + +That broadcast — "attach each cell's normal back onto every matching timestep" — +is exactly a relational **JOIN** on the grouping key. So the anomaly is a +climatology CTE joined back to the raw observations:: + + WITH clim AS (SELECT latitude, longitude, hour, AVG(T) ... GROUP BY ...) + SELECT a.T - c.clim_t AS anomaly + FROM era5 a JOIN clim c + ON (a.latitude, a.longitude, hour(a.time)) = (c.latitude, c.longitude, c.hour) + +Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days, +read once into memory (anomalies vs the diurnal cycle of that window). +""" + +from __future__ import annotations + +import xarray_sql as xql + +from _era5 import load_window, open_era5 +from _harness import check_close, run_case, show_sql, timed + +_TIME = slice("2020-06-01", "2020-06-03T23") +_LAT = slice(50.0, 25.0) +_LON = slice(235.0, 290.0) + + +def main() -> None: + full = open_era5() + with timed("read ERA5 window into memory"): + ds = load_window( + full, "2m_temperature", time=_TIME, latitude=_LAT, longitude=_LON + ) + print( + f" ERA5 2m_temperature window: {dict(ds.sizes)} " + f"(anomaly vs diurnal normals)" + ) + + ctx = xql.XarrayContext() + ctx.from_dataset("era5", ds, chunks={"time": 24}) + + sql = """ + WITH clim AS ( + SELECT latitude, longitude, + date_part('hour', time) AS hour, + AVG("2m_temperature") AS clim_t + FROM era5 + GROUP BY latitude, longitude, date_part('hour', time) + ) + SELECT a.time, a.latitude, a.longitude, + a."2m_temperature" - c.clim_t AS anomaly + FROM era5 a + JOIN clim c + ON a.latitude = c.latitude + AND a.longitude = c.longitude + AND date_part('hour', a.time) = c.hour + """ + show_sql(sql) + + with timed("SQL anomaly (climatology CTE self-join)"): + got = ctx.sql(sql).to_pandas() + + # Array reference: grouped broadcast-subtract. + with timed("xarray reference"): + clim = ds["2m_temperature"].groupby("time.hour").mean("time") + anom = ds["2m_temperature"].groupby("time.hour") - clim + ref = anom.to_dataframe(name="anomaly").reset_index() + + merged = got.merge( + ref, + on=["time", "latitude", "longitude"], + suffixes=("_sql", "_ref"), + validate="one_to_one", + ) + assert len(merged) == len(got) == len(ref), ( + f"row-count mismatch: sql={len(got)} ref={len(ref)} merged={len(merged)}" + ) + check_close( + "anomaly (T − diurnal climatology)", + merged["anomaly_sql"], + merged["anomaly_ref"], + rtol=1e-4, + atol=1e-3, + ) + + print(f"\n {len(got):,} hourly anomalies over the window.") + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Anomaly: climatology CTE self-JOIN (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/05_forecast_skill.py b/benchmarks/geospatial/05_forecast_skill.py new file mode 100644 index 0000000..4d14d00 --- /dev/null +++ b/benchmarks/geospatial/05_forecast_skill.py @@ -0,0 +1,183 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "pandas", +# "gcsfs", +# "zarr>=3", +# ] +# /// +"""Forecast skill — scoring ML weather models against ERA5 is a JOIN + aggregate. + +This is the real thing: scoring the **Pangu-Weather** and **GraphCast** machine- +learning forecast models against ERA5 ground truth, the headline workload of +[WeatherBench 2](https://weatherbench2.readthedocs.io/). A forecast is indexed by +*initialization time* and *lead time* (``prediction_timedelta``); the truth is +indexed by *valid time*. Evaluation aligns them by ``valid_time = init + lead`` +and reduces the error to RMSE as a function of lead — the classic "error grows +with forecast horizon" curve. + +That alignment is a relational **JOIN**, and ``valid_time = init + lead`` is just +timestamp + duration arithmetic the engine does natively:: + + SELECT f.prediction_timedelta AS lead, + SQRT(AVG(POWER(f.t - e.t, 2))) AS rmse + FROM forecast f + JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid_time = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude + GROUP BY f.prediction_timedelta + +The whole evaluation — temporal alignment, spatial matching, the score — is one +JOIN and one aggregate. We run it for both models and check each against a NumPy +reference. + +Datasets: WeatherBench 2 **Pangu**, **GraphCast**, and **ERA5** at a coarse +64×32 grid (so the demo is small and fast), read from the public ``gs:// +weatherbench2`` bucket. Requires network; skips cleanly offline. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +import xarray_sql as xql + +from _harness import CaseSkipped, check_close, run_case, show_sql, timed + +_SO = {"token": "anon"} +_GRID = "64x32_equiangular_conservative" +_ERA5 = f"gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-{_GRID}.zarr" +_PANGU = f"gs://weatherbench2/datasets/pangu/2018-2022_0012_{_GRID}.zarr" +_GRAPHCAST = ( + "gs://weatherbench2/datasets/graphcast/2020/" + f"date_range_2019-11-16_2021-02-01_12_hours-{_GRID}.zarr" +) +_VAR = "2m_temperature" +_INIT = slice("2020-01-01", "2020-01-10") # 20 init times (12-hourly) + + +def _open(url: str) -> xr.Dataset: + try: + import gcsfs # noqa: F401 + + # decode_timedelta=True: forecasts store prediction_timedelta as a + # real duration (and it silences xarray's decode-timedelta warning). + return xr.open_zarr( + url, chunks=None, storage_options=_SO, decode_timedelta=True + ) + except Exception as exc: # noqa: BLE001 + raise CaseSkipped(f"WeatherBench2 unavailable ({exc})") from exc + + +def _load_forecast(url: str, grid: xr.Dataset) -> xr.Dataset: + """Load a model's 2m-temperature over the init window, all lead times. + + Snaps lat/lon onto ERA5's exact coordinate arrays (identical 64×32 grid) so + the equality JOIN on coordinates is bit-safe across the two Zarr stores. + """ + da = _open(url)[_VAR].sel(time=_INIT) + da = da.assign_coords( + latitude=grid.latitude.values, longitude=grid.longitude.values + ) + return da.to_dataset().load() + + +def _rmse_sql(table: str) -> str: + """The forecast↔truth JOIN + RMSE-per-lead query for one model ``table``.""" + return f""" + SELECT f.prediction_timedelta AS lead, + SQRT(AVG(POWER( + CAST(f."{_VAR}" AS DOUBLE) - e."{_VAR}", 2))) AS rmse + FROM {table} f + JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude + GROUP BY f.prediction_timedelta + ORDER BY lead + """ + + +def _rmse_by_lead(ctx: xql.XarrayContext, table: str) -> pd.DataFrame: + return ctx.sql(_rmse_sql(table)).to_pandas() + + +def _reference_rmse(fc: xr.Dataset, truth: xr.Dataset) -> np.ndarray: + """NumPy reference: for each lead, align truth at valid_time and take RMSE.""" + f = fc[_VAR] + e = truth[_VAR] + out = [] + for lead in f.prediction_timedelta.values: + valid = f.time.values + lead + e_at_valid = e.sel(time=valid).values # (init, lat, lon) + diff = f.sel(prediction_timedelta=lead).values - e_at_valid + out.append(float(np.sqrt(np.mean(diff**2)))) + return np.array(out) + + +def main() -> None: + era5_full = _open(_ERA5) + + # Load the forecasts first; snap their grid to ERA5's exact coordinates. + with timed("read Pangu + GraphCast forecasts"): + pangu = _load_forecast(_PANGU, era5_full) + graphcast = _load_forecast(_GRAPHCAST, era5_full) + + # ERA5 truth must span every valid time: window start through the last + # init plus the longest lead, so the JOIN keeps every forecast pair. + valid_max = ( + pangu.time.values.max() + pangu.prediction_timedelta.values.max() + ) + with timed("read ERA5 truth window"): + truth = ( + era5_full[[_VAR]] + .sel(time=slice(_INIT.start, pd.Timestamp(valid_max))) + .load() + ) + + print( + f" 64×32 2m_temperature | init {_INIT.start}…{_INIT.stop} " + f"({pangu.sizes['time']} inits × {pangu.sizes['prediction_timedelta']} leads)" + ) + + ctx = xql.XarrayContext() + ctx.from_dataset("era5", truth, chunks={"time": 100}) + ctx.from_dataset("pangu", pangu, chunks={"time": 20}) + ctx.from_dataset("graphcast", graphcast, chunks={"time": 20}) + + print("\n Forecast↔ERA5 JOIN on valid_time = init + lead, RMSE per lead:") + show_sql(_rmse_sql("pangu")) + + results = {} + for name, fc in [("pangu", pangu), ("graphcast", graphcast)]: + with timed(f"SQL RMSE-by-lead: {name}"): + got = _rmse_by_lead(ctx, name) + with timed(f"NumPy reference: {name}"): + ref = _reference_rmse(fc, truth) + check_close( + f"{name} RMSE(lead)", got["rmse"], ref, rtol=1e-4, atol=1e-3 + ) + results[name] = got + + # Headline table: error growth with forecast horizon, both models. + leads = results["pangu"]["lead"].dt.total_seconds() / 86400.0 + print("\n 2m-temperature RMSE (K) vs lead time — lower is better:") + print(f" {'lead (days)':>12} {'Pangu':>9} {'GraphCast':>11}") + for i in range(0, len(leads), 4): + print( + f" {leads.iloc[i]:>12.2f} " + f"{results['pangu']['rmse'].iloc[i]:>9.3f} " + f"{results['graphcast']['rmse'].iloc[i]:>11.3f}" + ) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Forecast skill: Pangu vs GraphCast vs ERA5 (WB2)") + ) diff --git a/benchmarks/geospatial/06_zonal_vector.py b/benchmarks/geospatial/06_zonal_vector.py new file mode 100644 index 0000000..444e154 --- /dev/null +++ b/benchmarks/geospatial/06_zonal_vector.py @@ -0,0 +1,135 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "gcsfs", +# "zarr>=3", +# ] +# /// +"""Zonal statistics over regions — "rasterize the polygons, then mask" is a JOIN. + +"What is the average temperature inside each region?" is the canonical +*raster × vector* operation. The array paradigm rasterizes each region to a +mask and reduces the raster under it, one region at a time. But a region is +just a row in a table of bounds, and "pixel falls inside region" is a **range +predicate** — so zonal statistics is a JOIN between the raster table and the +regions table, plus a GROUP BY:: + + SELECT r.region, AVG(a."2m_temperature") - 273.15 AS avg_c + FROM era5.surface a JOIN regions r + ON a.latitude BETWEEN r.lat_min AND r.lat_max + AND a.longitude BETWEEN r.lon_min AND r.lon_max + GROUP BY r.region + +This is exactly the README's promise — *joining tabular data with raster data* — +made concrete: the raster is the full **ARCO-ERA5** archive (``WHERE time …`` +prunes it to one day), the regions are a second SQL table, and the spatial +relationship is an ordinary ``BETWEEN``. + +Dataset: full ARCO-ERA5 + a handful of continental-scale bounding boxes +(longitudes in ERA5's 0–360°E convention). +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +import xarray_sql as xql + +from _era5 import load_window, open_era5, register_era5 +from _harness import check_close, run_case, show_sql, timed + +_DAY = "2020-06-01" + +# Continental-scale boxes (name, lat_min, lat_max, lon_min, lon_max), lon 0–360°E. +_REGIONS = [ + ("Sahara", 18.0, 30.0, 0.0, 30.0), + ("Amazon", -10.0, 5.0, 290.0, 310.0), + ("Australia_Outback", -30.0, -20.0, 125.0, 140.0), + ("Greenland", 65.0, 80.0, 300.0, 340.0), + ("SE_Asia", 5.0, 20.0, 95.0, 110.0), +] + + +def _regions_dataset() -> xr.Dataset: + """A vector layer as an xarray Dataset: one row per region, bounds as vars.""" + bounds = np.array([r[1:] for r in _REGIONS], dtype="float64") + return xr.Dataset( + { + "lat_min": (["region"], bounds[:, 0]), + "lat_max": (["region"], bounds[:, 1]), + "lon_min": (["region"], bounds[:, 2]), + "lon_max": (["region"], bounds[:, 3]), + }, + coords={"region": np.arange(len(_REGIONS))}, + ).chunk({"region": len(_REGIONS)}) + + +def main() -> None: + full = open_era5() + print( + f" raster: full ARCO-ERA5 ({full.sizes['time']:,} timesteps, " + f"{full.sizes['latitude']}×{full.sizes['longitude']}) " + f"vector: {len(_REGIONS)} continental boxes" + ) + + ctx = xql.XarrayContext() + with timed("register full ERA5 + regions"): + register_era5(ctx, full) + ctx.from_dataset( + "regions", _regions_dataset(), chunks={"region": len(_REGIONS)} + ) + + sql = f""" + SELECT r.region AS region_id, + AVG(a."2m_temperature") - 273.15 AS avg_c, + COUNT(*) AS n_obs + FROM era5.surface a + JOIN regions r + ON a.latitude BETWEEN r.lat_min AND r.lat_max + AND a.longitude BETWEEN r.lon_min AND r.lon_max + WHERE a.time BETWEEN TIMESTAMP '{_DAY}' + AND TIMESTAMP '{_DAY} 23:00:00' + GROUP BY r.region + ORDER BY r.region + """ + show_sql(sql) + + with timed("SQL zonal stats (raster × vector range JOIN, WHERE-pruned)"): + got = ( + ctx.sql(sql) + .to_pandas() + .sort_values("region_id") + .reset_index(drop=True) + ) + + # Array reference: load the same day once, mask each region in memory. + with timed("xarray reference"): + day = load_window(full, "2m_temperature", time=_DAY)["2m_temperature"] + ref_vals = [] + for _, lat_min, lat_max, lon_min, lon_max in _REGIONS: + mask = ( + (day.latitude >= lat_min) + & (day.latitude <= lat_max) + & (day.longitude >= lon_min) + & (day.longitude <= lon_max) + ) + ref_vals.append(float(day.where(mask).mean()) - 273.15) + ref = np.array(ref_vals) + + check_close( + "zonal mean per region (°C)", got["avg_c"], ref, rtol=1e-4, atol=1e-2 + ) + + print("\n Region avg °C n_obs") + for (name, *_), row in zip(_REGIONS, got.itertuples()): + print(f" {name:<20} {row.avg_c:7.2f} {row.n_obs:>10,}") + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Zonal stats: raster × vector range JOIN (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py new file mode 100644 index 0000000..5a8ccd4 --- /dev/null +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -0,0 +1,166 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "pyproj", +# "pyarrow", +# ] +# /// +"""Reprojection — a per-pixel CRS transform is a scalar UDF (à la ST_Transform). + +Reprojection moves coordinates from one CRS to another (here UTM zone 32N, +EPSG:32632, → lon/lat, EPSG:4326). Crucially it is **row-independent**: each +pixel's new coordinate depends only on its own old coordinate. That is exactly +the shape of a SQL *scalar UDF*, and it is precisely how the geospatial SQL +world already does it — PostGIS ``ST_Transform`` and DuckDB-spatial +``ST_Transform`` are scalar PROJ wrappers. + +So we register a PROJ-backed scalar UDF and reproject in SQL:: + + SELECT x, y, utm_to_lon(x, y) AS lon, utm_to_lat(x, y) AS lat + FROM grid + +The UDF (built on ``pyproj``, vectorized over each Arrow batch) mirrors the +``cftime()`` UDF already shipped in xarray-sql (see ``xarray_sql/cftime.py``); +it could graduate into the package as ``xql.register_reproject_udfs``. + +**What this does *not* do:** it moves the coordinates, it does not resample onto +a regular target grid. Producing a gridded product in the new CRS still needs +interpolation — which is case 08, where regridding turns out to be a JOIN +against a weight table rather than a scalar UDF. + +**An honest caveat on threading:** PROJ's transformation context is not +thread-safe, and DataFusion evaluates independent projection expressions on +concurrent worker threads. Two separate PROJ UDFs in one ``SELECT`` (one for +lon, one for lat) race and segfault. The fix is to return *both* coordinates +from a **single** struct-returning UDF (so PROJ is touched once per row), and +to keep the source in one chunk so that single UDF runs serially. A +production-grade ``ST_Transform`` would additionally give each worker thread its +own PROJ context; the point here is the *shape* of the operation — a scalar +UDF — not the parallel execution. + +No network: a synthetic UTM grid over the extent of a Sentinel-2 tile. +""" + +from __future__ import annotations + +import numpy as np +import pyarrow as pa +import pyproj +import xarray as xr +from datafusion import udf + +import xarray_sql as xql + +from _harness import check_close, run_case, show_sql, timed + + +def register_reproject_udf( + ctx, src_crs: str, dst_crs: str, name: str = "reproject" +) -> None: + """Register a ``reproject(x, y) -> {lon, lat}`` PROJ scalar UDF. + + Mirrors ``xarray_sql.cftime.make_cftime_udf``: a vectorized scalar UDF over + Arrow arrays. ``always_xy=True`` keeps argument order (easting, northing) → + (lon, lat) regardless of CRS axis conventions. Like PostGIS/DuckDB + ``ST_Transform``, it returns *both* output coordinates from one call — here + as an Arrow struct, so callers write ``reproject(x, y)['lon']``. + + Returning a struct (rather than two separate UDFs) is deliberate: PROJ's + context is not thread-safe, and DataFusion evaluates independent projection + expressions concurrently — two PROJ UDFs in one SELECT race and crash. One + struct-returning UDF does the transform exactly once per row, on one thread. + """ + ret = pa.struct([("lon", pa.float64()), ("lat", pa.float64())]) + + def _fn(x: pa.Array, y: pa.Array) -> pa.Array: + # Build the Transformer inside the call so it lives on the worker + # thread that uses it (PROJ contexts are thread-bound). + transformer = pyproj.Transformer.from_crs( + src_crs, dst_crs, always_xy=True + ) + xs = np.asarray(x.to_numpy(zero_copy_only=False), dtype="float64") + ys = np.asarray(y.to_numpy(zero_copy_only=False), dtype="float64") + lon, lat = transformer.transform(xs, ys) + return pa.StructArray.from_arrays( + [ + pa.array(np.asarray(lon, "float64")), + pa.array(np.asarray(lat, "float64")), + ], + names=["lon", "lat"], + ) + + ctx.register_udf( + udf(_fn, [pa.float64(), pa.float64()], ret, "immutable", name) + ) + + +def _utm_grid() -> xr.Dataset: + """A synthetic field on a UTM grid matching a Sentinel-2 T32T tile extent.""" + # EPSG:32632 extent of tile T32TLQ (from the STAC proj:bbox), coarsened. + x = np.linspace(300_000.0, 409_800.0, 60, dtype="float64") + y = np.linspace(5_000_040.0, 4_890_240.0, 60, dtype="float64") + elevation = np.zeros( + (y.size, x.size), dtype="float64" + ) # value is irrelevant + # Single chunk → single partition → serial UDF (PROJ is not thread-safe). + return xr.Dataset( + {"elevation": (["y", "x"], elevation)}, + coords={"y": y, "x": x}, + ).chunk({"y": 60, "x": 60}) + + +def main() -> None: + src_crs, dst_crs = "EPSG:32632", "EPSG:4326" + ds = _utm_grid() + print( + f" UTM grid {dict(ds.sizes)} ({ds.sizes['y'] * ds.sizes['x']:,} points) " + f"{src_crs} → {dst_crs}" + ) + + ctx = xql.XarrayContext() + ctx.from_dataset("grid", ds, chunks={"y": 60, "x": 60}) + register_reproject_udf(ctx, src_crs, dst_crs) + + sql = """ + SELECT x, y, + reproject(x, y)['lon'] AS lon, + reproject(x, y)['lat'] AS lat + FROM grid + ORDER BY y, x + """ + show_sql(sql) + + with timed("SQL reprojection (PROJ scalar UDF)"): + got = ctx.sql(sql).to_pandas() + + # Array reference: the same PROJ transform applied to the full grid. + with timed("pyproj reference"): + transformer = pyproj.Transformer.from_crs( + src_crs, dst_crs, always_xy=True + ) + xx, yy = np.meshgrid(ds.x.values, ds.y.values) + ref_lon, ref_lat = transformer.transform(xx.ravel(), yy.ravel()) + + got = got.sort_values(["y", "x"]).reset_index(drop=True) + # The reference is built in (y, x) row-major order; align by sorting both. + order = np.lexsort((xx.ravel(), yy.ravel())) + check_close( + "reprojected longitude", got["lon"], ref_lon[order], rtol=0, atol=1e-9 + ) + check_close( + "reprojected latitude", got["lat"], ref_lat[order], rtol=0, atol=1e-9 + ) + + print( + f"\n Corner check: UTM ({got.x.iloc[0]:.0f}, {got.y.iloc[0]:.0f}) → " + f"lon {got.lon.iloc[0]:.4f}, lat {got.lat.iloc[0]:.4f}" + ) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Reprojection: PROJ scalar UDF (ST_Transform)") + ) diff --git a/benchmarks/geospatial/08_regrid_weights.py b/benchmarks/geospatial/08_regrid_weights.py new file mode 100644 index 0000000..4f94c40 --- /dev/null +++ b/benchmarks/geospatial/08_regrid_weights.py @@ -0,0 +1,158 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "scipy", +# ] +# /// +"""Regridding — interpolation to a new grid is a sparse matmul, i.e. a JOIN. + +Regridding (resampling a field from one grid onto another) is the operation we +most associate with the *array* paradigm — xESMF/ESMF, ``apply_ufunc``, +``.interp()``. But every linear regridding scheme (bilinear, conservative, +nearest) is mathematically a **sparse matrix–vector product**: each output cell +is a weighted sum of a few input cells. And a sparse matrix is just a table of +``(dst_id, src_id, weight)`` rows. So *applying* a regridding is:: + + SELECT w.dst_id, SUM(s.value * w.weight) AS regridded + FROM weights w JOIN src s ON s.cell_id = w.src_id + GROUP BY w.dst_id + +— a JOIN against the weight table plus a weighted GROUP BY. This is the most +relational the "array" paradigm ever gets: the operation we reach for xESMF to +do is a join. + +**Where the array paradigm still earns its keep:** *generating* the weights is +the genuinely geometric part (cell overlaps, interpolation stencils, spherical +coordinates). Here we build bilinear weights with a few lines of numpy; for +conservative remapping on real grids you would let ESMF/xESMF compute them once +and hand the resulting sparse matrix to SQL as a table. SQL *applies* the +weights; it does not invent the geometry. + +We validate the SQL matmul against xarray's own bilinear ``.interp()``. +No network; fully synthetic. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +import xarray_sql as xql + +from _harness import check_close, run_case, show_sql, timed + + +def _linear_weights( + src: np.ndarray, dst: np.ndarray +) -> list[tuple[int, int, float]]: + """1-D linear-interpolation weights: (dst_index, src_index, weight) triples. + + Each target point falls between two source points and borrows from both, + with weights summing to 1 — the 1-D building block of bilinear regridding. + """ + triples = [] + for t, x in enumerate(dst): + i = int(np.clip(np.searchsorted(src, x) - 1, 0, len(src) - 2)) + span = src[i + 1] - src[i] + hi = (x - src[i]) / span + triples.append((t, i, 1.0 - hi)) + triples.append((t, i + 1, hi)) + return triples + + +def _bilinear_weight_table( + slat: np.ndarray, slon: np.ndarray, tlat: np.ndarray, tlon: np.ndarray +) -> xr.Dataset: + """Build the sparse bilinear weight matrix as a (dst_id, src_id, weight) table. + + The 2-D weight is the outer product of the 1-D lat and lon weights; src_id + and dst_id are row-major flattenings of the source and target grids. + """ + nslon, ntlon = len(slon), len(tlon) + lat_w = _linear_weights(slat, tlat) + lon_w = _linear_weights(slon, tlon) + dst_ids, src_ids, weights = [], [], [] + for tj, si, wlat in lat_w: + for tk, sj, wlon in lon_w: + dst_ids.append(tj * ntlon + tk) + src_ids.append(si * nslon + sj) + weights.append(wlat * wlon) + n = len(weights) + return xr.Dataset( + { + "dst_id": (["pair"], np.array(dst_ids, dtype="int64")), + "src_id": (["pair"], np.array(src_ids, dtype="int64")), + "weight": (["pair"], np.array(weights, dtype="float64")), + }, + coords={"pair": np.arange(n)}, + ).chunk({"pair": n}) + + +def main() -> None: + # Coarse source grid with a smooth field; finer target grid strictly inside. + slat = np.linspace(0.0, 10.0, 11) + slon = np.linspace(0.0, 10.0, 11) + field = ( + np.sin(slat[:, None] / 3.0) * np.cos(slon[None, :] / 4.0) + + 0.1 * slat[:, None] + ) + src_da = xr.DataArray( + field, coords={"lat": slat, "lon": slon}, dims=["lat", "lon"] + ) + + tlat = np.linspace(0.3, 9.7, 25) + tlon = np.linspace(0.4, 9.6, 30) + print( + f" regrid {len(slat)}×{len(slon)} → {len(tlat)}×{len(tlon)} (bilinear)" + ) + + # Source field as a flat (cell_id, value) table. + src_table = xr.Dataset( + {"value": (["cell_id"], field.ravel())}, + coords={"cell_id": np.arange(field.size)}, + ).chunk({"cell_id": field.size}) + weights = _bilinear_weight_table(slat, slon, tlat, tlon) + print( + f" weight matrix: {weights.sizes['pair']:,} nonzeros " + f"({len(tlat) * len(tlon)} targets × 4 corners)" + ) + + ctx = xql.XarrayContext() + ctx.from_dataset("src", src_table, chunks={"cell_id": field.size}) + ctx.from_dataset("weights", weights, chunks={"pair": weights.sizes["pair"]}) + + sql = """ + SELECT w.dst_id, + SUM(s.value * w.weight) AS regridded + FROM weights w + JOIN src s ON s.cell_id = w.src_id + GROUP BY w.dst_id + ORDER BY w.dst_id + """ + show_sql(sql) + + with timed("SQL regrid (weight-table JOIN + weighted SUM)"): + got = ( + ctx.sql(sql) + .to_pandas() + .sort_values("dst_id") + .reset_index(drop=True) + ) + + # Array reference: xarray's own bilinear interpolation onto the target grid. + with timed("xarray .interp reference"): + ref = src_da.interp(lat=tlat, lon=tlon, method="linear") + + check_close("bilinear regrid", got["regridded"], ref, rtol=1e-9, atol=1e-9) + + print( + f"\n {len(got):,} target cells regridded; " + f"value range [{got.regridded.min():.3f}, {got.regridded.max():.3f}]." + ) + + +if __name__ == "__main__": + raise SystemExit(run_case(main, "Regridding: sparse weight-table JOIN")) diff --git a/benchmarks/geospatial/README.md b/benchmarks/geospatial/README.md new file mode 100644 index 0000000..82b691a --- /dev/null +++ b/benchmarks/geospatial/README.md @@ -0,0 +1,74 @@ +# Geospatial SQL benchmarks + +**Thesis:** the core geospatial operations we assume require an *array* paradigm +are, underneath, **relational** operations — `GROUP BY`, `JOIN`, window +functions, and `CASE`. Each script here takes one such operation, expresses it +in SQL against [`xarray-sql`](../../README.md), and **proves the SQL answer +matches an xarray/array reference implementation** (`numpy.assert_allclose`). +Wall-clock and peak memory are reported too, but the headline is correctness + +clarity of the SQL. + +This suite is *expressibility-first*: the point is that the SQL reads like the +plain-English definition of the operation, and computes the same numbers. + +## The cases + +| # | Case | Array mental model | Relational reality | +|---|------|--------------------|--------------------| +| 01 | `01_ndvi.py` | `apply_ufunc` over a raster | column algebra + `CASE` mask | +| 02 | `02_climatology.py` | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | +| 03 | `03_zonal_mean.py` | reduce over lon/time axes | `GROUP BY latitude` | +| 04 | `04_anomaly.py` | climatology broadcast-subtract | climatology CTE self-`JOIN` | +| 05 | `05_forecast_skill.py` | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` + aggregate | +| 06 | `06_zonal_vector.py` | rasterize + mask per region | range `JOIN` raster↔regions | +| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()`), à la PostGIS `ST_Transform` | +| 08 | `08_regrid_weights.py` | interpolation to a new grid | sparse-weight table `JOIN` + weighted `GROUP BY` | + +Cases 01–06 show operations that are *natively* relational. Cases 07–08 are the +"hardest" array operations — reprojection and regridding — and show where a UDF +fits (a per-row coordinate transform) versus where the operation is really a +sparse matrix multiply expressed as a `JOIN`. See +[`docs/geospatial.md`](../../docs/geospatial.md) for the full narrative, +including *where the array paradigm still earns its keep* (generating the +interpolation weights — the geometry — which SQL applies but does not compute). + +## Datasets + +- **01 NDVI** — a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample + service (bands B04/B08). Requires network; skips cleanly if offline. +- **02–06** — the full **[ARCO-ERA5](https://github.com/google-research/arco-era5)** + archive (0.25° global, ~1.3M hourly timesteps, 273 variables) read anonymously + from a public GCS bucket. Cases 03 and 06 register the *whole* archive and let + SQL `WHERE` prune it to a one-day window (the partition-pruning demo); cases + 02/04 read a bounded region/time window into memory once and compare + SQL against the same array. All require network (`gcsfs`); skip cleanly + offline. Each ERA5 case takes roughly a minute, dominated by the GCS read. +- **05 forecast skill** — the **[WeatherBench 2](https://weatherbench2.readthedocs.io/)** + Pangu-Weather, GraphCast, and ERA5 datasets at a coarse 64×32 grid, scoring + both ML models against ERA5 ground truth. Network-backed; runs in seconds + because the grid is small. +- **07–08** — small/synthetic grids plus precomputed regrid weights, so they run + without heavy geospatial dependencies (ESMF/ESMPy). + +## Running + +Inside the repo, use the project environment (so `import xarray_sql` resolves to +the locally built native extension): + +```shell +python benchmarks/geospatial/03_zonal_mean.py +``` + +Each script also carries [PEP 723 / `uv` inline metadata](https://docs.astral.sh/uv/guides/scripts/), +so it can be run standalone against the published `xarray-sql` wheel: + +```shell +uv run benchmarks/geospatial/03_zonal_mean.py +``` + +A passing case prints a `✅ … SQL matches array reference` line; a mismatch +raises `AssertionError` and exits non-zero. Cases that need an unavailable +dataset/dependency print `⏭ SKIPPED` and exit 0. + +Shared helpers (timing, peak memory, the `assert_allclose` wrapper, SQL echo) +live in [`_harness.py`](_harness.py). diff --git a/benchmarks/geospatial/_era5.py b/benchmarks/geospatial/_era5.py new file mode 100644 index 0000000..af73a0f --- /dev/null +++ b/benchmarks/geospatial/_era5.py @@ -0,0 +1,70 @@ +"""Shared ARCO-ERA5 access for the geospatial benchmarks. + +[ARCO-ERA5](https://github.com/google-research/arco-era5) is the ECMWF ERA5 +reanalysis re-published as analysis-ready, cloud-optimized Zarr on a public GCS +bucket: 273 variables, hourly, 0.25° global (721×1440), since 1940 — about 1.3 +million timesteps. We open the *whole* archive (no time/space slicing on the +xarray side) and let SQL ``WHERE`` clauses prune it. That is the demo: the +table is the full reanalysis; the query reads only the window it asks for. + +These cases require anonymous GCS access (``gcsfs``); they skip cleanly when +the network or bucket is unavailable. +""" + +from __future__ import annotations + +import xarray as xr + +import xarray_sql as xql + +from _harness import CaseSkipped + +URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" + +#: Friendly names for ERA5's two dimension groups (surface vs. atmosphere). +TABLE_NAMES = { + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", +} + + +def open_era5() -> xr.Dataset: + """Open the full ARCO-ERA5 archive (lazy, dask off), or skip if unreachable.""" + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + return xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + + +def register_era5(ctx: xql.XarrayContext, ds: xr.Dataset, *, chunks=None): + """Register ERA5 as ``era5.surface`` / ``era5.atmosphere`` tables.""" + ctx.from_dataset( + "era5", ds, chunks=chunks or {"time": 6}, table_names=TABLE_NAMES + ) + return ctx + + +def load_window( + ds: xr.Dataset, + var: str, + *, + time, + latitude=None, + longitude=None, +) -> xr.Dataset: + """Read a bounded ERA5 window for ``var`` into memory as a 1-variable Dataset. + + Multi-timestep cases (climatology, anomaly, forecast skill) read their window + *once* into memory here, then run both the SQL and its array reference + against the same in-memory data — keeping I/O bounded and the comparison + apples-to-apples. ERA5 latitude is descending, so pass + ``latitude=slice(north, south)``; longitude is 0–360°E ascending. + """ + da = ds[var].sel(time=time) + if latitude is not None: + da = da.sel(latitude=latitude) + if longitude is not None: + da = da.sel(longitude=longitude) + return da.to_dataset().load() diff --git a/benchmarks/geospatial/_harness.py b/benchmarks/geospatial/_harness.py new file mode 100644 index 0000000..f1e320e --- /dev/null +++ b/benchmarks/geospatial/_harness.py @@ -0,0 +1,120 @@ +"""Shared harness for the geospatial SQL benchmarks. + +The suite is *expressibility-first*: each case states a geospatial operation we +normally reach for an array library to perform, expresses it in SQL against +``xarray-sql``, and proves the SQL answer matches an xarray/array reference +implementation. Wall-clock and peak memory are reported too, but the headline +is correctness + clarity of the SQL. + +These helpers keep each case script short and uniform: + +* :func:`banner` / :func:`show_sql` — readable section headers and SQL echo. +* :func:`timed` — a context manager that reports elapsed time and peak memory. +* :func:`check_close` — assert a SQL result matches an array reference, then + print a PASS line. Raises ``AssertionError`` on mismatch (so a broken case + fails loudly rather than silently "passing"). +* :func:`run_case` — run a case's ``main()``, turning a raised + :class:`CaseSkipped` (e.g. an offline dataset) into a clean skip. +""" + +from __future__ import annotations + +import contextlib +import sys +import time +import tracemalloc +from collections.abc import Callable, Iterator +from typing import Any + +import numpy as np + +_WIDTH = 72 + + +class CaseSkipped(Exception): + """Raised by a case when it cannot run in this environment (e.g. offline).""" + + +def banner(text: str) -> None: + """Print a titled section divider.""" + print(f"\n{'─' * _WIDTH}") + print(f" {text}") + print(f"{'─' * _WIDTH}") + + +def show_sql(sql: str, *, label: str = "SQL") -> None: + """Echo a SQL statement so the reader sees exactly what ran.""" + print(f"\n {label}:") + for line in sql.strip("\n").splitlines(): + print(f" │ {line}") + print() + + +@contextlib.contextmanager +def timed(label: str) -> Iterator[None]: + """Time a block and report elapsed wall-clock and peak memory. + + Peak memory is the Python-allocator peak during the block (via + ``tracemalloc``); it captures the materialized result and intermediate + buffers, which is what we care about for "did this blow up memory". + """ + tracemalloc.start() + tracemalloc.reset_peak() + t0 = time.perf_counter() + try: + yield + finally: + elapsed = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + print(f" ⏱ {label}: {elapsed:.3f}s (peak {peak / 1e6:.1f} MB)") + + +def check_close( + name: str, + got: Any, + expected: Any, + *, + rtol: float = 1e-5, + atol: float = 1e-6, + equal_nan: bool = True, +) -> None: + """Assert ``got`` matches the array reference ``expected``, then print PASS. + + Both arguments are coerced to ``float64`` numpy arrays and flattened, so a + SQL result column (pandas Series / pyarrow) can be compared directly to an + xarray ``DataArray`` reference. Order-sensitive: sort both sides on a shared + key before calling if the SQL result order is not guaranteed. + """ + g = np.asarray(getattr(got, "values", got), dtype=np.float64).ravel() + e = np.asarray( + getattr(expected, "values", expected), dtype=np.float64 + ).ravel() + if g.shape != e.shape: + raise AssertionError( + f"{name}: shape mismatch — SQL {g.shape} vs reference {e.shape}" + ) + np.testing.assert_allclose(g, e, rtol=rtol, atol=atol, equal_nan=equal_nan) + print( + f" ✅ {name}: SQL matches array reference " + f"(n={g.size}, rtol={rtol:g}, atol={atol:g})" + ) + + +def run_case(main: Callable[[], None], title: str) -> int: + """Run a case ``main()``; turn :class:`CaseSkipped` into a clean skip. + + Returns a process exit code: 0 on success or skip, 1 on failure. Use as + ``if __name__ == '__main__': raise SystemExit(run_case(main, '...'))``. + """ + banner(title) + try: + main() + except CaseSkipped as exc: + print(f"\n ⏭ SKIPPED: {exc}") + return 0 + except Exception as exc: # noqa: BLE001 — surface any failure as exit 1 + print(f"\n ❌ FAILED: {type(exc).__name__}: {exc}", file=sys.stderr) + raise + print(f"\n 🎉 {title}: done.") + return 0 diff --git a/docs/geospatial.md b/docs/geospatial.md new file mode 100644 index 0000000..30af680 --- /dev/null +++ b/docs/geospatial.md @@ -0,0 +1,227 @@ +# Geospatial operations are relational operations + +A working hypothesis, and a slightly radical one: **the core operations of +geospatial and climate analysis — the ones we reach for an array library to +perform — are, underneath, relational operations.** Climatologies, anomalies, +zonal means, spectral indices, forecast skill, even regridding: each maps onto +ordinary SQL — `GROUP BY`, `JOIN`, window functions, `CASE`, and the occasional +scalar UDF. + +The array paradigm (NumPy, Xarray, Dask) is a wonderful *interface* for these +operations. But it is not the only one, and for a large and growing audience — +the people fluent in SQL rather than in `apply_ufunc` and rechunking — it is not +the most accessible one. [`xarray-sql`](../README.md) lets you pose these +questions in SQL and answers them with a real query engine (DataFusion), +complete with partition pruning and projection pushdown. + +This page makes the argument case by case. Every claim below is backed by a +runnable script in [`benchmarks/geospatial/`](../benchmarks/geospatial/) that +poses the operation in SQL and **asserts the answer matches an xarray/array +reference** to floating-point tolerance. The point is not that "SQL is faster"; +the point is that the SQL reads like the *definition* of the operation and +computes the same numbers — at ERA5's real 0.25° global resolution. + +## The mapping + +| Operation | The "array" framing | The relational reality | Script | +|-----------|---------------------|------------------------|--------| +| Spectral index (NDVI) | `apply_ufunc` over a raster | column algebra + `CASE` mask | [`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) | +| Climatology | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | [`02_climatology.py`](../benchmarks/geospatial/02_climatology.py) | +| Zonal mean | reduce over lon/time axes | `GROUP BY lat` | [`03_zonal_mean.py`](../benchmarks/geospatial/03_zonal_mean.py) | +| Anomaly | grouped broadcast-subtract | climatology CTE self-`JOIN` | [`04_anomaly.py`](../benchmarks/geospatial/04_anomaly.py) | +| Forecast skill (RMSE) | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` | [`05_forecast_skill.py`](../benchmarks/geospatial/05_forecast_skill.py) | +| Zonal stats over regions | rasterize polygons + mask | raster × vector range `JOIN` | [`06_zonal_vector.py`](../benchmarks/geospatial/06_zonal_vector.py) | +| Reprojection | per-pixel CRS transform | scalar **UDF** (`ST_Transform`-style) | [`07_reproject_udf.py`](../benchmarks/geospatial/07_reproject_udf.py) | +| Regridding | interpolation to a new grid | sparse-weight table `JOIN` | [`08_regrid_weights.py`](../benchmarks/geospatial/08_regrid_weights.py) | + +## 1. A pixel-wise formula is a column expression + +NDVI is `(NIR − Red) / (NIR + Red)`, per pixel, with clouds/nodata masked out. +The array idiom broadcasts a ufunc over the raster. But "one output per pixel, +computed from that pixel's bands" is the definition of a SQL projection, and +"drop invalid pixels" is `CASE WHEN`: + +```sql +SELECT x, y, + CASE WHEN red = 0 OR nir = 0 THEN NULL + ELSE (nir_refl - red_refl) / (nir_refl + red_refl) + END AS ndvi +FROM scene +``` + +[`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) runs this against a **real +Sentinel-2 L2A scene in Zarr** (ESA's EOPF sample service) and matches xarray's +`apply_ufunc`-style result over a million pixels. + +## 2. A climatology is a `GROUP BY` over the cycle + +A climatology is the average value for each time-of-cycle at each location. In +the array world this is the canonical painful workload — load native chunks, +*rechunk* so all of time lands in one chunk, reduce, rechunk back. The +rechunking serves the array layout, not the question. The question is: + +```sql +SELECT latitude, longitude, date_part('hour', time) AS hour, + AVG("2m_temperature") +FROM era5 GROUP BY latitude, longitude, date_part('hour', time) +``` + +The grouping keys are the dimensions you keep; everything else is reduced. No +layout to reason about. [`02_climatology.py`](../benchmarks/geospatial/02_climatology.py) +computes the **diurnal cycle** of ERA5 2m-temperature over a region — averaging +each cell by hour of day — and matches `da.groupby("time.hour").mean()` across +~500k cells. + +A **zonal mean** ([`03_zonal_mean.py`](../benchmarks/geospatial/03_zonal_mean.py)) +is the same idea with fewer keys: the axes you "reduce over" are simply the +columns you don't `GROUP BY`. + +## 3. Broadcasting a normal back onto observations is a `JOIN` + +An anomaly subtracts each cell's climatological normal from every matching +observation. Xarray expresses the realignment with grouped broadcasting +(`ds.groupby("time.hour") - climatology`). That realignment — *attach each +cell's normal to every timestep that shares its key* — is a JOIN on the +grouping key: + +```sql +WITH clim AS ( + SELECT latitude, longitude, date_part('hour', time) AS hour, + AVG("2m_temperature") AS clim_t + FROM era5 GROUP BY latitude, longitude, date_part('hour', time) +) +SELECT a.time, a.latitude, a.longitude, + a."2m_temperature" - c.clim_t AS anomaly +FROM era5 a JOIN clim c + ON a.latitude = c.latitude AND a.longitude = c.longitude + AND date_part('hour', a.time) = c.hour +``` + +[`04_anomaly.py`](../benchmarks/geospatial/04_anomaly.py) computes the +climatology once (the CTE) and joins it back to every observation. + +## 4. Forecast evaluation is a `JOIN` on valid time + aggregate + +This is the real workload of [WeatherBench 2](https://weatherbench2.readthedocs.io/): +scoring machine-learning weather models — **Pangu-Weather** and **GraphCast** — +against ERA5 ground truth. A forecast is indexed by *initialization time* and +*lead time* (`prediction_timedelta`); the truth is indexed by *valid time*. +Evaluation aligns them by `valid_time = init + lead` and reduces the error to +RMSE as a function of lead. + +That alignment is a relational JOIN, and `valid_time = init + lead` is just +timestamp + duration arithmetic the engine does natively: + +```sql +SELECT f.prediction_timedelta AS lead, + SQRT(AVG(POWER(f."2m_temperature" - e."2m_temperature", 2))) AS rmse +FROM forecast f +JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid_time = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude +GROUP BY f.prediction_timedelta +``` + +The entire evaluation — temporal alignment across three time axes, spatial +matching, and the score — is one JOIN and one aggregate. +[`05_forecast_skill.py`](../benchmarks/geospatial/05_forecast_skill.py) runs it +for both models, matches a NumPy reference, and reproduces the published result +that GraphCast edges out Pangu at every lead — the classic "error grows with +horizon" curve (≈0.3 K at 6 h rising to ≈2.5 K at 9 days): + +``` + lead (days) Pangu GraphCast + 0.25 0.336 0.296 + 5.25 1.469 1.228 + 9.25 2.814 2.380 +``` + +## 5. Raster × vector zonal statistics is a range `JOIN` + +"Average the raster inside each region" is the canonical raster-meets-vector +task. The array idiom rasterizes each polygon to a mask and reduces under it. But +a region is a row in a table of bounds, and "pixel inside region" is a range +predicate — so zonal statistics is a JOIN: + +```sql +SELECT r.region, AVG(a."2m_temperature") - 273.15 AS avg_c +FROM era5.surface a JOIN regions r + ON a.latitude BETWEEN r.lat_min AND r.lat_max + AND a.longitude BETWEEN r.lon_min AND r.lon_max +WHERE a.time BETWEEN TIMESTAMP '2020-06-01' AND TIMESTAMP '2020-06-01 23:00:00' +GROUP BY r.region +``` + +This is the README's promise — *joining tabular data with raster data* — made +literal: the raster is the full ERA5 archive (the `WHERE` prunes it to a day), +the regions are a second SQL table, and the spatial relationship is an ordinary +`BETWEEN`. See [`06_zonal_vector.py`](../benchmarks/geospatial/06_zonal_vector.py) +— it reports e.g. Sahara 33 °C vs Greenland −8 °C for a June day. (Rectangular +regions keep the demo dependency-free; arbitrary polygons are the natural next +step, via a point-in-polygon UDF — see below.) + +## 6. The hard cases: where a UDF fits, and where it doesn't + +Reprojection and regridding are the operations most wedded to the array +paradigm. They split cleanly along one line: **is the operation row-independent?** + +**Reprojection is.** Moving a coordinate from one CRS to another depends only on +that coordinate, so it is a *scalar function* — exactly what PostGIS and +DuckDB-spatial already ship as `ST_Transform`. We register a PROJ-backed scalar +UDF (mirroring the `cftime()` UDF already in `xarray_sql/cftime.py`) and +reproject in SQL: + +```sql +SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat +FROM grid +``` + +[`07_reproject_udf.py`](../benchmarks/geospatial/07_reproject_udf.py) matches +pyproj to 1e-9. Two honest caveats, both documented in the script: PROJ's +context is not thread-safe (so the UDF returns both coordinates from *one* call +and runs on a single partition), and reprojection moves coordinates without +resampling onto a grid — which is the next operation. + +**Regridding is not** row-independent: each output cell is a weighted blend of +several input cells. That is a *many-to-many* relationship — and a many-to-many +weighted blend is a sparse matrix–vector product, which is a `JOIN` against a +weight table plus a weighted `GROUP BY`: + +```sql +SELECT w.dst_id, SUM(s.value * w.weight) AS regridded +FROM weights w JOIN src s ON s.cell_id = w.src_id +GROUP BY w.dst_id +``` + +[`08_regrid_weights.py`](../benchmarks/geospatial/08_regrid_weights.py) matches +xarray's bilinear `.interp()` exactly. So regridding does not weaken the thesis — +it is the most relational operation of all. + +## Where the array paradigm still earns its keep + +The honest boundary is **weight generation**. Applying a regridding is a join; +*computing* the weights — cell overlaps for conservative remapping, stencils and +spherical geometry for bilinear, the whole machinery of xESMF/ESMF — is genuinely +geometric work that arrays (and specialized libraries) do well. The relational +view does not replace that; it consumes its output. The division of labor is +clean and, we think, the right one: + +> **Arrays compute the geometry (the weights). SQL applies it (the join).** + +Likewise, the array libraries remain the right tool for building the inputs in +the first place — opening Zarr, decoding CF metadata, the numerics of generating +a weight matrix. `xarray-sql` sits downstream of all that as a query front-end: +once the data is openable as an `xarray.Dataset`, these everyday operations are +expressible — and accessible — as SQL. + +## Running the suite + +```shell +python benchmarks/geospatial/02_climatology.py # inside the repo +uv run benchmarks/geospatial/02_climatology.py # standalone (PEP 723 deps) +``` + +Each script prints its SQL, runs the array reference, and asserts the two agree. +See [`benchmarks/geospatial/README.md`](../benchmarks/geospatial/README.md) for +the full list and dataset notes. diff --git a/zensical.toml b/zensical.toml index 21bce73..f182a70 100644 --- a/zensical.toml +++ b/zensical.toml @@ -9,6 +9,7 @@ edit_uri = "edit/main/docs/" nav = [ {"Home" = "index.md"}, {"Examples" = "examples.md"}, + {"Geospatial in SQL" = "geospatial.md"}, {"Contributing" = "contributing.md"}, {"Reference" = "reference/xarray_sql.md"} ] From f8dc42947ebfe412ff39412baa13ea6e29adb3f0 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 23 Jun 2026 20:00:36 +0300 Subject: [PATCH 02/26] NDVI: open Sentinel-2 idiomatically (pystac + open_datatree), drop manual scaling/CASE/sort Addresses PR review of case 01: - Discover the Sentinel-2 L2A product with pystac-client instead of a hard-coded object-store href. - Open it the canonical xarray way: xr.open_datatree on the EOPF Zarr, then the measurements/reflectance/r10m node. This yields B04/B08 already scaled to reflectance with x/y coordinates, so the manual zarr.open_array reads and proj:transform coordinate reconstruction are gone. - NDVI is now pure column arithmetic on both sides: xarray decodes the band _FillValue to NaN on open, so nodata masking is free and the CASE WHEN is dropped. The reference is a clean one-line xarray expression. - Use ORDER BY y, x in SQL and compare via coordinate-aligned xr.testing.assert_allclose (reindex_like) instead of manual .sortby() on both sides. Docs (suite README, docs/geospatial.md) updated to match. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/01_ndvi.py | 164 +++++++++++++++---------------- benchmarks/geospatial/README.md | 5 +- docs/geospatial.md | 24 ++--- 3 files changed, 96 insertions(+), 97 deletions(-) diff --git a/benchmarks/geospatial/01_ndvi.py b/benchmarks/geospatial/01_ndvi.py index 9c9e89e..42db282 100644 --- a/benchmarks/geospatial/01_ndvi.py +++ b/benchmarks/geospatial/01_ndvi.py @@ -4,139 +4,135 @@ # dependencies = [ # "xarray-sql", # "xarray", +# "pystac-client", # "zarr>=3", # "numpy", # ] # /// -"""NDVI + cloud/nodata masking — "apply_ufunc over a raster" is column algebra. +"""NDVI — "apply_ufunc over a raster" is just column arithmetic. The Normalized Difference Vegetation Index is the workhorse of optical remote -sensing: ``NDVI = (NIR - Red) / (NIR + Red)``, computed per pixel, with invalid -pixels (nodata / clouds) masked out. The array paradigm reaches for -``xarray.apply_ufunc`` (the coiled/benchmarks #1545 "vectorized operations" -case) to broadcast this over a whole scene. +sensing: ``NDVI = (NIR - Red) / (NIR + Red)``, computed per pixel. The array +paradigm reaches for ``xarray.apply_ufunc`` (the coiled/benchmarks #1545 +"vectorized operations" case) to broadcast this over a whole scene. But a per-pixel formula over two bands is just *column arithmetic over two -columns*, and "mask the invalid pixels" is just ``CASE WHEN``:: +columns*:: - SELECT x, y, - CASE WHEN red = 0 OR nir = 0 THEN NULL - ELSE (nir_refl - red_refl) / (nir_refl + red_refl) - END AS ndvi + SELECT x, y, (nir - red) / (nir + red) AS ndvi FROM scene + ORDER BY y, x -Each pixel is one row; the ufunc is the SELECT expression. No broadcasting -rules, no apply_ufunc signature — the SQL *is* the index definition. +Each pixel is one row; the ufunc is the SELECT expression. Invalid pixels are +already NaN (xarray decodes the band's ``_FillValue`` on open), and NaN +propagates through the arithmetic on both sides — so the masking is free, no +``CASE`` required. Dataset: a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample -service (bands B04=red, B08=NIR at 10 m, stored as uint16 with -``reflectance = DN * 0.0001 - 0.1``). We read one cloud-light window so the -case stays bounded. Requires network; skips cleanly if the store is offline. +service, discovered with ``pystac-client`` and opened the canonical way with +``xarray`` — ``xr.open_datatree`` yields the reflectance bands (B04=red, +B08=NIR at 10 m) already scaled to reflectance and carrying their ``x``/``y`` +coordinates. We read one window so the case stays bounded. Requires network; +skips cleanly if the service is offline. """ from __future__ import annotations -import numpy as np import xarray as xr import xarray_sql as xql -from _harness import CaseSkipped, check_close, run_case, show_sql, timed +from _harness import CaseSkipped, run_case, show_sql, timed -# A specific Sentinel-2 L2A product from the EOPF sample service (STAC -# collection ``sentinel-2-l2a``, tile T32TLQ over the Italian Alps). Hard-coded -# so the case needs no STAC client; if the sample service rotates this product -# out, the open below fails and the case skips. -_BASE = ( - "https://objectstore.eodc.eu:2222/e05ab01a9d56408d82ac32d69a5aae2a:" - "202505-s02msil2a/03/products/cpm_v256/" - "S2B_MSIL2A_20250503T102559_N0511_R108_T32TLQ_20250503T132157.zarr" -) -_R10M = _BASE + "/measurements/reflectance/r10m" +# EOPF sample-service STAC catalog; an agricultural AOI near Torino, Italy, in +# early May (peak spring growth). The search is deterministic — it resolves to +# a specific archived Sentinel-2 product. +_STAC = "https://stac.core.eopf.eodc.eu" +_BBOX = [7.2, 44.5, 7.4, 44.7] +_DATETIME = "2025-04-25/2025-05-05" -# proj:transform of the 10 m grid (EPSG:32632), from the STAC asset metadata. -_X0, _Y0, _RES = 300_000.0, 5_000_040.0, 10.0 +# A 1024×1024 (~105 km²) window over vegetated valley floor. +_Y0, _X0, _N = 4_000, 6_000, 1_024 -# Reflectance encoding (STAC raster:scale / raster:offset), DN -> reflectance. -_SCALE, _OFFSET = 0.0001, -0.1 -# A 1024×1024 window (~105 km²) offset into the tile to land on vegetated -# terrain (median NDVI ≈ 0.33 in early May) rather than the nodata border. -_ROW0, _COL0, _N = 2_500, 6_500, 1024 +def _load_scene() -> tuple[xr.Dataset, str]: + """Discover a Sentinel-2 L2A product and open its 10 m red/NIR bands. - -def _load_scene() -> xr.Dataset: - """Read the B04/B08 window from the EOPF Zarr store into an xr.Dataset.""" + Idiomatic end to end: ``pystac-client`` finds the product, ``open_datatree`` + opens the hierarchical EOPF Zarr, and the ``reflectance/r10m`` node already + carries B04/B08 scaled to reflectance (nodata decoded to NaN) with ``x``/``y`` + coordinates — no manual scaling or coordinate reconstruction. + """ try: - import zarr - - b04 = zarr.open_array(_R10M + "/b04", mode="r") - b08 = zarr.open_array(_R10M + "/b08", mode="r") - rows = slice(_ROW0, _ROW0 + _N) - cols = slice(_COL0, _COL0 + _N) - red = np.asarray(b04[rows, cols]) - nir = np.asarray(b08[rows, cols]) + from pystac_client import Client + + catalog = Client.open(_STAC) + search = catalog.search( + collections=["sentinel-2-l2a"], + bbox=_BBOX, + datetime=_DATETIME, + max_items=1, + ) + item = next(search.items()) + tree = xr.open_datatree( + item.assets["product"].href, engine="zarr", chunks={} + ) + except StopIteration as exc: + raise CaseSkipped("no Sentinel-2 product found for the query") from exc except Exception as exc: # noqa: BLE001 — any failure → skip, not crash - raise CaseSkipped(f"EOPF Sentinel-2 Zarr unavailable ({exc})") from exc - - # Pixel-center coordinates from the affine transform. - x = _X0 + (np.arange(_COL0, _COL0 + _N) + 0.5) * _RES - y = _Y0 - (np.arange(_ROW0, _ROW0 + _N) + 0.5) * _RES - return xr.Dataset( - {"red": (["y", "x"], red), "nir": (["y", "x"], nir)}, - coords={"y": y.astype("float64"), "x": x.astype("float64")}, - ).chunk({"y": 256, "x": 256}) + raise CaseSkipped(f"EOPF Sentinel-2 unavailable ({exc})") from exc + + r10m = tree["measurements/reflectance/r10m"].to_dataset() + scene = ( + r10m[["b04", "b08"]] + .rename(b04="red", b08="nir") + .isel(y=slice(_Y0, _Y0 + _N), x=slice(_X0, _X0 + _N)) + .load() + ) + return scene, item.id def main() -> None: - ds = _load_scene() - n = ds.sizes["y"] * ds.sizes["x"] + scene, item_id = _load_scene() + n = scene.sizes["y"] * scene.sizes["x"] + print(f" Sentinel-2 L2A {item_id}") print( - f" Sentinel-2 L2A scene window: {dict(ds.sizes)} " - f"({n:,} pixels, B04=red/B08=NIR, uint16 DN)" + f" scene window: {dict(scene.sizes)} ({n:,} pixels, B04=red/B08=NIR)" ) ctx = xql.XarrayContext() - ctx.from_dataset("scene", ds, chunks={"y": 256, "x": 256}) + ctx.from_dataset("scene", scene, chunks={"y": 256, "x": 256}) - # DN -> reflectance happens inline; CASE masks nodata (DN == 0). sql = """ - SELECT x, y, - CASE - WHEN red = 0 OR nir = 0 THEN NULL - ELSE ( (CAST(nir AS DOUBLE) * 0.0001 - 0.1) - - (CAST(red AS DOUBLE) * 0.0001 - 0.1) ) - / ( (CAST(nir AS DOUBLE) * 0.0001 - 0.1) - + (CAST(red AS DOUBLE) * 0.0001 - 0.1) ) - END AS ndvi + SELECT x, y, (nir - red) / (nir + red) AS ndvi FROM scene + ORDER BY y, x """ show_sql(sql) with timed("SQL NDVI"): got = ctx.sql(sql).to_dataset(dims=["y", "x"]).ndvi - # Array reference: the same formula via xarray broadcasting + .where mask. + # Array reference: the same formula in pure xarray. with timed("xarray reference"): - red = ds.red.astype("float64") * _SCALE + _OFFSET - nir = ds.nir.astype("float64") * _SCALE + _OFFSET - ref = ((nir - red) / (nir + red)).where((ds.red != 0) & (ds.nir != 0)) + ref = (scene.nir - scene.red) / (scene.nir + scene.red) - # Align both to the same (y, x) grid before comparing. - got = got.sortby(["y", "x"]) - ref = ref.sortby(["y", "x"]) - check_close("NDVI (per-pixel)", got, ref, rtol=1e-6, atol=1e-6) + # Compare by coordinate label (reindex_like), so neither side needs an + # explicit sort. NaNs (decoded nodata) in matching cells compare equal. + xr.testing.assert_allclose(got, ref.reindex_like(got), rtol=1e-6) + print( + f" ✅ NDVI (per-pixel): SQL matches xarray reference " + f"(n={got.size:,}, coordinate-aligned)" + ) - valid = np.isfinite(got.values) + valid = ref.notnull() print( - f"\n NDVI over {valid.sum():,} valid pixels: " - f"min {np.nanmin(got.values):.3f}, " - f"mean {np.nanmean(got.values):.3f}, " - f"max {np.nanmax(got.values):.3f}" + f"\n NDVI over {int(valid.sum()):,} valid pixels: " + f"min {float(ref.min()):.3f}, " + f"mean {float(ref.mean()):.3f}, " + f"max {float(ref.max()):.3f}" ) if __name__ == "__main__": - raise SystemExit( - run_case(main, "NDVI: per-pixel column algebra + CASE mask") - ) + raise SystemExit(run_case(main, "NDVI: per-pixel column arithmetic")) diff --git a/benchmarks/geospatial/README.md b/benchmarks/geospatial/README.md index 82b691a..96d0716 100644 --- a/benchmarks/geospatial/README.md +++ b/benchmarks/geospatial/README.md @@ -15,7 +15,7 @@ plain-English definition of the operation, and computes the same numbers. | # | Case | Array mental model | Relational reality | |---|------|--------------------|--------------------| -| 01 | `01_ndvi.py` | `apply_ufunc` over a raster | column algebra + `CASE` mask | +| 01 | `01_ndvi.py` | `apply_ufunc` over a raster | column arithmetic | | 02 | `02_climatology.py` | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | | 03 | `03_zonal_mean.py` | reduce over lon/time axes | `GROUP BY latitude` | | 04 | `04_anomaly.py` | climatology broadcast-subtract | climatology CTE self-`JOIN` | @@ -35,7 +35,8 @@ interpolation weights — the geometry — which SQL applies but does not comput ## Datasets - **01 NDVI** — a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample - service (bands B04/B08). Requires network; skips cleanly if offline. + service, discovered with `pystac-client` and opened with `xr.open_datatree` + (bands B04/B08). Requires network; skips cleanly if offline. - **02–06** — the full **[ARCO-ERA5](https://github.com/google-research/arco-era5)** archive (0.25° global, ~1.3M hourly timesteps, 273 variables) read anonymously from a public GCS bucket. Cases 03 and 06 register the *whole* archive and let diff --git a/docs/geospatial.md b/docs/geospatial.md index 30af680..ab8edbe 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -25,7 +25,7 @@ computes the same numbers — at ERA5's real 0.25° global resolution. | Operation | The "array" framing | The relational reality | Script | |-----------|---------------------|------------------------|--------| -| Spectral index (NDVI) | `apply_ufunc` over a raster | column algebra + `CASE` mask | [`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) | +| Spectral index (NDVI) | `apply_ufunc` over a raster | column arithmetic | [`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) | | Climatology | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | [`02_climatology.py`](../benchmarks/geospatial/02_climatology.py) | | Zonal mean | reduce over lon/time axes | `GROUP BY lat` | [`03_zonal_mean.py`](../benchmarks/geospatial/03_zonal_mean.py) | | Anomaly | grouped broadcast-subtract | climatology CTE self-`JOIN` | [`04_anomaly.py`](../benchmarks/geospatial/04_anomaly.py) | @@ -36,22 +36,24 @@ computes the same numbers — at ERA5's real 0.25° global resolution. ## 1. A pixel-wise formula is a column expression -NDVI is `(NIR − Red) / (NIR + Red)`, per pixel, with clouds/nodata masked out. -The array idiom broadcasts a ufunc over the raster. But "one output per pixel, -computed from that pixel's bands" is the definition of a SQL projection, and -"drop invalid pixels" is `CASE WHEN`: +NDVI is `(NIR − Red) / (NIR + Red)`, per pixel. The array idiom broadcasts a +ufunc over the raster. But "one output per pixel, computed from that pixel's +bands" is the definition of a SQL projection: ```sql -SELECT x, y, - CASE WHEN red = 0 OR nir = 0 THEN NULL - ELSE (nir_refl - red_refl) / (nir_refl + red_refl) - END AS ndvi +SELECT x, y, (nir - red) / (nir + red) AS ndvi FROM scene +ORDER BY y, x ``` +Invalid pixels need no special handling: xarray decodes the band's `_FillValue` +to `NaN` on open, and `NaN` propagates through the arithmetic on both sides, so +the masking is free. + [`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) runs this against a **real -Sentinel-2 L2A scene in Zarr** (ESA's EOPF sample service) and matches xarray's -`apply_ufunc`-style result over a million pixels. +Sentinel-2 L2A scene in Zarr** — discovered with `pystac-client` and opened the +canonical way with `xr.open_datatree` (ESA's EOPF sample service) — and matches +xarray's `apply_ufunc`-style result over a million pixels. ## 2. A climatology is a `GROUP BY` over the cycle From 85354c1d6ddbe1929c0dff2838dfa79c46138cfa Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 23 Jun 2026 22:44:10 +0300 Subject: [PATCH 03/26] Make benchmark cases idiomatic xarray: fluent opens, Dataset results, coordinate-aligned comparison MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses round-2 PR review (cases 02, 03) and applies the same patterns across the suite: - Remove the _era5.py helper module; each case now shows the xr.open_zarr(...).sel(...).load() fluent chain inline (and registration inline), since the repetition documents the canonical xarray usage. - Round-trip gridded SQL results back to xarray Datasets via to_dataset(...) instead of to_pandas() — a climatology/zonal-mean/anomaly/regrid is gridded data and is used as such. - Keep references in pure xarray (no to_dataframe); compare with a new assert_grid_close helper that aligns by coordinate label (xr.testing.assert_allclose), so no manual .sortby() is needed. - Add ORDER BY to the climatology and anomaly queries. - Apply the same Dataset-return + xarray-comparison treatment to the forecast skill (RMSE-by-lead as a DataArray over lead), reprojection (lon/lat grid), and regrid (reshaped to the target grid) cases. - Drop the now-unused check_close harness helper. docs/geospatial.md notes that every query round-trips its answer back to an xarray.Dataset. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/01_ndvi.py | 18 +++--- benchmarks/geospatial/02_climatology.py | 75 +++++++++++----------- benchmarks/geospatial/03_zonal_mean.py | 45 +++++++++---- benchmarks/geospatial/04_anomaly.py | 65 ++++++++++--------- benchmarks/geospatial/05_forecast_skill.py | 50 ++++++++------- benchmarks/geospatial/06_zonal_vector.py | 72 +++++++++++++-------- benchmarks/geospatial/07_reproject_udf.py | 30 +++++---- benchmarks/geospatial/08_regrid_weights.py | 20 +++--- benchmarks/geospatial/_era5.py | 70 -------------------- benchmarks/geospatial/_harness.py | 48 +++++++------- docs/geospatial.md | 5 +- 11 files changed, 241 insertions(+), 257 deletions(-) delete mode 100644 benchmarks/geospatial/_era5.py diff --git a/benchmarks/geospatial/01_ndvi.py b/benchmarks/geospatial/01_ndvi.py index 42db282..9f42497 100644 --- a/benchmarks/geospatial/01_ndvi.py +++ b/benchmarks/geospatial/01_ndvi.py @@ -42,7 +42,13 @@ import xarray_sql as xql -from _harness import CaseSkipped, run_case, show_sql, timed +from _harness import ( + CaseSkipped, + assert_grid_close, + run_case, + show_sql, + timed, +) # EOPF sample-service STAC catalog; an agricultural AOI near Torino, Italy, in # early May (peak spring growth). The search is deterministic — it resolves to @@ -117,13 +123,9 @@ def main() -> None: with timed("xarray reference"): ref = (scene.nir - scene.red) / (scene.nir + scene.red) - # Compare by coordinate label (reindex_like), so neither side needs an - # explicit sort. NaNs (decoded nodata) in matching cells compare equal. - xr.testing.assert_allclose(got, ref.reindex_like(got), rtol=1e-6) - print( - f" ✅ NDVI (per-pixel): SQL matches xarray reference " - f"(n={got.size:,}, coordinate-aligned)" - ) + # Compare the xarray way — aligned by coordinate label, so the ORDER BY + # above is enough and neither side needs an explicit sort. + assert_grid_close("NDVI (per-pixel)", got, ref, rtol=1e-6) valid = ref.notnull() print( diff --git a/benchmarks/geospatial/02_climatology.py b/benchmarks/geospatial/02_climatology.py index e294b74..002757b 100644 --- a/benchmarks/geospatial/02_climatology.py +++ b/benchmarks/geospatial/02_climatology.py @@ -3,8 +3,6 @@ # dependencies = [ # "xarray-sql", # "xarray", -# "numpy", -# "pandas", # "gcsfs", # "zarr>=3", # ] @@ -12,32 +10,34 @@ """Diurnal climatology — the "rechunk + grouped reduction" that is a GROUP BY. A *climatology* is the average value for each time-of-cycle, computed -independently at every location: "what is the typical temperature here at 06:00 -local cycle?" In the array paradigm (and in the coiled/benchmarks #1545 -write-up) this is the canonical painful workload — load native Zarr chunks, -*rechunk* to put all of time in one chunk ("pencils"), run a grouped reduction -over the calendar, then rechunk back to "pancakes" for output. +independently at every location: "what is the typical temperature here at +06:00?" In the array paradigm (and in the coiled/benchmarks #1545 write-up) +this is the canonical painful workload — load native Zarr chunks, *rechunk* to +put all of time in one chunk ("pencils"), run a grouped reduction over the +calendar, then rechunk back to "pancakes" for output. The rechunking exists only to serve the array layout. The *operation* is:: SELECT latitude, longitude, hour_of_day, AVG("2m_temperature") GROUP BY latitude, longitude, hour_of_day -Group by location and time-of-cycle, average the rest. Same answer as -``da.groupby("time.hour").mean()`` — here the **diurnal cycle** (mean -temperature by hour of day) over a region, at ERA5's 0.25° resolution. +Group by location and time-of-cycle, average the rest — the same answer as +``da.groupby("time.hour").mean()``. ERA5 is hourly, so grouping by hour of day +gives a clean 24-bin **diurnal cycle**, one sample per day in the window. Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days, -read once into memory (see ``_era5.load_window``). +opened and sliced with a single fluent xarray chain. """ from __future__ import annotations +import xarray as xr + import xarray_sql as xql -from _era5 import load_window, open_era5 -from _harness import check_close, run_case, show_sql, timed +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" # A few days over a CONUS-ish box (ERA5 latitude descends; lon is 0–360°E). _TIME = slice("2020-06-01", "2020-06-03T23") _LAT = slice(50.0, 25.0) @@ -45,11 +45,22 @@ def main() -> None: - full = open_era5() - with timed("read ERA5 window into memory"): - ds = load_window( - full, "2m_temperature", time=_TIME, latitude=_LAT, longitude=_LON - ) + # Open the full ARCO-ERA5 archive and slice to the window in one chain: + # pick the variable, select the box and days, and read it into memory. + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + with timed("open + slice ERA5 window"): + ds = ( + xr.open_zarr( + _URL, chunks=None, storage_options={"token": "anon"} + )[["2m_temperature"]] + .sel(time=_TIME, latitude=_LAT, longitude=_LON) + .load() + ) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + print( f" ERA5 2m_temperature window: {dict(ds.sizes)} " f"(diurnal climatology over {ds.sizes['latitude']}×{ds.sizes['longitude']} cells)" @@ -65,34 +76,24 @@ def main() -> None: AVG("2m_temperature") - 273.15 AS clim_c FROM era5 GROUP BY latitude, longitude, date_part('hour', time) + ORDER BY latitude DESC, longitude, hour """ show_sql(sql) + # A climatology is a gridded product, so round-trip the result back to an + # xarray Dataset keyed by (latitude, longitude, hour) — how it is used. with timed("SQL diurnal climatology"): - got = ctx.sql(sql).to_pandas() + got = ctx.sql(sql).to_dataset(dims=["latitude", "longitude", "hour"]) - # Array reference: the textbook groupby-over-the-cycle reduction. + # Array reference: the textbook groupby-over-the-cycle reduction, in °C. with timed("xarray reference"): - clim = ds["2m_temperature"].groupby("time.hour").mean("time") - 273.15 - ref = clim.to_dataframe(name="clim_c").reset_index() + ref = ds["2m_temperature"].groupby("time.hour").mean("time") - 273.15 - merged = got.merge( - ref, on=["latitude", "longitude", "hour"], validate="one_to_one" - ) - assert len(merged) == len(got) == len(ref), ( - f"row-count mismatch: sql={len(got)} ref={len(ref)} merged={len(merged)}" - ) - check_close( - "diurnal climatology (°C)", - merged["clim_c_x"], - merged["clim_c_y"], - rtol=1e-4, - atol=1e-3, + assert_grid_close( + "diurnal climatology (°C)", got.clim_c, ref, rtol=1e-4, atol=1e-2 ) - print( - f"\n {len(got):,} climatology cells (latitude × longitude × hour-of-day)." - ) + print(f"\n climatology Dataset: {dict(got.sizes)}") if __name__ == "__main__": diff --git a/benchmarks/geospatial/03_zonal_mean.py b/benchmarks/geospatial/03_zonal_mean.py index 06c95b4..a321d12 100644 --- a/benchmarks/geospatial/03_zonal_mean.py +++ b/benchmarks/geospatial/03_zonal_mean.py @@ -3,7 +3,6 @@ # dependencies = [ # "xarray-sql", # "xarray", -# "numpy", # "gcsfs", # "zarr>=3", # ] @@ -30,27 +29,46 @@ from __future__ import annotations +import xarray as xr import xarray_sql as xql -from _era5 import open_era5, register_era5 -from _harness import check_close, run_case, show_sql, timed +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" # One day of hourly data, global; the WHERE below prunes ERA5 to this window. _DAY = "2020-06-01" def main() -> None: - ds = open_era5() + # Open the full ARCO-ERA5 archive (lazy, dask off) — no slicing here; the + # SQL WHERE clause prunes it to the window we ask for. + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + print( f" ARCO-ERA5: {ds.sizes['time']:,} hourly timesteps, " f"{ds.sizes['latitude']}×{ds.sizes['longitude']} grid, " f"{len(ds.data_vars)} variables (no pre-slicing)" ) + # ERA5 mixes surface (time, lat, lon) and atmospheric (… level …) variables, + # so register it as two tables under an ``era5`` schema. ctx = xql.XarrayContext() with timed("register full ERA5"): - register_era5(ctx, ds) + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) sql = f""" SELECT latitude, @@ -63,26 +81,27 @@ def main() -> None: """ show_sql(sql) + # Round-trip the profile back to an xarray Dataset keyed by latitude. with timed("SQL zonal mean (WHERE-pruned to one day)"): - got = ctx.sql(sql).to_pandas() + got = ctx.sql(sql).to_dataset(dims=["latitude"]) # Array reference: reduce the same day over the two un-grouped axes. with timed("xarray reference"): - window = ds["2m_temperature"].sel(time=_DAY) - ref = (window.mean(dim=["longitude", "time"]) - 273.15).sortby( - "latitude", ascending=False + ref = ( + ds["2m_temperature"].sel(time=_DAY).mean(["longitude", "time"]) + - 273.15 ) - check_close( + assert_grid_close( "zonal mean (2m_temp vs latitude, °C)", - got["air_mean_c"], + got.air_mean_c, ref, rtol=1e-4, atol=1e-3, ) - print("\n Global temperature profile (every 10th parallel):") - print(got.iloc[::72].to_string(index=False)) + print("\n Global temperature profile (every 72nd parallel, °C):") + print(got.air_mean_c.isel(latitude=slice(None, None, 72)).to_series()) if __name__ == "__main__": diff --git a/benchmarks/geospatial/04_anomaly.py b/benchmarks/geospatial/04_anomaly.py index 96df0df..6824daa 100644 --- a/benchmarks/geospatial/04_anomaly.py +++ b/benchmarks/geospatial/04_anomaly.py @@ -3,8 +3,6 @@ # dependencies = [ # "xarray-sql", # "xarray", -# "numpy", -# "pandas", # "gcsfs", # "zarr>=3", # ] @@ -27,30 +25,41 @@ ON (a.latitude, a.longitude, hour(a.time)) = (c.latitude, c.longitude, c.hour) Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days, -read once into memory (anomalies vs the diurnal cycle of that window). +opened and sliced with a single fluent xarray chain (anomalies vs the diurnal +cycle of that window). """ from __future__ import annotations +import xarray as xr + import xarray_sql as xql -from _era5 import load_window, open_era5 -from _harness import check_close, run_case, show_sql, timed +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" _TIME = slice("2020-06-01", "2020-06-03T23") _LAT = slice(50.0, 25.0) _LON = slice(235.0, 290.0) def main() -> None: - full = open_era5() - with timed("read ERA5 window into memory"): - ds = load_window( - full, "2m_temperature", time=_TIME, latitude=_LAT, longitude=_LON - ) + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + with timed("open + slice ERA5 window"): + ds = ( + xr.open_zarr( + _URL, chunks=None, storage_options={"token": "anon"} + )[["2m_temperature"]] + .sel(time=_TIME, latitude=_LAT, longitude=_LON) + .load() + ) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + print( - f" ERA5 2m_temperature window: {dict(ds.sizes)} " - f"(anomaly vs diurnal normals)" + f" ERA5 2m_temperature window: {dict(ds.sizes)} (anomaly vs diurnal normals)" ) ctx = xql.XarrayContext() @@ -71,36 +80,28 @@ def main() -> None: ON a.latitude = c.latitude AND a.longitude = c.longitude AND date_part('hour', a.time) = c.hour + ORDER BY a.time, a.latitude DESC, a.longitude """ show_sql(sql) + # The anomaly is a gridded field; round-trip it to (time, lat, lon). with timed("SQL anomaly (climatology CTE self-join)"): - got = ctx.sql(sql).to_pandas() + got = ctx.sql(sql).to_dataset(dims=["time", "latitude", "longitude"]) - # Array reference: grouped broadcast-subtract. + # Array reference: grouped broadcast-subtract, in pure xarray. with timed("xarray reference"): - clim = ds["2m_temperature"].groupby("time.hour").mean("time") - anom = ds["2m_temperature"].groupby("time.hour") - clim - ref = anom.to_dataframe(name="anomaly").reset_index() + grouped = ds["2m_temperature"].groupby("time.hour") + ref = grouped - grouped.mean("time") - merged = got.merge( - ref, - on=["time", "latitude", "longitude"], - suffixes=("_sql", "_ref"), - validate="one_to_one", - ) - assert len(merged) == len(got) == len(ref), ( - f"row-count mismatch: sql={len(got)} ref={len(ref)} merged={len(merged)}" - ) - check_close( + assert_grid_close( "anomaly (T − diurnal climatology)", - merged["anomaly_sql"], - merged["anomaly_ref"], - rtol=1e-4, - atol=1e-3, + got.anomaly, + ref, + rtol=1e-3, + atol=1e-2, ) - print(f"\n {len(got):,} hourly anomalies over the window.") + print(f"\n anomaly Dataset: {dict(got.sizes)}") if __name__ == "__main__": diff --git a/benchmarks/geospatial/05_forecast_skill.py b/benchmarks/geospatial/05_forecast_skill.py index 4d14d00..2a007fb 100644 --- a/benchmarks/geospatial/05_forecast_skill.py +++ b/benchmarks/geospatial/05_forecast_skill.py @@ -32,8 +32,8 @@ GROUP BY f.prediction_timedelta The whole evaluation — temporal alignment, spatial matching, the score — is one -JOIN and one aggregate. We run it for both models and check each against a NumPy -reference. +JOIN and one aggregate. We run it for both models and check each against an +xarray reference. Datasets: WeatherBench 2 **Pangu**, **GraphCast**, and **ERA5** at a coarse 64×32 grid (so the demo is small and fast), read from the public ``gs:// @@ -48,7 +48,13 @@ import xarray_sql as xql -from _harness import CaseSkipped, check_close, run_case, show_sql, timed +from _harness import ( + CaseSkipped, + assert_grid_close, + run_case, + show_sql, + timed, +) _SO = {"token": "anon"} _GRID = "64x32_equiangular_conservative" @@ -104,21 +110,23 @@ def _rmse_sql(table: str) -> str: """ -def _rmse_by_lead(ctx: xql.XarrayContext, table: str) -> pd.DataFrame: - return ctx.sql(_rmse_sql(table)).to_pandas() +def _rmse_by_lead(ctx: xql.XarrayContext, table: str) -> xr.DataArray: + """Run the RMSE query, returning the skill curve as a DataArray over lead.""" + return ctx.sql(_rmse_sql(table)).to_dataset(dims=["lead"]).rmse -def _reference_rmse(fc: xr.Dataset, truth: xr.Dataset) -> np.ndarray: - """NumPy reference: for each lead, align truth at valid_time and take RMSE.""" +def _reference_rmse(fc: xr.Dataset, truth: xr.Dataset) -> xr.DataArray: + """xarray reference: for each lead, align truth at valid_time and take RMSE.""" f = fc[_VAR] e = truth[_VAR] + leads = f.prediction_timedelta out = [] - for lead in f.prediction_timedelta.values: - valid = f.time.values + lead - e_at_valid = e.sel(time=valid).values # (init, lat, lon) - diff = f.sel(prediction_timedelta=lead).values - e_at_valid - out.append(float(np.sqrt(np.mean(diff**2)))) - return np.array(out) + for lead in leads.values: + valid = f.time.values + lead # valid_time = init + lead + e_at_valid = e.sel(time=valid) + diff = f.sel(prediction_timedelta=lead) - e_at_valid.values + out.append(float(np.sqrt((diff**2).mean()))) + return xr.DataArray(out, coords={"lead": leads.values}, dims=["lead"]) def main() -> None: @@ -158,22 +166,20 @@ def main() -> None: for name, fc in [("pangu", pangu), ("graphcast", graphcast)]: with timed(f"SQL RMSE-by-lead: {name}"): got = _rmse_by_lead(ctx, name) - with timed(f"NumPy reference: {name}"): + with timed(f"xarray reference: {name}"): ref = _reference_rmse(fc, truth) - check_close( - f"{name} RMSE(lead)", got["rmse"], ref, rtol=1e-4, atol=1e-3 - ) + assert_grid_close(f"{name} RMSE(lead)", got, ref, rtol=1e-4, atol=1e-3) results[name] = got # Headline table: error growth with forecast horizon, both models. - leads = results["pangu"]["lead"].dt.total_seconds() / 86400.0 + lead_days = results["pangu"]["lead"].values / np.timedelta64(1, "D") print("\n 2m-temperature RMSE (K) vs lead time — lower is better:") print(f" {'lead (days)':>12} {'Pangu':>9} {'GraphCast':>11}") - for i in range(0, len(leads), 4): + for i in range(0, len(lead_days), 4): print( - f" {leads.iloc[i]:>12.2f} " - f"{results['pangu']['rmse'].iloc[i]:>9.3f} " - f"{results['graphcast']['rmse'].iloc[i]:>11.3f}" + f" {lead_days[i]:>12.2f} " + f"{float(results['pangu'][i]):>9.3f} " + f"{float(results['graphcast'][i]):>11.3f}" ) diff --git a/benchmarks/geospatial/06_zonal_vector.py b/benchmarks/geospatial/06_zonal_vector.py index 444e154..e517557 100644 --- a/benchmarks/geospatial/06_zonal_vector.py +++ b/benchmarks/geospatial/06_zonal_vector.py @@ -39,9 +39,9 @@ import xarray_sql as xql -from _era5 import load_window, open_era5, register_era5 -from _harness import check_close, run_case, show_sql, timed +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" _DAY = "2020-06-01" # Continental-scale boxes (name, lat_min, lat_max, lon_min, lon_max), lon 0–360°E. @@ -69,16 +69,30 @@ def _regions_dataset() -> xr.Dataset: def main() -> None: - full = open_era5() + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + print( - f" raster: full ARCO-ERA5 ({full.sizes['time']:,} timesteps, " - f"{full.sizes['latitude']}×{full.sizes['longitude']}) " + f" raster: full ARCO-ERA5 ({ds.sizes['time']:,} timesteps, " + f"{ds.sizes['latitude']}×{ds.sizes['longitude']}) " f"vector: {len(_REGIONS)} continental boxes" ) ctx = xql.XarrayContext() with timed("register full ERA5 + regions"): - register_era5(ctx, full) + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) ctx.from_dataset( "regions", _regions_dataset(), chunks={"region": len(_REGIONS)} ) @@ -99,34 +113,40 @@ def main() -> None: show_sql(sql) with timed("SQL zonal stats (raster × vector range JOIN, WHERE-pruned)"): - got = ( - ctx.sql(sql) - .to_pandas() - .sort_values("region_id") - .reset_index(drop=True) - ) + got = ctx.sql(sql).to_dataset(dims=["region_id"]) # Array reference: load the same day once, mask each region in memory. with timed("xarray reference"): - day = load_window(full, "2m_temperature", time=_DAY)["2m_temperature"] - ref_vals = [] - for _, lat_min, lat_max, lon_min, lon_max in _REGIONS: - mask = ( - (day.latitude >= lat_min) - & (day.latitude <= lat_max) - & (day.longitude >= lon_min) - & (day.longitude <= lon_max) + day = ( + xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})[ + "2m_temperature" + ] + .sel(time=_DAY) + .load() + ) + avg_c = [ + float( + day.where( + (day.latitude >= lat_min) + & (day.latitude <= lat_max) + & (day.longitude >= lon_min) + & (day.longitude <= lon_max) + ).mean() ) - ref_vals.append(float(day.where(mask).mean()) - 273.15) - ref = np.array(ref_vals) + - 273.15 + for _, lat_min, lat_max, lon_min, lon_max in _REGIONS + ] + ref = xr.DataArray( + avg_c, dims=["region_id"], coords={"region_id": got.region_id} + ) - check_close( - "zonal mean per region (°C)", got["avg_c"], ref, rtol=1e-4, atol=1e-2 + assert_grid_close( + "zonal mean per region (°C)", got.avg_c, ref, rtol=1e-4, atol=1e-2 ) print("\n Region avg °C n_obs") - for (name, *_), row in zip(_REGIONS, got.itertuples()): - print(f" {name:<20} {row.avg_c:7.2f} {row.n_obs:>10,}") + for (name, *_), avg, n in zip(_REGIONS, got.avg_c.values, got.n_obs.values): + print(f" {name:<20} {avg:7.2f} {int(n):>10,}") if __name__ == "__main__": diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py index 5a8ccd4..543f30c 100644 --- a/benchmarks/geospatial/07_reproject_udf.py +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -19,12 +19,12 @@ So we register a PROJ-backed scalar UDF and reproject in SQL:: - SELECT x, y, utm_to_lon(x, y) AS lon, utm_to_lat(x, y) AS lat + SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat FROM grid The UDF (built on ``pyproj``, vectorized over each Arrow batch) mirrors the ``cftime()`` UDF already shipped in xarray-sql (see ``xarray_sql/cftime.py``); -it could graduate into the package as ``xql.register_reproject_udfs``. +it could graduate into the package as ``xql.register_reproject_udf``. **What this does *not* do:** it moves the coordinates, it does not resample onto a regular target grid. Producing a gridded product in the new CRS still needs @@ -54,7 +54,7 @@ import xarray_sql as xql -from _harness import check_close, run_case, show_sql, timed +from _harness import assert_grid_close, run_case, show_sql, timed def register_reproject_udf( @@ -133,8 +133,9 @@ def main() -> None: """ show_sql(sql) + # Round-trip the reprojected coordinates back to an (y, x) grid. with timed("SQL reprojection (PROJ scalar UDF)"): - got = ctx.sql(sql).to_pandas() + got = ctx.sql(sql).to_dataset(dims=["y", "x"]) # Array reference: the same PROJ transform applied to the full grid. with timed("pyproj reference"): @@ -142,21 +143,22 @@ def main() -> None: src_crs, dst_crs, always_xy=True ) xx, yy = np.meshgrid(ds.x.values, ds.y.values) - ref_lon, ref_lat = transformer.transform(xx.ravel(), yy.ravel()) + lon, lat = transformer.transform(xx, yy) + coords = {"y": ds.y, "x": ds.x} + ref_lon = xr.DataArray(lon, dims=["y", "x"], coords=coords) + ref_lat = xr.DataArray(lat, dims=["y", "x"], coords=coords) - got = got.sort_values(["y", "x"]).reset_index(drop=True) - # The reference is built in (y, x) row-major order; align by sorting both. - order = np.lexsort((xx.ravel(), yy.ravel())) - check_close( - "reprojected longitude", got["lon"], ref_lon[order], rtol=0, atol=1e-9 + assert_grid_close( + "reprojected longitude", got.lon, ref_lon, rtol=0, atol=1e-9 ) - check_close( - "reprojected latitude", got["lat"], ref_lat[order], rtol=0, atol=1e-9 + assert_grid_close( + "reprojected latitude", got.lat, ref_lat, rtol=0, atol=1e-9 ) + corner = got.isel(x=0, y=0) print( - f"\n Corner check: UTM ({got.x.iloc[0]:.0f}, {got.y.iloc[0]:.0f}) → " - f"lon {got.lon.iloc[0]:.4f}, lat {got.lat.iloc[0]:.4f}" + f"\n Corner check: UTM ({float(corner.x):.0f}, {float(corner.y):.0f}) → " + f"lon {float(corner.lon):.4f}, lat {float(corner.lat):.4f}" ) diff --git a/benchmarks/geospatial/08_regrid_weights.py b/benchmarks/geospatial/08_regrid_weights.py index 4f94c40..1fa4b34 100644 --- a/benchmarks/geospatial/08_regrid_weights.py +++ b/benchmarks/geospatial/08_regrid_weights.py @@ -42,7 +42,7 @@ import xarray_sql as xql -from _harness import check_close, run_case, show_sql, timed +from _harness import assert_grid_close, run_case, show_sql, timed def _linear_weights( @@ -134,23 +134,25 @@ def main() -> None: """ show_sql(sql) + # The result is keyed by dst_id (row-major over the target grid); reshape + # it back to the (lat, lon) field it represents. with timed("SQL regrid (weight-table JOIN + weighted SUM)"): - got = ( - ctx.sql(sql) - .to_pandas() - .sort_values("dst_id") - .reset_index(drop=True) + flat = ctx.sql(sql).to_dataset(dims=["dst_id"]).regridded + got = xr.DataArray( + flat.values.reshape(len(tlat), len(tlon)), + dims=["lat", "lon"], + coords={"lat": tlat, "lon": tlon}, ) # Array reference: xarray's own bilinear interpolation onto the target grid. with timed("xarray .interp reference"): ref = src_da.interp(lat=tlat, lon=tlon, method="linear") - check_close("bilinear regrid", got["regridded"], ref, rtol=1e-9, atol=1e-9) + assert_grid_close("bilinear regrid", got, ref, rtol=1e-9, atol=1e-9) print( - f"\n {len(got):,} target cells regridded; " - f"value range [{got.regridded.min():.3f}, {got.regridded.max():.3f}]." + f"\n {got.size:,} target cells regridded; " + f"value range [{float(got.min()):.3f}, {float(got.max()):.3f}]." ) diff --git a/benchmarks/geospatial/_era5.py b/benchmarks/geospatial/_era5.py deleted file mode 100644 index af73a0f..0000000 --- a/benchmarks/geospatial/_era5.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Shared ARCO-ERA5 access for the geospatial benchmarks. - -[ARCO-ERA5](https://github.com/google-research/arco-era5) is the ECMWF ERA5 -reanalysis re-published as analysis-ready, cloud-optimized Zarr on a public GCS -bucket: 273 variables, hourly, 0.25° global (721×1440), since 1940 — about 1.3 -million timesteps. We open the *whole* archive (no time/space slicing on the -xarray side) and let SQL ``WHERE`` clauses prune it. That is the demo: the -table is the full reanalysis; the query reads only the window it asks for. - -These cases require anonymous GCS access (``gcsfs``); they skip cleanly when -the network or bucket is unavailable. -""" - -from __future__ import annotations - -import xarray as xr - -import xarray_sql as xql - -from _harness import CaseSkipped - -URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" - -#: Friendly names for ERA5's two dimension groups (surface vs. atmosphere). -TABLE_NAMES = { - ("time", "latitude", "longitude"): "surface", - ("time", "level", "latitude", "longitude"): "atmosphere", -} - - -def open_era5() -> xr.Dataset: - """Open the full ARCO-ERA5 archive (lazy, dask off), or skip if unreachable.""" - try: - import gcsfs # noqa: F401 — required by the gs:// protocol - - return xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) - except Exception as exc: # noqa: BLE001 — any failure → skip, not crash - raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc - - -def register_era5(ctx: xql.XarrayContext, ds: xr.Dataset, *, chunks=None): - """Register ERA5 as ``era5.surface`` / ``era5.atmosphere`` tables.""" - ctx.from_dataset( - "era5", ds, chunks=chunks or {"time": 6}, table_names=TABLE_NAMES - ) - return ctx - - -def load_window( - ds: xr.Dataset, - var: str, - *, - time, - latitude=None, - longitude=None, -) -> xr.Dataset: - """Read a bounded ERA5 window for ``var`` into memory as a 1-variable Dataset. - - Multi-timestep cases (climatology, anomaly, forecast skill) read their window - *once* into memory here, then run both the SQL and its array reference - against the same in-memory data — keeping I/O bounded and the comparison - apples-to-apples. ERA5 latitude is descending, so pass - ``latitude=slice(north, south)``; longitude is 0–360°E ascending. - """ - da = ds[var].sel(time=time) - if latitude is not None: - da = da.sel(latitude=latitude) - if longitude is not None: - da = da.sel(longitude=longitude) - return da.to_dataset().load() diff --git a/benchmarks/geospatial/_harness.py b/benchmarks/geospatial/_harness.py index f1e320e..44149ee 100644 --- a/benchmarks/geospatial/_harness.py +++ b/benchmarks/geospatial/_harness.py @@ -10,9 +10,10 @@ * :func:`banner` / :func:`show_sql` — readable section headers and SQL echo. * :func:`timed` — a context manager that reports elapsed time and peak memory. -* :func:`check_close` — assert a SQL result matches an array reference, then - print a PASS line. Raises ``AssertionError`` on mismatch (so a broken case - fails loudly rather than silently "passing"). +* :func:`assert_grid_close` — assert a SQL result (round-tripped to an + ``xr.DataArray``) matches an xarray reference, aligned by coordinate label. + Raises ``AssertionError`` on mismatch (so a broken case fails loudly rather + than silently "passing"). * :func:`run_case` — run a case's ``main()``, turning a raised :class:`CaseSkipped` (e.g. an offline dataset) into a clean skip. """ @@ -24,9 +25,8 @@ import time import tracemalloc from collections.abc import Callable, Iterator -from typing import Any -import numpy as np +import xarray as xr _WIDTH = 72 @@ -70,34 +70,32 @@ def timed(label: str) -> Iterator[None]: print(f" ⏱ {label}: {elapsed:.3f}s (peak {peak / 1e6:.1f} MB)") -def check_close( +def assert_grid_close( name: str, - got: Any, - expected: Any, + got: xr.DataArray, + ref: xr.DataArray, *, rtol: float = 1e-5, atol: float = 1e-6, - equal_nan: bool = True, ) -> None: - """Assert ``got`` matches the array reference ``expected``, then print PASS. + """Assert two gridded ``DataArray`` results match, then print PASS. - Both arguments are coerced to ``float64`` numpy arrays and flattened, so a - SQL result column (pandas Series / pyarrow) can be compared directly to an - xarray ``DataArray`` reference. Order-sensitive: sort both sides on a shared - key before calling if the SQL result order is not guaranteed. + For cases whose SQL result is round-tripped back to an ``xr.DataArray`` + (via ``XarrayDataFrame.to_dataset``), compare it to the array reference the + xarray way: align ``ref`` onto ``got``'s coordinates and dimension order, + then ``xr.testing.assert_allclose``. This aligns by *label*, so neither side + needs an explicit sort, and NaNs in matching cells compare equal. + + Helper coordinates xarray attaches along the way (e.g. the ``hour`` label a + ``groupby("time.hour")`` leaves behind) are dropped before comparing. """ - g = np.asarray(getattr(got, "values", got), dtype=np.float64).ravel() - e = np.asarray( - getattr(expected, "values", expected), dtype=np.float64 - ).ravel() - if g.shape != e.shape: - raise AssertionError( - f"{name}: shape mismatch — SQL {g.shape} vs reference {e.shape}" - ) - np.testing.assert_allclose(g, e, rtol=rtol, atol=atol, equal_nan=equal_nan) + aligned = ref.reindex_like(got).transpose(*got.dims) + extra = [c for c in aligned.coords if c not in got.coords] + aligned = aligned.drop_vars(extra) + xr.testing.assert_allclose(got, aligned, rtol=rtol, atol=atol) print( - f" ✅ {name}: SQL matches array reference " - f"(n={g.size}, rtol={rtol:g}, atol={atol:g})" + f" ✅ {name}: SQL matches xarray reference " + f"(n={got.size:,}, coordinate-aligned)" ) diff --git a/docs/geospatial.md b/docs/geospatial.md index ab8edbe..1c42f27 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -12,7 +12,10 @@ operations. But it is not the only one, and for a large and growing audience — the people fluent in SQL rather than in `apply_ufunc` and rechunking — it is not the most accessible one. [`xarray-sql`](../README.md) lets you pose these questions in SQL and answers them with a real query engine (DataFusion), -complete with partition pruning and projection pushdown. +complete with partition pruning and projection pushdown. And because a gridded +result is still gridded data, every query here round-trips its answer straight +back to an `xarray.Dataset` (via `to_dataset`) — SQL in, an array out, ready to +plot or save. This page makes the argument case by case. Every claim below is backed by a runnable script in [`benchmarks/geospatial/`](../benchmarks/geospatial/) that From 44a62fc0993396db17eab4645e5f3664af98c1ea Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 24 Jun 2026 08:28:22 +0300 Subject: [PATCH 04/26] 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 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/02_climatology.py | 77 ++++++---- benchmarks/geospatial/03_zonal_mean.py | 17 ++- benchmarks/geospatial/04_anomaly.py | 73 +++++---- benchmarks/geospatial/05_forecast_skill.py | 164 ++++++++++----------- benchmarks/geospatial/06_zonal_vector.py | 15 +- benchmarks/geospatial/README.md | 12 +- docs/geospatial.md | 25 ++-- 7 files changed, 220 insertions(+), 163 deletions(-) diff --git a/benchmarks/geospatial/02_climatology.py b/benchmarks/geospatial/02_climatology.py index 002757b..c3d722d 100644 --- a/benchmarks/geospatial/02_climatology.py +++ b/benchmarks/geospatial/02_climatology.py @@ -25,12 +25,17 @@ ``da.groupby("time.hour").mean()``. ERA5 is hourly, so grouping by hour of day gives a clean 24-bin **diurnal cycle**, one sample per day in the window. -Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days, -opened and sliced with a single fluent xarray chain. +The table is the *whole* lazily-opened ARCO-ERA5 archive — nothing is loaded or +pre-selected up front. The query reads only the ``2m_temperature`` column +(projection pushdown) for only the window the parameterized ``WHERE`` asks for +(partition pruning). The bounds are passed as **query parameters**, not +formatted into the SQL string. """ from __future__ import annotations +import datetime + import xarray as xr import xarray_sql as xql @@ -39,55 +44,73 @@ _URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" # A few days over a CONUS-ish box (ERA5 latitude descends; lon is 0–360°E). -_TIME = slice("2020-06-01", "2020-06-03T23") -_LAT = slice(50.0, 25.0) -_LON = slice(235.0, 290.0) +_START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23) +_LAT_N, _LAT_S = 50.0, 25.0 +_LON_W, _LON_E = 235.0, 290.0 +_PARAMS = { + "start": _START, + "end": _END, + "lat_s": _LAT_S, + "lat_n": _LAT_N, + "lon_w": _LON_W, + "lon_e": _LON_E, +} def main() -> None: - # Open the full ARCO-ERA5 archive and slice to the window in one chain: - # pick the variable, select the box and days, and read it into memory. + # Open the full ARCO-ERA5 archive lazily — dask off, nothing loaded or + # column-selected here. ERA5 mixes surface (time, lat, lon) and atmospheric + # (… level …) variables, so register it as two tables under an ``era5`` + # schema; the query below touches only the surface table's 2m_temperature. try: import gcsfs # noqa: F401 — required by the gs:// protocol - with timed("open + slice ERA5 window"): - ds = ( - xr.open_zarr( - _URL, chunks=None, storage_options={"token": "anon"} - )[["2m_temperature"]] - .sel(time=_TIME, latitude=_LAT, longitude=_LON) - .load() - ) + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) except Exception as exc: # noqa: BLE001 — any failure → skip, not crash raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc - print( - f" ERA5 2m_temperature window: {dict(ds.sizes)} " - f"(diurnal climatology over {ds.sizes['latitude']}×{ds.sizes['longitude']} cells)" - ) - ctx = xql.XarrayContext() - ctx.from_dataset("era5", ds, chunks={"time": 24}) + with timed("register full ERA5 (lazy)"): + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) sql = """ SELECT latitude, longitude, date_part('hour', time) AS hour, AVG("2m_temperature") - 273.15 AS clim_c - FROM era5 + FROM era5.surface + WHERE time BETWEEN $start AND $end + AND latitude BETWEEN $lat_s AND $lat_n + AND longitude BETWEEN $lon_w AND $lon_e GROUP BY latitude, longitude, date_part('hour', time) ORDER BY latitude DESC, longitude, hour """ show_sql(sql) - # A climatology is a gridded product, so round-trip the result back to an + # A climatology is a gridded product: round-trip the result back to an # xarray Dataset keyed by (latitude, longitude, hour) — how it is used. - with timed("SQL diurnal climatology"): - got = ctx.sql(sql).to_dataset(dims=["latitude", "longitude", "hour"]) + with timed("SQL diurnal climatology (lazy read, pushdown + pruning)"): + got = ctx.sql(sql, param_values=_PARAMS).to_dataset( + dims=["latitude", "longitude", "hour"] + ) - # Array reference: the textbook groupby-over-the-cycle reduction, in °C. + # Array reference: the textbook groupby-over-the-cycle reduction, in °C — + # the same lazy window, materialized only on demand. with timed("xarray reference"): - ref = ds["2m_temperature"].groupby("time.hour").mean("time") - 273.15 + window = ds["2m_temperature"].sel( + time=slice(_START, _END), + latitude=slice(_LAT_N, _LAT_S), + longitude=slice(_LON_W, _LON_E), + ) + ref = window.groupby("time.hour").mean("time") - 273.15 assert_grid_close( "diurnal climatology (°C)", got.clim_c, ref, rtol=1e-4, atol=1e-2 diff --git a/benchmarks/geospatial/03_zonal_mean.py b/benchmarks/geospatial/03_zonal_mean.py index a321d12..6924b34 100644 --- a/benchmarks/geospatial/03_zonal_mean.py +++ b/benchmarks/geospatial/03_zonal_mean.py @@ -29,6 +29,8 @@ from __future__ import annotations +import datetime + import xarray as xr import xarray_sql as xql @@ -38,6 +40,10 @@ _URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" # One day of hourly data, global; the WHERE below prunes ERA5 to this window. _DAY = "2020-06-01" +_START, _END = ( + datetime.datetime(2020, 6, 1, 0), + datetime.datetime(2020, 6, 1, 23), +) def main() -> None: @@ -70,12 +76,13 @@ def main() -> None: }, ) - sql = f""" + # The window bounds are passed as query parameters, not formatted into the + # SQL string; pruning still kicks in, so only one day is read. + sql = """ SELECT latitude, AVG("2m_temperature") - 273.15 AS air_mean_c FROM era5.surface - WHERE time BETWEEN TIMESTAMP '{_DAY}' - AND TIMESTAMP '{_DAY} 23:00:00' + WHERE time BETWEEN $start AND $end GROUP BY latitude ORDER BY latitude DESC """ @@ -83,7 +90,9 @@ def main() -> None: # Round-trip the profile back to an xarray Dataset keyed by latitude. with timed("SQL zonal mean (WHERE-pruned to one day)"): - got = ctx.sql(sql).to_dataset(dims=["latitude"]) + got = ctx.sql( + sql, param_values={"start": _START, "end": _END} + ).to_dataset(dims=["latitude"]) # Array reference: reduce the same day over the two un-grouped axes. with timed("xarray reference"): diff --git a/benchmarks/geospatial/04_anomaly.py b/benchmarks/geospatial/04_anomaly.py index 6824daa..ef920c9 100644 --- a/benchmarks/geospatial/04_anomaly.py +++ b/benchmarks/geospatial/04_anomaly.py @@ -24,13 +24,17 @@ FROM era5 a JOIN clim c ON (a.latitude, a.longitude, hour(a.time)) = (c.latitude, c.longitude, c.hour) -Dataset: **ARCO-ERA5** 2m-temperature over a North-American box for a few days, -opened and sliced with a single fluent xarray chain (anomalies vs the diurnal -cycle of that window). +The table is the *whole* lazily-opened ARCO-ERA5 archive — nothing loaded or +column-selected up front. Both the climatology CTE and the outer scan read only +``2m_temperature`` (projection pushdown) over only the window the parameterized +``WHERE`` asks for (partition pruning). Bounds are query parameters, not +string-formatted into the SQL. """ from __future__ import annotations +import datetime + import xarray as xr import xarray_sql as xql @@ -38,59 +42,78 @@ from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed _URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" -_TIME = slice("2020-06-01", "2020-06-03T23") -_LAT = slice(50.0, 25.0) -_LON = slice(235.0, 290.0) +_START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23) +_LAT_N, _LAT_S = 50.0, 25.0 +_LON_W, _LON_E = 235.0, 290.0 +_PARAMS = { + "start": _START, + "end": _END, + "lat_s": _LAT_S, + "lat_n": _LAT_N, + "lon_w": _LON_W, + "lon_e": _LON_E, +} def main() -> None: try: import gcsfs # noqa: F401 — required by the gs:// protocol - with timed("open + slice ERA5 window"): - ds = ( - xr.open_zarr( - _URL, chunks=None, storage_options={"token": "anon"} - )[["2m_temperature"]] - .sel(time=_TIME, latitude=_LAT, longitude=_LON) - .load() - ) + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) except Exception as exc: # noqa: BLE001 — any failure → skip, not crash raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc - print( - f" ERA5 2m_temperature window: {dict(ds.sizes)} (anomaly vs diurnal normals)" - ) - ctx = xql.XarrayContext() - ctx.from_dataset("era5", ds, chunks={"time": 24}) + with timed("register full ERA5 (lazy)"): + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) sql = """ WITH clim AS ( SELECT latitude, longitude, date_part('hour', time) AS hour, AVG("2m_temperature") AS clim_t - FROM era5 + FROM era5.surface + WHERE time BETWEEN $start AND $end + AND latitude BETWEEN $lat_s AND $lat_n + AND longitude BETWEEN $lon_w AND $lon_e GROUP BY latitude, longitude, date_part('hour', time) ) SELECT a.time, a.latitude, a.longitude, a."2m_temperature" - c.clim_t AS anomaly - FROM era5 a + FROM era5.surface a JOIN clim c ON a.latitude = c.latitude AND a.longitude = c.longitude AND date_part('hour', a.time) = c.hour + WHERE a.time BETWEEN $start AND $end + AND a.latitude BETWEEN $lat_s AND $lat_n + AND a.longitude BETWEEN $lon_w AND $lon_e ORDER BY a.time, a.latitude DESC, a.longitude """ show_sql(sql) # The anomaly is a gridded field; round-trip it to (time, lat, lon). - with timed("SQL anomaly (climatology CTE self-join)"): - got = ctx.sql(sql).to_dataset(dims=["time", "latitude", "longitude"]) + with timed("SQL anomaly (climatology CTE self-join, lazy read)"): + got = ctx.sql(sql, param_values=_PARAMS).to_dataset( + dims=["time", "latitude", "longitude"] + ) - # Array reference: grouped broadcast-subtract, in pure xarray. + # Array reference: grouped broadcast-subtract, in pure xarray (lazy window). with timed("xarray reference"): - grouped = ds["2m_temperature"].groupby("time.hour") + window = ds["2m_temperature"].sel( + time=slice(_START, _END), + latitude=slice(_LAT_N, _LAT_S), + longitude=slice(_LON_W, _LON_E), + ) + grouped = window.groupby("time.hour") ref = grouped - grouped.mean("time") assert_grid_close( diff --git a/benchmarks/geospatial/05_forecast_skill.py b/benchmarks/geospatial/05_forecast_skill.py index 2a007fb..9be7914 100644 --- a/benchmarks/geospatial/05_forecast_skill.py +++ b/benchmarks/geospatial/05_forecast_skill.py @@ -11,8 +11,8 @@ # /// """Forecast skill — scoring ML weather models against ERA5 is a JOIN + aggregate. -This is the real thing: scoring the **Pangu-Weather** and **GraphCast** machine- -learning forecast models against ERA5 ground truth, the headline workload of +Scoring the **Pangu-Weather** and **GraphCast** machine-learning forecast models +against ERA5 ground truth is the headline workload of [WeatherBench 2](https://weatherbench2.readthedocs.io/). A forecast is indexed by *initialization time* and *lead time* (``prediction_timedelta``); the truth is indexed by *valid time*. Evaluation aligns them by ``valid_time = init + lead`` @@ -22,18 +22,19 @@ That alignment is a relational **JOIN**, and ``valid_time = init + lead`` is just timestamp + duration arithmetic the engine does natively:: - SELECT f.prediction_timedelta AS lead, + SELECT f.model, f.prediction_timedelta AS lead, SQRT(AVG(POWER(f.t - e.t, 2))) AS rmse - FROM forecast f + FROM forecasts f JOIN era5 e ON e.time = f.time + f.prediction_timedelta -- valid_time = init + lead AND e.latitude = f.latitude AND e.longitude = f.longitude - GROUP BY f.prediction_timedelta + GROUP BY f.model, f.prediction_timedelta -The whole evaluation — temporal alignment, spatial matching, the score — is one -JOIN and one aggregate. We run it for both models and check each against an -xarray reference. +We stack the two models along a ``model`` dimension into a single forecast +table, so the query groups by a ``model`` *column* — no table name is formatted +into the SQL. Nothing is loaded up front either: the forecasts and ERA5 are +registered lazily, and the JOIN reads only what it needs at query time. Datasets: WeatherBench 2 **Pangu**, **GraphCast**, and **ERA5** at a coarse 64×32 grid (so the demo is small and fast), read from the public ``gs:// @@ -48,13 +49,7 @@ import xarray_sql as xql -from _harness import ( - CaseSkipped, - assert_grid_close, - run_case, - show_sql, - timed, -) +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed _SO = {"token": "anon"} _GRID = "64x32_equiangular_conservative" @@ -81,105 +76,98 @@ def _open(url: str) -> xr.Dataset: raise CaseSkipped(f"WeatherBench2 unavailable ({exc})") from exc -def _load_forecast(url: str, grid: xr.Dataset) -> xr.Dataset: - """Load a model's 2m-temperature over the init window, all lead times. - - Snaps lat/lon onto ERA5's exact coordinate arrays (identical 64×32 grid) so - the equality JOIN on coordinates is bit-safe across the two Zarr stores. - """ - da = _open(url)[_VAR].sel(time=_INIT) - da = da.assign_coords( - latitude=grid.latitude.values, longitude=grid.longitude.values - ) - return da.to_dataset().load() +def _reference_rmse(forecasts: xr.Dataset, truth: xr.Dataset) -> xr.DataArray: + """xarray reference: per (model, lead), align truth at valid_time, take RMSE. - -def _rmse_sql(table: str) -> str: - """The forecast↔truth JOIN + RMSE-per-lead query for one model ``table``.""" - return f""" - SELECT f.prediction_timedelta AS lead, - SQRT(AVG(POWER( - CAST(f."{_VAR}" AS DOUBLE) - e."{_VAR}", 2))) AS rmse - FROM {table} f - JOIN era5 e - ON e.time = f.time + f.prediction_timedelta -- valid = init + lead - AND e.latitude = f.latitude - AND e.longitude = f.longitude - GROUP BY f.prediction_timedelta - ORDER BY lead + The 64×32 windows are tiny, so the reference loads them and reduces in + memory; the SQL side above stays lazy. """ - - -def _rmse_by_lead(ctx: xql.XarrayContext, table: str) -> xr.DataArray: - """Run the RMSE query, returning the skill curve as a DataArray over lead.""" - return ctx.sql(_rmse_sql(table)).to_dataset(dims=["lead"]).rmse - - -def _reference_rmse(fc: xr.Dataset, truth: xr.Dataset) -> xr.DataArray: - """xarray reference: for each lead, align truth at valid_time and take RMSE.""" - f = fc[_VAR] - e = truth[_VAR] - leads = f.prediction_timedelta - out = [] - for lead in leads.values: - valid = f.time.values + lead # valid_time = init + lead - e_at_valid = e.sel(time=valid) + f = forecasts[_VAR].load() + e = truth[_VAR].load() + leads = f.prediction_timedelta.values + per_lead = [] + for lead in leads: + e_at_valid = e.sel(time=f.time.values + lead) # (init, lat, lon) diff = f.sel(prediction_timedelta=lead) - e_at_valid.values - out.append(float(np.sqrt((diff**2).mean()))) - return xr.DataArray(out, coords={"lead": leads.values}, dims=["lead"]) + per_lead.append( + np.sqrt((diff**2).mean(["time", "latitude", "longitude"])) + ) + return ( + xr.concat(per_lead, dim="lead") + .assign_coords(lead=leads) + .transpose("model", "lead") + ) def main() -> None: - era5_full = _open(_ERA5) - - # Load the forecasts first; snap their grid to ERA5's exact coordinates. - with timed("read Pangu + GraphCast forecasts"): - pangu = _load_forecast(_PANGU, era5_full) - graphcast = _load_forecast(_GRAPHCAST, era5_full) + # Open everything lazily — no .load() here. + era5 = _open(_ERA5) + + # The two models store different pressure-level sets (Pangu 13, GraphCast + # 37), so we keep the common surface field 2m_temperature and stack the + # models along a `model` dimension into one forecast table. Snap the grid + # onto ERA5's exact coordinates (same 64×32 grid) so the equality JOIN is + # bit-safe across the Zarr stores. + pangu = _open(_PANGU)[[_VAR]].sel(time=_INIT) + graphcast = _open(_GRAPHCAST)[[_VAR]].sel(time=_INIT) + forecasts = xr.concat([pangu, graphcast], dim="model").assign_coords( + model=["pangu", "graphcast"], + latitude=era5.latitude.values, + longitude=era5.longitude.values, + ) - # ERA5 truth must span every valid time: window start through the last - # init plus the longest lead, so the JOIN keeps every forecast pair. + # ERA5 truth must span every valid time (last init + longest lead); bound it + # lazily so the JOIN does not scan the whole 1959–2023 record. valid_max = ( pangu.time.values.max() + pangu.prediction_timedelta.values.max() ) - with timed("read ERA5 truth window"): - truth = ( - era5_full[[_VAR]] - .sel(time=slice(_INIT.start, pd.Timestamp(valid_max))) - .load() - ) + truth = era5[[_VAR]].sel(time=slice(_INIT.start, pd.Timestamp(valid_max))) print( f" 64×32 2m_temperature | init {_INIT.start}…{_INIT.stop} " - f"({pangu.sizes['time']} inits × {pangu.sizes['prediction_timedelta']} leads)" + f"({pangu.sizes['time']} inits × {pangu.sizes['prediction_timedelta']} " + f"leads × 2 models)" ) ctx = xql.XarrayContext() + ctx.from_dataset("forecasts", forecasts, chunks={"time": 6}) ctx.from_dataset("era5", truth, chunks={"time": 100}) - ctx.from_dataset("pangu", pangu, chunks={"time": 20}) - ctx.from_dataset("graphcast", graphcast, chunks={"time": 20}) - print("\n Forecast↔ERA5 JOIN on valid_time = init + lead, RMSE per lead:") - show_sql(_rmse_sql("pangu")) + sql = """ + SELECT f.model, + f.prediction_timedelta AS lead, + SQRT(AVG(POWER( + CAST(f."2m_temperature" AS DOUBLE) - e."2m_temperature", 2 + ))) AS rmse + FROM forecasts f + JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude + GROUP BY f.model, f.prediction_timedelta + ORDER BY f.model, lead + """ + show_sql(sql) + + with timed("SQL RMSE by (model, lead) — lazy JOIN"): + got = ctx.sql(sql).to_dataset(dims=["model", "lead"]).rmse + + with timed("xarray reference"): + ref = _reference_rmse(forecasts, truth) - results = {} - for name, fc in [("pangu", pangu), ("graphcast", graphcast)]: - with timed(f"SQL RMSE-by-lead: {name}"): - got = _rmse_by_lead(ctx, name) - with timed(f"xarray reference: {name}"): - ref = _reference_rmse(fc, truth) - assert_grid_close(f"{name} RMSE(lead)", got, ref, rtol=1e-4, atol=1e-3) - results[name] = got + assert_grid_close("RMSE(model, lead)", got, ref, rtol=1e-4, atol=1e-3) # Headline table: error growth with forecast horizon, both models. - lead_days = results["pangu"]["lead"].values / np.timedelta64(1, "D") + lead_days = got["lead"].values / np.timedelta64(1, "D") + pangu_rmse = got.sel(model="pangu") + graphcast_rmse = got.sel(model="graphcast") print("\n 2m-temperature RMSE (K) vs lead time — lower is better:") print(f" {'lead (days)':>12} {'Pangu':>9} {'GraphCast':>11}") for i in range(0, len(lead_days), 4): print( f" {lead_days[i]:>12.2f} " - f"{float(results['pangu'][i]):>9.3f} " - f"{float(results['graphcast'][i]):>11.3f}" + f"{float(pangu_rmse.isel(lead=i)):>9.3f} " + f"{float(graphcast_rmse.isel(lead=i)):>11.3f}" ) diff --git a/benchmarks/geospatial/06_zonal_vector.py b/benchmarks/geospatial/06_zonal_vector.py index e517557..8c7a931 100644 --- a/benchmarks/geospatial/06_zonal_vector.py +++ b/benchmarks/geospatial/06_zonal_vector.py @@ -34,6 +34,8 @@ from __future__ import annotations +import datetime + import numpy as np import xarray as xr @@ -43,6 +45,10 @@ _URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" _DAY = "2020-06-01" +_START, _END = ( + datetime.datetime(2020, 6, 1, 0), + datetime.datetime(2020, 6, 1, 23), +) # Continental-scale boxes (name, lat_min, lat_max, lon_min, lon_max), lon 0–360°E. _REGIONS = [ @@ -97,7 +103,7 @@ def main() -> None: "regions", _regions_dataset(), chunks={"region": len(_REGIONS)} ) - sql = f""" + sql = """ SELECT r.region AS region_id, AVG(a."2m_temperature") - 273.15 AS avg_c, COUNT(*) AS n_obs @@ -105,15 +111,16 @@ def main() -> None: JOIN regions r ON a.latitude BETWEEN r.lat_min AND r.lat_max AND a.longitude BETWEEN r.lon_min AND r.lon_max - WHERE a.time BETWEEN TIMESTAMP '{_DAY}' - AND TIMESTAMP '{_DAY} 23:00:00' + WHERE a.time BETWEEN $start AND $end GROUP BY r.region ORDER BY r.region """ show_sql(sql) with timed("SQL zonal stats (raster × vector range JOIN, WHERE-pruned)"): - got = ctx.sql(sql).to_dataset(dims=["region_id"]) + got = ctx.sql( + sql, param_values={"start": _START, "end": _END} + ).to_dataset(dims=["region_id"]) # Array reference: load the same day once, mask each region in memory. with timed("xarray reference"): diff --git a/benchmarks/geospatial/README.md b/benchmarks/geospatial/README.md index 96d0716..879f237 100644 --- a/benchmarks/geospatial/README.md +++ b/benchmarks/geospatial/README.md @@ -39,11 +39,13 @@ interpolation weights — the geometry — which SQL applies but does not comput (bands B04/B08). Requires network; skips cleanly if offline. - **02–06** — the full **[ARCO-ERA5](https://github.com/google-research/arco-era5)** archive (0.25° global, ~1.3M hourly timesteps, 273 variables) read anonymously - from a public GCS bucket. Cases 03 and 06 register the *whole* archive and let - SQL `WHERE` prune it to a one-day window (the partition-pruning demo); cases - 02/04 read a bounded region/time window into memory once and compare - SQL against the same array. All require network (`gcsfs`); skip cleanly - offline. Each ERA5 case takes roughly a minute, dominated by the GCS read. + from a public GCS bucket. Every case registers the *whole* archive **lazily** + (nothing loaded or column-selected up front) and filters it with a + **parameterized** `WHERE` (bounds bound as query parameters, not formatted into + the SQL); projection pushdown reads only `2m_temperature` and partition pruning + reads only the window asked for. All require network (`gcsfs`); skip cleanly + offline. ERA5 cases take roughly one to a few minutes, dominated by the GCS + read (the lazy reference re-reads the same window). - **05 forecast skill** — the **[WeatherBench 2](https://weatherbench2.readthedocs.io/)** Pangu-Weather, GraphCast, and ERA5 datasets at a coarse 64×32 grid, scoring both ML models against ERA5 ground truth. Network-backed; runs in seconds diff --git a/docs/geospatial.md b/docs/geospatial.md index 1c42f27..8fc61ef 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -12,10 +12,13 @@ operations. But it is not the only one, and for a large and growing audience — the people fluent in SQL rather than in `apply_ufunc` and rechunking — it is not the most accessible one. [`xarray-sql`](../README.md) lets you pose these questions in SQL and answers them with a real query engine (DataFusion), -complete with partition pruning and projection pushdown. And because a gridded -result is still gridded data, every query here round-trips its answer straight -back to an `xarray.Dataset` (via `to_dataset`) — SQL in, an array out, ready to -plot or save. +complete with partition pruning and projection pushdown. The datasets are +registered *lazily* — nothing is read or column-selected up front; each query +pulls only the variable and the partitions it needs, and value filters are +passed as bound **query parameters** rather than formatted into the SQL string. +And because a gridded result is still gridded data, every query here round-trips +its answer straight back to an `xarray.Dataset` (via `to_dataset`) — SQL in, an +array out, ready to plot or save. This page makes the argument case by case. Every claim below is backed by a runnable script in [`benchmarks/geospatial/`](../benchmarks/geospatial/) that @@ -118,20 +121,22 @@ That alignment is a relational JOIN, and `valid_time = init + lead` is just timestamp + duration arithmetic the engine does natively: ```sql -SELECT f.prediction_timedelta AS lead, +SELECT f.model, f.prediction_timedelta AS lead, SQRT(AVG(POWER(f."2m_temperature" - e."2m_temperature", 2))) AS rmse -FROM forecast f +FROM forecasts f JOIN era5 e ON e.time = f.time + f.prediction_timedelta -- valid_time = init + lead AND e.latitude = f.latitude AND e.longitude = f.longitude -GROUP BY f.prediction_timedelta +GROUP BY f.model, f.prediction_timedelta ``` -The entire evaluation — temporal alignment across three time axes, spatial -matching, and the score — is one JOIN and one aggregate. +Both models are stacked along a `model` dimension into one forecast table, so +the query scores them together by grouping on a `model` *column* — no table name +formatted into the SQL. The entire evaluation — temporal alignment across three +time axes, spatial matching, and the score — is one JOIN and one aggregate. [`05_forecast_skill.py`](../benchmarks/geospatial/05_forecast_skill.py) runs it -for both models, matches a NumPy reference, and reproduces the published result +for both models, matches an xarray reference, and reproduces the published result that GraphCast edges out Pangu at every lead — the classic "error grows with horizon" curve (≈0.3 K at 6 h rising to ≈2.5 K at 9 days): From 7dda833cabaf7696768438a392723bb6fcfacabe Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 24 Jun 2026 09:14:05 +0300 Subject: [PATCH 05/26] Cases 07/08: use Earth Engine via Xee (independent reproj reference; real SRTM regrid) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 07 reprojection: open a UTM grid through Xee carrying ee.Image.pixelLonLat(), so Earth Engine's own geodesy reports each pixel's true lon/lat. The SQL PROJ UDF is now validated against an *independent* implementation (EE), not against PROJ via pyproj. Grid built with xee.helpers.fit_geometry; compared at ~1e-5°. - 08 regridding: regrid real SRTM elevation (USGS/SRTMGL1_003, Sierra Nevada) opened via Xee, validated against xarray's bilinear .interp() on the same source field — a new (DEM) data type for the suite. - Both require Earth Engine access (earthengine authenticate + EARTHENGINE_PROJECT) and skip cleanly when EE is not initialized, like the other network cases. Deps add xee, earthengine-api, shapely. Verified offline: imports, the graceful EE-not-initialized skip, ruff/mypy, and the xee open_dataset / fit_geometry API surface. The live EE round-trip was not run here (no credentials in this environment). docs/geospatial.md and the suite README updated for the EE/Xee sources. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/07_reproject_udf.py | 135 +++++++++++++-------- benchmarks/geospatial/08_regrid_weights.py | 94 +++++++++++--- benchmarks/geospatial/README.md | 9 +- docs/geospatial.md | 18 ++- 4 files changed, 178 insertions(+), 78 deletions(-) diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py index 543f30c..761c140 100644 --- a/benchmarks/geospatial/07_reproject_udf.py +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -1,3 +1,4 @@ +#!/usr/bin/env python3 # /// script # requires-python = ">=3.11" # dependencies = [ @@ -6,12 +7,15 @@ # "numpy", # "pyproj", # "pyarrow", +# "xee", +# "earthengine-api", +# "shapely", # ] # /// """Reprojection — a per-pixel CRS transform is a scalar UDF (à la ST_Transform). -Reprojection moves coordinates from one CRS to another (here UTM zone 32N, -EPSG:32632, → lon/lat, EPSG:4326). Crucially it is **row-independent**: each +Reprojection moves coordinates from one CRS to another (here UTM zone 10N, +EPSG:32610, → lon/lat, EPSG:4326). Crucially it is **row-independent**: each pixel's new coordinate depends only on its own old coordinate. That is exactly the shape of a SQL *scalar UDF*, and it is precisely how the geospatial SQL world already does it — PostGIS ``ST_Transform`` and DuckDB-spatial @@ -22,30 +26,29 @@ SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat FROM grid +**The reference is Earth Engine itself.** We open a UTM grid through +[Xee](https://github.com/google/Xee) carrying ``ee.Image.pixelLonLat()`` — Earth +Engine's *own* geodesy engine computes the true lon/lat of every UTM pixel +centre. So this case validates our PROJ-in-SQL transform against a fully +**independent** reprojection implementation (EE), not against PROJ again. They +agree to sub-metre precision. + The UDF (built on ``pyproj``, vectorized over each Arrow batch) mirrors the ``cftime()`` UDF already shipped in xarray-sql (see ``xarray_sql/cftime.py``); it could graduate into the package as ``xql.register_reproject_udf``. -**What this does *not* do:** it moves the coordinates, it does not resample onto -a regular target grid. Producing a gridded product in the new CRS still needs -interpolation — which is case 08, where regridding turns out to be a JOIN -against a weight table rather than a scalar UDF. - -**An honest caveat on threading:** PROJ's transformation context is not -thread-safe, and DataFusion evaluates independent projection expressions on -concurrent worker threads. Two separate PROJ UDFs in one ``SELECT`` (one for -lon, one for lat) race and segfault. The fix is to return *both* coordinates -from a **single** struct-returning UDF (so PROJ is touched once per row), and -to keep the source in one chunk so that single UDF runs serially. A -production-grade ``ST_Transform`` would additionally give each worker thread its -own PROJ context; the point here is the *shape* of the operation — a scalar -UDF — not the parallel execution. - -No network: a synthetic UTM grid over the extent of a Sentinel-2 tile. +PROJ's context is not thread-safe and DataFusion evaluates projection +expressions concurrently, so we return *both* coordinates from one +struct-returning UDF and keep the source in a single chunk (one serial UDF). + +Requires Earth Engine access: ``earthengine authenticate`` once, then an +initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise. """ from __future__ import annotations +import os + import numpy as np import pyarrow as pa import pyproj @@ -54,7 +57,12 @@ import xarray_sql as xql -from _harness import assert_grid_close, run_case, show_sql, timed +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed + +_SRC_CRS, _DST_CRS = "EPSG:32610", "EPSG:4326" # UTM zone 10N → lon/lat +# A 1° box over the San Francisco Bay area, well inside UTM zone 10N. +_AOI = (-122.6, 37.4, -121.6, 38.4) +_SCALE_M = 2_000 # 2 km pixels → a ~50×60 grid def register_reproject_udf( @@ -97,32 +105,63 @@ def _fn(x: pa.Array, y: pa.Array) -> pa.Array: ) -def _utm_grid() -> xr.Dataset: - """A synthetic field on a UTM grid matching a Sentinel-2 T32T tile extent.""" - # EPSG:32632 extent of tile T32TLQ (from the STAC proj:bbox), coarsened. - x = np.linspace(300_000.0, 409_800.0, 60, dtype="float64") - y = np.linspace(5_000_040.0, 4_890_240.0, 60, dtype="float64") - elevation = np.zeros( - (y.size, x.size), dtype="float64" - ) # value is irrelevant - # Single chunk → single partition → serial UDF (PROJ is not thread-safe). - return xr.Dataset( - {"elevation": (["y", "x"], elevation)}, - coords={"y": y, "x": x}, - ).chunk({"y": 60, "x": 60}) +def _open_ee_lonlat_grid() -> xr.Dataset: + """Open ``ee.Image.pixelLonLat()`` on a UTM grid via Xee. + + Earth Engine evaluates ``pixelLonLat`` on the requested UTM grid, so each + pixel carries its UTM ``x``/``y`` (coordinates) and EE's own ``longitude`` / + ``latitude`` (data variables) — the independent reprojection reference. + """ + try: + import ee + import shapely.geometry as sgeom + from xee import helpers + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api xee`" + ) from exc + + try: + ee.Initialize( + project=os.environ.get("EARTHENGINE_PROJECT"), + opt_url="https://earthengine-highvolume.googleapis.com", + ) + except Exception as exc: # noqa: BLE001 — not authenticated → skip + raise CaseSkipped( + f"Earth Engine not initialized ({exc}); run " + "`earthengine authenticate` and set EARTHENGINE_PROJECT" + ) from exc + + # fit_geometry builds the pixel grid (crs, crs_transform, shape_2d) Xee's + # backend expects — here a UTM grid at _SCALE_M metres covering the AOI. + grid = helpers.fit_geometry( + sgeom.box(*_AOI), + geometry_crs="EPSG:4326", + grid_crs=_SRC_CRS, + grid_scale=(float(_SCALE_M), float(_SCALE_M)), + ) + ic = ee.ImageCollection([ee.Image.pixelLonLat()]) + ds = xr.open_dataset(ic, engine="ee", **grid) + # One image → a length-1 time axis; drop it. Xee names projected spatial + # coordinates "X"/"Y" (UTM metres); normalize to lower case for the SQL. + ds = ds.isel(time=0).rename({"X": "x", "Y": "y"}) + return ds.load() def main() -> None: - src_crs, dst_crs = "EPSG:32632", "EPSG:4326" - ds = _utm_grid() + ds = _open_ee_lonlat_grid() + n = ds.sizes["y"] * ds.sizes["x"] print( - f" UTM grid {dict(ds.sizes)} ({ds.sizes['y'] * ds.sizes['x']:,} points) " - f"{src_crs} → {dst_crs}" + f" EE pixelLonLat on UTM grid {dict(ds.sizes)} ({n:,} pixels) " + f"{_SRC_CRS} → {_DST_CRS}" ) ctx = xql.XarrayContext() - ctx.from_dataset("grid", ds, chunks={"y": 60, "x": 60}) - register_reproject_udf(ctx, src_crs, dst_crs) + # Single chunk → single partition → serial UDF (PROJ is not thread-safe). + ctx.from_dataset( + "grid", ds, chunks={"y": ds.sizes["y"], "x": ds.sizes["x"]} + ) + register_reproject_udf(ctx, _SRC_CRS, _DST_CRS) sql = """ SELECT x, y, @@ -133,26 +172,16 @@ def main() -> None: """ show_sql(sql) - # Round-trip the reprojected coordinates back to an (y, x) grid. with timed("SQL reprojection (PROJ scalar UDF)"): got = ctx.sql(sql).to_dataset(dims=["y", "x"]) - # Array reference: the same PROJ transform applied to the full grid. - with timed("pyproj reference"): - transformer = pyproj.Transformer.from_crs( - src_crs, dst_crs, always_xy=True - ) - xx, yy = np.meshgrid(ds.x.values, ds.y.values) - lon, lat = transformer.transform(xx, yy) - coords = {"y": ds.y, "x": ds.x} - ref_lon = xr.DataArray(lon, dims=["y", "x"], coords=coords) - ref_lat = xr.DataArray(lat, dims=["y", "x"], coords=coords) - + # Reference: Earth Engine's own per-pixel lon/lat (independent of PROJ). + # EE and PROJ are separate implementations, so compare at ~1e-5° (~1 m). assert_grid_close( - "reprojected longitude", got.lon, ref_lon, rtol=0, atol=1e-9 + "reprojected longitude", got.lon, ds.longitude, rtol=0, atol=1e-5 ) assert_grid_close( - "reprojected latitude", got.lat, ref_lat, rtol=0, atol=1e-9 + "reprojected latitude", got.lat, ds.latitude, rtol=0, atol=1e-5 ) corner = got.isel(x=0, y=0) @@ -164,5 +193,5 @@ def main() -> None: if __name__ == "__main__": raise SystemExit( - run_case(main, "Reprojection: PROJ scalar UDF (ST_Transform)") + run_case(main, "Reprojection: PROJ scalar UDF vs Earth Engine") ) diff --git a/benchmarks/geospatial/08_regrid_weights.py b/benchmarks/geospatial/08_regrid_weights.py index 1fa4b34..4191ead 100644 --- a/benchmarks/geospatial/08_regrid_weights.py +++ b/benchmarks/geospatial/08_regrid_weights.py @@ -5,6 +5,9 @@ # "xarray", # "numpy", # "scipy", +# "xee", +# "earthengine-api", +# "shapely", # ] # /// """Regridding — interpolation to a new grid is a sparse matmul, i.e. a JOIN. @@ -31,18 +34,29 @@ and hand the resulting sparse matrix to SQL as a table. SQL *applies* the weights; it does not invent the geometry. -We validate the SQL matmul against xarray's own bilinear ``.interp()``. -No network; fully synthetic. +The field is real **SRTM elevation** (terrain over the Sierra Nevada), opened +from the Earth Engine catalog through [Xee](https://github.com/google/Xee). We +regrid it coarse → fine in SQL and validate against xarray's own bilinear +``.interp()`` on the same source field. + +Requires Earth Engine access: ``earthengine authenticate`` once, then an +initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise. """ from __future__ import annotations +import os + import numpy as np import xarray as xr import xarray_sql as xql -from _harness import assert_grid_close, run_case, show_sql, timed +from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed + +# A 1° box over the Sierra Nevada — real terrain with strong relief. +_AOI = (-119.6, 37.0, -118.6, 38.0) +_SRC_SCALE_DEG = 0.02 # ~2 km source pixels (a coarse DEM to upsample) def _linear_weights( @@ -91,20 +105,64 @@ def _bilinear_weight_table( ).chunk({"pair": n}) -def main() -> None: - # Coarse source grid with a smooth field; finer target grid strictly inside. - slat = np.linspace(0.0, 10.0, 11) - slon = np.linspace(0.0, 10.0, 11) - field = ( - np.sin(slat[:, None] / 3.0) * np.cos(slon[None, :] / 4.0) - + 0.1 * slat[:, None] +def _open_srtm() -> xr.DataArray: + """Open SRTM elevation over the AOI as a coarse (lat, lon) field via Xee.""" + try: + import ee + import shapely.geometry as sgeom + from xee import helpers + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api xee`" + ) from exc + + try: + ee.Initialize( + project=os.environ.get("EARTHENGINE_PROJECT"), + opt_url="https://earthengine-highvolume.googleapis.com", + ) + except Exception as exc: # noqa: BLE001 — not authenticated → skip + raise CaseSkipped( + f"Earth Engine not initialized ({exc}); run " + "`earthengine authenticate` and set EARTHENGINE_PROJECT" + ) from exc + + # fit_geometry builds the pixel grid (crs, crs_transform, shape_2d) Xee's + # backend expects — here a geographic grid at _SRC_SCALE_DEG° over the AOI. + grid = helpers.fit_geometry( + sgeom.box(*_AOI), + grid_crs="EPSG:4326", + grid_scale=(_SRC_SCALE_DEG, _SRC_SCALE_DEG), ) - src_da = xr.DataArray( - field, coords={"lat": slat, "lon": slon}, dims=["lat", "lon"] + ic = ee.ImageCollection([ee.Image("USGS/SRTMGL1_003")]) # band: elevation + ds = xr.open_dataset(ic, engine="ee", **grid) + da = ds["elevation"].isel(time=0) + # Normalize Xee's spatial coordinate names to lat/lon and sort ascending so + # the 1-D weight construction (searchsorted) sees increasing coordinates. + rename = {} + for d in da.dims: + dl = d.lower() + if dl in ("y", "lat", "latitude"): + rename[d] = "lat" + elif dl in ("x", "lon", "longitude"): + rename[d] = "lon" + return da.rename(rename).sortby("lat").sortby("lon").load() + + +def main() -> None: + with timed("open SRTM via Xee"): + src_da = _open_srtm() + slat = src_da.lat.values + slon = src_da.lon.values + field = src_da.values.astype("float64") + print( + f" SRTM elevation source grid {len(slat)}×{len(slon)} " + f"({float(np.nanmin(field)):.0f}–{float(np.nanmax(field)):.0f} m)" ) - tlat = np.linspace(0.3, 9.7, 25) - tlon = np.linspace(0.4, 9.6, 30) + # Finer target grid strictly inside the source extent (bilinear upsampling). + tlat = np.linspace(slat[1], slat[-2], 60) + tlon = np.linspace(slon[1], slon[-2], 72) print( f" regrid {len(slat)}×{len(slon)} → {len(tlat)}×{len(tlon)} (bilinear)" ) @@ -144,7 +202,7 @@ def main() -> None: coords={"lat": tlat, "lon": tlon}, ) - # Array reference: xarray's own bilinear interpolation onto the target grid. + # Array reference: xarray's own bilinear interpolation of the same field. with timed("xarray .interp reference"): ref = src_da.interp(lat=tlat, lon=tlon, method="linear") @@ -152,9 +210,11 @@ def main() -> None: print( f"\n {got.size:,} target cells regridded; " - f"value range [{float(got.min()):.3f}, {float(got.max()):.3f}]." + f"elevation range [{float(got.min()):.0f}, {float(got.max()):.0f}] m." ) if __name__ == "__main__": - raise SystemExit(run_case(main, "Regridding: sparse weight-table JOIN")) + raise SystemExit( + run_case(main, "Regridding: sparse weight-table JOIN (SRTM)") + ) diff --git a/benchmarks/geospatial/README.md b/benchmarks/geospatial/README.md index 879f237..37df3e2 100644 --- a/benchmarks/geospatial/README.md +++ b/benchmarks/geospatial/README.md @@ -50,8 +50,13 @@ interpolation weights — the geometry — which SQL applies but does not comput Pangu-Weather, GraphCast, and ERA5 datasets at a coarse 64×32 grid, scoring both ML models against ERA5 ground truth. Network-backed; runs in seconds because the grid is small. -- **07–08** — small/synthetic grids plus precomputed regrid weights, so they run - without heavy geospatial dependencies (ESMF/ESMPy). +- **07–08** — the **Earth Engine** catalog via [Xee](https://github.com/google/Xee). + 07 reprojects a UTM grid and validates the SQL transform against Earth Engine's + *own* per-pixel lon/lat (`ee.Image.pixelLonLat()`) — an independent reprojection + reference, not PROJ-vs-PROJ. 08 regrids real **SRTM elevation** (Sierra Nevada) + and validates against xarray's bilinear `.interp()`. Both need EE access + (`earthengine authenticate` + an initialized project via `EARTHENGINE_PROJECT`) + and skip cleanly without it. ## Running diff --git a/docs/geospatial.md b/docs/geospatial.md index 8fc61ef..ca13306 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -187,11 +187,15 @@ SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat FROM grid ``` -[`07_reproject_udf.py`](../benchmarks/geospatial/07_reproject_udf.py) matches -pyproj to 1e-9. Two honest caveats, both documented in the script: PROJ's -context is not thread-safe (so the UDF returns both coordinates from *one* call -and runs on a single partition), and reprojection moves coordinates without -resampling onto a grid — which is the next operation. +[`07_reproject_udf.py`](../benchmarks/geospatial/07_reproject_udf.py) validates +this against **Earth Engine itself**: it opens a UTM grid through +[Xee](https://github.com/google/Xee) carrying `ee.Image.pixelLonLat()`, so EE's +own geodesy engine reports the true lon/lat of every pixel — an *independent* +reprojection reference, not PROJ-vs-PROJ. The SQL UDF and EE agree to sub-metre +precision. Two honest caveats, both documented in the script: PROJ's context is +not thread-safe (so the UDF returns both coordinates from *one* call and runs on +a single partition), and reprojection moves coordinates without resampling onto +a grid — which is the next operation. **Regridding is not** row-independent: each output cell is a weighted blend of several input cells. That is a *many-to-many* relationship — and a many-to-many @@ -204,7 +208,9 @@ FROM weights w JOIN src s ON s.cell_id = w.src_id GROUP BY w.dst_id ``` -[`08_regrid_weights.py`](../benchmarks/geospatial/08_regrid_weights.py) matches +[`08_regrid_weights.py`](../benchmarks/geospatial/08_regrid_weights.py) regrids +real **SRTM elevation** (Sierra Nevada terrain, opened from the Earth Engine +catalog through [Xee](https://github.com/google/Xee)) coarse → fine and matches xarray's bilinear `.interp()` exactly. So regridding does not weaken the thesis — it is the most relational operation of all. From 552a6db71a725aca1eb6b31827dcb118dad2d86a Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 24 Jun 2026 09:37:22 +0300 Subject: [PATCH 06/26] README: add "Does it work?" section (addresses #46); drop "honest" - Move the geospatial-suite pointer out of "What is this?" into a new "Does it work?" section, placed between "How does it work?" and "Why does this work?", reframed to answer that question with a findings summary (validated operations, real datasets, the GraphCast-beats-Pangu result). Addresses #46. - Reword the two remaining uses of "honest" in docs/geospatial.md. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- README.md | 27 +++++++++++++++++++++++++++ docs/geospatial.md | 10 +++++----- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index aeec14f..402eca7 100644 --- a/README.md +++ b/README.md @@ -190,6 +190,33 @@ that lets the DB engine translate the underlying Dataset arrays into DataFusion Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be translated into a 2D table -- underlies this performant query mechanism. +## Does it work? + +Yes. The recurring worry is that the SQL interface is a toy — fine for `SELECT`s, +but not for the operations geoscience actually runs. So we wrote a suite that +takes the staples of geospatial and climate analysis — the ones we assume *need* +an array library — and expresses each one in SQL, then **checks the SQL answer +against an xarray/array reference** to floating-point tolerance: + +* **Spectral indices** (NDVI) — column arithmetic over a real Sentinel-2 scene. +* **Climatology, anomalies, zonal means** — `GROUP BY` and self-`JOIN` over the + full 0.25° **ARCO-ERA5** archive (≈1.3M hourly steps), read lazily with + partition pruning and column pushdown. +* **Forecast skill** — scoring the **Pangu-Weather** and **GraphCast** ML models + against ERA5 (WeatherBench 2) as a `JOIN` on `valid_time = init + lead`; it + reproduces the published result that GraphCast beats Pangu at every lead. +* **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a + table of regions. +* **Reprojection and regridding** — a scalar PROJ UDF (validated against Earth + Engine's own geodesy via [Xee](https://github.com/google/Xee)) and a + sparse-weight-table `JOIN` (regridding real SRTM terrain). + +Every case matches its array reference. The headline finding: these operations +are not really "array" operations at all — they are `GROUP BY`, `JOIN`, window +functions, and `CASE` in disguise, and a query engine runs them at scale. See +[`benchmarks/geospatial/`](benchmarks/geospatial/) and the write-up, +[Geospatial operations are relational operations](docs/geospatial.md). + ## Why does this work? Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in diff --git a/docs/geospatial.md b/docs/geospatial.md index ca13306..794310a 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -192,10 +192,10 @@ this against **Earth Engine itself**: it opens a UTM grid through [Xee](https://github.com/google/Xee) carrying `ee.Image.pixelLonLat()`, so EE's own geodesy engine reports the true lon/lat of every pixel — an *independent* reprojection reference, not PROJ-vs-PROJ. The SQL UDF and EE agree to sub-metre -precision. Two honest caveats, both documented in the script: PROJ's context is -not thread-safe (so the UDF returns both coordinates from *one* call and runs on -a single partition), and reprojection moves coordinates without resampling onto -a grid — which is the next operation. +precision. Two caveats, both documented in the script: PROJ's context is not +thread-safe (so the UDF returns both coordinates from *one* call and runs on a +single partition), and reprojection moves coordinates without resampling onto a +grid — which is the next operation. **Regridding is not** row-independent: each output cell is a weighted blend of several input cells. That is a *many-to-many* relationship — and a many-to-many @@ -216,7 +216,7 @@ it is the most relational operation of all. ## Where the array paradigm still earns its keep -The honest boundary is **weight generation**. Applying a regridding is a join; +The boundary is **weight generation**. Applying a regridding is a join; *computing* the weights — cell overlaps for conservative remapping, stencils and spherical geometry for bilinear, the whole machinery of xESMF/ESMF — is genuinely geometric work that arrays (and specialized libraries) do well. The relational From fd4f45ae8a9420096f3aff4b84caaaa925007021 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 24 Jun 2026 10:19:40 +0300 Subject: [PATCH 07/26] =?UTF-8?q?EE=20cases:=20initialize=20via=20ADC=20(n?= =?UTF-8?q?o=20blocked=20OAuth);=20fix=2007=20coord=20names=20=E2=80=94=20?= =?UTF-8?q?both=20pass=20live?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add _harness.initialize_earth_engine(): initialize EE from Application Default Credentials (gcloud auth application-default login) with the Earth Engine scope, using the ADC/EARTHENGINE_PROJECT project. This avoids the separate `earthengine authenticate` OAuth flow (and the "this app is blocked" org-policy error), and needs no service account. - 07/08 use the helper instead of bare ee.Initialize. - 07: Xee names the projected grid coords x/y (not X/Y), so drop the erroneous rename. Verified end-to-end against Earth Engine: - 07 reprojection: PROJ UDF matches EE's own ee.Image.pixelLonLat() to ~1e-5° over 2,565 UTM pixels (independent geodesy engines agree; pixel centers align). - 08 regrid: real SRTM elevation (Sierra Nevada, 1206–3439 m) regridded 50×50 → 60×72 matches xarray bilinear .interp() exactly. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/07_reproject_udf.py | 30 +++++++---------- benchmarks/geospatial/08_regrid_weights.py | 23 +++++-------- benchmarks/geospatial/_harness.py | 39 ++++++++++++++++++++++ 3 files changed, 60 insertions(+), 32 deletions(-) diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py index 761c140..2578878 100644 --- a/benchmarks/geospatial/07_reproject_udf.py +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -47,8 +47,6 @@ from __future__ import annotations -import os - import numpy as np import pyarrow as pa import pyproj @@ -57,7 +55,14 @@ import xarray_sql as xql -from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed +from _harness import ( + CaseSkipped, + assert_grid_close, + initialize_earth_engine, + run_case, + show_sql, + timed, +) _SRC_CRS, _DST_CRS = "EPSG:32610", "EPSG:4326" # UTM zone 10N → lon/lat # A 1° box over the San Francisco Bay area, well inside UTM zone 10N. @@ -113,7 +118,6 @@ def _open_ee_lonlat_grid() -> xr.Dataset: ``latitude`` (data variables) — the independent reprojection reference. """ try: - import ee import shapely.geometry as sgeom from xee import helpers except ImportError as exc: # pragma: no cover @@ -121,16 +125,7 @@ def _open_ee_lonlat_grid() -> xr.Dataset: "Earth Engine support needs `pip install earthengine-api xee`" ) from exc - try: - ee.Initialize( - project=os.environ.get("EARTHENGINE_PROJECT"), - opt_url="https://earthengine-highvolume.googleapis.com", - ) - except Exception as exc: # noqa: BLE001 — not authenticated → skip - raise CaseSkipped( - f"Earth Engine not initialized ({exc}); run " - "`earthengine authenticate` and set EARTHENGINE_PROJECT" - ) from exc + ee = initialize_earth_engine() # fit_geometry builds the pixel grid (crs, crs_transform, shape_2d) Xee's # backend expects — here a UTM grid at _SCALE_M metres covering the AOI. @@ -142,10 +137,9 @@ def _open_ee_lonlat_grid() -> xr.Dataset: ) ic = ee.ImageCollection([ee.Image.pixelLonLat()]) ds = xr.open_dataset(ic, engine="ee", **grid) - # One image → a length-1 time axis; drop it. Xee names projected spatial - # coordinates "X"/"Y" (UTM metres); normalize to lower case for the SQL. - ds = ds.isel(time=0).rename({"X": "x", "Y": "y"}) - return ds.load() + # One image → a length-1 time axis; drop it. Xee gives x/y coordinates (UTM + # metres) and longitude/latitude data variables (EE's per-pixel geodesy). + return ds.isel(time=0).load() def main() -> None: diff --git a/benchmarks/geospatial/08_regrid_weights.py b/benchmarks/geospatial/08_regrid_weights.py index 4191ead..e2a825b 100644 --- a/benchmarks/geospatial/08_regrid_weights.py +++ b/benchmarks/geospatial/08_regrid_weights.py @@ -45,14 +45,19 @@ from __future__ import annotations -import os - import numpy as np import xarray as xr import xarray_sql as xql -from _harness import CaseSkipped, assert_grid_close, run_case, show_sql, timed +from _harness import ( + CaseSkipped, + assert_grid_close, + initialize_earth_engine, + run_case, + show_sql, + timed, +) # A 1° box over the Sierra Nevada — real terrain with strong relief. _AOI = (-119.6, 37.0, -118.6, 38.0) @@ -108,7 +113,6 @@ def _bilinear_weight_table( def _open_srtm() -> xr.DataArray: """Open SRTM elevation over the AOI as a coarse (lat, lon) field via Xee.""" try: - import ee import shapely.geometry as sgeom from xee import helpers except ImportError as exc: # pragma: no cover @@ -116,16 +120,7 @@ def _open_srtm() -> xr.DataArray: "Earth Engine support needs `pip install earthengine-api xee`" ) from exc - try: - ee.Initialize( - project=os.environ.get("EARTHENGINE_PROJECT"), - opt_url="https://earthengine-highvolume.googleapis.com", - ) - except Exception as exc: # noqa: BLE001 — not authenticated → skip - raise CaseSkipped( - f"Earth Engine not initialized ({exc}); run " - "`earthengine authenticate` and set EARTHENGINE_PROJECT" - ) from exc + ee = initialize_earth_engine() # fit_geometry builds the pixel grid (crs, crs_transform, shape_2d) Xee's # backend expects — here a geographic grid at _SRC_SCALE_DEG° over the AOI. diff --git a/benchmarks/geospatial/_harness.py b/benchmarks/geospatial/_harness.py index 44149ee..9925dea 100644 --- a/benchmarks/geospatial/_harness.py +++ b/benchmarks/geospatial/_harness.py @@ -21,20 +21,59 @@ from __future__ import annotations import contextlib +import os import sys import time import tracemalloc from collections.abc import Callable, Iterator +from typing import Any import xarray as xr _WIDTH = 72 +_EE_SCOPES = [ + "https://www.googleapis.com/auth/earthengine", + "https://www.googleapis.com/auth/cloud-platform", +] + class CaseSkipped(Exception): """Raised by a case when it cannot run in this environment (e.g. offline).""" +def initialize_earth_engine() -> Any: + """Initialize Earth Engine from Application Default Credentials, or skip. + + Uses the credentials from ``gcloud auth application-default login`` (with the + Earth Engine scope) and the ADC project — so no separate ``earthengine + authenticate`` OAuth flow is needed, which also sidesteps the "this app is + blocked" error some org policies raise. Override the project with the + ``EARTHENGINE_PROJECT`` environment variable. Returns the initialized ``ee`` + module; raises :class:`CaseSkipped` if EE is unavailable or unauthenticated. + """ + try: + import ee + import google.auth + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api`" + ) from exc + try: + credentials, adc_project = google.auth.default(scopes=_EE_SCOPES) + ee.Initialize( + credentials, + project=os.environ.get("EARTHENGINE_PROJECT") or adc_project, + opt_url="https://earthengine-highvolume.googleapis.com", + ) + except Exception as exc: # noqa: BLE001 — not authenticated → skip + raise CaseSkipped( + f"Earth Engine not initialized ({exc}); run " + "`gcloud auth application-default login` (or set EARTHENGINE_PROJECT)" + ) from exc + return ee + + def banner(text: str) -> None: """Print a titled section divider.""" print(f"\n{'─' * _WIDTH}") From 1de03b82fbfddb9cea31653dbe0e501c80dbb079 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 24 Jun 2026 11:28:58 +0300 Subject: [PATCH 08/26] Fix run_all.sh: run from any directory, run all cases, use the local xarray-sql MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Resolve the script's own directory (BASH_SOURCE) so it works from any cwd — the old `uv run ./01_nvdi.py` only worked from benchmarks/geospatial/ (and had a typo: 01_nvdi → 01_ndvi). - Glob all NN_*.py cases instead of an incomplete hand-written list; report a per-case PASS/FAIL summary and exit non-zero if any case fails. - Hand `uv` the local checkout via `--with-editable $REPO_ROOT`: these cases use xarray-sql features (XarrayDataFrame.to_dataset) newer than the latest PyPI release, so plain `uv run` would resolve an xarray-sql without `to_dataset`. - 01_ndvi.py: declare aiohttp + requests (needed by pystac-client / zarr-over- HTTPS under uv's isolated environment). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01AWvrZYAT2NbuETBqNAN3o9 --- benchmarks/geospatial/01_ndvi.py | 2 ++ benchmarks/geospatial/run_all.sh | 39 ++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) create mode 100755 benchmarks/geospatial/run_all.sh diff --git a/benchmarks/geospatial/01_ndvi.py b/benchmarks/geospatial/01_ndvi.py index 9f42497..b03d23b 100644 --- a/benchmarks/geospatial/01_ndvi.py +++ b/benchmarks/geospatial/01_ndvi.py @@ -4,6 +4,8 @@ # dependencies = [ # "xarray-sql", # "xarray", +# "aiohttp", +# "requests", # "pystac-client", # "zarr>=3", # "numpy", diff --git a/benchmarks/geospatial/run_all.sh b/benchmarks/geospatial/run_all.sh new file mode 100755 index 0000000..138fedc --- /dev/null +++ b/benchmarks/geospatial/run_all.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# +# Run every geospatial benchmark case with `uv run` (each script declares its +# own dependencies via PEP 723 inline metadata). Works from any directory: it +# resolves its own location, so the cases are found and the paths handed to +# `uv run` are absolute. +# +# ./run_all.sh # from anywhere +# bash benchmarks/geospatial/run_all.sh +# +# These cases use xarray-sql features (e.g. XarrayDataFrame.to_dataset) that may +# be newer than the latest PyPI release, so we hand `uv` the *local* checkout +# with --with-editable rather than letting it resolve xarray-sql from PyPI. Once +# a release ships those features, plain `uv run