diff --git a/README.md b/README.md index 37ca3dc..8bca886 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,9 @@ rel = con.sql('SELECT time, AVG("air") AS air FROM air GROUP BY time ORDER BY ti xql.to_dataset(rel, template=ds) # any engine's Arrow result round-trips ``` +`table_names` (below) works the same way on every engine, so a query written +against `era5.surface` is not tied to the engine it was written for. + See [Engines](https://xqlsystems.github.io/xarray-sql/latest/engines/) for the support matrix, DuckDB/Polars details, and the lazy chunked round-trip. diff --git a/benchmarks/geospatial/_engines.py b/benchmarks/geospatial/_engines.py index 66d769b..5de5b79 100644 --- a/benchmarks/geospatial/_engines.py +++ b/benchmarks/geospatial/_engines.py @@ -18,9 +18,11 @@ :meth:`EngineContext.sql_to_dataset`. On the ``datafusion`` path this is byte-for-byte the original behavior (``from_dataset`` + ``sql`` + ``XarrayDataFrame.to_dataset``); the other engines register one pyarrow -dataset per dimension group under flattened table names -(``era5.surface`` → ``era5_surface`` — rewritten in the SQL text) and the -result rows are round-tripped to an ``xr.Dataset`` through pandas. +dataset per dimension group under the flat table names +``xql.arrow_datasets`` returns (``era5.surface`` → ``era5_surface`` — +rewritten in the SQL text, since these paths register frames one at a +time rather than through a schema) and the result rows are +round-tripped to an ``xr.Dataset`` through pandas. The DataFusion-only UDF cases (07 and the UDF half of 09) build ``xql.XarrayContext`` directly rather than through this layer; the suite @@ -47,24 +49,6 @@ def engine_name() -> str: return engine -def _group_tables(name, ds, table_names): - """Split ``ds`` into per-dimension-group tables like XarrayContext does. - - Returns ``[(flat_name, dotted_name, sub_dataset)]``; a uniform dataset - keeps its plain name (flat == dotted == name). - """ - groups: dict[tuple, list] = {} - for var, v in ds.data_vars.items(): - groups.setdefault(tuple(v.dims), []).append(var) - if len(groups) == 1: - return [(name, name, ds)] - out = [] - for dims, variables in groups.items(): - sub = (table_names or {}).get(dims) or "_".join(dims) - out.append((f"{name}_{sub}", f"{name}.{sub}", ds[variables])) - return out - - def _literal(value: Any) -> str: """Render a parameter value as a SQL literal (for engines without binds).""" if isinstance(value, (datetime.datetime, pd.Timestamp, np.datetime64)): @@ -134,15 +118,16 @@ def from_dataset(self, name, ds, *, chunks=None, table_names=None): """Register ``ds`` as SQL table(s), mirroring XarrayContext naming.""" import xarray_sql as xql - for flat, dotted, sub in _group_tables(name, ds, table_names): - if dotted != flat: - self._renames[dotted] = flat - sub_chunks = ( - {d: c for d, c in chunks.items() if d in sub.dims} - if isinstance(chunks, dict) - else chunks - ) or None - self._register(flat, xql.arrow_dataset(sub, sub_chunks)) + tables = xql.arrow_datasets( + ds, name, chunks=chunks, table_names=table_names + ) + for flat, dataset in tables.items(): + if flat != name: + # `era5_surface` here is `era5.surface` in the case's SQL, + # which was written against XarrayContext's schema split. + group = flat[len(name) + 1 :] + self._renames[f"{name}.{group}"] = flat + self._register(flat, dataset) # -- querying ---------------------------------------------------------- diff --git a/docs/engines.md b/docs/engines.md index 11dacac..04201f9 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -90,6 +90,27 @@ exact expression via pyarrow — pruning is only an optimization on top. `XarrayArrowStream`, the dependency-light re-scannable C-stream wrapper without pushdown, remains available as a fallback. +As is standard in for all Xarray-SQL engines, users may provide a mapping of +groups of dimensions to their preferred table names, like so: + +```python +xql.register(con, "era5", ds, table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", +}) + +con.sql("SELECT AVG(temperature) FROM era5.atmosphere WHERE level = 500") +``` + +Mixed-dimension Datasets split as they do everywhere else, with one +DuckDB-specific wrinkle: `con.register` can only place an object in +DuckDB's temporary namespace, so each group is registered flat as +`era5_surface` and mirrored as a view `era5.surface` in a schema of its +own. Both spellings hit the same scan — pushdown and projection travel +through the view — and the dotted one is what keeps the SQL portable. +A read-only connection cannot create the schema; registration then +warns and leaves the flat tables. + Details that matter in production: - **Finely partitioned axes** (e.g. hourly-chunked reanalysis time with @@ -146,6 +167,19 @@ out = ( xql.to_dataset(out, template=ds) # polars frames speak Arrow PyCapsule ``` +`arrow_dataset` wants a Dataset whose variables share one set of +dimensions; `xql.arrow_datasets(ds, "era5", table_names=...)` splits a +mixed-dimension one and hands back the tables named, reading the shared +dimension coordinates once for all of them. + +```python +tables = xql.arrow_datasets(ds, "era5", table_names={...}) + +ctx = pl.SQLContext() +for table, dataset in tables.items(): # 'era5_surface', ... + ctx.register(table, pl.scan_pyarrow_dataset(dataset)) +``` + Polars pushes its predicate and column selection into the dataset scan (verified: a filtered group-by read 1 of 20 chunks and 3 of 5 columns), and its results round-trip through `xql.to_dataset` unchanged. The @@ -166,7 +200,8 @@ What each integration provides. Known issues and constraints live on | Eager round-trip (`xql.to_dataset`) | yes | yes | yes | | Chunked round-trip (`chunks=`) | re-execution | `spill=True` [^spill-only] | re-execution (streaming engine) | | `geometry` column ([geospatial](geospatial.md#geoarrow-point-geometry-columns)) | annotated WKB passes through | native `GEOMETRY` (`"wkb"` encoding) | plain binary/struct | -| Mixed-dimension datasets | one schema, `name.group` tables | `_` tables | filter `data_vars` before `arrow_dataset` | +| Mixed-dimension datasets | one schema, `name.group` tables | `name.group` views over `name_group` tables | `xql.arrow_datasets(ds, name)`, one per group | +| Naming those tables (`table_names=`) | yes | yes | yes | | Version floor | bundled (core dependency) | `duckdb >= 1.4` (tested on 1.5) | tested on `polars 1.42` | [^spill-only]: Why DuckDB relations do not re-execute — and two other diff --git a/docs/examples.md b/docs/examples.md index 483fb2a..f2b925b 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -107,7 +107,11 @@ ctx.sql(''' 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`. +`era5.time_level_latitude_longitude`. Keys naming a dimension group the Dataset +does not have are ignored, so one naming map can be reused across Datasets +holding different subsets of the same variables. The keyword is not +DataFusion-only — `xql.register` takes it on every engine (see +[the same tables on DuckDB and Polars](#the-same-tables-on-duckdb-and-polars)). ## GOES satellite imagery (scalar variables) @@ -154,5 +158,44 @@ not DataFusion-specific: `xql.register(con, name, ds)` attaches the same lazy, pushdown-scanned table to a DuckDB connection, and `pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` serves Polars — same splitting rules for mixed-dimension Datasets, same round-trip through -`xql.to_dataset(result, template=ds)`. See [Engines](engines.md) for the -support matrix and per-engine details. +`xql.to_dataset(result, template=ds)`. + +`table_names` travels with them, so the ERA5 example above changes engine +without changing a word of its SQL: + +```python +import duckdb + +con = duckdb.connect() +xql.register(con, 'era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) + +rel = con.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 +''') +xql.to_dataset(rel, template=ds, dims=['level']) +``` + +Polars registers table by table, so it takes the split datasets by name: + +```python +import polars as pl + +tables = xql.arrow_datasets(ds, 'era5', table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) + +ctx = pl.SQLContext() +for table, dataset in tables.items(): # 'era5_surface', 'era5_atmosphere' + ctx.register(table, pl.scan_pyarrow_dataset(dataset)) +``` + +See [Engines](engines.md) for the support matrix and per-engine details. diff --git a/docs/limitations.md b/docs/limitations.md index 8aab000..b7e5d60 100644 --- a/docs/limitations.md +++ b/docs/limitations.md @@ -144,12 +144,14 @@ A materialized table or bare C-stream has no query behind it, so the re-execution form of `chunks=` cannot serve it — `spill=True` (one pass to a temporary Parquet file) is the chunked path for these. -### Mixed-dimension datasets split into one DuckDB table per dim group - -DuckDB registration has no schema namespace, so variables with -different dims land in suffixed tables (`_`), sharing one -set of coordinate reads. DataFusion registers the same layout as -`name.group` tables inside one schema. +### A read-only DuckDB connection loses the dotted table names + +Mixed-dimension Datasets split into one table per dimension group on +every engine, addressed as `name.group`. On DuckDB that dotted spelling +is a view in a schema, because `con.register` reaches only the +temporary namespace — and creating a schema needs a writable catalog. +Registering on a read-only connection therefore warns and leaves the +flat `name_group` tables, which are always registered and always work. ### Pointwise indexers on lazy round-trip arrays are slower diff --git a/tests/test_table_names.py b/tests/test_table_names.py new file mode 100644 index 0000000..fca4b66 --- /dev/null +++ b/tests/test_table_names.py @@ -0,0 +1,262 @@ +"""Naming the tables a mixed-dimension Dataset splits into, on every engine. + +A Dataset whose variables sit on different dimensions (ARCO-ERA5: 262 +surface fields on ``(time, latitude, longitude)``, 11 atmospheric ones +on ``(time, level, latitude, longitude)``) becomes one table per +dimension group. ``table_names`` is what lets the user call those +tables ``surface`` and ``atmosphere`` instead of +``time_latitude_longitude``. + +The contract these tests hold: + +1. ``table_names`` means the same thing wherever a Dataset is + registered — ``XarrayContext``, a plain ``SessionContext``, DuckDB, + or the pyarrow-dataset path Polars scans. +2. ``name.group`` is the portable spelling: the same SQL text runs on + DataFusion and DuckDB. DuckDB additionally keeps the flat + ``name_group`` tables it has always registered. +3. Naming changes only names. Pushdown, projection, and results are + what they were. +""" + +import numpy as np +import pandas as pd +import pytest +import xarray as xr +from datafusion import SessionContext + +import xarray_sql as xql + +duckdb = pytest.importorskip("duckdb") + +# The ARCO-ERA5 shape in miniature: a surface group, an atmospheric +# group on an extra `level` dim, and a scalar metadata variable. +NAMES = { + ("time", "lat", "lon"): "surface", + ("time", "level", "lat", "lon"): "atmosphere", + (): "meta", +} + +SURFACE_QUERY = "SELECT AVG(t2m) AS avg_t FROM era5.surface" + + +@pytest.fixture +def ds() -> xr.Dataset: + np.random.seed(11) + return xr.Dataset( + { + "t2m": (["time", "lat", "lon"], np.random.rand(6, 3, 4)), + "temperature": ( + ["time", "level", "lat", "lon"], + np.random.rand(6, 2, 3, 4), + ), + "projection": ((), 42.0), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=6, freq="D"), + "lat": np.linspace(-90, 90, 3), + "lon": np.linspace(-180, 180, 4), + "level": [500, 1000], + }, + ).chunk({"time": 1}) + + +def expected_avg(ds: xr.Dataset) -> float: + return float(ds["t2m"].mean().compute()) + + +# 1. The same naming map on every engine ------------------------------ + + +def test_xarray_context_names_tables(ds): + ctx = xql.XarrayContext() + xql.register(ctx, "era5", ds, table_names=NAMES) + + result = ctx.sql(SURFACE_QUERY).to_pandas()["avg_t"][0] + assert result == pytest.approx(expected_avg(ds)) + + +def test_plain_session_context_splits_and_names(ds): + # A bare SessionContext used to register the whole Dataset as one + # table, which a mixed-dimension Dataset cannot be. + con = SessionContext() + xql.register(con, "era5", ds, table_names=NAMES) + + result = con.sql(SURFACE_QUERY).to_pandas()["avg_t"][0] + assert result == pytest.approx(expected_avg(ds)) + + +def test_duckdb_names_tables(ds): + con = duckdb.connect() + xql.register(con, "era5", ds, table_names=NAMES) + + result = con.sql(SURFACE_QUERY).fetchone()[0] + assert result == pytest.approx(expected_avg(ds)) + + +def test_arrow_datasets_names_each_group(ds): + tables = xql.arrow_datasets(ds, "era5", table_names=NAMES) + + assert set(tables) == {"era5_surface", "era5_atmosphere", "era5_meta"} + assert set(tables["era5_surface"].schema.names) == { + "time", + "lat", + "lon", + "t2m", + } + + +def test_polars_queries_named_tables(ds): + pl = pytest.importorskip("polars") + ctx = pl.SQLContext() + tables = xql.arrow_datasets(ds, "era5", table_names=NAMES) + for table, dataset in tables.items(): + ctx.register(table, pl.scan_pyarrow_dataset(dataset)) + + result = ctx.execute( + "SELECT AVG(t2m) AS avg_t FROM era5_surface", eager=True + ) + assert result["avg_t"][0] == pytest.approx(expected_avg(ds)) + + +def test_every_engine_answers_the_same_query(ds): + """One SQL string, three engines, one number.""" + xarray_ctx = xql.XarrayContext() + xql.register(xarray_ctx, "era5", ds, table_names=NAMES) + session_ctx = SessionContext() + xql.register(session_ctx, "era5", ds, table_names=NAMES) + con = duckdb.connect() + xql.register(con, "era5", ds, table_names=NAMES) + + answers = [ + xarray_ctx.sql(SURFACE_QUERY).to_pandas()["avg_t"][0], + session_ctx.sql(SURFACE_QUERY).to_pandas()["avg_t"][0], + con.sql(SURFACE_QUERY).fetchone()[0], + ] + assert answers == pytest.approx([expected_avg(ds)] * 3) + + +# 2. Spellings: dotted everywhere, flat still on DuckDB ---------------- + + +def test_duckdb_keeps_the_flat_spelling(ds): + con = duckdb.connect() + xql.register(con, "era5", ds, table_names=NAMES) + + dotted = con.sql("SELECT COUNT(*) FROM era5.atmosphere").fetchone()[0] + flat = con.sql("SELECT COUNT(*) FROM era5_atmosphere").fetchone()[0] + assert dotted == flat == 6 * 2 * 3 * 4 + + +def test_duckdb_defaults_to_joined_dim_names(ds): + con = duckdb.connect() + xql.register(con, "era5", ds) + + assert con.sql("SELECT COUNT(*) FROM era5.time_lat_lon").fetchone()[0] == ( + 6 * 3 * 4 + ) + assert con.sql("SELECT COUNT(*) FROM era5_time_lat_lon").fetchone()[0] == ( + 6 * 3 * 4 + ) + + +def test_scalar_group_can_be_named(ds): + con = duckdb.connect() + xql.register(con, "era5", ds, table_names=NAMES) + + assert con.sql("SELECT projection FROM era5.meta").fetchall() == [(42.0,)] + + +def test_uniform_dataset_keeps_the_bare_name(ds): + # One dimension group is one table, named `name` — there is no + # group to name, on any engine. + surface_only = ds[["t2m"]] + con = duckdb.connect() + xql.register(con, "era5", surface_only, table_names=NAMES) + + assert con.sql("SELECT COUNT(*) FROM era5").fetchone()[0] == 6 * 3 * 4 + assert set(xql.arrow_datasets(surface_only, "era5")) == {"era5"} + + +def test_arrow_datasets_without_a_prefix(ds): + # Engines registered table-by-table (Polars) may not want the + # namespace prefix at all. + tables = xql.arrow_datasets(ds, table_names=NAMES) + + assert set(tables) == {"surface", "atmosphere", "meta"} + + +# 3. Naming changes names, nothing else ------------------------------- + + +def test_pushdown_survives_the_duckdb_schema_view(ds): + # The dotted spelling is a view over the registered table; a view + # that blocked pushdown would turn a pruned scan into a full one. + reads: list = [] + projections: list = [] + con = duckdb.connect() + xql.register( + con, + "era5", + ds, + table_names=NAMES, + _iteration_callback=lambda block, proj: ( + reads.append(block), + projections.append(proj), + ), + ) + + con.sql( + "SELECT AVG(t2m) FROM era5.surface WHERE time = TIMESTAMP '2020-01-03'" + ).fetchall() + + assert len(reads) == 1 # 1 of 6 daily chunks + assert projections[0] == ["time", "t2m"] + + +def test_named_result_round_trips_to_xarray(ds): + con = duckdb.connect() + xql.register(con, "era5", ds, table_names=NAMES) + + rel = con.sql( + "SELECT time, lat, lon, t2m FROM era5.surface ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + xr.testing.assert_allclose(out, ds[["t2m"]].compute()) + + +def test_duplicate_names_are_rejected(ds): + duplicated = {**NAMES, ("time", "level", "lat", "lon"): "surface"} + con = duckdb.connect() + + with pytest.raises(ValueError, match="same table name 'surface'"): + xql.register(con, "era5", ds, table_names=duplicated) + + +def test_names_for_absent_groups_are_ignored(ds): + # One canonical naming map, reused over Datasets holding different + # subsets of the same variables. + con = duckdb.connect() + xql.register(con, "era5", ds[["t2m", "projection"]], table_names=NAMES) + + assert con.sql("SELECT COUNT(*) FROM era5.surface").fetchone()[0] == ( + 6 * 3 * 4 + ) + views = con.sql("SELECT view_name FROM duckdb_views()").fetchall() + assert "atmosphere" not in {row[0] for row in views} + + +def test_read_only_duckdb_warns_but_still_registers(ds, tmp_path): + # Mirroring the groups as `era5.` needs a writable catalog. + # A read-only connection loses the dotted spelling, not the tables. + path = tmp_path / "ro.db" + duckdb.connect(str(path)).close() + con = duckdb.connect(str(path), read_only=True) + + with pytest.warns(RuntimeWarning, match="could not create the 'era5'"): + xql.register(con, "era5", ds, table_names=NAMES) + + assert con.sql("SELECT COUNT(*) FROM era5_surface").fetchone()[0] == ( + 6 * 3 * 4 + ) diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py index 0f98e1e..1bbfe56 100644 --- a/xarray_sql/__init__.py +++ b/xarray_sql/__init__.py @@ -1,5 +1,5 @@ from . import cftime -from .backends import arrow_dataset, register +from .backends import arrow_dataset, arrow_datasets, register from .geometry import bbox_conjuncts from .df import from_map from .reader import read_xarray, read_xarray_table @@ -12,6 +12,7 @@ "read_xarray_table", "read_xarray", "arrow_dataset", + "arrow_datasets", "bbox_conjuncts", "register", "to_dataset", diff --git a/xarray_sql/backends/__init__.py b/xarray_sql/backends/__init__.py index 82aad8e..af0334b 100644 --- a/xarray_sql/backends/__init__.py +++ b/xarray_sql/backends/__init__.py @@ -21,6 +21,7 @@ XarrayArrowStream, XarrayPushdownDataset, arrow_dataset, + arrow_datasets, ) __all__ = [ @@ -28,6 +29,7 @@ "XarrayArrowStream", "XarrayPushdownDataset", "arrow_dataset", + "arrow_datasets", "get_adapter", "register", "register_adapter", diff --git a/xarray_sql/backends/base.py b/xarray_sql/backends/base.py index 59eb488..1446bef 100644 --- a/xarray_sql/backends/base.py +++ b/xarray_sql/backends/base.py @@ -18,7 +18,7 @@ import xarray as xr -from ..df import Chunks +from ..df import Chunks, TableNames ConT = TypeVar("ConT") """An engine's native connection type (e.g. ``duckdb.DuckDBPyConnection``).""" @@ -39,6 +39,7 @@ def register( ds: xr.Dataset, *, chunks: Chunks = None, + table_names: TableNames = None, **kwargs: Any, ) -> ConT: """Register *ds* as table *name* on *con*; returns *con*.""" @@ -74,6 +75,7 @@ def register( ds: xr.Dataset, *, chunks: Chunks = None, + table_names: TableNames = None, **kwargs: Any, ) -> ConT: """Register a lazy xarray Dataset as a table on an engine connection. @@ -94,20 +96,37 @@ def register( rel = con.sql("SELECT time, AVG(t2m) AS t2m FROM era5 GROUP BY time") result = xql.to_dataset(rel, template=ds) + A Dataset whose variables sit on different dimensions is split into + one table per dimension group. Name those tables with + ``table_names``, and the same SQL runs on every engine:: + + xql.register(con, "era5", ds, table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }) + con.sql("SELECT AVG(temperature) FROM era5.atmosphere") + Args: con: An engine connection: a ``datafusion.SessionContext`` (or [xarray_sql.XarrayContext][]) or a ``duckdb.DuckDBPyConnection``. name: The table name to register the Dataset under. Datasets whose variables have differing dimensions are split into one - table per dimension group (a SQL schema ``name.group`` on - DataFusion; ``name_group`` tables on DuckDB). + table per dimension group, addressed as ``name.group`` on + every engine (DuckDB also keeps the flat ``name_group`` + spelling, since its registration namespace is flat). ds: An xarray Dataset. chunks: Xarray-like chunks specification controlling partition granularity. Defaults to the Dataset's existing chunks. + table_names: Maps a dimension group's exact dim tuple to the name + its table takes. Groups left unnamed take their dimensions + joined by underscores (``time_latitude_longitude``); the + group holding scalar variables, if any, takes ``scalar``. + Keys matching no group in ``ds`` are ignored, so one naming + map can be reused across Datasets holding different subsets + of the same variables. **kwargs: Adapter-specific options, forwarded as-is — e.g. - ``table_names`` on DataFusion, ``batch_size`` / ``prefetch`` - on DuckDB. + ``batch_size`` / ``prefetch`` on DuckDB. Returns: The connection, to allow chaining. @@ -115,4 +134,9 @@ def register( # The connection type is erased by the runtime dispatch; every adapter # returns the connection it was given. adapter: Any = get_adapter(con) - return cast(ConT, adapter.register(con, name, ds, chunks=chunks, **kwargs)) + return cast( + ConT, + adapter.register( + con, name, ds, chunks=chunks, table_names=table_names, **kwargs + ), + ) diff --git a/xarray_sql/backends/datafusion.py b/xarray_sql/backends/datafusion.py index 3d7c29e..4fc652f 100644 --- a/xarray_sql/backends/datafusion.py +++ b/xarray_sql/backends/datafusion.py @@ -13,8 +13,15 @@ import xarray as xr from datafusion import SessionContext +from datafusion.catalog import Schema -from ..df import Chunks +from ..df import ( + Chunks, + TableNames, + group_vars_by_dims, + resolve_table_names, + shared_coord_arrays, +) from ..reader import read_xarray_table from ..sql import XarrayContext from .base import register_adapter @@ -35,12 +42,37 @@ def register( ds: xr.Dataset, *, chunks: Chunks = None, + table_names: TableNames = None, **kwargs: Any, ) -> SessionContext: - # XarrayContext.from_dataset adds dim-group splitting, cftime UDF - # registration, and round-trip metadata tracking on top of the - # plain table registration; use it when available. + # XarrayContext.from_dataset adds cftime UDF registration and + # round-trip metadata tracking on top of the split below; use it + # when available. if isinstance(con, XarrayContext): - return con.from_dataset(name, ds, chunks=chunks, **kwargs) - con.register_table(name, read_xarray_table(ds, chunks, **kwargs)) + return con.from_dataset( + name, ds, chunks=chunks, table_names=table_names, **kwargs + ) + + groups = group_vars_by_dims(ds) + names = resolve_table_names(ds, table_names) + if len(groups) <= 1: + con.register_table(name, read_xarray_table(ds, chunks, **kwargs)) + return con + + # A mixed-dimension Dataset becomes one table per dimension group + # in a schema named after the Dataset, exactly as XarrayContext + # registers it — so ``name.group`` is the same SQL either way. + coord_arrays = shared_coord_arrays(ds) + schema = Schema.memory_schema(con) + con.catalog().register_schema(name, schema) + for dims, var_names in groups.items(): + schema.register_table( + names[dims], + read_xarray_table( + ds[var_names], + chunks, + coord_arrays=coord_arrays, + **kwargs, + ), + ) return con diff --git a/xarray_sql/backends/duckdb.py b/xarray_sql/backends/duckdb.py index 1651ac0..c602cc2 100644 --- a/xarray_sql/backends/duckdb.py +++ b/xarray_sql/backends/duckdb.py @@ -21,11 +21,18 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING, Any, TypeGuard import xarray as xr -from ..df import Chunks, group_vars_by_dims +from ..df import ( + Chunks, + TableNames, + group_vars_by_dims, + resolve_table_names, + shared_coord_arrays, +) from .base import register_adapter from .pyarrow import XarrayArrowStream, XarrayPushdownDataset @@ -35,6 +42,44 @@ __all__ = ["DuckDBAdapter", "XarrayArrowStream", "XarrayPushdownDataset"] +def _quote(identifier: str) -> str: + """Render *identifier* as a quoted SQL identifier.""" + escaped = identifier.replace('"', '""') + return f'"{escaped}"' + + +def _mirror_as_schema( + con: duckdb.DuckDBPyConnection, name: str, tables: dict[str, str] +) -> None: + """Expose flat tables as views in a DuckDB schema named *name*. + + This translates a view to `name_group` can be queried as `name.group` + in DuckDB. + + *tables* maps each group's table name to the flat name it was + registered under. + """ + try: + con.execute(f"CREATE SCHEMA IF NOT EXISTS {_quote(name)}") + for group, flat in tables.items(): + con.execute( + f"CREATE OR REPLACE VIEW {_quote(name)}.{_quote(group)} " + f"AS SELECT * FROM {_quote(flat)}" + ) + except Exception as exc: # noqa: BLE001 — degrade, don't fail the register + # Creating a schema needs a writable catalog; a read-only + # connection (or a name already taken by something else) is not a + # reason to lose the registration. + warnings.warn( + f"Registered the dimension groups of {name!r} as " + f"{', '.join(sorted(tables.values()))}, but could not create " + f"the {name!r} schema mirroring them as {name}. " + f"({exc}). Query the flat table names instead.", + RuntimeWarning, + stacklevel=3, + ) + + @register_adapter class DuckDBAdapter: """Registers Datasets on ``duckdb.DuckDBPyConnection`` connections.""" @@ -53,32 +98,43 @@ def register( ds: xr.Dataset, *, chunks: Chunks = None, + table_names: TableNames = None, **kwargs: Any, ) -> duckdb.DuckDBPyConnection: """Register ``ds`` on a DuckDB connection. Datasets whose variables all share the same dimensions become a single table named ``name``. Mixed-dimension datasets are split - into one table per dimension group, named - ``___...`` (DuckDB registration has no schema - namespace to mirror the DataFusion adapter's ``name.group`` - layout). Extra keyword arguments (``batch_size``, ``prefetch``) - are forwarded to [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. + into one table per dimension group, named after the group's + dimensions joined by underscores unless ``table_names`` gives the + group a name of its own:: + + xql.register(con, 'era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', + }) + con.sql('SELECT ... FROM era5.surface') # or era5_surface + + Each group is registered under the flat name ``_`` + and mirrored as a view ``.`` in a DuckDB schema, so + the dotted spelling the DataFusion adapter uses queries the same + table here. Extra keyword arguments (``batch_size``, + ``prefetch``) are forwarded to + [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. """ groups = group_vars_by_dims(ds) + names = resolve_table_names(ds, table_names) if len(groups) <= 1: con.register(name, XarrayPushdownDataset(ds, chunks, **kwargs)) return con # Materialise dim coordinates once and share across sub-tables. - coord_arrays = { - str(dim): ds.coords[dim].values - for dim in ds.dims - if dim in ds.coords - } + coord_arrays = shared_coord_arrays(ds) + flat_names = {} for dims, var_names in groups.items(): - suffix = "_".join(dims) or "scalar" + group = names[dims] + flat_names[group] = f"{name}_{group}" con.register( - f"{name}_{suffix}", + flat_names[group], XarrayPushdownDataset( ds[var_names], chunks, @@ -86,4 +142,5 @@ def register( **kwargs, ), ) + _mirror_as_schema(con, name, flat_names) return con diff --git a/xarray_sql/backends/pyarrow.py b/xarray_sql/backends/pyarrow.py index dd0f254..86283fe 100644 --- a/xarray_sql/backends/pyarrow.py +++ b/xarray_sql/backends/pyarrow.py @@ -45,10 +45,14 @@ Block, Chunks, DEFAULT_BATCH_SIZE, + TableNames, _ensure_default_indexes, _parse_schema, + group_vars_by_dims, iter_record_batches, resolve_chunks, + resolve_table_names, + shared_coord_arrays, ) from ..geometry import GEOMETRY_COLUMN, build_geometry, geometry_field from ..reader import XarrayRecordBatchReader @@ -1141,3 +1145,82 @@ def arrow_dataset( geometry_encoding=geometry_encoding, geometry_crs=geometry_crs, ) + + +def arrow_datasets( + ds: xr.Dataset, + name: str | None = None, + *, + chunks: Chunks = None, + table_names: TableNames = None, + **kwargs: Any, +) -> dict[str, XarrayPushdownDataset]: + """Named pushdown datasets, one per dimension group of ``ds``. + + [arrow_dataset][xarray_sql.backends.pyarrow.arrow_dataset] needs a + Dataset whose data variables all share one set of dimensions. This + applies the same split the DataFusion and DuckDB adapters do — one + table per dimension group — and hands back the tables *named*, ready + to register on an engine that has no connection object to dispatch + on:: + + tables = xql.arrow_datasets(ds, 'era5', table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', + }) + + ctx = pl.SQLContext() + for table, dataset in tables.items(): # 'era5_surface', ... + ctx.register(table, pl.scan_pyarrow_dataset(dataset)) + + ctx.execute('SELECT AVG("2m_temperature") FROM era5_surface') + + Dimension coordinates are read once and shared across the returned + tables, which is a network round-trip saved per dimension per group + on Zarr-backed stores. + + Args: + ds: An xarray Dataset, with variables on any mix of dimensions. + name: Prefix for the table names, as in ``era5_surface``. Omit it + to get the bare group names (``surface``). + chunks: Xarray-like chunks specification controlling partition + granularity. Keys naming dimensions a group does not have are + ignored for that group. Defaults to the Dataset's existing + chunks. + table_names: Maps a dimension group's exact dim tuple to the name + its table takes, e.g. + ``{('time', 'latitude', 'longitude'): 'surface'}``. Groups + left unnamed take their dimensions joined by underscores + (``time_latitude_longitude``); the group holding scalar + variables, if any, takes ``scalar``. + **kwargs: Forwarded to + [arrow_dataset][xarray_sql.backends.pyarrow.arrow_dataset] + (``batch_size``, ``prefetch``, ``geometry``, ...), applied to + every returned table. + + Returns: + Table name to + [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. + A Dataset with a single dimension group yields one entry, keyed + by ``name`` when given. + """ + groups = group_vars_by_dims(ds) + names = resolve_table_names(ds, table_names) + + def table_name(dims: tuple[str, ...]) -> str: + group = names[dims] + return f"{name}_{group}" if name else group + + if len(groups) <= 1: + # One group is one table, named `name` — there is no group to + # tell apart. Without a `name`, it takes the group's own. + only = name if name else next(iter(names.values()), "scalar") + return {only: arrow_dataset(ds, chunks, **kwargs)} + + coord_arrays = shared_coord_arrays(ds) + return { + table_name(dims): XarrayPushdownDataset( + ds[var_names], chunks, coord_arrays=coord_arrays, **kwargs + ) + for dims, var_names in groups.items() + } diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 83879f7..9a05ade 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -13,6 +13,8 @@ Block = dict[Hashable, slice] Chunks = dict[str, int] | None +TableNames = Mapping[tuple[str, ...], str] | None +"""Maps a dimension group's exact dim tuple to the table name it takes.""" # Borrowed from Xarray @@ -149,6 +151,59 @@ def group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: return groups +def default_table_name(dims: tuple[str, ...]) -> str: + """The table name a dimension group gets when the user names none.""" + # Scalar variables group under empty dims, where "_".join(()) is the + # empty string; fall back to a valid default table name. + return "_".join(dims) or "scalar" + + +def resolve_table_names( + ds: xr.Dataset, table_names: TableNames = None +) -> dict[tuple[str, ...], str]: + """Name every dimension group of ``ds``, honouring user overrides. + + ``table_names`` maps a group's exact dimension tuple to the name its + table should carry; groups it does not mention keep + [default_table_name][xarray_sql.df.default_table_name]. Keys that + match no group in ``ds`` are ignored, so one naming map can be reused + across Datasets that hold different subsets of the same variables. + + Every engine adapter routes through here, which is what makes + ``table_names={('time', 'lat', 'lon'): 'surface'}`` mean the same + thing on DataFusion, DuckDB, and the pyarrow-dataset engines. + + Raises: + ValueError: if two groups would end up with the same name, which + would silently register one table over the other. + """ + overrides = table_names or {} + names = { + dims: overrides.get(dims) or default_table_name(dims) + for dims in group_vars_by_dims(ds) + } + taken: dict[str, tuple[str, ...]] = {} + for dims, name in names.items(): + if name in taken: + raise ValueError( + f"table_names maps two dimension groups to the same table " + f"name {name!r}: {taken[name]} and {dims}. Give each group " + f"a distinct name." + ) + taken[name] = dims + return names + + +def shared_coord_arrays(ds: xr.Dataset) -> dict[str, np.ndarray]: + """Materialise ``ds``'s dimension coordinates once, to share. + + Splitting a Dataset into per-dimension-group tables otherwise reads + every shared dim coordinate once per table — a network round-trip + apiece for Zarr-backed parents like ARCO-ERA5. + """ + return {str(dim): ds.coords[dim].values for dim in ds.dims} + + def _block_len(block: Block) -> int: return int(np.prod([v.stop - v.start for v in block.values()])) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index a5df4fb..c9911c1 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -4,7 +4,13 @@ from types import ModuleType from . import cftime as cft -from .df import Chunks, group_vars_by_dims +from .df import ( + Chunks, + TableNames, + group_vars_by_dims, + resolve_table_names, + shared_coord_arrays, +) from .ds import XarrayDataFrame from .reader import read_xarray_table @@ -39,7 +45,7 @@ def from_dataset( name: str, input_table: xr.Dataset, *, - table_names: dict[tuple[str, ...], str] | None = None, + table_names: TableNames = None, chunks: Chunks = None, ): """Register an xarray Dataset as one or more queryable SQL tables. @@ -100,13 +106,12 @@ def from_dataset( self, to allow chaining. """ groups = group_vars_by_dims(input_table) + names = resolve_table_names(input_table, table_names) # Materialise dim coordinates once and share across every sub-table. # For Zarr-backed parents (e.g. ARCO-ERA5 on GCS) this saves one # network round-trip per dim per dim-group. - coord_arrays = { - str(dim): input_table.coords[dim].values for dim in input_table.dims - } + coord_arrays = shared_coord_arrays(input_table) if len(groups) <= 1: self._registered_datasets[name] = input_table @@ -114,14 +119,11 @@ def from_dataset( name, input_table, chunks, coord_arrays=coord_arrays ) - table_names = table_names or {} schema = Schema.memory_schema(self) self.catalog().register_schema(name, schema) for dims, var_names in groups.items(): - # Scalar variables group under empty dims, where "_".join(()) is - # the empty string; fall back to a valid default table name. - sub_name = table_names.get(dims, "_".join(dims) or "scalar") + sub_name = names[dims] sub_ds = input_table[var_names] self._from_dataset( sub_name,