From 2363d2dfdea6adf1152633925384093416c660e0 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 5 Sep 2026 12:41:13 -0700 Subject: [PATCH 1/3] Name your tables on every engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `table_names` gave ARCO-ERA5's two dimension groups the names `surface` and `atmosphere` instead of `time_latitude_longitude` and `time_level_latitude_longitude`. It only ever reached DataFusion. Route it through the register seam so it means the same thing wherever a Dataset is registered. One naming rule, `df.resolve_table_names`, now serves every adapter: * DuckDB takes `table_names`, keeps registering the flat `era5_surface` tables, and mirrors the groups as views in a schema so `era5.surface` resolves too — the DataFusion spelling, so query text moves between engines unchanged. Pushdown and projection travel through the view. A read-only connection cannot create a schema; that path warns and leaves the flat tables rather than losing the registration. * A plain `SessionContext` splits mixed-dimension Datasets into a schema the way `XarrayContext` does, instead of registering them as one table they cannot be. * `arrow_datasets(ds, name, table_names=...)` returns the split tables named, for Polars and anything else scanning the pyarrow dataset protocol, reading the shared dimension coordinates once. Groups left unnamed keep their joined dimension names. Keys naming a group the Dataset does not have are ignored, so one map can be reused across variable subsets; two groups claiming one name is an error rather than a silent clobber. The geospatial suite's private copy of the split is gone — the harness calls `arrow_datasets`, so benchmark naming cannot drift from the library's. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013LfwhbX9VLmvoN6ka7py5w --- README.md | 8 + benchmarks/geospatial/_engines.py | 45 ++--- docs/engines.md | 56 ++++++- docs/examples.md | 49 +++++- docs/limitations.md | 14 +- tests/test_table_names.py | 262 ++++++++++++++++++++++++++++++ xarray_sql/__init__.py | 3 +- xarray_sql/backends/__init__.py | 2 + xarray_sql/backends/base.py | 44 ++++- xarray_sql/backends/datafusion.py | 44 ++++- xarray_sql/backends/duckdb.py | 87 ++++++++-- xarray_sql/backends/pyarrow.py | 83 ++++++++++ xarray_sql/df.py | 55 +++++++ xarray_sql/sql.py | 20 ++- 14 files changed, 696 insertions(+), 76 deletions(-) create mode 100644 tests/test_table_names.py diff --git a/README.md b/README.md index 37ca3dc2..34181e41 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. @@ -102,6 +105,11 @@ ds = xr.open_zarr( ctx = xql.XarrayContext() # Make sure to pass `chunks`! +# ERA5's variables sit on two different grids, so it registers as two tables. +# `table_names` is what lets you call them `surface` and `atmosphere` instead +# of `time_latitude_longitude` and `time_level_latitude_longitude`. The same +# kwarg works on `xql.register(con, ...)` for DuckDB and on +# `xql.arrow_datasets(ds, ...)` for Polars. ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', diff --git a/benchmarks/geospatial/_engines.py b/benchmarks/geospatial/_engines.py index 66d769bc..5de5b797 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 11dacac7..fba6f291 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -16,6 +16,45 @@ builds for itself: Everything between the seams — geometry functions, dialects, optimizers — belongs to the engine. +## Naming the tables + +A Dataset whose variables sit on different dimensions is registered as +one table per dimension group, because a table has one shape. +ARCO-ERA5 splits into a surface group on +`(time, latitude, longitude)` and an atmospheric group on +`(time, level, latitude, longitude)`; unnamed, those tables are called +`era5.time_latitude_longitude` and `era5.time_level_latitude_longitude`. +`table_names` maps a group's dimensions to the name you would rather +call it: + +```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") +``` + +This is the same keyword `XarrayContext.from_dataset` takes, and the +`name.group` spelling resolves on every engine that has a connection to +register into — so the query text above moves between DataFusion and +DuckDB unchanged. Groups you do not name keep their joined dimension +names, and keys naming a group the Dataset does not have are ignored, +so one naming map can be reused across Datasets holding different +subsets of the same variables. + +Polars has no connection object to dispatch on, so it takes the tables +directly: + +```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)) +``` + ## DataFusion (default) DataFusion is the built-in engine, wrapped in a session: @@ -90,6 +129,15 @@ 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. +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 +194,11 @@ 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. + 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 +219,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 483fb2a5..f2b925b7 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 8aab000c..b7e5d609 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 00000000..fca4b66e --- /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 0f98e1e0..1bbfe565 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 82aad8eb..af0334be 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 59eb488c..3a461005 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,9 +39,16 @@ def register( ds: xr.Dataset, *, chunks: Chunks = None, + table_names: TableNames = None, **kwargs: Any, ) -> ConT: - """Register *ds* as table *name* on *con*; returns *con*.""" + """Register *ds* as table *name* on *con*; returns *con*. + + ``table_names`` names the per-dimension-group tables a + mixed-dimension Dataset splits into (see + [resolve_table_names][xarray_sql.df.resolve_table_names]); + adapters that split must honour it. + """ ... @@ -74,6 +81,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 +102,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 +140,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 3d7c29e1..4fc652f0 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 1651ac05..bcd6325c 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,48 @@ __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*. + + ``con.register`` can only place an object in DuckDB's temporary + namespace, so a dim-group split lands as flat ``name_group`` tables. + A schema of views over them makes the dotted ``name.group`` spelling + resolve too — the same SQL text that queries the split on + DataFusion, which is the point: the table names are the user's, and + they should not have to rewrite queries to change engine. + + *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 +102,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 +146,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 dd0f2543..86283fed 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 83879f77..9a05ade5 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 a5df4fb9..c9911c12 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, From 1a9cb1bd9f588c7d010d3ca56e0252939c9e0fe2 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 5 Sep 2026 12:57:36 -0700 Subject: [PATCH 2/3] More succinct docs. --- docs/engines.md | 59 +++++++++++++++++-------------------------------- 1 file changed, 20 insertions(+), 39 deletions(-) diff --git a/docs/engines.md b/docs/engines.md index fba6f291..04201f92 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -16,45 +16,6 @@ builds for itself: Everything between the seams — geometry functions, dialects, optimizers — belongs to the engine. -## Naming the tables - -A Dataset whose variables sit on different dimensions is registered as -one table per dimension group, because a table has one shape. -ARCO-ERA5 splits into a surface group on -`(time, latitude, longitude)` and an atmospheric group on -`(time, level, latitude, longitude)`; unnamed, those tables are called -`era5.time_latitude_longitude` and `era5.time_level_latitude_longitude`. -`table_names` maps a group's dimensions to the name you would rather -call it: - -```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") -``` - -This is the same keyword `XarrayContext.from_dataset` takes, and the -`name.group` spelling resolves on every engine that has a connection to -register into — so the query text above moves between DataFusion and -DuckDB unchanged. Groups you do not name keep their joined dimension -names, and keys naming a group the Dataset does not have are ignored, -so one naming map can be reused across Datasets holding different -subsets of the same variables. - -Polars has no connection object to dispatch on, so it takes the tables -directly: - -```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)) -``` - ## DataFusion (default) DataFusion is the built-in engine, wrapped in a session: @@ -129,6 +90,18 @@ 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 @@ -199,6 +172,14 @@ 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 From a5ce9aac02aaa590fe73f84427330645b93d01db Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 5 Sep 2026 13:08:23 -0700 Subject: [PATCH 3/3] Fix Claudisms. --- README.md | 5 ----- xarray_sql/backends/base.py | 8 +------- xarray_sql/backends/duckdb.py | 8 ++------ 3 files changed, 3 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 34181e41..8bca8865 100644 --- a/README.md +++ b/README.md @@ -105,11 +105,6 @@ ds = xr.open_zarr( ctx = xql.XarrayContext() # Make sure to pass `chunks`! -# ERA5's variables sit on two different grids, so it registers as two tables. -# `table_names` is what lets you call them `surface` and `atmosphere` instead -# of `time_latitude_longitude` and `time_level_latitude_longitude`. The same -# kwarg works on `xql.register(con, ...)` for DuckDB and on -# `xql.arrow_datasets(ds, ...)` for Polars. ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', diff --git a/xarray_sql/backends/base.py b/xarray_sql/backends/base.py index 3a461005..1446bef8 100644 --- a/xarray_sql/backends/base.py +++ b/xarray_sql/backends/base.py @@ -42,13 +42,7 @@ def register( table_names: TableNames = None, **kwargs: Any, ) -> ConT: - """Register *ds* as table *name* on *con*; returns *con*. - - ``table_names`` names the per-dimension-group tables a - mixed-dimension Dataset splits into (see - [resolve_table_names][xarray_sql.df.resolve_table_names]); - adapters that split must honour it. - """ + """Register *ds* as table *name* on *con*; returns *con*.""" ... diff --git a/xarray_sql/backends/duckdb.py b/xarray_sql/backends/duckdb.py index bcd6325c..c602cc25 100644 --- a/xarray_sql/backends/duckdb.py +++ b/xarray_sql/backends/duckdb.py @@ -53,12 +53,8 @@ def _mirror_as_schema( ) -> None: """Expose flat tables as views in a DuckDB schema named *name*. - ``con.register`` can only place an object in DuckDB's temporary - namespace, so a dim-group split lands as flat ``name_group`` tables. - A schema of views over them makes the dotted ``name.group`` spelling - resolve too — the same SQL text that queries the split on - DataFusion, which is the point: the table names are the user's, and - they should not have to rewrite queries to change engine. + 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.