Geospatial SQL benchmark suite: prove core geo/climate ops are relational

[alxmrs]

What & why

A suite of worked examples in benchmarks/geospatial/ arguing a deliberately radical thesis: 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 reprojection and regridding map onto ordinary SQL (GROUP BY, JOIN, window functions, CASE, and the occasional scalar UDF).

Every case poses the operation in SQL against xarray-sql, round-trips the result back to an xarray.Dataset/DataArray (via to_dataset), and asserts it matches an xarray reference (coordinate-aligned xr.testing.assert_allclose). SQL in, an array out, verified — and each run prints the result's standard xarray repr. The headline is correctness + clarity of the SQL, not raw speed.

The README gains a "Does it work?" section summarizing the findings (addresses #46), and a full write-up lives in docs/geospatial.md (wired into the docs nav).

The cases (all validated end-to-end on real data)

#	Operation	Relational form	Data + reference
01	NDVI	column arithmetic (nodata auto-masked on decode)	real Sentinel-2 L2A Zarr (ESA EOPF, via pystac-client + xr.open_datatree)
02	Diurnal climatology	GROUP BY lat, lon, hour	ARCO-ERA5
03	Zonal mean	GROUP BY latitude	full ARCO-ERA5, WHERE-pruned
04	Anomaly	climatology CTE self-JOIN	ARCO-ERA5
05	Forecast skill	forecast↔truth JOIN on valid_time = init + lead	WeatherBench 2 Pangu & GraphCast vs ERA5
06	Zonal stats over regions	raster × vector range JOIN	full ARCO-ERA5 + continental boxes
07	Reprojection	PROJ scalar UDF (ST_Transform-style)	Earth Engine via Xee, validated against EE's own pixelLonLat
08	Regridding	sparse weight-table JOIN + weighted GROUP BY	real SRTM elevation via Xee, vs xarray .interp()

Highlights

• It reads naturally to the people who would use it. The array side is idiomatic Xarray; the query side is idiomatic, lazy SQL. You can point a query at the entire ERA5 archive — decades of global hourly data — and it reads only the slice the query asks for.
• Real models, checked against the literature. Case 05 scores the Pangu-Weather and GraphCast forecast models against ERA5 and reproduces the published WeatherBench 2 result that GraphCast leads at every forecast horizon — written as a single JOIN that lines each forecast up with its valid time.
• Validated against independent implementations. Case 07 checks its SQL reprojection against Earth Engine's own coordinate geometry (through Xee) — a genuinely different implementation, not PROJ compared with itself. Case 08 regrids real SRTM terrain and matches Xarray's own interpolation.
• Honest about the limits. Reprojection and regridding are the operations most tied to arrays; the write-up shows where SQL expresses them cleanly and where arrays still earn their keep — computing the interpolation weights (the geometry), which SQL then applies as a join.

Running

Run a single case with uv run benchmarks/geospatial/<case>.py, or the whole suite with benchmarks/geospatial/run_all.sh — from any directory. Cases that need data or credentials you do not have skip cleanly with a clear message: the cloud datasets (ERA5, WeatherBench 2) are read anonymously, and the Earth Engine cases reuse your existing gcloud login. (Until the next xarray-sql release, the scripts run against the version in this branch.)

Performance

Profiled timings — 8 cases on an e2-standard-8 VM in us-central1

Measured cold vs cold with benchmarks/geospatial/run_perf.sh: each case runs
once per fresh process, with no warmup, repeated 5 times. Fairness took some care
— the trap is caching. A reference that calls .load() caches its data in place
on the very object the SQL table also reads from, so a later read (even just running
the reference after the SQL query in the same process) can be served warm. Case 05
loaded shared objects and now uses .compute() instead (fresh array, inputs stay
lazy); the other references either reopen their data or recompute eagerly with
chunks=None (NumPy, no graph to keep warm). Verified directly: reading a window
repeatedly in one process stays flat, and neither side warms the other. Reported are
the median/spread of wall-clock and peak Python-allocator memory across the cold
runs. Hardware: Google Compute Engine e2-standard-8 (8 vCPU, 32 GB), us-central1
— in-region with the ARCO-ERA5 and WeatherBench 2 buckets.

Case	Step	reps	median (s)	stdev (s)	min (s)	max (s)	peak (MB)
01 · NDVI (per-pixel arithmetic)	SQL	5	3.093	0.050	2.998	3.128	98.5
	xarray reference	5	0.402	0.060	0.372	0.507	42.0
02 · Climatology (GROUP BY lat, lon, hour)	SQL	5	5.987	1.328	5.901	8.942	513.8
	xarray reference	5	2.449	0.401	2.332	3.312	43.7
03 · Zonal mean (GROUP BY latitude)	SQL	5	3.224	0.048	3.175	3.288	245.3
	xarray reference	5	0.552	0.020	0.529	0.576	249.5
04 · Anomaly (climatology self-JOIN)	SQL	5	9.479	0.413	9.382	10.280	524.4
	xarray reference	5	3.013	0.584	2.875	4.199	80.5
05 · Forecast skill (forecast↔truth JOIN)	SQL	5	23.241	0.101	23.114	23.348	34.2
	xarray reference	5	0.336	0.014	0.314	0.352	2.2
06 · Zonal stats (raster × vector JOIN)	SQL	5	6.285	0.052	6.191	6.317	509.9
	xarray reference	5	2.235	0.216	2.195	2.703	359.9
07 · Reprojection (PROJ scalar UDF)	SQL	1	0.039	0.000	0.039	0.039	0.3
08 · Regridding (weight-table JOIN)	SQL	5	0.061	0.001	0.060	0.063	0.8
	xarray reference	5	0.018	0.001	0.018	0.020	0.2

• SQL is ctx.sql(...).to_dataset(...) — the cold lazy read from cloud storage plus the relational compute; xarray reference is the array-paradigm computation each case is validated against. SQL is slower because the paradigm has real cost: the grid is exploded to rows, joined with a hash join, and shuttled through Arrow — work the array path does in-place over contiguous buffers.
• Case 05's SQL time is CPU-bound (the join + GIL-held row production) and the most hardware-sensitive number here — a second e2-standard-8 run measured ≈12 s rather than ≈23 s. The reference is read-bound and stable, so read the 05 ratio as "the relational form costs real CPU," not a fixed multiplier.
• Case 01 reads Sentinel-2 from EODC (Europe), the only non-US source, so its time includes a cross-region read.
• Cases 07–08 load their Earth Engine inputs into memory once and then compute, so they are methodology-agnostic (cold ≡ warm) and are reported from a single in-memory run. Case 07's PROJ UDF is not re-entrant across repeated in-process executions, so it is measured once (reps=1).

Notes for review

• ruff and mypy pass via pre-commit. The suite lives outside the package (no __init__.py), so it is not collected by pytest or shipped in the wheel.
• Each script declares its own dependencies inline (PEP 723); nothing is added to pyproject.toml.
• Adds a "Does it work?" section to the README (addresses Add a "Does it work?" section to the README #46) and a longer write-up at docs/geospatial.md.

Follow-ups

• Promote the reprojection UDF (case 07) into the package, alongside the existing cftime UDF.
• Once a release includes the SQL→Xarray round-trip, the cases run straight from PyPI.

🤖 Generated with Claude Code

[alxmrs]

Review of first example

[alxmrs]

Is this really the canonical way to open the dataset in Xarray? Please review the docs and sentinel example to see if we can make this more typical. Is this fixed if we use pystack?

[alxmrs]

We should be able to calculate NDVI in pure xarray after opening the dataset.

[alxmrs]

I belive we can omit this if we use ORDER BY in the SQL statement.

[alxmrs]

Thanks for the review of case 01 — all four points addressed in 6bc41aa:

1. Use pystac for discovery — now uses pystac-client to search the EOPF STAC catalog and resolve the product, instead of a hard-coded object-store href. Added to the script's PEP 723 deps.
2. Idiomatic xarray open — replaced zarr.open_array + manual proj:transform coordinate reconstruction with xr.open_datatree(item.assets["product"].href, engine="zarr") and the measurements/reflectance/r10m node. xarray decodes the bands to reflectance and brings x/y coordinates for free.
3. NDVI in pure xarray — the reference is now a one-liner: (scene.nir - scene.red) / (scene.nir + scene.red). Because xarray decodes the band _FillValue to NaN on open, nodata masking is automatic and the CASE WHEN is gone on both sides.
4. ORDER BY instead of sortby — the SQL uses ORDER BY y, x, and the comparison aligns by coordinate label via xr.testing.assert_allclose(got, ref.reindex_like(got)), so the manual .sortby() on both sides is removed.

Verified end-to-end: 1,048,576 pixels match the xarray reference; mean NDVI 0.235 over the window. Docs (suite README + docs/geospatial.md) updated to match.

One note: xarray-eopf's eopf-zarr engine would let you open the group directly, but it pins zarr<3, which conflicts with this project's zarr 3 — so I kept the plain engine="zarr" + open_datatree path, which needs no extra engine.

[alxmrs]

I think repeating the sel command here is instructive.

[alxmrs]

I recommend a fluent interface

[alxmrs]

In fact, the open_zarr call and the slicing could be one fluent chain.

[alxmrs]

This should be loaded to a dataset. That's how climatologies are typically used.

[alxmrs]

Omit transforming this to a dataframe.

[alxmrs]

I don't like this helper either since repetition here serves as useful documentation.

[alxmrs]

I believe to_pandas() invokes a copy. Is it needed?

[alxmrs]

Round-2 feedback addressed in e1cbc8b (cases 02/03 explicitly, plus the same treatment swept across the suite).

Your comments on 02 / 03:

• Inline the open/slice; the helper hides it — deleted _era5.py. Each case now shows the fluent chain inline, e.g.

  ds = (
      xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})[["2m_temperature"]]
      .sel(time=_TIME, latitude=_LAT, longitude=_LON)
      .load()
  )

  and registration is inline too (03/06 show the full from_dataset(..., table_names=...)).
• ORDER BY latitude DESC, longitude, … — added to the climatology (… , hour) and anomaly queries.
• Time is already hourly — right; the docstring now states ERA5 is hourly, so grouping by hour-of-day gives a clean 24-bin diurnal cycle (one sample per day in the window).
• Load the climatology to a Dataset — the result now round-trips via to_dataset(dims=["latitude","longitude","hour"]) and is returned/used as a gridded xr.Dataset.
• Keep the reference in pure xarray (no to_dataframe) — done: ds["2m_temperature"].groupby("time.hour").mean("time"), compared by coordinate label (no merge, no sortby).
• to_pandas() copy in zonal mean — gone; it returns to_dataset(dims=["latitude"]).

Swept across the suite (anticipating the same notes):

• All gridded results (NDVI, climatology, zonal mean, anomaly, forecast RMSE-by-lead, reprojection, regrid) now round-trip to xr.Dataset/DataArray and are compared with a new assert_grid_close helper that aligns by coordinate label via xr.testing.assert_allclose — so no case sorts manually.
• Forecast skill returns each model's skill curve as a DataArray over lead (a real timedelta coordinate), with an xarray reference.
• Dropped the now-unused check_close harness helper.

Full suite re-run is green (NDVI 1.05M px; climatology 535,704 cells; anomaly 1.6M cells; GraphCast still beats Pangu at every lead). Background note: I evaluated xarray-eopf for case 01's open but it pins zarr<3 (conflicts with this project), so case 01 stays on pystac-client + xr.open_datatree.

Thanks for the careful review — the cases read much more like idiomatic xarray now.

[alxmrs]

I think we should omit the load calls to test out the lazy property of xql. What do you think?

[alxmrs]

I don't think we need this. Xarray-SQL with have a pushdown opt to only select columns we care about.

[alxmrs]

I don't like "this is that real thing"

[alxmrs]

I think at can omit the var selection here too, and elsewhere.

[alxmrs]

This is vulnerable to a SQL injection IIUC. Can we demo this in a way that puts forth best SQL practices?

[alxmrs]

If we don't load up front, then I don't think these timed blocks for measuring loading are necessary. This is good, it demos laziness.

[alxmrs]

Round-3 feedback addressed in a971f2b (cases 04/05 explicitly, extended across the suite).

Your comments:

• 04:54 / 05:90 — don't pre-select the variable; pushdown handles columns. The ERA5 cases (02, 03, 04, 06) now register the whole archive lazily with no variable selection; SELECT "2m_temperature" FROM era5.surface lets projection pushdown read just that column. (05 is the one exception, see below.)
• 05:136 — don't load up front; demo laziness. Nothing is .load()-ed up front anymore — datasets are registered lazily and the query pulls only the partitions/columns it needs at read time. The "read into memory" timed blocks are gone; the references use lazy .sel().
• 05:97 — FROM {table} is a SQL-injection smell. Eliminated the dynamic identifier entirely: the two models are stacked along a model dimension into one forecast table, and the query does GROUP BY f.model, f.prediction_timedelta — no table name formatted into SQL. More broadly, all value filters (dates, region bounds) are now passed as bound query parameters (param_values=...) instead of f-string interpolation. I verified partition pruning still fires with bound parameters (a parameterized one-day WHERE reads one day, not the 1.3M-step archive).
• 05:14 — "this is the real thing". Reworded.

One deliberate exception: 05 still selects 2m_temperature in xarray. Pangu stores 13 pressure levels and GraphCast stores 37, so the full datasets can't be concatenated over model; selecting the common surface field is semantically required to align them, not just an optimization. Noted in a code comment.

Full suite re-run is green (NDVI 1.05M px; climatology 535,704 cells; anomaly 1.6M cells; Sahara 33 °C vs Greenland −8 °C; GraphCast still beats Pangu at every lead). The lazy cases are slower now (each ERA5 case reads its window for the SQL and again for the reference), which is the expected cost of demonstrating laziness rather than pre-loading.

Background note (carried over): xarray-eopf's eopf-zarr engine for case 01 pins zarr<3, conflicting with this project, so case 01 stays on pystac-client + xr.open_datatree.

[alxmrs]

A few last thoughts. The geospatial docs are excellent.

[alxmrs]

I think this paragraph should live in a new section below called "does it work?" and be framed to answer that question. Then add can close this issue: #46

[alxmrs]

I think "does it work" should be between "how does this work" and "why does this work".

[alxmrs]

Omit the word "honest".

[alxmrs]

Round-4 feedback addressed in df67d98:

• README:116 — "Does it work?" section. Moved 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?", and reframed it to answer that question — a findings summary (which operations were validated, on which real datasets, and the headline GraphCast-beats-Pangu result). This addresses Add a "Does it work?" section to the README #46; I'll leave the issue for you to close on merge.
• 07:34 — drop "honest". The 07 docstring already lost that wording in the Xee rewrite; I also reworded the two remaining uses ("Two honest caveats" → "Two caveats", "The honest boundary" → "The boundary") in docs/geospatial.md. No "honest" left in the suite or docs.

README section order is now: How does it work? → Does it work? → Why does this work?

[ahuang11]

Skimmed the PR; wondering if the benchmark timings can be documented (e.g. how long it took on the last run)

[ghostiee-11]

Read through the suite end to end, really like the relational framing and that every case round-trips back to xarray and asserts against a reference. Two small things on the harness and one idea below.

[ghostiee-11]

nit: if the body raises, we skip tracemalloc.stop() and leave it running. small try/finally like timed():

Suggested change

	t0 = time.perf_counter()
	yield
	elapsed = time.perf_counter() - t0
	_, peak = tracemalloc.get_traced_memory()
	tracemalloc.stop()
	t0 = time.perf_counter()
	try:
	yield
	finally:
	elapsed = time.perf_counter() - t0
	_, peak = tracemalloc.get_traced_memory()
	tracemalloc.stop()

[alxmrs]

Latest round of feedback

[alxmrs]

Why do we have this load op here still?

[alxmrs]

Does this storage options value have to be extracted into a constant if we are already extracting the open call to a function?

[alxmrs]

So we need to reorganize the script as a data array? Why do we do this copy? Isn't the output + da access enough?

[ghostiee-11]

built the lazy round-trip case I mentioned and opened it against this branch as #178 (also carries the two harness tweaks). same slab is 1,325 rows / ~0.7 MB lazy vs 3.86M / ~161 MB eager, asserted equal.

[ghostiee-11]

@alxmrs a few cases I'd like to add next, all pulled from the #1545 workflow list so they're community-asked rather than SQL-friendly cherry-picks:

• TEM diagnostic: zonal mean + eddy covariances (U'V', V'T', U'W'), which is just GROUP BY (time, level, lat) plus AVG(x*y) - AVG(x)*AVG(y). dcherian posted the exact xarray for this in #1545; validates against aostools/pyzome on ERA5.
• forecast skill, ACC + MAE/bias: extends case 05 with anomaly correlation, which needs a JOIN to a climatology table to form the anomalies. reuses the gs://weatherbench2 data already in the suite.
• conservative region aggregation: raster onto polygons as a sparse weight-table JOIN + weighted GROUP BY (the carbonplan/extreme-heat workflow). same shape as the regrid case, different weights.
• building-footprint spatial join: bbox range-join + ST_Intersects refinement on Overture. the one #1545 category the suite doesn't cover yet, though it needs spatial UDFs since DataFusion has none natively.
• spectral-index time series (NDWI/NDSI): scales the NDVI case from one scene to a GROUP BY over time on Sentinel-2.

[alxmrs]

Let's not check in this script.

[alxmrs]

Done — removed perf_summary.py and folded its aggregation into run_perf.sh as an inline step, so the perf tooling is one self-contained driver (0941e0b).

[alxmrs]

Is it really this many hourly steps?

[alxmrs]

Did we really do this broad of a query?

[alxmrs]

Good catch — I dropped the “≈1.3M hourly steps” figure; it overstated what the queries touch. Reworded to say the archive is multi-decade and registered as a lazy table, while each query reads only its small window (0941e0b).

[alxmrs]

Right — the queries are bounded to a small window (a few days over a region), not a scan of the archive. Reworded so the README says exactly that (0941e0b).

[alxmrs]

This doesn't appear to be the whole ERA5 archive, because we have a time range. What is the truth?

[alxmrs]

Agreed. Reworded: we register the whole archive as a lazy table, but this climatology is computed over a bounded window — the WHERE prunes the read, so the query never scans the full record (0941e0b).

[alxmrs]

This is not the whole ERA5 archive.

[alxmrs]

Reworded the same way — the table spans the archive lazily; the anomaly is computed over a bounded window the WHERE prunes to (0941e0b).

[alxmrs]

If we use the dataframe code above, then let's print that repr here.

[alxmrs]

Done — the doc now shows the got.to_pandas() repr (RMSE by lead × model), matching the script's new headline (1cd2c93).

[alxmrs]

Let's put this section above results.

[alxmrs]

Done — “Running the suite” now sits above Results (1cd2c93).

[alxmrs]

This doesn't make sense to me. This ERA5 dataset is a 6 hourly step version at a small grid designed to be comparable to the forecasts. By throw away data, I mean we load 100 timesteps in the chunk and then discard the remainder that doesn't match the forecast chunk (as they divide together).

Can you write a small prototype (that's discarded) that tries 6 and 6 chunks ve 6 and 100 chunks to see what the empirical answer is? I'm still not convinced.

[alxmrs]

Since we're not using dask anyway, let's get rid of this load. (Is this essentially for any reason?)

[alxmrs]

I want this to be a valid perf comparison test of laziness for the xr ref too, so let's get rid of the load here.

[alxmrs]

At this point, fine, let's make this a second script.

[alxmrs]

Done — moved the aggregation back into perf_summary.py as a second script (the inline heredoc was worse to read) (7bfd56f).