Made new xarray context behavior to register multiple tables from a dataset

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?
 

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

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()