xarray_sql 0.3.0__tar.gz → 0.3.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (44) hide show
  1. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/Cargo.lock +1 -1
  2. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/Cargo.toml +6 -2
  3. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/PKG-INFO +81 -4
  4. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/README.md +80 -3
  5. xarray_sql-0.3.1/benchmarks/geospatial/01_ndvi.py +149 -0
  6. xarray_sql-0.3.1/benchmarks/geospatial/02_climatology.py +137 -0
  7. xarray_sql-0.3.1/benchmarks/geospatial/03_zonal_mean.py +132 -0
  8. xarray_sql-0.3.1/benchmarks/geospatial/04_anomaly.py +144 -0
  9. xarray_sql-0.3.1/benchmarks/geospatial/05_forecast_skill.py +200 -0
  10. xarray_sql-0.3.1/benchmarks/geospatial/06_zonal_vector.py +175 -0
  11. xarray_sql-0.3.1/benchmarks/geospatial/07_reproject_udf.py +196 -0
  12. xarray_sql-0.3.1/benchmarks/geospatial/08_regrid_weights.py +228 -0
  13. xarray_sql-0.3.1/benchmarks/geospatial/09_warp.py +292 -0
  14. xarray_sql-0.3.1/benchmarks/geospatial/README.md +106 -0
  15. xarray_sql-0.3.1/benchmarks/geospatial/_harness.py +276 -0
  16. xarray_sql-0.3.1/benchmarks/geospatial/perf_summary.py +103 -0
  17. xarray_sql-0.3.1/benchmarks/geospatial/run_all.sh +37 -0
  18. xarray_sql-0.3.1/benchmarks/geospatial/run_perf.sh +44 -0
  19. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/docs/examples.md +26 -6
  20. xarray_sql-0.3.1/docs/geospatial.md +474 -0
  21. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/test_ds.py +104 -0
  22. xarray_sql-0.3.1/tests/test_to_dataset_perf.py +115 -0
  23. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/ds.py +102 -14
  24. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/zensical.toml +1 -0
  25. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/.gitignore +0 -0
  26. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/LICENSE +0 -0
  27. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/docs/assets/logo.svg +0 -0
  28. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/docs/contributing.md +0 -0
  29. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/docs/index.md +0 -0
  30. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/docs/reference/xarray_sql.md +0 -0
  31. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/pyproject.toml +0 -0
  32. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/src/lib.rs +0 -0
  33. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/__init__.py +0 -0
  34. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/conftest.py +0 -0
  35. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/test_cft.py +0 -0
  36. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/test_df.py +0 -0
  37. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/test_reader.py +0 -0
  38. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/tests/test_sql.py +0 -0
  39. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/__init__.py +0 -0
  40. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/cftime.py +0 -0
  41. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/core.py +0 -0
  42. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/df.py +0 -0
  43. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/reader.py +0 -0
  44. {xarray_sql-0.3.0 → xarray_sql-0.3.1}/xarray_sql/sql.py +0 -0
@@ -3375,7 +3375,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb"
3375
3375
 
3376
3376
  [[package]]
3377
3377
  name = "xarray_sql"
3378
- version = "0.3.0"
3378
+ version = "0.3.1"
3379
3379
  dependencies = [
3380
3380
  "arrow",
3381
3381
  "async-stream",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "xarray_sql"
3
- version = "0.3.0"
3
+ version = "0.3.1"
4
4
  authors = ["Alex Merose"]
5
5
  edition = "2021"
6
6
  exclude = [
@@ -25,7 +25,11 @@ async-trait = "0.1"
25
25
  datafusion = { version = "52.0.0" }
26
26
  datafusion-ffi = { version = "52.0.0" }
27
27
  futures = { version = "0.3" }
28
- pyo3 = { version = "0.26.0", features = ["extension-module"] }
28
+ # `abi3-py310` builds against CPython's stable ABI, so a single wheel per
29
+ # platform works on all CPython >= 3.10 (matching `requires-python`). This
30
+ # lets the release workflow ship pre-built wheels for every interpreter
31
+ # without compiling per-version, avoiding local rebuilds on install.
32
+ pyo3 = { version = "0.26.0", features = ["extension-module", "abi3-py310"] }
29
33
  tokio = { version = "1.46.1", features = ["rt"] }
30
34
 
31
35
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: xarray_sql
3
- Version: 0.3.0
3
+ Version: 0.3.1
4
4
  Classifier: Development Status :: 4 - Beta
5
5
  Classifier: Intended Audience :: Science/Research
6
6
  Classifier: Intended Audience :: Developers
@@ -57,6 +57,57 @@ pip install xarray-sql
57
57
  ## What is this?
58
58
 
59
59
  This is an experiment to provide a SQL interface for array datasets.
60
+ Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run
61
+ SQL queries against them.
62
+
63
+ ## Quickstart
64
+
65
+ Open a Dataset, register it as a table with `from_dataset`, compute a
66
+ climatology in SQL, then write the result back to Xarray and plot it:
67
+
68
+ > **Note:** this example also needs `pooch` and a netCDF backend (for the
69
+ > tutorial download) and `matplotlib` (for the plot):
70
+ > `pip install pooch netCDF4 matplotlib`.
71
+
72
+ ```python
73
+ import xarray as xr
74
+ import xarray_sql as xql
75
+
76
+ # 4x-daily surface air temperature on a lat/lon grid, 2013-2014.
77
+ ds = xr.tutorial.open_dataset('air_temperature')
78
+
79
+ ctx = xql.XarrayContext()
80
+ ctx.from_dataset('air', ds, chunks=dict(time=100))
81
+
82
+ # A climatology — the mean annual cycle — computed in SQL: average air
83
+ # temperature for each month of the year, over all grid cells and years.
84
+ clim = ctx.sql('''
85
+ SELECT
86
+ CAST(date_part('month', "time") AS INTEGER) AS month,
87
+ AVG("air") AS air
88
+ FROM "air"
89
+ GROUP BY CAST(date_part('month', "time") AS INTEGER)
90
+ ORDER BY month
91
+ ''')
92
+
93
+ # Write the SQL result back to an Xarray Dataset. `month` is a derived
94
+ # column, so name it as the dimension; the variable's units are recovered
95
+ # from the registered table. The result is one value per month: air(month).
96
+ clim_ds = clim.to_dataset(dims=["month"])
97
+
98
+ # Plot the annual cycle as a time series.
99
+ clim_ds["air"].plot() # in a script, call matplotlib.pyplot.show() to display
100
+ ```
101
+
102
+ That's the round trip — Xarray in, SQL in the middle, Xarray (and a plot) back
103
+ out.
104
+
105
+ ## A bigger example: ARCO-ERA5
106
+
107
+ The same interface scales to cloud-native datasets with hundreds of variables,
108
+ like [ARCO-ERA5](https://github.com/google-research/arco-era5).
109
+
110
+ > **Note:** reading from `gs://` requires `gcsfs` (`pip install gcsfs`).
60
111
 
61
112
  ```python
62
113
  import xarray as xr
@@ -148,9 +199,6 @@ ctx.sql('''
148
199
  _(A runnable version of this example lives at
149
200
  [`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_
150
201
 
151
- Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run
152
- SQL queries against them.
153
-
154
202
  ## Why build this?
155
203
 
156
204
  A few reasons:
@@ -185,6 +233,35 @@ that lets the DB engine translate the underlying Dataset arrays into DataFusion
185
233
  Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be
186
234
  translated into a 2D table -- underlies this performant query mechanism.
187
235
 
236
+ ## Does it work?
237
+
238
+ Yes. The recurring worry is that the SQL interface is a toy — fine for `SELECT`s,
239
+ but not for the operations geoscience actually runs. So we wrote a suite that
240
+ takes the staples of geospatial and climate analysis — the ones we assume *need*
241
+ an array library — and expresses each one in SQL, then **checks the SQL answer
242
+ against an xarray/array reference** to floating-point tolerance:
243
+
244
+ * **Spectral indices** (NDVI) — column arithmetic over a real Sentinel-2 scene.
245
+ * **Climatology, anomalies, zonal means** — `GROUP BY` and self-`JOIN` against
246
+ the 0.25° **ARCO-ERA5** archive registered as a lazy table. Each query is
247
+ bounded to a small window (a few days over a region) and reads only that
248
+ slice — the point is that you can aim a query at a multi-decade archive and
249
+ pay only for the data it asks for, not that the query scans the whole record.
250
+ * **Forecast skill** — scoring the **Pangu-Weather** and **GraphCast** ML models
251
+ against ERA5 (WeatherBench 2) as a `JOIN` on `valid_time = init + lead`; it
252
+ reproduces the published result that GraphCast beats Pangu at every lead.
253
+ * **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a
254
+ table of regions.
255
+ * **Reprojection and regridding** — a scalar PROJ UDF (validated against Earth
256
+ Engine's own geodesy via [Xee](https://github.com/google/Xee)) and a
257
+ sparse-weight-table `JOIN` (regridding real SRTM terrain).
258
+
259
+ Every case matches its array reference. The headline finding: these operations
260
+ are not really "array" operations at all — they are `GROUP BY`, `JOIN`, window
261
+ functions, and `CASE` in disguise, and a query engine runs them at scale. See
262
+ [`benchmarks/geospatial/`](benchmarks/geospatial/) and the write-up,
263
+ [Geospatial operations are relational operations](docs/geospatial.md).
264
+
188
265
  ## Why does this work?
189
266
 
190
267
  Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in
@@ -14,6 +14,57 @@ pip install xarray-sql
14
14
  ## What is this?
15
15
 
16
16
  This is an experiment to provide a SQL interface for array datasets.
17
+ Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run
18
+ SQL queries against them.
19
+
20
+ ## Quickstart
21
+
22
+ Open a Dataset, register it as a table with `from_dataset`, compute a
23
+ climatology in SQL, then write the result back to Xarray and plot it:
24
+
25
+ > **Note:** this example also needs `pooch` and a netCDF backend (for the
26
+ > tutorial download) and `matplotlib` (for the plot):
27
+ > `pip install pooch netCDF4 matplotlib`.
28
+
29
+ ```python
30
+ import xarray as xr
31
+ import xarray_sql as xql
32
+
33
+ # 4x-daily surface air temperature on a lat/lon grid, 2013-2014.
34
+ ds = xr.tutorial.open_dataset('air_temperature')
35
+
36
+ ctx = xql.XarrayContext()
37
+ ctx.from_dataset('air', ds, chunks=dict(time=100))
38
+
39
+ # A climatology — the mean annual cycle — computed in SQL: average air
40
+ # temperature for each month of the year, over all grid cells and years.
41
+ clim = ctx.sql('''
42
+ SELECT
43
+ CAST(date_part('month', "time") AS INTEGER) AS month,
44
+ AVG("air") AS air
45
+ FROM "air"
46
+ GROUP BY CAST(date_part('month', "time") AS INTEGER)
47
+ ORDER BY month
48
+ ''')
49
+
50
+ # Write the SQL result back to an Xarray Dataset. `month` is a derived
51
+ # column, so name it as the dimension; the variable's units are recovered
52
+ # from the registered table. The result is one value per month: air(month).
53
+ clim_ds = clim.to_dataset(dims=["month"])
54
+
55
+ # Plot the annual cycle as a time series.
56
+ clim_ds["air"].plot() # in a script, call matplotlib.pyplot.show() to display
57
+ ```
58
+
59
+ That's the round trip — Xarray in, SQL in the middle, Xarray (and a plot) back
60
+ out.
61
+
62
+ ## A bigger example: ARCO-ERA5
63
+
64
+ The same interface scales to cloud-native datasets with hundreds of variables,
65
+ like [ARCO-ERA5](https://github.com/google-research/arco-era5).
66
+
67
+ > **Note:** reading from `gs://` requires `gcsfs` (`pip install gcsfs`).
17
68
 
18
69
  ```python
19
70
  import xarray as xr
@@ -105,9 +156,6 @@ ctx.sql('''
105
156
  _(A runnable version of this example lives at
106
157
  [`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_
107
158
 
108
- Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run
109
- SQL queries against them.
110
-
111
159
  ## Why build this?
112
160
 
113
161
  A few reasons:
@@ -142,6 +190,35 @@ that lets the DB engine translate the underlying Dataset arrays into DataFusion
142
190
  Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be
143
191
  translated into a 2D table -- underlies this performant query mechanism.
144
192
 
193
+ ## Does it work?
194
+
195
+ Yes. The recurring worry is that the SQL interface is a toy — fine for `SELECT`s,
196
+ but not for the operations geoscience actually runs. So we wrote a suite that
197
+ takes the staples of geospatial and climate analysis — the ones we assume *need*
198
+ an array library — and expresses each one in SQL, then **checks the SQL answer
199
+ against an xarray/array reference** to floating-point tolerance:
200
+
201
+ * **Spectral indices** (NDVI) — column arithmetic over a real Sentinel-2 scene.
202
+ * **Climatology, anomalies, zonal means** — `GROUP BY` and self-`JOIN` against
203
+ the 0.25° **ARCO-ERA5** archive registered as a lazy table. Each query is
204
+ bounded to a small window (a few days over a region) and reads only that
205
+ slice — the point is that you can aim a query at a multi-decade archive and
206
+ pay only for the data it asks for, not that the query scans the whole record.
207
+ * **Forecast skill** — scoring the **Pangu-Weather** and **GraphCast** ML models
208
+ against ERA5 (WeatherBench 2) as a `JOIN` on `valid_time = init + lead`; it
209
+ reproduces the published result that GraphCast beats Pangu at every lead.
210
+ * **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a
211
+ table of regions.
212
+ * **Reprojection and regridding** — a scalar PROJ UDF (validated against Earth
213
+ Engine's own geodesy via [Xee](https://github.com/google/Xee)) and a
214
+ sparse-weight-table `JOIN` (regridding real SRTM terrain).
215
+
216
+ Every case matches its array reference. The headline finding: these operations
217
+ are not really "array" operations at all — they are `GROUP BY`, `JOIN`, window
218
+ functions, and `CASE` in disguise, and a query engine runs them at scale. See
219
+ [`benchmarks/geospatial/`](benchmarks/geospatial/) and the write-up,
220
+ [Geospatial operations are relational operations](docs/geospatial.md).
221
+
145
222
  ## Why does this work?
146
223
 
147
224
  Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in
@@ -0,0 +1,149 @@
1
+ #!/usr/bin/env python3
2
+ # /// script
3
+ # requires-python = ">=3.11"
4
+ # dependencies = [
5
+ # "xarray-sql",
6
+ # "xarray",
7
+ # "aiohttp",
8
+ # "requests",
9
+ # "pystac-client",
10
+ # "zarr>=3",
11
+ # "numpy",
12
+ # ]
13
+ #
14
+ # [tool.uv.sources]
15
+ # xarray-sql = { path = "../../", editable = true }
16
+ # ///
17
+ """NDVI — "apply_ufunc over a raster" is just column arithmetic.
18
+
19
+ The Normalized Difference Vegetation Index is the workhorse of optical remote
20
+ sensing: ``NDVI = (NIR - Red) / (NIR + Red)``, computed per pixel. The array
21
+ paradigm reaches for ``xarray.apply_ufunc`` (the coiled/benchmarks #1545
22
+ "vectorized operations" case) to broadcast this over a whole scene.
23
+
24
+ But a per-pixel formula over two bands is just *column arithmetic over two
25
+ columns*::
26
+
27
+ SELECT x, y, (nir - red) / (nir + red) AS ndvi
28
+ FROM scene
29
+ ORDER BY y, x
30
+
31
+ Each pixel is one row; the ufunc is the SELECT expression. Invalid pixels are
32
+ already NaN (xarray decodes the band's ``_FillValue`` on open), and NaN
33
+ propagates through the arithmetic on both sides — so the masking is free, no
34
+ ``CASE`` required.
35
+
36
+ Dataset: a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample
37
+ service, discovered with ``pystac-client`` and opened the canonical way with
38
+ ``xarray`` — ``xr.open_datatree`` yields the reflectance bands (B04=red,
39
+ B08=NIR at 10 m) already scaled to reflectance and carrying their ``x``/``y``
40
+ coordinates. We read one window so the case stays bounded. Requires network;
41
+ skips cleanly if the service is offline.
42
+ """
43
+
44
+ from __future__ import annotations
45
+
46
+ import xarray as xr
47
+
48
+ import xarray_sql as xql
49
+
50
+ from _harness import (
51
+ CaseSkipped,
52
+ assert_grid_close,
53
+ measured,
54
+ run_case,
55
+ show_result,
56
+ show_sql,
57
+ )
58
+
59
+ # EOPF sample-service STAC catalog; an agricultural AOI near Torino, Italy, in
60
+ # early May (peak spring growth). The search is deterministic — it resolves to
61
+ # a specific archived Sentinel-2 product.
62
+ _STAC = "https://stac.core.eopf.eodc.eu"
63
+ _BBOX = [7.2, 44.5, 7.4, 44.7]
64
+ _DATETIME = "2025-04-25/2025-05-05"
65
+
66
+ # A 1024×1024 (~105 km²) window over vegetated valley floor.
67
+ _Y0, _X0, _N = 4_000, 6_000, 1_024
68
+
69
+
70
+ def _load_scene() -> tuple[xr.Dataset, str]:
71
+ """Discover a Sentinel-2 L2A product and open its 10 m red/NIR bands.
72
+
73
+ Idiomatic end to end: ``pystac-client`` finds the product, ``open_datatree``
74
+ opens the hierarchical EOPF Zarr, and the ``reflectance/r10m`` node already
75
+ carries B04/B08 scaled to reflectance (nodata decoded to NaN) with
76
+ ``x``/``y`` coordinates — no manual scaling or coordinate reconstruction.
77
+ """
78
+ try:
79
+ from pystac_client import Client
80
+
81
+ catalog = Client.open(_STAC)
82
+ search = catalog.search(
83
+ collections=["sentinel-2-l2a"],
84
+ bbox=_BBOX,
85
+ datetime=_DATETIME,
86
+ max_items=1,
87
+ )
88
+ item = next(search.items())
89
+ tree = xr.open_datatree(
90
+ item.assets["product"].href, engine="zarr", chunks={}
91
+ )
92
+ except StopIteration as exc:
93
+ raise CaseSkipped("no Sentinel-2 product found for the query") from exc
94
+ except Exception as exc: # noqa: BLE001 — any failure → skip, not crash
95
+ raise CaseSkipped(f"EOPF Sentinel-2 unavailable ({exc})") from exc
96
+
97
+ r10m = tree["measurements/reflectance/r10m"].to_dataset()
98
+ scene = (
99
+ r10m[["b04", "b08"]]
100
+ .rename(b04="red", b08="nir")
101
+ .isel(y=slice(_Y0, _Y0 + _N), x=slice(_X0, _X0 + _N))
102
+ )
103
+ return scene, item.id
104
+
105
+
106
+ def main() -> None:
107
+ scene, item_id = _load_scene()
108
+ n = scene.sizes["y"] * scene.sizes["x"]
109
+ print(f" Sentinel-2 L2A {item_id}")
110
+ print(
111
+ f" scene window: {dict(scene.sizes)} ({n:,} pixels, B04=red/B08=NIR)"
112
+ )
113
+
114
+ ctx = xql.XarrayContext()
115
+ ctx.from_dataset("scene", scene, chunks={"y": 256, "x": 256})
116
+
117
+ sql = """
118
+ SELECT x, y, (nir - red) / (nir + red) AS ndvi
119
+ FROM scene
120
+ ORDER BY y, x
121
+ """
122
+ show_sql(sql)
123
+
124
+ for _ in measured("SQL NDVI"):
125
+ got = ctx.sql(sql).to_dataset(dims=["y", "x"]).ndvi
126
+
127
+ # Array reference: the same formula in pure xarray. ``.compute()`` reads the
128
+ # window and evaluates it here (the scene is lazy), so this measures the same
129
+ # read-and-compute the SQL side does — not just graph construction.
130
+ for _ in measured("xarray reference"):
131
+ ref = ((scene.nir - scene.red) / (scene.nir + scene.red)).compute()
132
+
133
+ # Compare the xarray way — aligned by coordinate label, so the ORDER BY
134
+ # above is enough and neither side needs an explicit sort.
135
+ assert_grid_close("NDVI (per-pixel)", got, ref, rtol=1e-6)
136
+
137
+ show_result(got)
138
+
139
+ valid = ref.notnull()
140
+ print(
141
+ f"\n NDVI over {int(valid.sum()):,} valid pixels: "
142
+ f"min {float(ref.min()):.3f}, "
143
+ f"mean {float(ref.mean()):.3f}, "
144
+ f"max {float(ref.max()):.3f}"
145
+ )
146
+
147
+
148
+ if __name__ == "__main__":
149
+ raise SystemExit(run_case(main, "NDVI: per-pixel column arithmetic"))
@@ -0,0 +1,137 @@
1
+ # /// script
2
+ # requires-python = ">=3.11"
3
+ # dependencies = [
4
+ # "xarray-sql",
5
+ # "xarray",
6
+ # "gcsfs",
7
+ # "zarr>=3",
8
+ # ]
9
+ #
10
+ # [tool.uv.sources]
11
+ # xarray-sql = { path = "../../", editable = true }
12
+ # ///
13
+ """Diurnal climatology — the "rechunk + grouped reduction" that is a GROUP BY.
14
+
15
+ A *climatology* is the average value for each time-of-cycle, computed
16
+ independently at every location: "what is the typical temperature here at
17
+ 06:00?" In the array paradigm (and in the coiled/benchmarks #1545 write-up)
18
+ this is the canonical painful workload — load native Zarr chunks, *rechunk* to
19
+ put all of time in one chunk ("pencils"), run a grouped reduction over the
20
+ calendar, then rechunk back to "pancakes" for output.
21
+
22
+ The rechunking exists only to serve the array layout. The *operation* is::
23
+
24
+ SELECT latitude, longitude, hour_of_day, AVG("2m_temperature")
25
+ GROUP BY latitude, longitude, hour_of_day
26
+
27
+ Group by location and time-of-cycle, average the rest — the same answer as
28
+ ``da.groupby("time.hour").mean()``. ERA5 is hourly, so grouping by hour of day
29
+ gives a clean 24-bin **diurnal cycle**, one sample per day in the window.
30
+
31
+ We register the full ARCO-ERA5 archive as a lazy table, but the climatology here
32
+ is computed over a *bounded window* — a few summer days over a CONUS-ish box. The
33
+ ``WHERE`` prunes the read, so the query touches only ``2m_temperature`` over that
34
+ window and never scans the rest of the archive. The point is not that we reduce
35
+ the whole record; it is that you can aim a query at a multi-decade archive and pay
36
+ only for the slice it asks for.
37
+ """
38
+
39
+ from __future__ import annotations
40
+
41
+ import datetime
42
+
43
+ import xarray as xr
44
+
45
+ import xarray_sql as xql
46
+
47
+ from _harness import (
48
+ CaseSkipped,
49
+ assert_grid_close,
50
+ measured,
51
+ run_case,
52
+ show_result,
53
+ show_sql,
54
+ timed,
55
+ )
56
+
57
+ _URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
58
+ # A few days over a CONUS-ish box (ERA5 latitude descends; lon is 0–360°E).
59
+ _START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23)
60
+ _LAT_N, _LAT_S = 50.0, 25.0
61
+ _LON_W, _LON_E = 235.0, 290.0
62
+ _PARAMS = {
63
+ "start": _START,
64
+ "end": _END,
65
+ "lat_s": _LAT_S,
66
+ "lat_n": _LAT_N,
67
+ "lon_w": _LON_W,
68
+ "lon_e": _LON_E,
69
+ }
70
+
71
+
72
+ def main() -> None:
73
+ # Open the full ARCO-ERA5 archive lazily — no data is read here. ERA5 mixes
74
+ # surface (time, lat, lon) and atmospheric (… level …) variables, so register
75
+ # it as two tables under an ``era5`` schema; the query below touches only the
76
+ # surface table's 2m_temperature.
77
+ try:
78
+ import gcsfs # noqa: F401 — required by the gs:// protocol
79
+
80
+ ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})
81
+ except Exception as exc: # noqa: BLE001 — any failure → skip, not crash
82
+ raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc
83
+
84
+ ctx = xql.XarrayContext()
85
+ with timed("register full ERA5 (lazy)"):
86
+ ctx.from_dataset(
87
+ "era5",
88
+ ds,
89
+ chunks={"time": 6},
90
+ table_names={
91
+ ("time", "latitude", "longitude"): "surface",
92
+ ("time", "level", "latitude", "longitude"): "atmosphere",
93
+ },
94
+ )
95
+
96
+ sql = """
97
+ SELECT latitude,
98
+ longitude,
99
+ date_part('hour', time) AS hour,
100
+ AVG("2m_temperature") - 273.15 AS clim_c
101
+ FROM era5.surface
102
+ WHERE time BETWEEN $start AND $end
103
+ AND latitude BETWEEN $lat_s AND $lat_n
104
+ AND longitude BETWEEN $lon_w AND $lon_e
105
+ GROUP BY latitude, longitude, date_part('hour', time)
106
+ ORDER BY latitude DESC, longitude, hour
107
+ """
108
+ show_sql(sql)
109
+
110
+ # A climatology is a gridded product: round-trip the result back to an
111
+ # xarray Dataset keyed by (latitude, longitude, hour) — how it is used.
112
+ for _ in measured("SQL diurnal climatology (lazy read)"):
113
+ got = ctx.sql(sql, param_values=_PARAMS).to_dataset(
114
+ dims=["latitude", "longitude", "hour"]
115
+ )
116
+
117
+ # Array reference: the textbook groupby-over-the-cycle reduction, in °C —
118
+ # the same lazy window, materialized only on demand.
119
+ for _ in measured("xarray reference"):
120
+ window = ds["2m_temperature"].sel(
121
+ time=slice(_START, _END),
122
+ latitude=slice(_LAT_N, _LAT_S),
123
+ longitude=slice(_LON_W, _LON_E),
124
+ )
125
+ ref = window.groupby("time.hour").mean("time") - 273.15
126
+
127
+ assert_grid_close(
128
+ "diurnal climatology (°C)", got.clim_c, ref, rtol=1e-4, atol=1e-2
129
+ )
130
+
131
+ show_result(got)
132
+
133
+
134
+ if __name__ == "__main__":
135
+ raise SystemExit(
136
+ run_case(main, "Climatology: GROUP BY lat, lon, hour (ARCO-ERA5)")
137
+ )
@@ -0,0 +1,132 @@
1
+ # /// script
2
+ # requires-python = ">=3.11"
3
+ # dependencies = [
4
+ # "xarray-sql",
5
+ # "xarray",
6
+ # "gcsfs",
7
+ # "zarr>=3",
8
+ # ]
9
+ #
10
+ # [tool.uv.sources]
11
+ # xarray-sql = { path = "../../", editable = true }
12
+ # ///
13
+ """Zonal mean — the array reduction that is secretly a GROUP BY.
14
+
15
+ A *zonal mean* averages a field around each circle of latitude (over all
16
+ longitudes, and here over a day of hours too), collapsing a 3-D field to a 1-D
17
+ profile of value-vs-latitude — the classic pole-to-pole temperature curve. In
18
+ the array paradigm this is ``da.mean(dim=["longitude", "time"])``, a reduction
19
+ over two axes.
20
+
21
+ Relationally it is nothing more than::
22
+
23
+ SELECT latitude, AVG("2m_temperature") GROUP BY latitude
24
+
25
+ The "axes" we reduce over are just the columns we *don't* group by. Same answer,
26
+ and the SQL reads like the plain-English definition of a zonal mean.
27
+
28
+ Dataset: the full **ARCO-ERA5** archive (0.25° global, 1.3M hourly timesteps).
29
+ The table is the whole reanalysis; ``WHERE time …`` prunes it to one day, and
30
+ the GROUP BY produces a 721-point global temperature profile.
31
+ """
32
+
33
+ from __future__ import annotations
34
+
35
+ import datetime
36
+
37
+ import xarray as xr
38
+
39
+ import xarray_sql as xql
40
+
41
+ from _harness import (
42
+ CaseSkipped,
43
+ assert_grid_close,
44
+ measured,
45
+ run_case,
46
+ show_result,
47
+ show_sql,
48
+ timed,
49
+ )
50
+
51
+ _URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
52
+ # One day of hourly data, global; the WHERE below prunes ERA5 to this window.
53
+ _DAY = "2020-06-01"
54
+ _START, _END = (
55
+ datetime.datetime(2020, 6, 1, 0),
56
+ datetime.datetime(2020, 6, 1, 23),
57
+ )
58
+
59
+
60
+ def main() -> None:
61
+ # Open the full ARCO-ERA5 archive (lazy, dask off) — no slicing here; the
62
+ # SQL WHERE clause prunes it to the window we ask for.
63
+ try:
64
+ import gcsfs # noqa: F401 — required by the gs:// protocol
65
+
66
+ ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})
67
+ except Exception as exc: # noqa: BLE001 — any failure → skip, not crash
68
+ raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc
69
+
70
+ print(
71
+ f" ARCO-ERA5: {ds.sizes['time']:,} hourly timesteps, "
72
+ f"{ds.sizes['latitude']}×{ds.sizes['longitude']} grid, "
73
+ f"{len(ds.data_vars)} variables (no pre-slicing)"
74
+ )
75
+
76
+ # ERA5 mixes surface (time, lat, lon) and atmospheric (… level …) variables,
77
+ # so register it as two tables under an ``era5`` schema.
78
+ ctx = xql.XarrayContext()
79
+ with timed("register full ERA5"):
80
+ ctx.from_dataset(
81
+ "era5",
82
+ ds,
83
+ chunks={"time": 6},
84
+ table_names={
85
+ ("time", "latitude", "longitude"): "surface",
86
+ ("time", "level", "latitude", "longitude"): "atmosphere",
87
+ },
88
+ )
89
+
90
+ # Pass the day's bounds as query parameters; the query still reads only that
91
+ # one day out of the whole archive.
92
+ sql = """
93
+ SELECT latitude,
94
+ AVG("2m_temperature") - 273.15 AS air_mean_c
95
+ FROM era5.surface
96
+ WHERE time BETWEEN $start AND $end
97
+ GROUP BY latitude
98
+ ORDER BY latitude DESC
99
+ """
100
+ show_sql(sql)
101
+
102
+ # Round-trip the profile back to an xarray Dataset keyed by latitude.
103
+ for _ in measured("SQL zonal mean (reads one day)"):
104
+ got = ctx.sql(
105
+ sql, param_values={"start": _START, "end": _END}
106
+ ).to_dataset(dims=["latitude"])
107
+
108
+ # Array reference: reduce the same day over the two un-grouped axes.
109
+ for _ in measured("xarray reference"):
110
+ ref = (
111
+ ds["2m_temperature"].sel(time=_DAY).mean(["longitude", "time"])
112
+ - 273.15
113
+ )
114
+
115
+ assert_grid_close(
116
+ "zonal mean (2m_temp vs latitude, °C)",
117
+ got.air_mean_c,
118
+ ref,
119
+ rtol=1e-4,
120
+ atol=1e-3,
121
+ )
122
+
123
+ show_result(got)
124
+
125
+ print("\n Global temperature profile (every 72nd parallel, °C):")
126
+ print(got.air_mean_c.isel(latitude=slice(None, None, 72)).to_series())
127
+
128
+
129
+ if __name__ == "__main__":
130
+ raise SystemExit(
131
+ run_case(main, "Zonal mean: GROUP BY latitude (ARCO-ERA5)")
132
+ )