diff --git a/docs/configuration/sources.md b/docs/configuration/sources.md index 73633721b..bef27b49d 100644 --- a/docs/configuration/sources.md +++ b/docs/configuration/sources.md @@ -37,6 +37,7 @@ ui.servable() | Oracle | Oracle via SQLAlchemy | | MSSQL | Microsoft SQL Server via SQLAlchemy | | Intake | Data catalogs | +| Xarray | NetCDF and Zarr scientific datasets | ## Database connections @@ -173,6 +174,28 @@ source = DuckDBSource( 1. Required for HTTP/S3 access +### Xarray datasets + +Use `XarraySource` when your data already lives in NetCDF or Zarr and you want +to preserve coordinate-aware filtering before flattening to a DataFrame. + +``` py title="Xarray NetCDF source" +from lumen.sources.xarray import XarraySource +import lumen.ai as lmai + +source = XarraySource( + uri="air_temperature.nc", + filterable_coords=["time", "lat", "lon"], + max_rows=5000, +) + +ui = lmai.ExplorerUI(data=source) +ui.servable() +``` + +Set `dataset_format="zarr"` for suffixless stores, or leave the default `auto` +format detection for normal NetCDF and local Zarr directories. + ### Multiple sources ``` py title="Mix sources" diff --git a/docs/configuration/spec/sources.md b/docs/configuration/spec/sources.md index 95c042afe..11a5ef73c 100644 --- a/docs/configuration/spec/sources.md +++ b/docs/configuration/spec/sources.md @@ -26,9 +26,57 @@ sources: | `intake` | Intake catalog entries | Data catalogs | | `rest` | REST API endpoints | API data | | `live` | Live website status checks | Website monitoring | +| `xarray` | NetCDF and Zarr datasets | Labeled n-dimensional scientific data | This guide focuses on `file` sources (most common). See Lumen's source reference for other types. +## Xarray sources + +`xarray` sources expose each `Dataset.data_var` as a logical table. Queries filter +coordinates with xarray selection semantics, and results are flattened to pandas +DataFrames so they can flow through normal Lumen pipelines and views. + +### Basic xarray spec + +```yaml +sources: + forecast: + type: xarray + uri: ./air_temperature.nc + filterable_coords: [time, lat, lon] + max_rows: 5000 + +pipelines: + air: + source: forecast + table: air + filters: + - type: widget + field: time + - type: widget + field: lat +``` + +Use `uri` for NetCDF or Zarr-backed datasets, or construct `XarraySource(dataset=...)` +directly in Python when you already have an in-memory `xarray.Dataset`. + +### Key parameters + +- `uri`: Local path, `file://` URI, or remote URI to an xarray-readable dataset +- `dataset_format`: `auto`, `netcdf`, or `zarr`; set explicitly for suffixless stores +- `load_kwargs`: Passed through to `xr.open_dataset()` or `xr.open_zarr()` +- `filterable_coords`: Restricts which 1D coordinates appear in schema and queries +- `max_rows`: Rejects oversized flattened results before DataFrame materialization + +### Current behavior + +- Each data variable becomes its own table +- Only 1D coordinates are queryable +- `get_schema()` includes `__len__` and supports sampled schemas via `limit` and `shuffle` +- Zarr stores can be selected explicitly with `dataset_format: zarr` or inferred from local store markers + +See [examples/xarray_air_temperature.yaml](../../../examples/xarray_air_temperature.yaml) for a minimal review-friendly spec. + ## File sources FileSource loads data from files in various formats: diff --git a/examples/xarray_air_temperature.yaml b/examples/xarray_air_temperature.yaml new file mode 100644 index 000000000..40ababb52 --- /dev/null +++ b/examples/xarray_air_temperature.yaml @@ -0,0 +1,29 @@ +config: + title: Xarray Air Temperature + reloadable: false + +sources: + air_temperature: + type: xarray + uri: ./air_temperature.nc + filterable_coords: [time, lat, lon] + max_rows: 5000 + +pipelines: + air_temperature: + source: air_temperature + table: air + filters: + - type: widget + field: time + - type: widget + field: lat + +layouts: + - title: Air Temperature + pipeline: air_temperature + views: + - type: table + page_size: 20 + sizing_mode: stretch_width + sizing_mode: stretch_width diff --git a/examples/xarray_air_temperature_demo.py b/examples/xarray_air_temperature_demo.py new file mode 100644 index 000000000..77a320559 --- /dev/null +++ b/examples/xarray_air_temperature_demo.py @@ -0,0 +1,51 @@ +""" +Small manual demo for the experimental XarraySource. + +This script uses xarray's tutorial dataset, so it requires: +- xarray +- pooch +- network access on first run to download the dataset +""" + +from __future__ import annotations + +import sys + +from pathlib import Path + +import xarray as xr + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from lumen.sources.xarray import XarraySource + + +def main() -> None: + ds = xr.tutorial.open_dataset("air_temperature") + source = XarraySource(dataset=ds) + + try: + table = source.get_tables()[0] + + print("Tables:") + print(source.get_tables()) + + print("\nSchema keys:") + print(list(source.get_schema(table).keys())) + + print("\nMetadata:") + print(source.get_metadata(table)) + + print("\nFiltered result (lat=(30.0, 60.0), time='2013-01-01'):") + result = source.get(table, lat=(30.0, 60.0), time="2013-01-01") + print(result.head()) + print(f"\nReturned {len(result)} rows") + finally: + ds.close() + source.close() + + +if __name__ == "__main__": + main() diff --git a/lumen/sources/__init__.py b/lumen/sources/__init__.py index 9b5ed21c9..a4fcef621 100644 --- a/lumen/sources/__init__.py +++ b/lumen/sources/__init__.py @@ -1 +1,16 @@ +from importlib import import_module +from importlib.util import find_spec + from .base import * +from .base import __all__ as _base_all + +__all__ = list(_base_all) + +if find_spec("xarray") is not None: + __all__.append("XarraySource") + + +def __getattr__(name): + if name == "XarraySource" and find_spec("xarray") is not None: + return import_module(f"{__name__}.xarray").XarraySource + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/lumen/sources/xarray.py b/lumen/sources/xarray.py new file mode 100644 index 000000000..764aa7e9e --- /dev/null +++ b/lumen/sources/xarray.py @@ -0,0 +1,417 @@ +from __future__ import annotations + +import datetime as dt +import hashlib + +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar +from urllib.parse import unquote, urlparse +from urllib.request import url2pathname + +import numpy as np +import pandas as pd +import param + +from ..util import get_dataframe_schema +from ..validation import ValidationError, match_suggestion_message +from .base import Source, cached, cached_schema + +try: + import xarray as xr +except ImportError: + xr = None # type: ignore + +if TYPE_CHECKING: + import xarray as xr + + +class XarraySource(Source): + """ + Minimal xarray-backed Source implementation. + + This first draft is intentionally conservative: + - accepts an in-memory xarray.Dataset + - exposes Dataset.data_vars as logical tables + - supports coordinate-based filtering via .sel(...) + - returns pandas DataFrames for downstream compatibility + """ + + dataset = param.Parameter(doc=""" + An xarray.Dataset providing the data variables exposed by the source. + """) + + uri = param.String(default=None, doc=""" + Optional path or URI to an xarray-readable dataset. Used when no in-memory + dataset is supplied. + """) + + dataset_format = param.Selector(default="auto", objects=["auto", "netcdf", "zarr"], doc=""" + Dataset storage format. By default Lumen will infer the format from + the URI, local store markers, or load kwargs. Set explicitly for + suffixless or non-standard stores. + """) + + load_kwargs = param.Dict(default={}, doc=""" + Additional keyword arguments forwarded to xarray.open_dataset/open_zarr. + """) + + filterable_coords = param.List(default=None, allow_None=True, doc=""" + Optional list of coordinate names that should be exposed in schema and + allowed in queries. By default all 1D coordinates on a variable are + queryable. + """) + + max_rows = param.Integer(default=0, bounds=(0, None), doc=""" + Optional safety guard on the number of rows returned after flattening. + A value of 0 disables the guard. + """) + + source_type: ClassVar[str] = "xarray" + + def __init__(self, **params): + super().__init__(**params) + self._loaded_dataset = None + self._validate_dataset() + + def _validate_dataset(self) -> None: + if xr is None: + raise ImportError("XarraySource requires the 'xarray' package to be installed.") + if self.dataset is None and not self.uri: + raise ValueError("XarraySource requires either a dataset or a uri.") + if self.dataset is not None and not isinstance(self.dataset, xr.Dataset): + raise TypeError("XarraySource 'dataset' must be an xarray.Dataset.") + + def _is_windows_path(self, uri: str) -> bool: + parsed = urlparse(uri) + return len(parsed.scheme) == 1 and uri[1:2] == ":" + + def _is_remote_uri(self, uri: str) -> bool: + parsed = urlparse(uri) + return bool( + parsed.scheme + and parsed.scheme not in ("file",) + and not self._is_windows_path(uri) + ) + + def _uri_to_path(self, uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme == "file": + location = f"//{parsed.netloc}{parsed.path}" if parsed.netloc else parsed.path + local_path = url2pathname(unquote(location)) + if local_path.startswith("/") and len(local_path) > 2 and local_path[2] == ":": + local_path = local_path[1:] + return Path(local_path) + return Path(uri) + + def _resolve_uri(self) -> str: + assert self.uri is not None + if self._is_remote_uri(self.uri): + return self.uri + if self._is_windows_path(self.uri): + return self.uri + path = self._uri_to_path(self.uri) + if not path.is_absolute(): + path = self.root / path + return str(path) + + def _is_local_zarr_store(self, uri: str) -> bool: + path = Path(uri) + if not path.exists() or not path.is_dir(): + return False + return (path / ".zgroup").exists() or (path / "zarr.json").exists() + + def _get_dataset_format(self, uri: str) -> str: + if self.dataset_format != "auto": + return self.dataset_format + if self.load_kwargs.get("engine") == "zarr": + return "zarr" + parsed = urlparse(uri) + if parsed.path.endswith(".zarr"): + return "zarr" + if not self._is_remote_uri(uri) and self._is_local_zarr_store(uri): + return "zarr" + return "netcdf" + + def _open_dataset(self) -> xr.Dataset: + uri = self._resolve_uri() + if self._get_dataset_format(uri) == "zarr": + return xr.open_zarr(uri, **self.load_kwargs) + return xr.open_dataset(uri, **self.load_kwargs) + + def _get_dataset(self) -> xr.Dataset: + if self.dataset is not None: + return self.dataset + if self._loaded_dataset is None: + self._loaded_dataset = self._open_dataset() + return self._loaded_dataset + + def _get_source_hash(self): + sha = hashlib.sha256() + for k, v in self.param.values().items(): + if k in ("root",): + continue + sha.update(k.encode("utf-8")) + if k == "uri" and isinstance(v, str) and not self._is_remote_uri(v): + path = Path(self._resolve_uri()) + if path.exists(): + sha.update(str(path.stat().st_mtime_ns).encode("utf-8")) + else: + sha.update(str(v).encode("utf-8")) + return sha.hexdigest() + + def close(self) -> None: + if self._loaded_dataset is not None and hasattr(self._loaded_dataset, "close"): + self._loaded_dataset.close() + self._loaded_dataset = None + + def clear_cache(self, *events: param.parameterized.Event): + self.close() + super().clear_cache(*events) + + def get_tables(self) -> list[str]: + return list(self._get_dataset().data_vars) + + def _get_array(self, table: str) -> xr.DataArray: + tables = self.get_tables() + dataset = self._get_dataset() + if table not in dataset.data_vars: + raise KeyError(f"Table {table!r} not found. Available tables: {tables}.") + arr = dataset[table] + return arr if arr.name == table else arr.rename(table) + + def _iter_queryable_coords(self, arr: xr.DataArray) -> list[tuple[str, xr.DataArray]]: + coords = [] + allowed = set(self.filterable_coords) if self.filterable_coords is not None else None + for coord_name, coord in arr.coords.items(): + if coord.ndim != 1: + continue + if allowed is not None and coord_name not in allowed: + continue + coords.append((coord_name, coord)) + return coords + + def _iter_auxiliary_coords(self, arr: xr.DataArray) -> list[tuple[str, xr.DataArray]]: + queryable = {name for name, _ in self._iter_queryable_coords(arr)} + return [(coord_name, coord) for coord_name, coord in arr.coords.items() if coord_name not in queryable] + + def _dtype_to_schema(self, dtype: np.dtype) -> dict[str, Any]: + if np.issubdtype(dtype, np.bool_): + return {"type": "boolean"} + if np.issubdtype(dtype, np.integer): + return {"type": "integer"} + if np.issubdtype(dtype, np.floating): + return {"type": "number"} + if np.issubdtype(dtype, np.datetime64): + return {"type": "string", "format": "datetime"} + return {"type": "string"} + + def _coord_schema(self, coord: xr.DataArray) -> dict[str, Any]: + schema = self._dtype_to_schema(coord.dtype) + values = coord.values + if values.ndim == 1 and len(values) and np.issubdtype(coord.dtype, np.number): + schema["inclusiveMinimum"] = values.min().item() + schema["inclusiveMaximum"] = values.max().item() + return schema + + def _get_sampled_schema(self, arr: xr.DataArray, table: str, limit: int | None, shuffle: bool) -> dict[str, Any]: + df = arr.to_dataframe(name=table).reset_index() + ordered_columns = [*(name for name, _ in self._iter_queryable_coords(arr)), table] + df = df[[col for col in ordered_columns if col in df.columns]] + if shuffle and not df.empty: + sample_size = min(limit or len(df), len(df)) + df = df.sample(n=sample_size, random_state=0) + elif limit is not None: + df = df.head(limit) + return get_dataframe_schema(df)["items"]["properties"] + + @cached_schema + def get_schema( + self, table: str | None = None, limit: int | None = None, shuffle: bool = False + ) -> dict[str, dict[str, Any]] | dict[str, Any]: + tables = [table] if table is not None else self.get_tables() + schemas: dict[str, dict[str, Any]] = {} + names = self.get_tables() + + for name in tables: + try: + arr = self._get_array(name) + except KeyError as e: + msg = f"{type(self).__name__} does not contain {name!r}." + msg = match_suggestion_message(name, names, msg) + raise ValidationError(msg) from e + if limit is not None or shuffle: + schema = self._get_sampled_schema(arr, name, limit, shuffle) + else: + schema = {} + for coord_name, coord in self._iter_queryable_coords(arr): + schema[coord_name] = self._coord_schema(coord) + schema[name] = self._dtype_to_schema(arr.dtype) + schema["__len__"] = self._get_result_length(arr) + schemas[name] = schema + + return schemas if table is None else schemas[table] + + def _normalize_query_value(self, value: Any) -> Any: + if isinstance(value, tuple): + if len(value) != 2: + raise ValueError( + "Tuple query values must contain exactly two elements for range selection." + ) + return slice(value[0], value[1]) + return value + + def _normalize_scalar_for_coord(self, coord: xr.DataArray, value: Any) -> Any: + if np.issubdtype(coord.dtype, np.datetime64): + if isinstance(value, dt.date) and not isinstance(value, dt.datetime): + return np.datetime64(value) + if isinstance(value, str | dt.datetime | np.datetime64): + return np.datetime64(value) + return value + + def _normalize_selector(self, coord: xr.DataArray, value: Any) -> Any: + if isinstance(value, slice): + return slice( + self._normalize_scalar_for_coord(coord, value.start), + self._normalize_scalar_for_coord(coord, value.stop), + value.step, + ) + if isinstance(value, list): + return [self._normalize_scalar_for_coord(coord, v) for v in value] + return self._normalize_scalar_for_coord(coord, value) + + def _select_range(self, arr: xr.DataArray, key: str, selector: slice) -> xr.DataArray: + coord = arr.coords[key] + if coord.ndim != 1: + raise ValueError( + f"Filtering on non-1D coordinate {key!r} is not supported; got {coord.ndim}D coordinate." + ) + values = coord.values + if len(values) < 2: + return arr.sel({key: selector}) + + start = selector.start + stop = selector.stop + if start is None or stop is None: + return arr.sel({key: selector}) + + increasing = bool(values[0] <= values[-1]) + if increasing: + ordered = slice(start, stop) + else: + ordered = slice(stop, start) if start < stop else slice(start, stop) + return arr.sel({key: ordered}) + + def _is_range_list(self, value: Any) -> bool: + return ( + isinstance(value, list) + and bool(value) + and all(isinstance(v, tuple) and len(v) == 2 for v in value) + ) + + def _select_range_list(self, arr: xr.DataArray, key: str, ranges: list[tuple[Any, Any]]) -> xr.DataArray: + coord = arr.coords[key] + if coord.ndim != 1: + raise ValueError( + f"Range-list filtering is only supported for 1D coordinates, got {coord.ndim}D for {key!r}." + ) + values = coord.values + mask = np.zeros(len(values), dtype=bool) + for start, end in ranges: + start = self._normalize_scalar_for_coord(coord, start) + end = self._normalize_scalar_for_coord(coord, end) + mask |= (values >= start) & (values <= end) + return arr.sel({key: values[mask]}) + + def _apply_query(self, arr: xr.DataArray, query: dict[str, Any]) -> xr.DataArray: + queryable = dict(self._iter_queryable_coords(arr)) + for key, value in query.items(): + if key.startswith("__"): + continue + if key not in queryable: + if key in arr.coords and arr.coords[key].ndim != 1: + raise ValueError( + f"Filtering on non-1D coordinate {key!r} is not supported; got {arr.coords[key].ndim}D coordinate." + ) + raise ValueError( + f"Unsupported query key {key!r} for table {arr.name!r}. " + f"Expected one of: {list(queryable)}." + ) + coord = queryable[key] + if self._is_range_list(value): + arr = self._select_range_list(arr, key, value) + continue + normalized = self._normalize_query_value(value) + normalized = self._normalize_selector(coord, normalized) + try: + if isinstance(normalized, slice): + arr = self._select_range(arr, key, normalized) + else: + arr = arr.sel({key: normalized}) + except Exception as e: + raise ValueError( + f"Could not apply query {key}={value!r} to table {arr.name!r}." + ) from e + return arr + + def _get_result_length(self, arr: xr.DataArray) -> int: + if arr.ndim == 0: + return 1 + return int(np.prod(arr.shape, dtype=np.int64)) + + def _enforce_max_rows(self, arr: xr.DataArray, table: str) -> None: + if not self.max_rows: + return + rows = self._get_result_length(arr) + if rows > self.max_rows: + raise ValueError( + f"Query result for table {table!r} produced {rows} rows, " + f"which exceeds max_rows={self.max_rows}." + ) + + def _to_dataframe(self, arr: xr.DataArray, table: str) -> pd.DataFrame: + arr = arr if arr.name == table else arr.rename(table) + self._enforce_max_rows(arr, table) + if arr.ndim == 0: + return pd.DataFrame({table: [arr.item()]}) + df = arr.to_dataframe(name=table).reset_index() + ordered_columns = [*(name for name, _ in self._iter_queryable_coords(self._get_array(table))), table] + df = df[[col for col in ordered_columns if col in df.columns]] + return df + + def _get_table_metadata(self, tables: list[str]) -> dict[str, dict[str, Any]]: + metadata: dict[str, dict[str, Any]] = {} + for table in tables: + arr = self._get_array(table) + columns: dict[str, dict[str, Any]] = {} + queryable_coords = self._iter_queryable_coords(arr) + for coord_name, coord in queryable_coords: + coord_meta = { + "description": coord.attrs.get("long_name") or coord.attrs.get("description") or f"{coord_name} coordinate", + "data_type": str(coord.dtype), + } + if "units" in coord.attrs: + coord_meta["unit"] = coord.attrs["units"] + columns[coord_name] = coord_meta + + value_meta = { + "description": arr.attrs.get("long_name") or arr.attrs.get("description") or table, + "data_type": str(arr.dtype), + } + if "units" in arr.attrs: + value_meta["unit"] = arr.attrs["units"] + columns[table] = value_meta + + metadata[table] = { + "description": arr.attrs.get("long_name") or arr.attrs.get("description") or table, + "dims": list(arr.dims), + "queryable_coords": [coord_name for coord_name, _ in queryable_coords], + "auxiliary_coords": [coord_name for coord_name, _ in self._iter_auxiliary_coords(arr)], + "columns": columns, + } + return metadata + + @cached + def get(self, table: str, **query) -> pd.DataFrame: + arr = self._get_array(table) + arr = self._apply_query(arr, query) + return self._to_dataframe(arr, table) diff --git a/lumen/tests/sources/test_xarray.py b/lumen/tests/sources/test_xarray.py new file mode 100644 index 000000000..475ec1d03 --- /dev/null +++ b/lumen/tests/sources/test_xarray.py @@ -0,0 +1,512 @@ +import datetime as dt + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import yaml + +xr = pytest.importorskip("xarray") + +from lumen.dashboard import Dashboard +from lumen.sources.xarray import XarraySource +from lumen.validation import ValidationError + + +def assert_df_equal(df1, df2): + pd.testing.assert_frame_equal(df1, df2, check_dtype=False) + + +@pytest.fixture +def xr_dataset(): + time = np.array(["2024-01-01", "2024-01-02"], dtype="datetime64[ns]") + lat = np.array([10.0, 20.0]) + lon = np.array([70.0, 80.0]) + + temperature = np.array([ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + ]) + pressure = np.array([ + [[101.0, 102.0], [103.0, 104.0]], + [[105.0, 106.0], [107.0, 108.0]], + ]) + + ds = xr.Dataset( + data_vars={ + "temperature": (("time", "lat", "lon"), temperature), + "pressure": (("time", "lat", "lon"), pressure), + }, + coords={ + "time": time, + "lat": lat, + "lon": lon, + }, + ) + ds["temperature"].attrs["units"] = "degC" + ds["temperature"].attrs["long_name"] = "Air temperature" + ds["lat"].attrs["long_name"] = "Latitude" + return ds + + +def test_xarray_source_get_tables(xr_dataset): + source = XarraySource(dataset=xr_dataset) + assert source.get_tables() == ["temperature", "pressure"] + + +def test_xarray_source_loads_from_uri(tmp_path, xr_dataset): + path = tmp_path / "sample.nc" + xr_dataset.to_netcdf(path) + + source = XarraySource(uri=str(path)) + + assert source.get_tables() == ["temperature", "pressure"] + result = source.get("temperature", time="2024-01-01") + assert len(result) == 4 + + +def test_xarray_source_explicit_zarr_format_uses_open_zarr(monkeypatch): + calls = [] + expected = object() + + def fake_open_zarr(uri, **kwargs): + calls.append(("zarr", uri, kwargs)) + return expected + + def fake_open_dataset(uri, **kwargs): + calls.append(("dataset", uri, kwargs)) + raise AssertionError("open_dataset should not be called for explicit zarr format") + + monkeypatch.setattr(xr, "open_zarr", fake_open_zarr) + monkeypatch.setattr(xr, "open_dataset", fake_open_dataset) + + source = XarraySource(uri="memory-store", dataset_format="zarr", load_kwargs={"chunks": {}}) + resolved = source._resolve_uri() + + assert source._open_dataset() is expected + assert calls == [("zarr", resolved, {"chunks": {}})] + + +def test_xarray_source_auto_detects_local_zarr_store_without_suffix(tmp_path, monkeypatch): + store = tmp_path / "temperature_store" + store.mkdir() + (store / "zarr.json").write_text("{}", encoding="utf-8") + calls = [] + expected = object() + + def fake_open_zarr(uri, **kwargs): + calls.append(("zarr", uri, kwargs)) + return expected + + def fake_open_dataset(uri, **kwargs): + calls.append(("dataset", uri, kwargs)) + raise AssertionError("open_dataset should not be called for a detected local zarr store") + + monkeypatch.setattr(xr, "open_zarr", fake_open_zarr) + monkeypatch.setattr(xr, "open_dataset", fake_open_dataset) + + source = XarraySource(uri=str(store)) + + assert source._open_dataset() is expected + assert calls == [("zarr", str(store), {})] + + +def test_xarray_source_windows_uri_treated_as_local_path(tmp_path, xr_dataset): + path = tmp_path / "sample.nc" + xr_dataset.to_netcdf(path) + + source = XarraySource(uri=r"C:\data\sample.nc", root=tmp_path) + + resolved = source._resolve_uri() + + assert resolved == r"C:\data\sample.nc" + assert not source._is_remote_uri(r"C:\data\sample.nc") + + +def test_xarray_example_spec_renders_dashboard(tmp_path): + example_path = Path(__file__).resolve().parents[3] / "examples" / "xarray_air_temperature.yaml" + spec = yaml.safe_load(example_path.read_text(encoding="utf-8")) + + ds = xr.Dataset( + data_vars={ + "air": ( + ("time", "lat", "lon"), + np.array([ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + ]), + ), + }, + coords={ + "time": np.array(["2024-01-01", "2024-01-02"], dtype="datetime64[ns]"), + "lat": np.array([10.0, 20.0]), + "lon": np.array([70.0, 80.0]), + }, + ) + ds.to_netcdf(tmp_path / "air_temperature.nc") + + dashboard = Dashboard(spec, root=str(tmp_path)) + dashboard._render_dashboard() + + layout = dashboard.layouts[0] + pipeline = next(iter(layout._pipelines.values())) + + assert isinstance(pipeline.source, XarraySource) + assert [filt.field for filt in pipeline.filters] == ["time", "lat"] + assert list(pipeline.data.columns) == ["time", "lat", "lon", "air"] + assert len(pipeline.data) == 8 + + +def test_xarray_source_get_schema_single_table(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + schema = source.get_schema("temperature") + + assert set(schema) == {"time", "lat", "lon", "temperature", "__len__"} + assert schema["time"]["type"] == "string" + assert schema["time"]["format"] == "datetime" + assert schema["lat"]["type"] == "number" + assert schema["lon"]["type"] == "number" + assert schema["temperature"]["type"] == "number" + assert schema["__len__"] == 8 + + +def test_xarray_source_get_schema_only_exposes_queryable_1d_coords(): + ds = xr.Dataset( + data_vars={ + "temperature": ( + ("x", "y"), + np.array([[1.0, 2.0], [3.0, 4.0]]), + ), + }, + coords={ + "x": np.array([0, 1]), + "y": np.array([0, 1]), + "mesh": (("x", "y"), np.array([[10, 20], [30, 40]])), + }, + ) + source = XarraySource(dataset=ds) + + schema = source.get_schema("temperature") + + assert set(schema) == {"x", "y", "temperature", "__len__"} + + +def test_xarray_source_get_schema_all_tables(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + schemas = source.get_schema() + + assert set(schemas) == {"temperature", "pressure"} + assert set(schemas["temperature"]) == {"time", "lat", "lon", "temperature", "__len__"} + assert set(schemas["pressure"]) == {"time", "lat", "lon", "pressure", "__len__"} + assert schemas["temperature"]["__len__"] == 8 + assert schemas["pressure"]["__len__"] == 8 + + +def test_xarray_source_get_schema_with_limit(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + schema = source.get_schema("temperature", limit=1) + + assert schema == { + "time": { + "type": "string", + "format": "datetime", + "inclusiveMinimum": "2024-01-01T00:00:00", + "inclusiveMaximum": "2024-01-01T00:00:00", + }, + "lat": {"type": "number", "inclusiveMinimum": 10.0, "inclusiveMaximum": 10.0}, + "lon": {"type": "number", "inclusiveMinimum": 70.0, "inclusiveMaximum": 70.0}, + "temperature": {"type": "number", "inclusiveMinimum": 1.0, "inclusiveMaximum": 1.0}, + "__len__": 8, + } + + +def test_xarray_source_get_schema_with_shuffle(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + schema = source.get_schema("temperature", limit=1, shuffle=True) + + assert schema == { + "time": { + "type": "string", + "format": "datetime", + "inclusiveMinimum": "2024-01-02T00:00:00", + "inclusiveMaximum": "2024-01-02T00:00:00", + }, + "lat": {"type": "number", "inclusiveMinimum": 20.0, "inclusiveMaximum": 20.0}, + "lon": {"type": "number", "inclusiveMinimum": 70.0, "inclusiveMaximum": 70.0}, + "temperature": {"type": "number", "inclusiveMinimum": 7.0, "inclusiveMaximum": 7.0}, + "__len__": 8, + } + + +def test_xarray_source_get_returns_dataframe(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature") + + assert isinstance(result, pd.DataFrame) + assert list(result.columns) == ["time", "lat", "lon", "temperature"] + assert len(result) == 8 + + +def test_xarray_source_get_scalar_coord_filter(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", time=np.datetime64("2024-01-01")) + expected = pd.DataFrame( + { + "time": np.array(["2024-01-01"] * 4, dtype="datetime64[ns]"), + "lat": [10.0, 10.0, 20.0, 20.0], + "lon": [70.0, 80.0, 70.0, 80.0], + "temperature": [1.0, 2.0, 3.0, 4.0], + } + ) + + assert_df_equal(result, expected) + + +def test_xarray_source_get_scalar_date_coord_filter(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", time=dt.date(2024, 1, 1)) + + assert len(result) == 4 + assert result["time"].nunique() == 1 + assert pd.Timestamp(result["time"].iloc[0]) == pd.Timestamp("2024-01-01") + + +def test_xarray_source_get_scalar_string_datetime_coord_filter(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", time="2024-01-01") + + assert len(result) == 4 + assert result["time"].nunique() == 1 + assert pd.Timestamp(result["time"].iloc[0]) == pd.Timestamp("2024-01-01") + + +def test_xarray_source_get_range_coord_filter(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", lat=(10.0, 10.0)) + expected = pd.DataFrame( + { + "time": np.array(["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"], dtype="datetime64[ns]"), + "lat": [10.0, 10.0, 10.0, 10.0], + "lon": [70.0, 80.0, 70.0, 80.0], + "temperature": [1.0, 2.0, 5.0, 6.0], + } + ) + + assert_df_equal(result, expected) + + +def test_xarray_source_get_date_range_coord_filter(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", time=(dt.date(2024, 1, 1), dt.date(2024, 1, 1))) + + assert len(result) == 4 + assert result["time"].nunique() == 1 + assert pd.Timestamp(result["time"].iloc[0]) == pd.Timestamp("2024-01-01") + + +def test_xarray_source_get_range_coord_filter_descending_coordinate(): + ds = xr.Dataset( + data_vars={ + "temperature": ( + ("lat", "lon"), + np.array([ + [1.0, 2.0], + [3.0, 4.0], + [5.0, 6.0], + ]), + ), + }, + coords={ + "lat": np.array([30.0, 20.0, 10.0]), + "lon": np.array([70.0, 80.0]), + }, + ) + source = XarraySource(dataset=ds) + + result = source.get("temperature", lat=(10.0, 20.0)) + + assert list(result["lat"].unique()) == [20.0, 10.0] + assert len(result) == 4 + + +def test_xarray_source_get_range_list_coord_filter(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", time=[(np.datetime64("2024-01-01"), np.datetime64("2024-01-01"))]) + expected = pd.DataFrame( + { + "time": np.array(["2024-01-01"] * 4, dtype="datetime64[ns]"), + "lat": [10.0, 10.0, 20.0, 20.0], + "lon": [70.0, 80.0, 70.0, 80.0], + "temperature": [1.0, 2.0, 3.0, 4.0], + } + ) + + assert_df_equal(result, expected) + + +def test_xarray_source_metadata_preserves_variable_attrs(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + metadata = source.get_metadata("temperature") + + assert metadata["description"] == "Air temperature" + assert metadata["dims"] == ["time", "lat", "lon"] + assert metadata["columns"]["temperature"]["unit"] == "degC" + assert metadata["columns"]["lat"]["description"] == "Latitude" + + +def test_xarray_source_metadata_reports_auxiliary_coords(): + ds = xr.Dataset( + data_vars={ + "temperature": ( + ("x", "y"), + np.array([[1.0, 2.0], [3.0, 4.0]]), + ), + }, + coords={ + "x": np.array([0, 1]), + "y": np.array([0, 1]), + "mesh": (("x", "y"), np.array([[10, 20], [30, 40]])), + }, + ) + source = XarraySource(dataset=ds) + + metadata = source.get_metadata("temperature") + + assert metadata["queryable_coords"] == ["x", "y"] + assert metadata["auxiliary_coords"] == ["mesh"] + + +def test_xarray_source_unknown_table_raises(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + with pytest.raises(KeyError): + source.get("not_a_table") + + +def test_xarray_source_unknown_schema_table_raises_validation_error(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + with pytest.raises(ValidationError, match="does not contain"): + source.get_schema("not_a_table") + + +def test_xarray_source_unknown_query_field_raises(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + with pytest.raises(ValueError, match="Unsupported query key"): + source.get("temperature", station=1) + + +def test_xarray_source_filterable_coords_restrict_query_surface(xr_dataset): + source = XarraySource(dataset=xr_dataset, filterable_coords=["time"]) + + schema = source.get_schema("temperature") + + assert set(schema) == {"time", "temperature", "__len__"} + with pytest.raises(ValueError, match="Unsupported query key"): + source.get("temperature", lat=(10.0, 20.0)) + + +def test_xarray_source_invalid_range_tuple_raises(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + with pytest.raises(ValueError, match="exactly two elements"): + source.get("temperature", lat=(10.0, 20.0, 30.0)) + + +def test_xarray_source_empty_selection_returns_empty_dataframe(xr_dataset): + source = XarraySource(dataset=xr_dataset) + + result = source.get("temperature", time=[]) + + assert isinstance(result, pd.DataFrame) + assert list(result.columns) == ["time", "lat", "lon", "temperature"] + assert result.empty + + +def test_xarray_source_handles_variable_with_different_dims(): + ds = xr.Dataset( + data_vars={ + "temperature": ( + ("time", "lat", "lon"), + np.array([ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + ]), + ), + "station_bias": ( + ("lat", "lon"), + np.array([ + [0.1, 0.2], + [0.3, 0.4], + ]), + ), + }, + coords={ + "time": np.array(["2024-01-01", "2024-01-02"], dtype="datetime64[ns]"), + "lat": np.array([10.0, 20.0]), + "lon": np.array([70.0, 80.0]), + }, + ) + source = XarraySource(dataset=ds) + + schema = source.get_schema("station_bias") + result = source.get("station_bias") + + assert set(source.get_tables()) == {"temperature", "station_bias"} + assert set(schema) == {"lat", "lon", "station_bias", "__len__"} + assert list(result.columns) == ["lat", "lon", "station_bias"] + assert len(result) == 4 + assert schema["__len__"] == 4 + + +def test_xarray_source_max_rows_guard_raises(xr_dataset): + source = XarraySource(dataset=xr_dataset, max_rows=2) + + with pytest.raises(ValueError, match="exceeds max_rows=2"): + source.get("temperature") + + +def test_xarray_source_max_rows_guard_raises_before_materializing_dataframe(xr_dataset, monkeypatch): + source = XarraySource(dataset=xr_dataset, max_rows=2) + + def fail_to_dataframe(*args, **kwargs): + raise AssertionError("to_dataframe should not be called when max_rows is exceeded") + + monkeypatch.setattr(xr.DataArray, "to_dataframe", fail_to_dataframe) + + with pytest.raises(ValueError, match="exceeds max_rows=2"): + source.get("temperature") + + +def test_xarray_source_non_1d_coordinate_filter_raises(): + ds = xr.Dataset( + data_vars={ + "temperature": ( + ("x", "y"), + np.array([[1.0, 2.0], [3.0, 4.0]]), + ), + }, + coords={ + "x": np.array([0, 1]), + "y": np.array([0, 1]), + "mesh": (("x", "y"), np.array([[10, 20], [30, 40]])), + }, + ) + source = XarraySource(dataset=ds) + + with pytest.raises(ValueError, match="non-1D coordinate"): + source.get("temperature", mesh=(10, 20)) diff --git a/pixi.toml b/pixi.toml index 06a9d438d..203f2b588 100644 --- a/pixi.toml +++ b/pixi.toml @@ -80,6 +80,8 @@ pytest-rerunfailures = "*" pytest-xdist = "*" pytest-asyncio = "*" dask = "<=2024.10" # Temporary workaround for pyarrow string issue in tests +xarray = "*" +scipy = "*" [feature.test-core.tasks] test-unit = 'pytest lumen/tests -n logical --dist loadgroup' diff --git a/pyproject.toml b/pyproject.toml index 288607d17..063f8121e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ HoloViz = "https://holoviz.org/" [project.optional-dependencies] tests = ['pytest', 'pytest-rerunfailures', 'pytest-asyncio'] +xarray = ['xarray'] sql = ['sqlalchemy'] ai-local = ['huggingface_hub', 'hf_xet'] ai-openai = ['openai']