xarray_sql 0.2.2__tar.gz → 0.2.3__tar.gz

{xarray_sql-0.2.2 → xarray_sql-0.2.3}/Cargo.lock

@@ -3375,7 +3375,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb"
 
 [[package]]
 name = "xarray_sql"
-version = "0.2.2"
+version = "0.2.3"
 dependencies = [
  "arrow",
  "async-stream",

{xarray_sql-0.2.2 → xarray_sql-0.2.3}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "xarray_sql"
-version = "0.2.2"
+version = "0.2.3"
 authors = ["Alex Merose"]
 edition = "2021"
 exclude = [

{xarray_sql-0.2.2 → xarray_sql-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xarray_sql
-Version: 0.2.2
+Version: 0.2.3
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: Intended Audience :: Developers
@@ -62,52 +62,97 @@ 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=dict(time=1),
+               storage_options={'token': 'anon'})  # Anonymous read from the public GCS bucket — no auth required.
+  .sel(time='2020')
+)
+
 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?
 

{xarray_sql-0.2.2 → xarray_sql-0.2.3}/README.md

@@ -19,52 +19,97 @@ 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=dict(time=1),
+               storage_options={'token': 'anon'})  # Anonymous read from the public GCS bucket — no auth required.
+  .sel(time='2020')
+)
+
 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?
 

xarray_sql-0.2.3/docs/examples.md

@@ -0,0 +1,93 @@
+# Examples
+
+```python
+import xarray as xr
+import xarray_sql as xql
+
+ds = xr.tutorial.open_dataset('air_temperature')
+
+ctx = xql.XarrayContext()
+ctx.from_dataset('air', ds, chunks=dict(time=24))
+
+result = ctx.sql('''
+  SELECT
+    "lat", "lon", AVG("air") as air_avg
+  FROM
+    "air"
+  GROUP BY
+   "lat", "lon"
+''')
+
+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

{xarray_sql-0.2.2 → xarray_sql-0.2.3}/pyproject.toml

@@ -64,11 +64,13 @@ module-name = "xarray_sql._native"
 [tool.setuptools.packages.find]
 exclude = ["demo", "perf_tests", "tests", "tests.*"]
 
-[tool.pyink]
+[tool.ruff]
 line-length = 80
-preview = true
-pyink-indentation = 2
-pyink-use-majority-quotes = true
+indent-width = 4
+
+[tool.ruff.format]
+indent-style = "space"
+quote-style = "double"
 
 [tool.mypy]
 python_version = "3.11"
@@ -98,7 +100,7 @@ dev = [
     "xarray_sql[test]",
     "xarray_sql[docs]",
     "py-spy>=0.4.0",
-    "pyink>=24.10.1",
+    "ruff>=0.15.10",
     "maturin>=1.9.1",
 ]
 

xarray_sql-0.2.3/tests/conftest.py

@@ -0,0 +1,150 @@
+import pytest
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def rand_wx(start: str, end: str) -> xr.Dataset:
+    np.random.seed(42)
+    lat = np.linspace(-90, 90, num=720)
+    lon = np.linspace(-180, 180, num=1440)
+    time = pd.date_range(start, end, freq="h")
+    level = np.array([1000, 500], dtype=np.int32)
+    reference_time = pd.Timestamp(start)
+    temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
+    precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
+    return xr.Dataset(
+        data_vars=dict(
+            temperature=(["lat", "lon", "time", "level"], temperature),
+            precipitation=(["lat", "lon", "time", "level"], precipitation),
+        ),
+        coords=dict(
+            lat=lat,
+            lon=lon,
+            time=time,
+            level=level,
+            reference_time=reference_time,
+        ),
+        attrs=dict(description="Random weather."),
+    )
+
+
+def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
+    """Create a large xarray dataset for memory testing."""
+    np.random.seed(42)
+
+    time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
+    lat = np.linspace(-90, 90, lat_points)
+    lon = np.linspace(-180, 180, lon_points)
+
+    temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
+    precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
+
+    return xr.Dataset(
+        {
+            "temperature": (["time", "lat", "lon"], temp_data),
+            "precipitation": (["time", "lat", "lon"], precip_data),
+        },
+        coords={"time": time, "lat": lat, "lon": lon},
+    )
+
+
+@pytest.fixture
+def air():
+    ds = xr.tutorial.open_dataset("air_temperature")
+    chunks = {"time": 240}
+    return ds.chunk(chunks)
+
+
+@pytest.fixture
+def air_small(air):
+    return air.isel(
+        time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
+    ).chunk({"time": 240})
+
+
+@pytest.fixture
+def randwx():
+    return rand_wx("1995-01-13T00", "1995-01-13T01")
+
+
+@pytest.fixture
+def large_ds():
+    return create_large_dataset().chunk({"time": 25})
+
+
+@pytest.fixture
+def air_dataset_small():
+    ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
+    return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
+
+
+@pytest.fixture
+def air_dataset_large():
+    return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
+
+
+@pytest.fixture
+def rasm_ds():
+    """rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time."""
+    return xr.tutorial.open_dataset("rasm")
+
+
+@pytest.fixture
+def weather_dataset():
+    ds = rand_wx("2023-01-01T00", "2023-01-01T12")
+    return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
+        {"time": 3}
+    )
+
+
+@pytest.fixture
+def synthetic_dataset():
+    return create_large_dataset(
+        time_steps=50, lat_points=20, lon_points=20
+    ).chunk({"time": 25})
+
+
+@pytest.fixture
+def station_dataset():
+    return xr.Dataset(
+        {
+            "station_id": (["station"], [1, 2, 3, 4, 5]),
+            "elevation": (["station"], [100, 250, 500, 750, 1000]),
+            "name": (
+                ["station"],
+                [
+                    "Station_A",
+                    "Station_B",
+                    "Station_C",
+                    "Station_D",
+                    "Station_E",
+                ],
+            ),
+        }
+    ).chunk({"station": 5})
+
+
+@pytest.fixture
+def air_and_stations():
+    air = (
+        xr.tutorial.open_dataset("air_temperature")
+        .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
+        .chunk({"time": 6})
+    )
+    stations = xr.Dataset(
+        {
+            "station_id": (["station"], [101, 102, 103]),
+            "lat": (
+                ["station"],
+                [air.lat.values[0], air.lat.values[2], air.lat.values[4]],
+            ),
+            "lon": (
+                ["station"],
+                [air.lon.values[1], air.lon.values[3], air.lon.values[5]],
+            ),
+            "elevation": (["station"], [100, 250, 500]),
+        }
+    ).chunk({"station": 3})
+    return air, stations