diff --git a/docs/configuration/sources.md b/docs/configuration/sources.md index 73633721b..5305e3b1e 100644 --- a/docs/configuration/sources.md +++ b/docs/configuration/sources.md @@ -29,6 +29,7 @@ ui.servable() |--------|---------| | Files | CSV, Parquet, JSON (local or URL) | | DuckDB | Local SQL queries on files | +| xarray | N-dimensional scientific data (NetCDF, Zarr, HDF5) | | Snowflake | Cloud data warehouse | | BigQuery | Google's data warehouse | | PostgreSQL | PostgreSQL via SQLAlchemy | @@ -173,6 +174,45 @@ source = DuckDBSource( 1. Required for HTTP/S3 access +### xarray for scientific data + +Query N-dimensional scientific datasets (NetCDF, Zarr, HDF5) with SQL via Apache DataFusion: + +``` bash title="Install dependencies" +pip install lumen[xarray] +``` + +``` py title="Query a NetCDF file" +from lumen.sources.xarray_sql import XArraySQLSource +import lumen.ai as lmai + +source = XArraySQLSource(uri='air_temperature.nc') + +ui = lmai.ExplorerUI(data=source) +ui.servable() +``` + +Each data variable in the dataset becomes a separate SQL table with coordinate columns: + +``` py title="Direct SQL queries" +source = XArraySQLSource(uri='air_temperature.nc') +source.get_tables() # ['air'] + +df = source.execute( + "SELECT lat, lon, AVG(air) as avg_temp " + "FROM air GROUP BY lat, lon" +) +``` + +**Select specific variables:** + +``` py hl_lines="3" +source = XArraySQLSource( + uri='climate_data.nc', + variables=['temperature', 'pressure'] +) +``` + ### Multiple sources ``` py title="Mix sources" diff --git a/lumen/sources/xarray_sql.py b/lumen/sources/xarray_sql.py new file mode 100644 index 000000000..1b51d7886 --- /dev/null +++ b/lumen/sources/xarray_sql.py @@ -0,0 +1,351 @@ +""" +XArraySQLSource -- SQL-backed xarray source for Lumen. + +Uses xarray-sql (Apache DataFusion) to provide SQL query capabilities +over N-dimensional xarray datasets. Each data variable in the dataset +is exposed as a separate SQL table with coordinate columns. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, ClassVar + +import param + +from ..transforms.sql import SQLFilter +from .base import BaseSQLSource, cached + + +def _check_xarray_available(): + """Check xarray and xarray-sql are installed, import lazily.""" + try: + import xarray # noqa + import xarray_sql # noqa + return True + except ImportError: + return False + +class XArraySQLSource(BaseSQLSource): + """ + SQL-queryable xarray data source for Lumen. + + Wraps xarray datasets with Apache DataFusion (via xarray-sql) to + provide full SQL query support. Each data variable in the dataset + is exposed as a separate SQL table with coordinate columns. + + To use this source, install the optional xarray dependencies:: + + pip install lumen[xarray] + + Parameters + ---------- + uri : str + Path or URL to a NetCDF/Zarr/HDF5/GRIB file. + engine : str or None + xarray engine to use. Auto-detected from file extension if None. + chunks : dict or str or None + Chunking specification for dask-backed lazy loading. + Use 'auto' for automatic chunking, or a dict like {'time': 100}. + open_kwargs : dict + Additional keyword arguments passed to xr.open_dataset(). + variables : list or None + Subset of data variables to expose. None means all variables. + + Examples + -------- + >>> source = XArraySQLSource(uri="air_temperature.nc") + >>> source.get_tables() + ['air'] + >>> df = source.execute("SELECT lat, lon, AVG(air) FROM air GROUP BY lat, lon") + >>> schema = source.get_schema("air") + """ + + source_type: ClassVar[str | None] = "xarray" + + _xarray_engines: ClassVar[dict[str, str]] = { + ".nc": "netcdf4", + ".nc4": "netcdf4", + ".netcdf": "netcdf4", + ".zarr": "zarr", + ".grib": "cfgrib", + ".grib2": "cfgrib", + ".grb": "cfgrib", + ".grb2": "cfgrib", + } + + uri = param.String(default=None, allow_None=True, doc=""" + Path or URL to the xarray-compatible data file. + Supports NetCDF (.nc), Zarr (.zarr), HDF5 (.h5), GRIB (.grib).""") + + engine = param.String(default=None, allow_None=True, doc=""" + xarray backend engine. Auto-detected from file extension if None. + Options: netcdf4, zarr, cfgrib, h5netcdf.""") + + chunks = param.Parameter(default="auto", doc=""" + Dask chunk specification for lazy loading. + Use 'auto', a dict like {'time': 100}, or None for eager loading.""") + + open_kwargs = param.Dict(default={}, doc=""" + Additional keyword arguments for xr.open_dataset().""") + + variables = param.List(default=None, allow_None=True, doc=""" + Subset of data variables to expose as tables. + None means all variables are exposed.""") + + tables = param.ClassSelector(class_=(list, dict), default=None, allow_None=True, doc=""" + List or dict of tables. Populated automatically from dataset variables. + When a dict, maps logical table names to SQL expressions.""") + + sql_expr = param.String(default='SELECT * FROM {table}', doc=""" + The SQL expression template for table queries.""") + + def __init__(self, _dataset=None, _ctx=None, **params): + if not _check_xarray_available(): + raise ImportError( + "xarray and xarray-sql are required for XArraySQLSource. " + "Install them with: pip install lumen[xarray]" + ) + super().__init__(**params) + self._dataset = _dataset + self._ctx = _ctx + self._registered_tables = set() + if self._ctx is None: + self._initialize() + else: + if self._dataset is not None: + self._registered_tables = set(self._dataset.data_vars) + if self.tables is None: + self.tables = list(self._registered_tables) + + @classmethod + def _detect_engine(cls, path: str) -> str | None: + """Detect xarray engine from file extension.""" + suffix = Path(path).suffix.lower() + return cls._xarray_engines.get(suffix) + + def _initialize(self): + """Load the xarray dataset and register variables with DataFusion.""" + self._dataset = self._load_dataset( + dataset=self._dataset, + uri=self.uri, + engine=self.engine, + chunks=self.chunks, + open_kwargs=self.open_kwargs, + variables=self.variables, + ) + + from xarray_sql import XarrayContext + self._ctx = XarrayContext() + chunk_spec = self._resolve_chunks(self._dataset, self.chunks) + + for var_name in self._dataset.data_vars: + var_ds = self._dataset[[var_name]] + self._ctx.from_dataset(var_name, var_ds, chunks=chunk_spec) + self._registered_tables.add(var_name) + + if self.tables is None: + self.tables = list(self._registered_tables) + + @staticmethod + def _resolve_chunks(dataset, chunks): + """Resolve chunk specification for DataFusion registration.""" + if isinstance(chunks, dict): + return chunks + if dataset.chunks: + return {dim: sizes[0] for dim, sizes in dataset.chunks.items()} + return dict(dataset.sizes) + + @classmethod + def _load_dataset(cls, dataset, uri, engine, chunks, open_kwargs, variables): + """Load an xarray dataset from URI or use a provided one.""" + if dataset is not None: + ds = dataset + elif uri is not None: + resolved_engine = engine or cls._detect_engine(uri) + kw = dict(open_kwargs) + if resolved_engine: + kw["engine"] = resolved_engine + if chunks is not None: + kw["chunks"] = chunks + import xarray as xr + ds = xr.open_dataset(uri, **kw) + else: + raise ValueError("Either 'uri' or '_dataset' must be provided.") + + if variables: + available = list(ds.data_vars) + missing = [v for v in variables if v not in available] + if missing: + raise ValueError( + f"Variables {missing} not found in dataset. " + f"Available: {available}" + ) + ds = ds[variables] + return ds + + @classmethod + def from_dataset(cls, dataset, chunks='auto', **kwargs): + """ + Creates an in-memory XArraySQLSource from an xarray Dataset. + + Arguments + --------- + dataset: xr.Dataset + The xarray Dataset to query. + chunks: dict or str + Chunk specification for DataFusion registration. + kwargs: any + Additional keyword arguments for the source. + + Returns + ------- + source: XArraySQLSource + """ + return cls(_dataset=dataset, chunks=chunks, **kwargs) + + @property + def dataset(self): + """Access the underlying xarray Dataset.""" + return self._dataset + + def get_schema(self, table=None, limit=None, shuffle=False): + # DataFusion supports neither TABLESAMPLE nor ORDER BY RAND(), + # so force shuffle=False to make the parent use SQLLimit instead. + return super().get_schema(table, limit, shuffle=False) + + # ---- BaseSQLSource required overrides ---- + + def get_tables(self) -> list[str]: + if isinstance(self.tables, (dict, list)): + return sorted(self.tables) + return sorted(self._registered_tables) + + def normalize_table(self, table: str) -> str: + table = table.strip('"').strip("'").strip("`") + known = set(self._registered_tables) + if isinstance(self.tables, dict): + known.update(self.tables.keys()) + elif isinstance(self.tables, list): + known.update(self.tables) + if table in known: + return table + # Case-insensitive fallback (DataFusion lowercases unquoted identifiers) + lower_map = {t.lower(): t for t in known} + return lower_map.get(table.lower(), table) + + def execute(self, sql_query, params=None, *args, **kwargs): + """Execute a SQL query against the DataFusion context.""" + result = self._ctx.sql(sql_query) + return result.to_pandas() + + @cached + def get(self, table, **query): + query.pop('__dask', None) + sql_expr = self.get_sql_expr(table) + sql_transforms = query.pop('sql_transforms', []) + conditions = [ + (col, (val.start, val.stop) if isinstance(val, slice) else val) + for col, val in query.items() + if not col.startswith("__") + ] + sql_transforms = [SQLFilter(conditions=conditions, read=self.dialect)] + sql_transforms + for st in sql_transforms: + sql_expr = st.apply(sql_expr) + return self.execute(sql_expr) + + def create_sql_expr_source( + self, + tables: dict[str, str], + params: dict[str, list | dict] | None = None, + **kwargs, + ): + """Create a new XArraySQLSource sharing the same DataFusion context.""" + all_tables = {} + if isinstance(self.tables, dict): + all_tables.update(self.tables) + elif isinstance(self.tables, (list, tuple)): + for t in self.tables: + all_tables[t] = self.sql_expr.format(table=t) + all_tables.update(tables) + + all_params = dict(self.table_params) + if params: + all_params.update(params) + + return XArraySQLSource( + _dataset=self._dataset, + _ctx=self._ctx, + uri=self.uri, + engine=self.engine, + chunks=self.chunks, + variables=self.variables, + open_kwargs=self.open_kwargs, + tables=all_tables, + table_params=all_params, + ) + + def _build_column_metadata( + self, var: Any, table: str, + ) -> dict[str, dict[str, str]]: + """Build per-column metadata for a single data variable.""" + ds = self._dataset + columns = {} + for dim in var.dims: + if dim not in ds.coords: + columns[dim] = { + "data_type": "int64", + "description": f"index along {dim} dimension", + } + continue + coord_attrs = dict(ds.coords[dim].attrs) + columns[dim] = { + "data_type": str(ds.coords[dim].dtype), + "description": coord_attrs.get("long_name", ""), + } + columns[table] = { + "data_type": str(var.dtype), + "description": var.attrs.get("long_name", ""), + } + return columns + + def _format_aux_coords(self, var: Any) -> str: + """Format auxiliary coordinate labels for a data variable.""" + ds = self._dataset + aux_coords = [c for c in var.coords if c not in var.dims and c != var.name] + if not aux_coords: + return "" + labels = [] + for ac in aux_coords: + long_name = ds.coords[ac].attrs.get("long_name") + labels.append(f"{ac} [{long_name}]" if long_name else ac) + return f" (auxiliary coords: {', '.join(labels)})" + + def _get_table_metadata(self, tables: list[str]) -> dict[str, Any]: + """Build metadata from xarray attributes for AI table discovery.""" + ds = self._dataset + metadata = {} + for table in tables: + if table not in ds.data_vars: + continue + var = ds[table] + attrs = dict(var.attrs) + + description = attrs.get("long_name", table) + if "units" in attrs: + description += f" [{attrs['units']}]" + description += self._format_aux_coords(var) + + metadata[table] = { + "description": description, + "columns": self._build_column_metadata(var, table), + "rows": var.size, + } + return metadata + + def close(self): + """Clean up resources.""" + if self._dataset is not None: + self._dataset.close() + self._ctx = None + self._registered_tables.clear() diff --git a/lumen/tests/sources/test_xarray_sql.py b/lumen/tests/sources/test_xarray_sql.py new file mode 100644 index 000000000..70b136a94 --- /dev/null +++ b/lumen/tests/sources/test_xarray_sql.py @@ -0,0 +1,529 @@ +""" +Tests for XArraySQLSource -- SQL-backed xarray source via DataFusion. + +Tests cover: +- Construction from dataset, file path, and Zarr store +- SQL query execution via DataFusion +- Table listing, schema, and metadata generation +- Lumen integration (get, get_sql_expr, normalize_table) +- Serialization (to_spec / from_spec) +- Async operations +- Edge cases and error handling +""" + +import warnings + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +xr = pytest.importorskip("xarray") +pytest.importorskip("xarray_sql") + +from lumen.sources.xarray_sql import XArraySQLSource + +# ---- Fixtures ---- + +@pytest.fixture +def synthetic_dataset(): + """ + Synthetic 3D dataset with known values for deterministic testing. + Dimensions: time(10) x lat(5) x lon(4) + Variables: temperature, pressure + """ + times = pd.date_range("2020-01-01", periods=10, freq="D") + lats = np.array([10.0, 20.0, 30.0, 40.0, 50.0]) + lons = np.array([100.0, 110.0, 120.0, 130.0]) + + rng = np.random.default_rng(42) + temp = rng.uniform(250, 310, (10, 5, 4)) + pressure = rng.uniform(950, 1050, (10, 5, 4)) + + return xr.Dataset( + { + "temperature": (["time", "lat", "lon"], temp, { + "units": "K", + "long_name": "Air Temperature", + }), + "pressure": (["time", "lat", "lon"], pressure, { + "units": "hPa", + "long_name": "Surface Pressure", + }), + }, + coords={ + "time": times, + "lat": ("lat", lats, {"units": "degrees_north"}), + "lon": ("lon", lons, {"units": "degrees_east"}), + }, + attrs={ + "title": "Synthetic Test Dataset", + "source": "lumen-xarray tests", + }, + ) + + +@pytest.fixture +def nc_file(synthetic_dataset, tmp_path): + """Write synthetic dataset to a temporary NetCDF file.""" + pytest.importorskip("netCDF4") + path = tmp_path / "test_data.nc" + synthetic_dataset.to_netcdf(path) + return str(path) + + +@pytest.fixture +def zarr_path(synthetic_dataset, tmp_path): + """Write synthetic dataset to a temporary Zarr store.""" + zarr = pytest.importorskip("zarr") # noqa: F841 + path = tmp_path / "test_data.zarr" + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + synthetic_dataset.to_zarr(path) + return str(path) + + +# ---- Construction ---- + +class TestConstruction: + + def test_from_dataset(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.get_tables() == ["pressure", "temperature"] + + def test_from_netcdf(self, nc_file): + source = XArraySQLSource(uri=nc_file) + assert "temperature" in source.get_tables() + assert "pressure" in source.get_tables() + + def test_from_zarr(self, zarr_path): + source = XArraySQLSource(uri=zarr_path, engine="zarr") + assert "temperature" in source.get_tables() + + def test_auto_engine_detection(self, nc_file): + source = XArraySQLSource(uri=nc_file) + assert source.dataset is not None + + def test_variable_filtering(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset, variables=["temperature"]) + assert source.get_tables() == ["temperature"] + + def test_invalid_variable_raises(self, synthetic_dataset): + with pytest.raises(ValueError, match="not found"): + XArraySQLSource(_dataset=synthetic_dataset, variables=["nonexistent"]) + + def test_no_uri_no_dataset_raises(self): + with pytest.raises(ValueError, match="Either 'uri' or '_dataset'"): + XArraySQLSource() + + def test_dialect(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.dialect == "any" + + def test_from_dataset_classmethod(self, synthetic_dataset): + source = XArraySQLSource.from_dataset(synthetic_dataset) + assert source.get_tables() == ["pressure", "temperature"] + df = source.execute("SELECT * FROM temperature LIMIT 3") + assert len(df) == 3 + + def test_supports_sql(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source._supports_sql is True + + def test_source_type(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.source_type == "xarray" + + def test_engine_detection_all_formats(self): + assert XArraySQLSource._detect_engine("/data/file.nc") == "netcdf4" + assert XArraySQLSource._detect_engine("/data/file.nc4") == "netcdf4" + assert XArraySQLSource._detect_engine("/data/file.zarr") == "zarr" + assert XArraySQLSource._detect_engine("/data/file.grib") == "cfgrib" + assert XArraySQLSource._detect_engine("/data/file.grib2") == "cfgrib" + assert XArraySQLSource._detect_engine("/data/file.grb") == "cfgrib" + assert XArraySQLSource._detect_engine("/data/file.h5") is None + assert XArraySQLSource._detect_engine("/data/file.txt") is None + + def test_from_grib(self): + """GRIB file I/O via cfgrib engine.""" + pytest.importorskip("cfgrib") + grib_path = Path(xr.__file__).parent / "tests" / "data" / "example.grib" + if not grib_path.exists(): + pytest.skip("xarray example.grib not found") + source = XArraySQLSource(uri=str(grib_path), engine="cfgrib") + tables = source.get_tables() + assert "z" in tables + assert "t" in tables + + df = source.execute("SELECT * FROM z LIMIT 5") + assert len(df) == 5 + assert "z" in df.columns + assert "latitude" in df.columns + assert "longitude" in df.columns + + +# ---- SQL Execution ---- + +class TestSQLExecution: + + def test_select_all(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute("SELECT * FROM temperature LIMIT 5") + assert len(df) == 5 + assert "temperature" in df.columns + assert "lat" in df.columns + assert "lon" in df.columns + assert "time" in df.columns + + def test_select_columns(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute("SELECT lat, lon, temperature FROM temperature LIMIT 3") + assert list(df.columns) == ["lat", "lon", "temperature"] + assert len(df) == 3 + + def test_where_filter(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute("SELECT * FROM temperature WHERE lat > 30") + assert all(df["lat"] > 30) + assert len(df) > 0 + + def test_aggregation(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute("SELECT AVG(temperature) as avg_temp FROM temperature") + assert len(df) == 1 + assert "avg_temp" in df.columns + assert 250 < df["avg_temp"].iloc[0] < 310 + + def test_group_by(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute( + "SELECT lat, AVG(temperature) as avg_temp FROM temperature GROUP BY lat ORDER BY lat" + ) + assert len(df) == 5 # 5 lat values + assert df["lat"].is_monotonic_increasing + + def test_count(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute("SELECT COUNT(*) as cnt FROM temperature") + assert df["cnt"].iloc[0] == 10 * 5 * 4 # time * lat * lon = 200 + + def test_multiple_conditions(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute( + "SELECT * FROM temperature WHERE lat >= 20 AND lat <= 40 AND lon > 110" + ) + assert all(df["lat"] >= 20) + assert all(df["lat"] <= 40) + assert all(df["lon"] > 110) + + def test_order_by(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.execute( + "SELECT lat, temperature FROM temperature ORDER BY temperature DESC LIMIT 10" + ) + assert df["temperature"].is_monotonic_decreasing + + def test_cross_table_independence(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df_t = source.execute("SELECT COUNT(*) as cnt FROM temperature") + df_p = source.execute("SELECT COUNT(*) as cnt FROM pressure") + assert df_t["cnt"].iloc[0] == df_p["cnt"].iloc[0] + + +# ---- Schema & Metadata ---- + +class TestSchemaMetadata: + + def test_get_schema_single_table(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + schema = source.get_schema("temperature") + assert "temperature" in schema + assert "lat" in schema + assert "lon" in schema + assert "time" in schema + assert "__len__" in schema + + def test_get_schema_all_tables(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + schemas = source.get_schema() + assert "temperature" in schemas + assert "pressure" in schemas + + def test_schema_has_types(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + schema = source.get_schema("temperature") + assert "type" in schema["temperature"] + assert "type" in schema["lat"] + + def test_schema_has_minmax(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + schema = source.get_schema("temperature") + assert "inclusiveMinimum" in schema.get("lat", {}) or "inclusiveMinimum" in schema.get("temperature", {}) + + def test_get_schema_with_limit(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + schema = source.get_schema("temperature", limit=5) + assert "__len__" in schema + + def test_get_schema_with_shuffle(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + schema = source.get_schema("temperature", shuffle=True) + assert "temperature" in schema + + def test_get_metadata(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + meta = source.get_metadata("temperature") + assert "description" in meta + assert "columns" in meta + assert "rows" in meta + assert meta["rows"] == 10 * 5 * 4 # time * lat * lon + + def test_metadata_has_descriptions(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + meta = source.get_metadata("temperature") + assert "Air Temperature" in meta["description"] + cols = meta["columns"] + assert cols["temperature"]["description"] == "Air Temperature" + assert cols["lat"]["data_type"] == "float64" + + def test_metadata_all_tables(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + meta = source.get_metadata() + assert "temperature" in meta + assert "pressure" in meta + + def test_metadata_2d_auxiliary_coords(self): + """Metadata correctly handles datasets with 2D auxiliary coordinates (RASM-like).""" + y = np.arange(5) + x = np.arange(4) + yy, xx = np.meshgrid(y, x, indexing="ij") + ds = xr.Dataset( + {"Tair": (["time", "y", "x"], np.random.randn(3, 5, 4).astype("float32"))}, + coords={ + "time": np.arange(3), + "y": y, + "x": x, + "xc": (["y", "x"], -120.0 + xx * 0.5), + "yc": (["y", "x"], 50.0 + yy * 0.5), + }, + ) + ds["Tair"].attrs = {"long_name": "Air Temperature", "units": "K"} + ds.coords["xc"].attrs = {"long_name": "longitude"} + ds.coords["yc"].attrs = {"long_name": "latitude"} + source = XArraySQLSource(_dataset=ds) + + meta = source._get_table_metadata(["Tair"]) + assert "Tair" in meta + + # SQL columns = dims + data var (not 2D auxiliary coords) + assert set(meta["Tair"]["columns"]) == {"time", "y", "x", "Tair"} + assert "xc" not in meta["Tair"]["columns"] + assert "yc" not in meta["Tair"]["columns"] + + # Auxiliary coords noted in description with long_name + desc = meta["Tair"]["description"] + assert "auxiliary coords" in desc + assert "xc [longitude]" in desc + assert "yc [latitude]" in desc + + # Row count correct + assert meta["Tair"]["rows"] == 3 * 5 * 4 + + # SQL queries still work with dimension coords + df = source.execute("SELECT * FROM Tair LIMIT 5") + assert set(df.columns) == {"time", "y", "x", "Tair"} + assert len(df) == 5 + + +# ---- Lumen Integration (get, get_sql_expr) ---- + +class TestLumenIntegration: + + def test_get_returns_dataframe(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.get("temperature") + assert isinstance(df, pd.DataFrame) + assert "temperature" in df.columns + assert len(df) > 0 + + def test_get_with_filter(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.get("temperature", lat=30.0) + assert all(df["lat"] == 30.0) + + def test_get_sql_expr(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + expr = source.get_sql_expr("temperature") + assert isinstance(expr, str) + assert "temperature" in expr.lower() + + def test_get_sql_expr_dict_tables(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + derived = {"avg_temp": "SELECT lat, AVG(temperature) FROM temperature GROUP BY lat"} + new_source = source.create_sql_expr_source(derived) + expr = new_source.get_sql_expr("avg_temp") + assert "AVG" in expr + + def test_get_sql_expr_missing_table_raises(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + derived = {"avg_temp": "SELECT 1"} + new_source = source.create_sql_expr_source(derived) + with pytest.raises(KeyError, match="not found"): + new_source.get_sql_expr("nonexistent") + + def test_get_skips_dunder_keys(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.get("temperature", __limit__=10) + assert len(df) > 0 + + def test_get_with_list_filter(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.get("temperature", lat=[10.0, 30.0]) + assert set(df["lat"].unique()) <= {10.0, 30.0} + + def test_get_with_slice_filter(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = source.get("temperature", lat=slice(20.0, 40.0)) + assert all(df["lat"] >= 20.0) + assert all(df["lat"] <= 40.0) + + def test_get_with_timestamp_filter(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + ts = pd.Timestamp("2020-01-03") + df = source.get("temperature", time=ts) + assert len(df) > 0 + + +# ---- Normalize Table ---- + +class TestNormalizeTable: + + def test_exact_match(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.normalize_table("temperature") == "temperature" + + def test_case_insensitive(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.normalize_table("TEMPERATURE") == "temperature" + assert source.normalize_table("Temperature") == "temperature" + + def test_strip_quotes(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.normalize_table('"temperature"') == "temperature" + assert source.normalize_table("'temperature'") == "temperature" + + def test_unknown_returns_as_is(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.normalize_table("nonexistent") == "nonexistent" + + def test_normalize_derived_table(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + derived = {"avg_temp": "SELECT 1"} + new_source = source.create_sql_expr_source(derived) + assert new_source.normalize_table("AVG_TEMP") == "avg_temp" + + +# ---- get_tables ---- + +class TestGetTables: + + def test_includes_registered_tables(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert "temperature" in source.get_tables() + assert "pressure" in source.get_tables() + + def test_includes_derived_tables(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + derived = {"avg_temp": "SELECT lat, AVG(temperature) FROM temperature GROUP BY lat"} + new_source = source.create_sql_expr_source(derived) + tables = new_source.get_tables() + assert "avg_temp" in tables + assert "temperature" in tables + assert "pressure" in tables + + +# ---- Serialization ---- + +class TestSerialization: + + def test_to_spec(self, nc_file): + source = XArraySQLSource(uri=nc_file) + spec = source.to_spec() + assert spec["type"] == "xarray" + assert spec["uri"] == nc_file + + def test_to_spec_with_options(self, nc_file): + source = XArraySQLSource( + uri=nc_file, engine="netcdf4", variables=["temperature"] + ) + spec = source.to_spec() + assert spec["engine"] == "netcdf4" + assert spec["variables"] == ["temperature"] + + def test_roundtrip_spec(self, nc_file): + source1 = XArraySQLSource(uri=nc_file) + spec = source1.to_spec() + source2 = XArraySQLSource.from_spec(spec) + assert source1.get_tables() == source2.get_tables() + + +# ---- Resource Management ---- + +class TestResourceManagement: + + def test_close(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + source.close() + assert source._ctx is None + + def test_clear_cache(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + _ = source.get_schema("temperature") + source.clear_cache() + schema = source.get_schema("temperature") + assert "temperature" in schema + + def test_dataset_property(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + assert source.dataset is synthetic_dataset + + def test_create_sql_expr_source(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + new_tables = {"avg_temp": "SELECT lat, AVG(temperature) FROM temperature GROUP BY lat"} + new_source = source.create_sql_expr_source(new_tables) + assert new_source._ctx is source._ctx # shared context + assert new_source._dataset is source._dataset + + def test_create_sql_expr_source_with_params(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + new_tables = {"filtered": "SELECT * FROM temperature WHERE lat > :min_lat"} + params = {"filtered": {"min_lat": 30.0}} + new_source = source.create_sql_expr_source(new_tables, params=params) + assert new_source.table_params["filtered"] == {"min_lat": 30.0} + + def test_create_sql_expr_source_merges_tables(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + derived = {"avg_temp": "SELECT lat, AVG(temperature) FROM temperature GROUP BY lat"} + new_source = source.create_sql_expr_source(derived) + tables = new_source.get_tables() + assert "temperature" in tables + assert "pressure" in tables + assert "avg_temp" in tables + + +# ---- Async ---- + +class TestAsync: + + @pytest.mark.asyncio + async def test_execute_async(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = await source.execute_async("SELECT * FROM temperature LIMIT 5") + assert len(df) == 5 + + @pytest.mark.asyncio + async def test_get_async(self, synthetic_dataset): + source = XArraySQLSource(_dataset=synthetic_dataset) + df = await source.get_async("temperature") + assert len(df) > 0 diff --git a/pixi.toml b/pixi.toml index 8ee0ae7e8..18b60b578 100644 --- a/pixi.toml +++ b/pixi.toml @@ -10,8 +10,8 @@ install = 'python -m pip install --no-deps --disable-pip-version-check -e .' PYTHONIOENCODING = "utf-8" [environments] -test-312 = ["py312", "test-core", "test", "ai", "sql"] -test-313 = ["py313", "test-core", "test", "ai", "sql"] +test-312 = ["py312", "test-core", "test", "ai", "sql", "xarray"] +test-313 = ["py313", "test-core", "test", "ai", "sql", "xarray"] test-core = ["py313", "test-core"] docs = ["py312", "doc", "ai", "sql"] build = ["py312", "build"] @@ -68,6 +68,14 @@ duckdb = "*" intake-sql = "*" sqlalchemy = "*" +[feature.xarray.dependencies] +xarray = "*" +netcdf4 = "*" +zarr = "*" + +[feature.xarray.pypi-dependencies] +xarray-sql = "*" + # ============================================= # =================== TESTS =================== # ============================================= diff --git a/pyproject.toml b/pyproject.toml index ead041c71..d6c3d6fc9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,8 @@ ai-litellm = ['litellm'] ai-ollama = ['ollama'] chromadb = ['chromadb'] bigquery = ['google-cloud-bigquery', 'sqlalchemy-bigquery'] +xarray = ['xarray', 'xarray-sql', 'netCDF4', 'zarr'] +xarray-grib = ['lumen[xarray]', 'cfgrib', 'eccodes'] snowflake = ['snowflake-connector-python', 'cryptography'] ae5 = ['ae5-tools'] mcp = ['fastmcp']