xarray_sql 0.2.3__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.
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/Cargo.lock +1 -1
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/Cargo.toml +6 -2
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/PKG-INFO +132 -56
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/README.md +131 -55
- xarray_sql-0.3.1/benchmarks/geospatial/01_ndvi.py +149 -0
- xarray_sql-0.3.1/benchmarks/geospatial/02_climatology.py +137 -0
- xarray_sql-0.3.1/benchmarks/geospatial/03_zonal_mean.py +132 -0
- xarray_sql-0.3.1/benchmarks/geospatial/04_anomaly.py +144 -0
- xarray_sql-0.3.1/benchmarks/geospatial/05_forecast_skill.py +200 -0
- xarray_sql-0.3.1/benchmarks/geospatial/06_zonal_vector.py +175 -0
- xarray_sql-0.3.1/benchmarks/geospatial/07_reproject_udf.py +196 -0
- xarray_sql-0.3.1/benchmarks/geospatial/08_regrid_weights.py +228 -0
- xarray_sql-0.3.1/benchmarks/geospatial/09_warp.py +292 -0
- xarray_sql-0.3.1/benchmarks/geospatial/README.md +106 -0
- xarray_sql-0.3.1/benchmarks/geospatial/_harness.py +276 -0
- xarray_sql-0.3.1/benchmarks/geospatial/perf_summary.py +103 -0
- xarray_sql-0.3.1/benchmarks/geospatial/run_all.sh +37 -0
- xarray_sql-0.3.1/benchmarks/geospatial/run_perf.sh +44 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/docs/examples.md +59 -7
- xarray_sql-0.3.1/docs/geospatial.md +474 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/pyproject.toml +1 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/tests/test_df.py +91 -0
- xarray_sql-0.3.1/tests/test_ds.py +675 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/tests/test_reader.py +38 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/tests/test_sql.py +51 -0
- xarray_sql-0.3.1/tests/test_to_dataset_perf.py +115 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/xarray_sql/df.py +80 -19
- xarray_sql-0.3.1/xarray_sql/ds.py +926 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/xarray_sql/reader.py +31 -4
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/xarray_sql/sql.py +69 -7
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/zensical.toml +1 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/.gitignore +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/LICENSE +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/docs/assets/logo.svg +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/docs/contributing.md +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/docs/index.md +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/docs/reference/xarray_sql.md +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/src/lib.rs +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/tests/__init__.py +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/tests/conftest.py +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/tests/test_cft.py +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/xarray_sql/__init__.py +1 -1
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/xarray_sql/cftime.py +0 -0
- {xarray_sql-0.2.3 → xarray_sql-0.3.1}/xarray_sql/core.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[package]
|
|
2
2
|
name = "xarray_sql"
|
|
3
|
-
version = "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
|
-
|
|
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
|
+
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,28 +57,78 @@ 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`.
|
|
60
71
|
|
|
61
72
|
```python
|
|
62
73
|
import xarray as xr
|
|
63
74
|
import xarray_sql as xql
|
|
64
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`).
|
|
65
111
|
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
112
|
+
```python
|
|
113
|
+
import xarray as xr
|
|
114
|
+
import xarray_sql as xql
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
# Open ARCO-ERA5 — a weather dataset with 273 variables since 1940.
|
|
118
|
+
# Turning off dask means we don't have to wait to construct a task graph.
|
|
119
|
+
ds = xr.open_zarr(
|
|
120
|
+
'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3',
|
|
121
|
+
chunks=None, # Turn dask off
|
|
122
|
+
storage_options={'token': 'anon'} # Anonymous read from the public GCS bucket — no auth required.
|
|
73
123
|
)
|
|
74
124
|
|
|
75
125
|
ctx = xql.XarrayContext()
|
|
76
|
-
|
|
126
|
+
# Make sure to pass `chunks`!
|
|
127
|
+
ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={
|
|
77
128
|
('time', 'latitude', 'longitude'): 'surface',
|
|
78
129
|
('time', 'level', 'latitude', 'longitude'): 'atmosphere',
|
|
79
130
|
})
|
|
80
|
-
# Registration
|
|
81
|
-
|
|
131
|
+
# Registration takes ~10s on my machine.
|
|
82
132
|
|
|
83
133
|
# Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library
|
|
84
134
|
# pushes column projection down to Zarr, so SELECT only fetches what you ask
|
|
@@ -100,60 +150,55 @@ ctx.sql('''
|
|
|
100
150
|
# 0 8.640069
|
|
101
151
|
|
|
102
152
|
# Average temperature per pressure level, globally.
|
|
103
|
-
ctx.sql('''
|
|
153
|
+
result = ctx.sql('''
|
|
104
154
|
SELECT level, AVG(temperature) - 273.15 AS avg_c
|
|
105
155
|
FROM era5.atmosphere
|
|
106
156
|
WHERE time BETWEEN TIMESTAMP '2020-01-01'
|
|
107
157
|
AND TIMESTAMP '2020-01-01 05:00:00'
|
|
108
158
|
GROUP BY level
|
|
109
159
|
ORDER BY level DESC
|
|
110
|
-
''')
|
|
111
|
-
#
|
|
112
|
-
#
|
|
113
|
-
#
|
|
114
|
-
#
|
|
115
|
-
#
|
|
116
|
-
#
|
|
117
|
-
#
|
|
118
|
-
#
|
|
119
|
-
#
|
|
120
|
-
#
|
|
121
|
-
#
|
|
122
|
-
#
|
|
123
|
-
#
|
|
124
|
-
#
|
|
125
|
-
#
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
#
|
|
136
|
-
#
|
|
137
|
-
#
|
|
138
|
-
#
|
|
139
|
-
#
|
|
140
|
-
#
|
|
141
|
-
#
|
|
142
|
-
#
|
|
143
|
-
#
|
|
144
|
-
#
|
|
145
|
-
#
|
|
146
|
-
#
|
|
147
|
-
# 35 2 -13.355764
|
|
148
|
-
# 36 1 -9.020513 ← top of atmosphere
|
|
160
|
+
''')
|
|
161
|
+
# DataFrame()
|
|
162
|
+
# +-------+----------------------+
|
|
163
|
+
# | level | avg_c |
|
|
164
|
+
# +-------+----------------------+
|
|
165
|
+
# | 1000 | 6.6210120796502565 |
|
|
166
|
+
# | 975 | 5.185637919348153 |
|
|
167
|
+
# | 950 | 4.028428657263021 |
|
|
168
|
+
# | 925 | 3.0828117974912743 |
|
|
169
|
+
# | 900 | 2.2109172992531967 |
|
|
170
|
+
# | 875 | 1.395017610194202 |
|
|
171
|
+
# | 850 | 0.6342670572626616 |
|
|
172
|
+
# | 825 | -0.21037158786759846 |
|
|
173
|
+
# | 800 | -1.1810754318269687 |
|
|
174
|
+
# | 775 | -2.3064649711534457 |
|
|
175
|
+
# +-------+----------------------+
|
|
176
|
+
|
|
177
|
+
ctx.sql('''
|
|
178
|
+
SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c
|
|
179
|
+
FROM era5.surface
|
|
180
|
+
WHERE time BETWEEN TIMESTAMP '2020-01-01'
|
|
181
|
+
AND TIMESTAMP '2020-01-01 05:00:00'
|
|
182
|
+
GROUP BY latitude, longitude
|
|
183
|
+
ORDER BY latitude DESC, longitude
|
|
184
|
+
''').to_dataset(dims=['latitude', 'longitude'], template=ds)
|
|
185
|
+
# <xarray.Dataset> Size: 8MB
|
|
186
|
+
# Dimensions: (latitude: 721, longitude: 1440)
|
|
187
|
+
# Coordinates:
|
|
188
|
+
# * latitude (latitude) float32 3kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0
|
|
189
|
+
# * longitude (longitude) float32 6kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8
|
|
190
|
+
# Data variables:
|
|
191
|
+
# avg_c (latitude, longitude) float64 8MB -26.84 -26.84 ... -27.38 -27.38
|
|
192
|
+
# Attributes:
|
|
193
|
+
# last_updated: 2026-06-20 02:33:34.265980+00:00
|
|
194
|
+
# valid_time_start: 1940-01-01
|
|
195
|
+
# valid_time_stop: 2025-12-31
|
|
196
|
+
# valid_time_stop_era5t: 2026-06-14
|
|
149
197
|
```
|
|
150
198
|
|
|
151
199
|
_(A runnable version of this example lives at
|
|
152
200
|
[`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_
|
|
153
201
|
|
|
154
|
-
Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run
|
|
155
|
-
SQL queries against them.
|
|
156
|
-
|
|
157
202
|
## Why build this?
|
|
158
203
|
|
|
159
204
|
A few reasons:
|
|
@@ -188,6 +233,35 @@ that lets the DB engine translate the underlying Dataset arrays into DataFusion
|
|
|
188
233
|
Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be
|
|
189
234
|
translated into a 2D table -- underlies this performant query mechanism.
|
|
190
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
|
+
|
|
191
265
|
## Why does this work?
|
|
192
266
|
|
|
193
267
|
Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in
|
|
@@ -225,14 +299,14 @@ _2025 update_: Something like this is being built across a few projects! The one
|
|
|
225
299
|
_2026 update_: A colleague and I are experimenting with native Zarr RDBMS engines. Check out:
|
|
226
300
|
|
|
227
301
|
- [Zarr-Datafusion](https://lib.rs/crates/zarr-datafusion)
|
|
228
|
-
- [DuckDB-Zarr](https://github.com/
|
|
302
|
+
- [DuckDB-Zarr](https://github.com/alxmrs/duckdb-zarr)
|
|
229
303
|
|
|
230
304
|
## Roadmap
|
|
231
305
|
|
|
232
306
|
- [x] ~Lazy evaluation via the pyarrow Dataset interface [#93](https://github.com/alxmrs/xarray-sql/issues/93).~ _Implemented in [#100](https://github.com/alxmrs/xarray-sql/pull/100)_
|
|
233
307
|
- [x] Support proper parallelism via proper partition handling on the rust/datafusion side. [#106](https://github.com/alxmrs/xarray-sql/issues/106)
|
|
234
308
|
- [x] Support core datafusion optimizations to scan less data, like [104](https://github.com/alxmrs/xarray-sql/issues/104), ...
|
|
235
|
-
- [
|
|
309
|
+
- [x] Translate a single Zarr to a collection of tables [#85](https://github.com/alxmrs/xarray-sql/issues/85).
|
|
236
310
|
- [ ] Distributed beyond a single node through the DataFusion integration with Ray Datasets [#68](https://github.com/alxmrs/xarray-sql/issues/68) or Apache Ballista [#98](https://github.com/alxmrs/xarray-sql/issues/98).
|
|
237
311
|
- [ ] Demo: calculate Sea Surface Temperature from 1940 - Present in SQL [#36](https://github.com/alxmrs/xarray-sql/issues/36).
|
|
238
312
|
- [ ] Provide an option to integrate DataFusion directly to Zarr via Rust [#4](https://github.com/alxmrs/xarray-sql/issues/4).
|
|
@@ -254,6 +328,8 @@ I want to give a special thanks to the following folks and institutions:
|
|
|
254
328
|
who are working to make this library better.
|
|
255
329
|
- Andrew Huang for the sense of taste he brings to the project and consummate code
|
|
256
330
|
changes.
|
|
331
|
+
- Aman Kumar for spending a considerable amount of his GSoC internship
|
|
332
|
+
contributing to this project.
|
|
257
333
|
|
|
258
334
|
|
|
259
335
|
## License
|
|
@@ -14,28 +14,78 @@ 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`.
|
|
17
28
|
|
|
18
29
|
```python
|
|
19
30
|
import xarray as xr
|
|
20
31
|
import xarray_sql as xql
|
|
21
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`).
|
|
22
68
|
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
69
|
+
```python
|
|
70
|
+
import xarray as xr
|
|
71
|
+
import xarray_sql as xql
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
# Open ARCO-ERA5 — a weather dataset with 273 variables since 1940.
|
|
75
|
+
# Turning off dask means we don't have to wait to construct a task graph.
|
|
76
|
+
ds = xr.open_zarr(
|
|
77
|
+
'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3',
|
|
78
|
+
chunks=None, # Turn dask off
|
|
79
|
+
storage_options={'token': 'anon'} # Anonymous read from the public GCS bucket — no auth required.
|
|
30
80
|
)
|
|
31
81
|
|
|
32
82
|
ctx = xql.XarrayContext()
|
|
33
|
-
|
|
83
|
+
# Make sure to pass `chunks`!
|
|
84
|
+
ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={
|
|
34
85
|
('time', 'latitude', 'longitude'): 'surface',
|
|
35
86
|
('time', 'level', 'latitude', 'longitude'): 'atmosphere',
|
|
36
87
|
})
|
|
37
|
-
# Registration
|
|
38
|
-
|
|
88
|
+
# Registration takes ~10s on my machine.
|
|
39
89
|
|
|
40
90
|
# Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library
|
|
41
91
|
# pushes column projection down to Zarr, so SELECT only fetches what you ask
|
|
@@ -57,60 +107,55 @@ ctx.sql('''
|
|
|
57
107
|
# 0 8.640069
|
|
58
108
|
|
|
59
109
|
# Average temperature per pressure level, globally.
|
|
60
|
-
ctx.sql('''
|
|
110
|
+
result = ctx.sql('''
|
|
61
111
|
SELECT level, AVG(temperature) - 273.15 AS avg_c
|
|
62
112
|
FROM era5.atmosphere
|
|
63
113
|
WHERE time BETWEEN TIMESTAMP '2020-01-01'
|
|
64
114
|
AND TIMESTAMP '2020-01-01 05:00:00'
|
|
65
115
|
GROUP BY level
|
|
66
116
|
ORDER BY level DESC
|
|
67
|
-
''')
|
|
68
|
-
#
|
|
69
|
-
#
|
|
70
|
-
#
|
|
71
|
-
#
|
|
72
|
-
#
|
|
73
|
-
#
|
|
74
|
-
#
|
|
75
|
-
#
|
|
76
|
-
#
|
|
77
|
-
#
|
|
78
|
-
#
|
|
79
|
-
#
|
|
80
|
-
#
|
|
81
|
-
#
|
|
82
|
-
#
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
#
|
|
93
|
-
#
|
|
94
|
-
#
|
|
95
|
-
#
|
|
96
|
-
#
|
|
97
|
-
#
|
|
98
|
-
#
|
|
99
|
-
#
|
|
100
|
-
#
|
|
101
|
-
#
|
|
102
|
-
#
|
|
103
|
-
#
|
|
104
|
-
# 35 2 -13.355764
|
|
105
|
-
# 36 1 -9.020513 ← top of atmosphere
|
|
117
|
+
''')
|
|
118
|
+
# DataFrame()
|
|
119
|
+
# +-------+----------------------+
|
|
120
|
+
# | level | avg_c |
|
|
121
|
+
# +-------+----------------------+
|
|
122
|
+
# | 1000 | 6.6210120796502565 |
|
|
123
|
+
# | 975 | 5.185637919348153 |
|
|
124
|
+
# | 950 | 4.028428657263021 |
|
|
125
|
+
# | 925 | 3.0828117974912743 |
|
|
126
|
+
# | 900 | 2.2109172992531967 |
|
|
127
|
+
# | 875 | 1.395017610194202 |
|
|
128
|
+
# | 850 | 0.6342670572626616 |
|
|
129
|
+
# | 825 | -0.21037158786759846 |
|
|
130
|
+
# | 800 | -1.1810754318269687 |
|
|
131
|
+
# | 775 | -2.3064649711534457 |
|
|
132
|
+
# +-------+----------------------+
|
|
133
|
+
|
|
134
|
+
ctx.sql('''
|
|
135
|
+
SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c
|
|
136
|
+
FROM era5.surface
|
|
137
|
+
WHERE time BETWEEN TIMESTAMP '2020-01-01'
|
|
138
|
+
AND TIMESTAMP '2020-01-01 05:00:00'
|
|
139
|
+
GROUP BY latitude, longitude
|
|
140
|
+
ORDER BY latitude DESC, longitude
|
|
141
|
+
''').to_dataset(dims=['latitude', 'longitude'], template=ds)
|
|
142
|
+
# <xarray.Dataset> Size: 8MB
|
|
143
|
+
# Dimensions: (latitude: 721, longitude: 1440)
|
|
144
|
+
# Coordinates:
|
|
145
|
+
# * latitude (latitude) float32 3kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0
|
|
146
|
+
# * longitude (longitude) float32 6kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8
|
|
147
|
+
# Data variables:
|
|
148
|
+
# avg_c (latitude, longitude) float64 8MB -26.84 -26.84 ... -27.38 -27.38
|
|
149
|
+
# Attributes:
|
|
150
|
+
# last_updated: 2026-06-20 02:33:34.265980+00:00
|
|
151
|
+
# valid_time_start: 1940-01-01
|
|
152
|
+
# valid_time_stop: 2025-12-31
|
|
153
|
+
# valid_time_stop_era5t: 2026-06-14
|
|
106
154
|
```
|
|
107
155
|
|
|
108
156
|
_(A runnable version of this example lives at
|
|
109
157
|
[`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_
|
|
110
158
|
|
|
111
|
-
Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run
|
|
112
|
-
SQL queries against them.
|
|
113
|
-
|
|
114
159
|
## Why build this?
|
|
115
160
|
|
|
116
161
|
A few reasons:
|
|
@@ -145,6 +190,35 @@ that lets the DB engine translate the underlying Dataset arrays into DataFusion
|
|
|
145
190
|
Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be
|
|
146
191
|
translated into a 2D table -- underlies this performant query mechanism.
|
|
147
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
|
+
|
|
148
222
|
## Why does this work?
|
|
149
223
|
|
|
150
224
|
Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in
|
|
@@ -182,14 +256,14 @@ _2025 update_: Something like this is being built across a few projects! The one
|
|
|
182
256
|
_2026 update_: A colleague and I are experimenting with native Zarr RDBMS engines. Check out:
|
|
183
257
|
|
|
184
258
|
- [Zarr-Datafusion](https://lib.rs/crates/zarr-datafusion)
|
|
185
|
-
- [DuckDB-Zarr](https://github.com/
|
|
259
|
+
- [DuckDB-Zarr](https://github.com/alxmrs/duckdb-zarr)
|
|
186
260
|
|
|
187
261
|
## Roadmap
|
|
188
262
|
|
|
189
263
|
- [x] ~Lazy evaluation via the pyarrow Dataset interface [#93](https://github.com/alxmrs/xarray-sql/issues/93).~ _Implemented in [#100](https://github.com/alxmrs/xarray-sql/pull/100)_
|
|
190
264
|
- [x] Support proper parallelism via proper partition handling on the rust/datafusion side. [#106](https://github.com/alxmrs/xarray-sql/issues/106)
|
|
191
265
|
- [x] Support core datafusion optimizations to scan less data, like [104](https://github.com/alxmrs/xarray-sql/issues/104), ...
|
|
192
|
-
- [
|
|
266
|
+
- [x] Translate a single Zarr to a collection of tables [#85](https://github.com/alxmrs/xarray-sql/issues/85).
|
|
193
267
|
- [ ] Distributed beyond a single node through the DataFusion integration with Ray Datasets [#68](https://github.com/alxmrs/xarray-sql/issues/68) or Apache Ballista [#98](https://github.com/alxmrs/xarray-sql/issues/98).
|
|
194
268
|
- [ ] Demo: calculate Sea Surface Temperature from 1940 - Present in SQL [#36](https://github.com/alxmrs/xarray-sql/issues/36).
|
|
195
269
|
- [ ] Provide an option to integrate DataFusion directly to Zarr via Rust [#4](https://github.com/alxmrs/xarray-sql/issues/4).
|
|
@@ -211,6 +285,8 @@ I want to give a special thanks to the following folks and institutions:
|
|
|
211
285
|
who are working to make this library better.
|
|
212
286
|
- Andrew Huang for the sense of taste he brings to the project and consummate code
|
|
213
287
|
changes.
|
|
288
|
+
- Aman Kumar for spending a considerable amount of his GSoC internship
|
|
289
|
+
contributing to this project.
|
|
214
290
|
|
|
215
291
|
|
|
216
292
|
## License
|