diff --git a/Cargo.lock b/Cargo.lock index 41ef28b..1d53519 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3375,7 +3375,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] name = "xarray_sql" -version = "0.2.1" +version = "0.2.2" dependencies = [ "arrow", "async-stream", diff --git a/README.md b/README.md index 5b03ea5..58afc7c 100644 --- a/README.md +++ b/README.md @@ -19,52 +19,98 @@ This is an experiment to provide a SQL interface for array datasets. import xarray as xr import xarray_sql as xql -ds = xr.tutorial.open_dataset('air_temperature') -# The same as a dask-sql Context; i.e. an Apache DataFusion Context. +# Open a year of ARCO-ERA5 — all 273 variables. Selecting a year up front +# keeps Dask's partition setup cheap before any chunks are read from GCS. +ds = ( + xr.open_zarr('gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', + chunks=None, + storage_options={'token': 'anon'}) # Anonymous read from the public GCS bucket — no auth required. + .sel(time='2020') + .chunk({'time': 1}) +) + ctx = xql.XarrayContext() -ctx.from_dataset('air', ds, chunks=dict(time=24)) # the dataset needs to be chunked! -# data is only materialized when we make a query. - -result = ctx.sql(''' - SELECT - "lat", "lon", AVG("air") as air_avg - FROM - "air" - GROUP BY - "lat", "lon" -''') -# DataFrame() -# +------+-------+--------------------+ -# | lat | lon | air_avg | -# +------+-------+--------------------+ -# | 75.0 | 205.0 | 259.88662671232834 | -# | 75.0 | 207.5 | 259.48268150684896 | -# | 75.0 | 230.0 | 258.9192123287667 | -# | 75.0 | 275.0 | 257.07574315068456 | -# | 75.0 | 322.5 | 250.11792123287654 | -# | 75.0 | 325.0 | 250.81590068493134 | -# | 72.5 | 205.0 | 262.74933904109537 | -# | 72.5 | 207.5 | 262.5384315068488 | -# | 72.5 | 230.0 | 260.82879452054743 | -# | 72.5 | 275.0 | 257.3063321917804 | -# +------+-------+--------------------+ -# Data truncated. - -# The full query is only made when we call `collect()`, or, in this case, -# `to_pandas()`. -df = result.to_pandas() -df.head() -# lat lon air_avg -# 0 75.0 232.5 258.836188 -# 1 75.0 247.5 257.716171 -# 2 75.0 262.5 257.347959 -# 3 75.0 277.5 257.671308 -# 4 72.5 232.5 260.654401 +ctx.from_dataset('era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) +# Registration: ~0.5s for a full year of hourly ERA5, all variables. + + +# Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library +# pushes column projection down to Zarr, so SELECT only fetches what you ask +# for — but `SELECT * FROM era5.surface` would try to pull every variable +# across the year (terabytes from GCS). +# ---> Always SELECT specific columns. <--- + +# Average 2m-temperature over NYC on the morning of 2020-01-01. The library +# pushes WHERE clauses on dimension columns down to partition pruning. +ctx.sql(''' + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 -- ERA5 uses 0-360 longitudes +''').to_pandas() +# avg_c +# 0 8.640069 + +# Average temperature per pressure level, globally. +ctx.sql(''' + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC +''').to_pandas() +# level avg_c +# 0 1000 6.621012 ← surface +# 1 975 5.185638 +# 2 950 4.028429 +# 3 925 3.082812 +# 4 900 2.210917 +# 5 875 1.395018 +# 6 850 0.634267 +# 7 825 -0.210372 +# 8 800 -1.181075 +# 9 775 -2.306465 +# 10 750 -3.535534 +# 11 700 -6.241685 +# 12 650 -9.236364 +# 13 600 -12.580938 +# 14 550 -16.335386 +# 15 500 -20.643604 +# 16 450 -25.573401 +# 17 400 -31.156920 +# 18 350 -37.400552 +# 19 300 -43.852607 +# 20 250 -49.322132 +# 21 225 -51.569113 +# 22 200 -53.693248 +# 23 175 -55.890484 +# 24 150 -58.382290 +# 25 125 -61.091916 +# 26 100 -63.624885 ← tropopause +# 27 70 -63.182300 +# 28 50 -60.124845 +# 29 30 -55.986327 +# 30 20 -52.433089 +# 31 10 -44.140750 +# 32 7 -38.707350 +# 33 5 -32.621999 +# 34 3 -21.509175 +# 35 2 -13.355764 +# 36 1 -9.020513 ← top of atmosphere ``` -Succinctly, we "pivot" Xarray Datasets (with consistent dimensions) to treat them like tables so we can run -SQL queries against them. +_(A runnable version of this example lives at +[`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_ + +Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run +SQL queries against them. ## Why build this? diff --git a/docs/examples.md b/docs/examples.md index cff0f33..0a3066c 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -21,3 +21,73 @@ result = ctx.sql(''' df = result.to_pandas() df.head() ``` + +## Mixed-dimension datasets: ARCO-ERA5 + +When a Dataset has variables with differing dimensions (e.g. surface fields on +`(time, latitude, longitude)` and atmospheric fields on +`(time, level, latitude, longitude)`), `from_dataset` splits them into one +table per dimension group, registered together under a SQL schema named after +the first argument. [ARCO-ERA5][arco-era5] is a good example: 262 of its +variables are surface fields and 11 are atmospheric. + +Open a year of ARCO-ERA5 and let SQL `WHERE` clauses do the filtering — the +library prunes time partitions and pushes dimension-column filters down. Use +the `table_names` kwarg to give each dimension group a friendly name: + +```python +import xarray as xr +import xarray_sql as xql + +# Open ARCO-ERA5 directly from GCS (anonymous read). +url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' +full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) + +# A full year of hourly ERA5 — all 273 variables. No spatial slicing on the +# xarray side; SQL WHERE clauses below express the filters. `chunks={'time': 1}` +# aligns Dask chunks to native Zarr chunks of shape (1, 37, 721, 1440) so +# chunk reads from GCS happen concurrently. +# +# Heads up: 262 of those variables are surface and 11 are atmospheric. The +# library pushes column projection down, so SELECT only fetches what you ask +# for — but `SELECT * FROM era5.surface` would try to pull every variable +# across the year (terabytes from GCS). Always SELECT specific columns. +ds = full.sel(time='2020').chunk({'time': 1}) + +ctx = xql.XarrayContext() +ctx.from_dataset('era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) +# Registers two tables under a SQL schema named 'era5': 'surface' and 'atmosphere'. + +# Average 2m-temperature over the NYC area on the morning of 2020-01-01. +ctx.sql(''' + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 +''').to_pandas() + +# Average temperature per pressure level, globally — the standard +# atmospheric temperature profile. Scans ~230M rows. +ctx.sql(''' + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC -- surface (1000 hPa) first +''').to_pandas() +``` + +If you omit `table_names`, each table is named by joining its dimension names +with underscores, e.g. `era5.time_latitude_longitude` and +`era5.time_level_latitude_longitude`. + +A runnable version of this example lives at +[`perf_tests/era5_temp_profile.py`](../perf_tests/era5_temp_profile.py). + +[arco-era5]: https://github.com/google-research/arco-era5 diff --git a/perf_tests/era5_temp_profile.py b/perf_tests/era5_temp_profile.py new file mode 100644 index 0000000..28275a9 --- /dev/null +++ b/perf_tests/era5_temp_profile.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Surface and global-atmospheric temperatures on 2020-01-01, in SQL. + +Two queries against ARCO-ERA5 on the morning of January 1, 2020: + + * **Surface (local).** Average 2m-temperature over a small grid covering the + New York City area for the first six hours. + * **Atmosphere (global).** Average temperature per pressure level, computed + over the entire planet for the same six hours — a classic atmospheric + temperature profile (surface around 1000 hPa is warmest, tropopause near + 100 hPa is coldest). + +Both queries express their filters entirely in SQL: ``xr.open_zarr`` is given +a single calendar year and no spatial slicing. The library's table provider +prunes time partitions for ``WHERE time …`` filters, and pushes ``WHERE +latitude/longitude …`` down to dimension columns. + +ARCO-ERA5's atmospheric variables are stored in native Zarr chunks of shape +``(1, 37, 721, 1440)`` — about 150 MB per hour. We align Dask chunks to that +shape with ``chunks={'time': 1}`` so chunks fetch from GCS concurrently. The +global atmospheric query scans ~230M rows after pruning. + +The Zarr is read anonymously from the public GCS bucket — no auth required. +""" + +import time + +import xarray as xr + +import xarray_sql as xql + + +URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" + + +def main() -> None: + full = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) + + # Open a full calendar year — all 273 variables. No spatial slicing on + # the xarray side; SQL WHERE clauses below express the filters. + # + # Heads up: the library pushes column projection down to Zarr, so SELECT + # only fetches what you ask for — but `SELECT * FROM era5.surface` would + # try to read every variable across the year (terabytes from GCS). + # Always SELECT specific columns. + ds = full.sel(time="2020").chunk({"time": 1}) + print( + "ARCO-ERA5 opened: year 2020, " + f"{ds.sizes['time']:,} hourly time steps, " + f"{len(ds.data_vars)} variables (no spatial pre-slicing)." + ) + + ctx = xql.XarrayContext() + t0 = time.perf_counter() + ctx.from_dataset( + "era5", + ds, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + print(f"Registration: {time.perf_counter() - t0:.2f}s") + ctx.sql("SELECT 1").to_pandas() # warm the planner + + print("\nAverage 2m-temperature over NYC, 2020-01-01 00:00-05:00 UTC (°C):") + t0 = time.perf_counter() + surface = ctx.sql( + """ + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 -- ERA5 uses 0-360 longitudes + """ + ).to_pandas() + print(surface) + print(f" ({time.perf_counter() - t0:.2f}s)") + + print( + "\nAverage temperature per pressure level, globally, " + "2020-01-01 00:00-05:00 UTC (°C):" + ) + t0 = time.perf_counter() + profile = ctx.sql( + """ + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC -- surface (1000 hPa) first, top of atmosphere last + """ + ).to_pandas() + print(profile.to_string(index=False)) + print(f" ({time.perf_counter() - t0:.2f}s, ~230M rows scanned)") + + +if __name__ == "__main__": + main() diff --git a/tests/test_sql.py b/tests/test_sql.py index 01929f0..74f4c66 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -6,6 +6,7 @@ import xarray as xr from xarray_sql import XarrayContext +from xarray_sql.sql import _group_vars_by_dims def test_sanity(air_dataset_small): @@ -322,3 +323,117 @@ def test_gregorian_like_no_cftime_udf(self): ctx.sql( "SELECT COUNT(*) FROM rasm WHERE time >= cftime('1980-01-01')" ).collect() + + +class TestGroupVarsByDims: + """Unit tests for the private _group_vars_by_dims helper.""" + + def test_single_dim_group(self): + ds = xr.Dataset( + { + "a": (["x", "y"], np.zeros((2, 3))), + "b": (["x", "y"], np.ones((2, 3))), + } + ) + groups = _group_vars_by_dims(ds) + assert groups == {("x", "y"): ["a", "b"]} + + def test_multiple_dim_groups(self): + ds = xr.Dataset( + { + "surface": (["time", "lat", "lon"], np.zeros((2, 3, 4))), + "upper": ( + ["time", "lat", "lon", "level"], + np.zeros((2, 3, 4, 5)), + ), + } + ) + groups = _group_vars_by_dims(ds) + assert set(groups.keys()) == { + ("time", "lat", "lon"), + ("time", "lat", "lon", "level"), + } + assert groups[("time", "lat", "lon")] == ["surface"] + assert groups[("time", "lat", "lon", "level")] == ["upper"] + + def test_empty_dataset(self): + assert _group_vars_by_dims(xr.Dataset()) == {} + + def test_ignores_coords(self): + """Coordinate variables shouldn't be returned as groups.""" + ds = xr.Dataset( + {"v": (["x"], np.arange(3))}, + coords={"x": np.arange(3), "label": ("x", ["a", "b", "c"])}, + ) + groups = _group_vars_by_dims(ds) + assert groups == {("x",): ["v"]} + + +class TestFromDatasetMultiDims: + """from_dataset should split datasets with mixed dims into multiple tables.""" + + @pytest.fixture + def mixed_ds(self): + np.random.seed(0) + return xr.Dataset( + { + "temperature_2m": ( + ["time", "lat", "lon"], + np.random.rand(2, 3, 4), + ), + "pressure": ( + ["time", "lat", "lon", "level"], + np.random.rand(2, 3, 4, 2), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=2), + "lat": np.linspace(-90, 90, 3), + "lon": np.linspace(-180, 180, 4), + "level": [500, 1000], + }, + ).chunk({"time": 1}) + + def test_registers_multiple_tables(self, mixed_ds): + ctx = XarrayContext() + ctx.from_dataset("era5", mixed_ds) + surface = ctx.sql("SELECT * FROM era5.time_lat_lon").to_pandas() + upper = ctx.sql("SELECT * FROM era5.time_lat_lon_level").to_pandas() + assert "temperature_2m" in surface.columns + assert "pressure" in upper.columns + assert len(surface) == 2 * 3 * 4 + assert len(upper) == 2 * 3 * 4 * 2 + + def test_table_names_override(self, mixed_ds): + ctx = XarrayContext() + ctx.from_dataset( + "era5", + mixed_ds, + table_names={("time", "lat", "lon"): "surface"}, + ) + result = ctx.sql("SELECT * FROM era5.surface").to_pandas() + assert "temperature_2m" in result.columns + # Non-aliased group falls back to the default joined-dim table name. + upper = ctx.sql("SELECT * FROM era5.time_lat_lon_level").to_pandas() + assert "pressure" in upper.columns + + def test_schema_registered_in_catalog(self, mixed_ds): + """Mixed-dim datasets should create a SQL schema under the catalog.""" + ctx = XarrayContext() + ctx.from_dataset("era5", mixed_ds) + assert "era5" in ctx.catalog().schema_names() + tables = set(ctx.catalog().schema("era5").table_names()) + assert tables == {"time_lat_lon", "time_lat_lon_level"} + + def test_uniform_dims_uses_name_directly(self, mixed_ds): + """A single dim group should register under the bare name.""" + ds = mixed_ds[["temperature_2m"]] + ctx = XarrayContext() + ctx.from_dataset("surface", ds) + result = ctx.sql("SELECT * FROM surface").to_pandas() + assert "temperature_2m" in result.columns + + def test_table_names_is_keyword_only(self, mixed_ds): + ctx = XarrayContext() + with pytest.raises(TypeError): + ctx.from_dataset("era5", mixed_ds, {("time",): "x"}) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index f30f5b8..3881d98 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,5 +1,7 @@ import xarray as xr from datafusion import SessionContext +from datafusion.catalog import Schema +from collections import defaultdict from . import cftime as cft from .df import Chunks @@ -11,11 +13,34 @@ class XarrayContext(SessionContext): def from_dataset( self, - table_name: str, + name: str, input_table: xr.Dataset, + *, + table_names: dict[tuple[str, ...], str] | None = None, chunks: Chunks = None, ): - """Register an xarray Dataset as a queryable SQL table. + """Register an xarray Dataset as one or more queryable SQL tables. + + When all data variables share the same dimensions, the dataset is + registered as a single table named ``name``. When variables have + differing dimensions (e.g. some on a 3D grid and others on a 4D + grid), the dataset is split into one table per dimension group. + The tables are registered under a SQL schema (namespace) named + ``name`` and named ``__...`` by default:: + + ctx.from_dataset('era5', ds, chunks={'time': 24}) + # registers tables: 'era5.time_lat_lon' and + # 'era5.time_lat_lon_level' + ctx.sql('SELECT AVG(temperature_2m) FROM era5.time_lat_lon') + + Use ``table_names`` to override the name for specific dimension + tuples:: + + ctx.from_dataset( + 'era5', ds, + table_names={('time', 'lat', 'lon'): 'surface'}, + ) + ctx.sql('SELECT * FROM era5.surface') For datasets with non-Gregorian cftime coordinates (e.g. 360_day, julian), a ``cftime()`` scalar UDF is automatically registered so @@ -24,9 +49,6 @@ def from_dataset( ctx.from_dataset("ds360", ds, chunks={"time": 6}) ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')") - The UDF converts a date string to the int64 offset used to store - that calendar's time axis. - .. note:: Only one ``cftime()`` UDF is registered per context, using the @@ -39,25 +61,69 @@ def from_dataset( for each calendar. Args: - table_name: The SQL table name to register the dataset under. - input_table: An xarray Dataset. All data_vars must share the - same dimensions. + name: The SQL identifier under which the dataset is registered. + For datasets with uniform dimensions, this is the table + name. For datasets with mixed dimensions, this is the name + of a SQL schema (namespace) containing one table per + dimension group. + input_table: An xarray Dataset. + table_names: Optional mapping from dimension tuples to custom + table names within the schema, used when the dataset has + variables with differing dimensions. chunks: Xarray-like chunks specification. If not provided, uses the Dataset's existing chunks. Returns: self, to allow chaining. """ + groups = _group_vars_by_dims(input_table) + + if len(groups) <= 1: + return self._from_dataset(name, input_table, chunks) + + table_names = table_names or {} + schema = Schema.memory_schema(self) + self.catalog().register_schema(name, schema) + + for dims, var_names in groups.items(): + sub_name = table_names.get(dims, "_".join(dims)) + sub_ds = input_table[var_names] + schema.register_table(sub_name, read_xarray_table(sub_ds, chunks)) + self._maybe_register_cftime_udf(sub_ds) + + return self + + def _from_dataset( + self, + table_name: str, + input_table: xr.Dataset, + chunks: Chunks = None, + ): + """Register a uniform-dimension Dataset as a single SQL table.""" + table = read_xarray_table(input_table, chunks) self.register_table(table_name, table) + self._maybe_register_cftime_udf(input_table) + return self - # Auto-register a cftime() UDF for non-Gregorian cftime coordinates - # so users can write: WHERE time > cftime('0500-01-01') - for coord_name in input_table.dims: - if cft.is_cftime_index(input_table, coord_name): - units, cal = cft.encoding(input_table, coord_name) + def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None: + """Auto-register a cftime() UDF for non-Gregorian cftime coordinates.""" + for coord_name in ds.dims: + if cft.is_cftime_index(ds, coord_name): + units, cal = cft.encoding(ds, coord_name) if not cft.is_gregorian_like(cal): self.register_udf(cft.make_cftime_udf(units, cal)) break # One UDF per context is enough. - return self + +def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: + """Group variables in the dataset based on shared dims. + + ("time", "lat", "lon"): ["temperature_2m", "wind_speed"], + ("time", "lat", "lon", "level"): ["pressure", "humidity"] + """ + groups = defaultdict(list) + for var_name, var in ds.data_vars.items(): + dims = var.dims + groups[dims].append(var_name) + return groups