diff --git a/Cargo.lock b/Cargo.lock index 1d53519..8b58e75 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3375,7 +3375,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] name = "xarray_sql" -version = "0.2.2" +version = "0.2.3" dependencies = [ "arrow", "async-stream", diff --git a/README.md b/README.md index ee512dc..e4094da 100644 --- a/README.md +++ b/README.md @@ -20,21 +20,19 @@ import xarray as xr import xarray_sql as xql -# Open a year of ARCO-ERA5 — all 273 variables. Selecting a year up front -# keeps Dask's partition setup cheap before any chunks are read from GCS. -ds = ( - xr.open_zarr('gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', - chunks=dict(time=1), - storage_options={'token': 'anon'}) # Anonymous read from the public GCS bucket — no auth required. - .sel(time='2020') -) - +# Open ARCO-ERA5 (273 variables) and turn off Dask. +ds = xr.open_zarr( + 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', + chunks=None, # no dask graph + storage_options={'token': 'anon'}, # Anonymous read from the public GCS bucket — no auth required. +) + +# Register the dataset to a SQL context -- `chunks` are required. ctx = xql.XarrayContext() -ctx.from_dataset('era5', ds, table_names={ +ctx.from_dataset('era5', ds, chunks=dict(time=1), table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', }) -# Registration: ~0.5s for a full year of hourly ERA5, all variables. # Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library @@ -64,45 +62,36 @@ ctx.sql(''' AND TIMESTAMP '2020-01-01 05:00:00' GROUP BY level ORDER BY level DESC -''').to_pandas() -# level avg_c -# 0 1000 6.621012 ← surface -# 1 975 5.185638 -# 2 950 4.028429 -# 3 925 3.082812 -# 4 900 2.210917 -# 5 875 1.395018 -# 6 850 0.634267 -# 7 825 -0.210372 -# 8 800 -1.181075 -# 9 775 -2.306465 -# 10 750 -3.535534 -# 11 700 -6.241685 -# 12 650 -9.236364 -# 13 600 -12.580938 -# 14 550 -16.335386 -# 15 500 -20.643604 -# 16 450 -25.573401 -# 17 400 -31.156920 -# 18 350 -37.400552 -# 19 300 -43.852607 -# 20 250 -49.322132 -# 21 225 -51.569113 -# 22 200 -53.693248 -# 23 175 -55.890484 -# 24 150 -58.382290 -# 25 125 -61.091916 -# 26 100 -63.624885 ← tropopause -# 27 70 -63.182300 -# 28 50 -60.124845 -# 29 30 -55.986327 -# 30 20 -52.433089 -# 31 10 -44.140750 -# 32 7 -38.707350 -# 33 5 -32.621999 -# 34 3 -21.509175 -# 35 2 -13.355764 -# 36 1 -9.020513 ← top of atmosphere +''') +# DataFrame() +# +-------+----------------------+ +# | level | avg_c | +# +-------+----------------------+ +# | 1000 | 6.6210120796502565 | +# | 975 | 5.185637919348153 | +# | 950 | 4.028428657263021 | +# | 925 | 3.0828117974912743 | +# | 900 | 2.2109172992531967 | +# | 875 | 1.395017610194202 | +# | 850 | 0.6342670572626616 | +# | 825 | -0.21037158786759846 | +# | 800 | -1.1810754318269687 | +# | 775 | -2.3064649711534457 | +# +-------+----------------------+ + +# Average temperature global 2m temperature around the morning of 2020-01-01. +( + ctx.sql(''' + SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY latitude, longitude + ORDER BY latitude DESC, longitude + ''') + .to_dataset(dims=["latitude", "longitude"], template=ds) +) +# ... ``` _(A runnable version of this example lives at diff --git a/perf_tests/era5_temp_profile.py b/perf_tests/era5_temp_profile.py index 28275a9..652e75e 100644 --- a/perf_tests/era5_temp_profile.py +++ b/perf_tests/era5_temp_profile.py @@ -1,24 +1,29 @@ #!/usr/bin/env python3 -"""Surface and global-atmospheric temperatures on 2020-01-01, in SQL. - -Two queries against ARCO-ERA5 on the morning of January 1, 2020: - - * **Surface (local).** Average 2m-temperature over a small grid covering the - New York City area for the first six hours. - * **Atmosphere (global).** Average temperature per pressure level, computed - over the entire planet for the same six hours — a classic atmospheric - temperature profile (surface around 1000 hPa is warmest, tropopause near - 100 hPa is coldest). - -Both queries express their filters entirely in SQL: ``xr.open_zarr`` is given -a single calendar year and no spatial slicing. The library's table provider -prunes time partitions for ``WHERE time …`` filters, and pushes ``WHERE -latitude/longitude …`` down to dimension columns. +"""Surface and atmospheric temperatures on 2020-01-01, in SQL. + +Three queries against the full ARCO-ERA5 archive, all filtered to the morning +of January 1, 2020 entirely in SQL: + + * **Surface point (local).** Average 2m-temperature over a small grid box + covering the New York City area for the first six hours — a scalar. + * **Atmospheric profile (global).** Average temperature per pressure level + over the whole planet — a classic profile (surface ~1000 hPa warmest, + tropopause coldest), returned as a DataFrame. + * **Surface map (global).** Average 2m-temperature per grid point, + reconstructed back into an ``xr.Dataset`` (a 721x1440 raster) via + ``.to_dataset()``. + +Nothing is sliced on the xarray side: ``xr.open_zarr`` opens the entire archive +with Dask turned off (``chunks=None``), and ``from_dataset`` partitions it with +``chunks={'time': 1}``. Every filter is expressed in SQL — the table provider +prunes time partitions for ``WHERE time …`` and pushes ``WHERE +latitude/longitude …`` down to the dimension columns — so each query reads only +the partitions it needs. ARCO-ERA5's atmospheric variables are stored in native Zarr chunks of shape -``(1, 37, 721, 1440)`` — about 150 MB per hour. We align Dask chunks to that -shape with ``chunks={'time': 1}`` so chunks fetch from GCS concurrently. The -global atmospheric query scans ~230M rows after pruning. +``(1, 37, 721, 1440)`` (~150 MB/hour); ``chunks={'time': 1}`` aligns partitions +to that shape so they fetch from GCS concurrently. The global atmospheric query +scans ~230M rows after pruning. The Zarr is read anonymously from the public GCS bucket — no auth required. """ @@ -34,27 +39,27 @@ def main() -> None: - full = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) - - # Open a full calendar year — all 273 variables. No spatial slicing on - # the xarray side; SQL WHERE clauses below express the filters. + # Open the whole archive with Dask turned off. SQL WHERE clauses below + # express every filter; nothing is pre-sliced on the xarray side. # # Heads up: the library pushes column projection down to Zarr, so SELECT # only fetches what you ask for — but `SELECT * FROM era5.surface` would - # try to read every variable across the year (terabytes from GCS). + # try to read every variable across the archive (terabytes from GCS). # Always SELECT specific columns. - ds = full.sel(time="2020").chunk({"time": 1}) + ds = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) print( - "ARCO-ERA5 opened: year 2020, " - f"{ds.sizes['time']:,} hourly time steps, " - f"{len(ds.data_vars)} variables (no spatial pre-slicing)." + f"ARCO-ERA5 opened: {ds.sizes['time']:,} hourly time steps, " + f"{len(ds.data_vars)} variables (no Dask, no pre-slicing)." ) ctx = xql.XarrayContext() t0 = time.perf_counter() + # `chunks` is required here: `ds` is not Dask-backed, so the partition grid + # is given explicitly (one partition per hourly time step). ctx.from_dataset( "era5", ds, + chunks={"time": 1}, table_names={ ("time", "latitude", "longitude"): "surface", ("time", "level", "latitude", "longitude"): "atmosphere", @@ -63,6 +68,8 @@ def main() -> None: print(f"Registration: {time.perf_counter() - t0:.2f}s") ctx.sql("SELECT 1").to_pandas() # warm the planner + # 1. Surface point: average 2m-temperature over NYC. WHERE clauses on + # dimension columns prune partitions and push down to Zarr. print("\nAverage 2m-temperature over NYC, 2020-01-01 00:00-05:00 UTC (°C):") t0 = time.perf_counter() surface = ctx.sql( @@ -78,6 +85,7 @@ def main() -> None: print(surface) print(f" ({time.perf_counter() - t0:.2f}s)") + # 2. Atmospheric profile: average temperature per pressure level, globally. print( "\nAverage temperature per pressure level, globally, " "2020-01-01 00:00-05:00 UTC (°C):" @@ -96,6 +104,27 @@ def main() -> None: print(profile.to_string(index=False)) print(f" ({time.perf_counter() - t0:.2f}s, ~230M rows scanned)") + # 3. Surface map: average 2m-temperature per grid point, reconstructed back + # into an xr.Dataset (a 721x1440 global raster). `template=ds` recovers + # coordinate metadata; the aggregation result is materialized eagerly. + print( + "\nAverage 2m-temperature per grid point, globally, " + "2020-01-01 00:00-05:00 UTC -> xr.Dataset:" + ) + t0 = time.perf_counter() + temp_map = ctx.sql( + """ + SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY latitude, longitude + ORDER BY latitude DESC, longitude + """ + ).to_dataset(dims=["latitude", "longitude"], template=ds) + print(temp_map) + print(f" ({time.perf_counter() - t0:.2f}s, ~6M rows scanned)") + if __name__ == "__main__": main() diff --git a/pyproject.toml b/pyproject.toml index d9253bd..7246548 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,6 @@ classifiers = [ "Topic :: Database :: Front-Ends", ] dependencies = [ - "dask>=2024.8.0", "datafusion==52.0.0", # This needs to match the cargo datafusion version!! "xarray>=2024.7.0", ] @@ -41,6 +40,7 @@ test = [ "pytest", "xarray[io]", "gcsfs", + "dask", ] docs = [ "zensical", @@ -90,6 +90,7 @@ module = [ "pyarrow.*", "datafusion.*", "xarray.*", + "pandas.*", ] ignore_missing_imports = true diff --git a/tests/test_df.py b/tests/test_df.py index a2b443f..e7e2b1e 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -90,6 +90,27 @@ def test_block_slices_dimensional_unchunked_raises(): list(block_slices(ds)) +def test_block_slices_dask_free_on_eager_dataset(): + # The chunk spec alone partitions an eager (non-dask) Dataset -- no dask + # backing required and no transient dask graph built -- tiling each + # dimension exactly, with a short final chunk for an uneven split. + ds = xr.Dataset( + {"v": (("time", "x"), np.arange(10 * 3).reshape(10, 3))}, + coords={"time": np.arange(10), "x": np.arange(3)}, + ) + assert not ds.chunks # genuinely not dask-backed + + blocks = list(block_slices(ds, chunks={"time": 4})) + # 10 / 4 -> partitions of length (4, 4, 2). + assert [(b["time"].start, b["time"].stop) for b in blocks] == [ + (0, 4), + (4, 8), + (8, 10), + ] + # The unchunked dimension spans its full extent in every block. + assert all(b["x"] == slice(0, 3) for b in blocks) + + def test_from_map_basic(): def make_df(x): return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]}) diff --git a/tests/test_ds.py b/tests/test_ds.py new file mode 100644 index 0000000..3dcc079 --- /dev/null +++ b/tests/test_ds.py @@ -0,0 +1,595 @@ +"""Tests for the SQL -> xarray reverse path. + +Covers the user-facing contract of ``ctx.sql(...).to_dataset(...)``: + +* Wrapper behavior on the object returned by ``ctx.sql`` (method passthrough, + ``to_pandas`` equivalence). +* Round-trip identity across the eager and chunked paths (one parametrized + ``assert_identical`` test). +* Aggregation behavior: dim reduction, single-scan execution, ``ORDER BY`` + direction, and the ``chunks`` argument (eager / inherit / ``"auto"``). +* ``dims`` inference and ``template`` resolution (name or Dataset), with error + paths and metadata recovery. +* Indexing the chunked backend and filtered-query coordinate handling. + +The tests favor the user-visible contract (values, dims, attrs) over the +implementation path, so the suite stays useful as the backend evolves. +""" + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from xarray_sql import XarrayContext +from xarray_sql.ds import XarrayDataFrame + + +# --------------------------------------------------------------------------- +# Wrapper: ctx.sql(...) returns XarrayDataFrame +# --------------------------------------------------------------------------- + + +def test_sql_returns_wrapper_that_forwards_methods(air_dataset_small): + """``ctx.sql`` returns an ``XarrayDataFrame`` that forwards un-overridden + DataFusion methods (e.g. ``schema()``) via ``__getattr__``.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql("SELECT * FROM air LIMIT 5") + assert isinstance(result, XarrayDataFrame) + names = [f.name for f in result.schema()] + assert {"lat", "lon", "time", "air"}.issubset(set(names)) + + +def test_to_pandas_unchanged_behavior(air_dataset_small): + """Wrapped ``.to_pandas()`` is bit-for-bit equal to the un-wrapped path.""" + from datafusion import SessionContext + + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + wrapped = ctx.sql("SELECT * FROM air LIMIT 7").to_pandas() + raw = SessionContext.sql(ctx, "SELECT * FROM air LIMIT 7").to_pandas() + pd.testing.assert_frame_equal(wrapped, raw) + + +def test_xarray_dataframe_satisfies_datafusion_contract(air_dataset_small): + """``XarrayDataFrame`` must expose the full DataFusion ``DataFrame`` API it + stands in for -- public methods *and* the functional dunders (``df[col]``, + the Arrow C-stream export) that Python resolves on the type, so a wrapper + cannot silently drop them. Subclassing inherits all of it; this guards + against regressing to a composition wrapper (where ``__getattr__`` cannot + forward special methods invoked via syntax). + """ + from datafusion import DataFrame + + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + wrapped = ctx.sql("SELECT * FROM air LIMIT 2") + + # Every public attribute of DataFusion's DataFrame is present. + expected = {n for n in dir(DataFrame) if not n.startswith("_")} + missing = sorted(n for n in expected if not hasattr(wrapped, n)) + assert not missing, f"missing DataFusion API: {missing}" + + # Functional dunders resolve via syntax (not attribute forwarding) and work. + assert wrapped["air"].to_pandas() is not None # df[col] selection + assert wrapped.__arrow_c_stream__() is not None # zero-copy Arrow export + + +# --------------------------------------------------------------------------- +# Round-trip identity (parametrized over local + tutorial datasets) +# --------------------------------------------------------------------------- + + +def _clear_encoding(ds: xr.Dataset) -> xr.Dataset: + """Strip ``encoding`` from a Dataset and all its variables. + + Round-trip identity tests should not be coupled to encoding choices, + since template-recovery deliberately drops dtype-bound keys. + """ + ds = ds.copy() + for v in ds.variables.values(): + v.encoding.clear() + ds.encoding.clear() + return ds + + +@pytest.mark.parametrize( + "fixture_name", + # ``air`` exercises the eager path (single-chunk source); ``weather`` + # exercises the chunked path and adds datetime + non-dim coordinates. + ["air_dataset_small", "weather_dataset"], +) +def test_round_trip_identity(request, fixture_name): + """``SELECT *`` round-trips to a Dataset that is ``assert_identical`` + to the source: values, dims, coord values, dtypes, non-dim coords, + and attrs all match (modulo coord ordering, normalized on both + sides). + """ + source = request.getfixturevalue(fixture_name).copy() + source.attrs["round_trip_marker"] = "yes" + first_var = next(iter(source.data_vars)) + source[first_var].attrs["units"] = "test_units" + + ctx = XarrayContext() + ctx.from_dataset("t", source) + # ORDER BY each dimension in its native direction: this gives a single, + # deterministically-ordered stream so the round-trip reproduces the + # source coordinate order exactly -- no re-sorting. (Without ORDER BY, + # row order across partitions is unspecified, so coord order would be.) + src = source.compute() + order = ", ".join( + f'"{d}" {"DESC" if src[d].values[0] > src[d].values[-1] else "ASC"}' + for d in src.dims + ) + out = ctx.sql(f"SELECT * FROM t ORDER BY {order}").to_dataset().compute() + + actual = _clear_encoding(out) + expected = _clear_encoding(src) + xr.testing.assert_identical(actual, expected) + + +def test_aggregation_drops_dim(air_dataset_small): + """``GROUP BY lat, lon`` over time -> 2D Dataset with the alias.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql( + """SELECT lat, lon, AVG(air) AS air_avg + FROM air + GROUP BY lat, lon + ORDER BY lat DESC, lon""" + ).to_dataset(dims=["lat", "lon"]) + assert set(out.dims) == {"lat", "lon"} + assert "air_avg" in out.data_vars + assert "air" not in out.data_vars + expected = ( + air_dataset_small.compute() + .mean(dim="time")["air"] + .values + ) + actual = out["air_avg"].values + np.testing.assert_allclose(actual, expected) + + +def test_barrier_query_scans_source_once(air_dataset_small): + """A barrier plan (aggregation) executes the source exactly once. + + The lazy scan path re-runs the whole upstream plan for every coordinate + discovery and every variable access; for an aggregation -- which cannot push + an indexer filter below the GROUP BY -- that is pure re-computation of an + expensive scan. ``to_dataset()`` on a barrier plan must instead make a + single streamed pass over the source, and ``.compute()`` must trigger no + further reads. + """ + from xarray_sql.df import block_slices + from xarray_sql.reader import read_xarray_table + + reads: list = [] + table = read_xarray_table( + air_dataset_small, + chunks={"time": 6}, + _iteration_callback=lambda block, proj: reads.append(block), + ) + n_partitions = len(list(block_slices(air_dataset_small, {"time": 6}))) + + ctx = XarrayContext() + ctx.register_table("air", table) + ctx._registered_datasets["air"] = air_dataset_small + + out = ctx.sql( + """ + SELECT lat, lon, AVG(air) AS air_avg + FROM air + GROUP BY lat, lon + ORDER BY lat DESC, lon + """ + ).to_dataset(dims=["lat", "lon"]) + reads_after_construct = len(reads) + out.compute() + reads_after_compute = len(reads) + + # Exactly one pass over the source (each partition read once) ... + assert reads_after_construct == n_partitions + # ... and computing the materialized result re-reads nothing. + assert reads_after_compute == reads_after_construct + + +def test_order_by_direction_sets_dim_order(air_dataset_small): + """A barrier query's ORDER BY direction carries through to the Dataset + dimension order, rather than being force-sorted ascending. + + ``ORDER BY lat DESC`` must yield a strictly descending ``lat`` dimension, + with data still correctly aligned to those (descending) coordinates. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql( + "SELECT lat, AVG(air) AS air_avg FROM air GROUP BY lat ORDER BY lat DESC" + ).to_dataset(dims=["lat"]) + + lat = out["lat"].values + assert (np.diff(lat) < 0).all(), f"expected descending lat, got {lat}" + + # air's native latitude is already descending, matching ORDER BY lat + # DESC, so the result aligns with the source mean without re-sorting. + expected = air_dataset_small.compute().mean(dim=["time", "lon"])["air"] + np.testing.assert_allclose(out["air_avg"].values, expected.values) + + +def test_chunks_argument_controls_partitioning(synthetic_dataset): + """``chunks`` controls eager-vs-chunked and inherits the source grid. + + The default ``"inherit"`` reuses the source's genuinely multi-chunk + dimensions, so the output chunk grid maps onto the source partitions; + ``chunks=None`` forces an eager, in-memory result. Both reproduce the source. + """ + import dask.array as da + + ctx = XarrayContext() + ctx.from_dataset("t", synthetic_dataset) + var = next(iter(synthetic_dataset.data_vars)) + + inherited = ctx.sql("SELECT * FROM t").to_dataset() + assert isinstance(inherited[var].data, da.Array) + # Output time chunks align to the source's time partitions. + assert ( + inherited.chunksizes["time"] == synthetic_dataset.chunksizes["time"] + ) + + eager = ctx.sql("SELECT * FROM t").to_dataset(chunks=None) + assert not isinstance(eager[var].data, da.Array) + + # No ORDER BY, so coordinate order across partitions is unspecified; + # normalize both sides before comparing values. + xr.testing.assert_allclose( + inherited.compute().sortby(["time", "lat", "lon"]), + synthetic_dataset.compute().sortby(["time", "lat", "lon"]), + ) + + +def test_chunks_auto_snaps_to_source_partitions(): + """``chunks="auto"`` coarsens to the byte budget but snaps chunk boundaries + to whole source partitions (so no chunk splits a source partition).""" + import dask + + # 12 source partitions of size 2 along time. + ds = xr.Dataset( + {"v": (("time", "x"), np.arange(24 * 4, dtype="float64").reshape(24, 4))}, + coords={"time": np.arange(24), "x": np.arange(4)}, + ).chunk({"time": 2}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + + # block bytes = 8 * 2(time) * 4(x) = 64; target 192 -> merge 3 partitions. + with dask.config.set({"array.chunk-size": "192B"}): + out = ctx.sql("SELECT * FROM t").to_dataset(chunks="auto") + + time_chunks = out.chunksizes["time"] + assert all(c % 2 == 0 for c in time_chunks) # aligned to source size 2 + assert time_chunks[0] > 2 # genuinely coarsened + assert len(time_chunks) < 12 # fewer chunks than source partitions + + xr.testing.assert_allclose( + out.compute().sortby(["time", "x"]), ds.compute().sortby(["time", "x"]) + ) + + +# --------------------------------------------------------------------------- +# dimension_columns / template resolution rules +# --------------------------------------------------------------------------- + + +def test_to_dataset_infer_fails_when_no_template_fits(air_dataset_small): + """If no registered Dataset's dims fit the result -> clear error.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises( + ValueError, match="dims cannot be inferred" + ): + ctx.sql( + """ + SELECT lat, lon, AVG(air) AS air_avg + FROM air + GROUP BY lat, lon + ORDER BY lat DESC, lon + """ + ).to_dataset() + + +def test_template_accepts_name_or_dataset(air_dataset_small): + """``template=`` accepts either a registered table name or a Dataset + object, with equivalent metadata recovery.""" + other = air_dataset_small.copy() + other.attrs = {"flag": "other"} + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + ctx.from_dataset("other", other) + + by_name = ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template="other" + ) + by_object = ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template=other + ) + assert by_name.attrs == {"flag": "other"} + assert by_object.attrs == {"flag": "other"} + + +def test_template_unknown_name_raises(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(ValueError, match="not a registered table"): + ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template="missing" + ) + + +def test_template_recovers_var_encoding_strips_dtype(air_dataset_small): + """``zlib`` survives; dtype-bound keys are stripped (SQL may have cast).""" + ds = air_dataset_small.copy() + ds["air"].encoding = { + "zlib": True, + "dtype": "int16", + "_FillValue": -999, + "missing_value": -999, + } + ctx = XarrayContext() + ctx.from_dataset("air", ds) + out = ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"] + ) + assert out["air"].encoding.get("zlib") is True + assert "dtype" not in out["air"].encoding + assert "_FillValue" not in out["air"].encoding + assert "missing_value" not in out["air"].encoding + + +def test_template_aggregation_alias_no_attrs(air_dataset_small): + """``air_avg`` from ``AVG(air)`` does NOT inherit attrs from ``air``.""" + ds = air_dataset_small.copy() + ds["air"].attrs = {"units": "K"} + ctx = XarrayContext() + ctx.from_dataset("air", ds) + out = ctx.sql( + """ + SELECT lat, lon, AVG(air) AS air_avg + FROM air + GROUP BY lat, lon + ORDER BY lat DESC, lon + """ + ).to_dataset(dims=["lat", "lon"]) + assert "air_avg" in out.data_vars + assert out["air_avg"].attrs == {} + + +# --------------------------------------------------------------------------- +# Chunked (lazy SQLBackendArray) backend: value-level contract +# --------------------------------------------------------------------------- + + +def test_chunked_backend_indexing_matches_eager(air_dataset_small): + """Indexing a chunked result (the lazy ``SQLBackendArray`` path, reached by + passing ``chunks=``) matches the eager equivalent across every indexer kind. + + ``chunks=`` forces the dask-wrapped backend whose chunks read their + coordinate range via DataFusion filter pushdown; this exercises the int, + slice, outer-array, and vectorized indexer translations -- the last via + xarray's ``IndexingSupport.OUTER`` adapter (outer reads + numpy gather). + """ + import dask.array as da + + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + chunked = ctx.sql("SELECT * FROM air ORDER BY lat DESC, lon").to_dataset(chunks={"time": 4}) + assert isinstance(chunked["air"].data, da.Array) # genuinely lazy/chunked + # Compare lazy indexing against computing-then-indexing the SAME Dataset, so + # both sides share one coordinate order (positional indexers stay aligned). + eager = chunked.compute() + + # int indexer + np.testing.assert_array_equal( + chunked["air"].isel(time=0), eager["air"].isel(time=0) + ) + # slice indexer + np.testing.assert_array_equal( + chunked["air"].isel(time=slice(0, 3)), + eager["air"].isel(time=slice(0, 3)), + ) + # outer (fancy) array indexer + np.testing.assert_array_equal( + chunked["air"].isel(lat=[0, 3, 5]).values, + eager["air"].isel(lat=[0, 3, 5]).values, + ) + # vectorized indexer (xarray adapter -> outer + gather) + pt = xr.DataArray([0, 3, 1], dims="point") + pl = xr.DataArray([2, 0, 5], dims="point") + np.testing.assert_array_equal( + chunked["air"].isel(time=pt, lat=pl).values, + eager["air"].isel(time=pt, lat=pl).values, + ) + + +# --------------------------------------------------------------------------- +# Filtered queries +# --------------------------------------------------------------------------- + + +def test_filtered_query_keeps_present_coords(air_dataset_small): + """A filtered query yields a Dataset whose coordinates are exactly the + values present in the result (a smaller grid than the source). Users who + want the full grid back reindex the result themselves.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + threshold = float(air_dataset_small["lat"].values[5]) + out = ctx.sql(f"SELECT * FROM air WHERE lat > {threshold}").to_dataset() + assert (out["lat"].values > threshold).all() + assert out.sizes["lat"] < air_dataset_small.sizes["lat"] + + +# --------------------------------------------------------------------------- +# Filtered/lazy path, empty results, and malformed-grid guards +# --------------------------------------------------------------------------- + + +def test_filtered_query_through_chunked_path(): + """A WHERE filter combined with chunks= exercises the lazy SQLBackendArray + pushdown (indexer -> df.filter -> searchsorted scatter), not just eager.""" + import dask.array as da + + ds = xr.Dataset( + {"v": (("time", "lat"), np.arange(8 * 5, dtype="float64").reshape(8, 5))}, + coords={"time": np.arange(8), "lat": np.arange(5) * 1.0}, + ).chunk({"time": 2}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + + chunked = ctx.sql( + "SELECT * FROM t WHERE lat >= 2 ORDER BY time, lat" + ).to_dataset(chunks={"time": 2}) + assert isinstance(chunked["v"].data, da.Array) # genuinely lazy/chunked + assert list(chunked["lat"].values) == [2.0, 3.0, 4.0] # filtered subset + + np.testing.assert_array_equal( + chunked.compute()["v"].values, + ds.sel(lat=ds["lat"] >= 2)["v"].values, + ) + + +def test_empty_result_eager_and_chunked(): + """A query that returns zero rows yields a well-formed empty Dataset on + both the eager and chunked paths (no exception).""" + ds = xr.Dataset( + {"v": (("time", "lat"), np.arange(8 * 5, dtype="float64").reshape(8, 5))}, + coords={"time": np.arange(8), "lat": np.arange(5) * 1.0}, + ).chunk({"time": 2}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + q = "SELECT * FROM t WHERE lat > 9999" + + eager = ctx.sql(q).to_dataset(chunks=None) + assert eager.sizes["time"] == 0 and eager.sizes["lat"] == 0 + + chunked = ctx.sql(q).to_dataset(chunks={"time": 2}).compute() + assert chunked.sizes["time"] == 0 and chunked.sizes["lat"] == 0 + + +def test_integer_var_with_missing_cells_raises(): + """An integer data variable that doesn't fill the grid raises (integers + have no missing-value sentinel) rather than emit silent garbage.""" + ds = xr.Dataset( + {"n": (("lat", "lon"), np.arange(6, dtype=np.int64).reshape(3, 2))}, + coords={"lat": [0, 1, 2], "lon": [10, 11]}, + ).chunk({"lat": 3}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + # Drop a single cell -> the lat/lon grid has a hole, but n is int64. + with pytest.raises(ValueError, match="integer dtype"): + ctx.sql( + "SELECT lat, lon, n FROM t WHERE NOT (lat = 1 AND lon = 11)" + ).to_dataset(dims=["lat", "lon"]) + + +def test_duplicate_dimension_tuples_raise(): + """Dropping a dimension without aggregating yields duplicate dim tuples + (more rows than cells) -> ValueError, not silent last-write-wins.""" + ds = xr.Dataset( + {"v": (("time", "lat"), np.arange(6, dtype="float64").reshape(3, 2))}, + coords={"time": [0, 1, 2], "lat": [0.0, 1.0]}, + ).chunk({"time": 3}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + with pytest.raises(ValueError, match="duplicate dimension tuples"): + ctx.sql("SELECT lat, v FROM t").to_dataset(dims=["lat"]) + + +# --------------------------------------------------------------------------- +# dims inference / multi-table / chunks inherit edge cases +# --------------------------------------------------------------------------- + + +def test_inherit_reduction_drops_chunked_dim_is_eager(synthetic_dataset): + """The default chunks='inherit' over a reduction that drops the chunked + dimension yields an eager (non-dask) result -- the stated default.""" + import dask.array as da + + ctx = XarrayContext() + ctx.from_dataset("t", synthetic_dataset) + out = ctx.sql( + "SELECT lat, lon, AVG(temperature) AS m FROM t " + "GROUP BY lat, lon ORDER BY lat, lon" + ).to_dataset(dims=["lat", "lon"]) + assert not isinstance(out["m"].data, da.Array) + + +def test_multi_table_schema_round_trip(): + """A mixed-dimension Dataset registers as namespaced tables; one sub-table + round-trips, inferring dims from (and recovering metadata via) its + registered-name template.""" + ds = xr.Dataset( + { + "t2m": (("time", "x"), np.arange(12, dtype="float32").reshape(4, 3)), + "temp": ( + ("time", "lev", "x"), + np.arange(24, dtype="float32").reshape(4, 2, 3), + ), + }, + coords={"time": np.arange(4), "lev": [100, 200], "x": [0, 1, 2]}, + ).chunk({"time": 2}) + ctx = XarrayContext() + ctx.from_dataset( + "era5", + ds, + table_names={ + ("time", "x"): "surface", + ("time", "lev", "x"): "atmosphere", + }, + ) + out = ( + ctx.sql("SELECT * FROM era5.surface ORDER BY time, x") + .to_dataset(template="era5.surface") + .compute() + ) + assert set(out.dims) == {"time", "x"} # dims inferred from the surface table + np.testing.assert_array_equal(out["t2m"].values, ds["t2m"].values) + + +def test_dims_inference_ambiguous_raises(): + """When several registered Datasets' dims all fit the result, dims cannot + be inferred unambiguously.""" + a = xr.Dataset( + {"u": (("lat", "lon"), np.zeros((2, 2)))}, + coords={"lat": [0, 1], "lon": [0, 1]}, + ).chunk({"lat": 2}) + b = a.rename({"u": "w"}) + ctx = XarrayContext() + ctx.from_dataset("a", a) + ctx.from_dataset("b", b) + with pytest.raises(ValueError, match="cannot be inferred"): + ctx.sql("SELECT lat, lon, u FROM a").to_dataset() + + +def test_cftime_gregorian_like_round_trips_values(): + """A noleap (Gregorian-like cftime) dataset round-trips its DATA correctly. + + Known limitation: the calendar label is NOT restored on the reverse path -- + the time coordinate comes back as datetime64, not cftime. The forward path + encodes Gregorian-like cftime as Arrow timestamps; to_dataset does not + re-decode them to cftime. This test pins the data round-trip and documents + the gap (it deliberately does not assert the time dtype). + """ + times = xr.date_range( + "2000-01-01", periods=6, freq="D", calendar="noleap", use_cftime=True + ) + ds = xr.Dataset( + {"v": (("time", "x"), np.arange(18, dtype="float64").reshape(6, 3))}, + coords={"time": times, "x": [0, 1, 2]}, + ).chunk({"time": 3}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + out = ( + ctx.sql("SELECT * FROM t ORDER BY time, x") + .to_dataset(template=ds) + .compute() + ) + assert out.sizes == {"time": 6, "x": 3} + np.testing.assert_array_equal(out["v"].values, ds["v"].values) diff --git a/uv.lock b/uv.lock index fa753dc..ae366f8 100644 --- a/uv.lock +++ b/uv.lock @@ -168,40 +168,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/50/cd/30110dc0ffcf3b131156077b90e9f60ed75711223f306da4db08eff8403b/beautifulsoup4-4.13.4-py3-none-any.whl", hash = "sha256:9bbbb14bfde9d79f38b8cd5f8c7c85f4b8f2523190ebed90e950a8dea4cb1c4b", size = 187285, upload-time = "2025-04-15T17:05:12.221Z" }, ] -[[package]] -name = "black" -version = "24.10.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "pathspec" }, - { name = "platformdirs" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d8/0d/cc2fb42b8c50d80143221515dd7e4766995bd07c56c9a3ed30baf080b6dc/black-24.10.0.tar.gz", hash = "sha256:846ea64c97afe3bc677b761787993be4991810ecc7a4a937816dd6bddedc4875", size = 645813, upload-time = "2024-10-07T19:20:50.361Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/f3/465c0eb5cddf7dbbfe1fecd9b875d1dcf51b88923cd2c1d7e9ab95c6336b/black-24.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6668650ea4b685440857138e5fe40cde4d652633b1bdffc62933d0db4ed9812", size = 1623211, upload-time = "2024-10-07T19:26:12.43Z" }, - { url = "https://files.pythonhosted.org/packages/df/57/b6d2da7d200773fdfcc224ffb87052cf283cec4d7102fab450b4a05996d8/black-24.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1c536fcf674217e87b8cc3657b81809d3c085d7bf3ef262ead700da345bfa6ea", size = 1457139, upload-time = "2024-10-07T19:25:06.453Z" }, - { url = "https://files.pythonhosted.org/packages/6e/c5/9023b7673904a5188f9be81f5e129fff69f51f5515655fbd1d5a4e80a47b/black-24.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:649fff99a20bd06c6f727d2a27f401331dc0cc861fb69cde910fe95b01b5928f", size = 1753774, upload-time = "2024-10-07T19:23:58.47Z" }, - { url = "https://files.pythonhosted.org/packages/e1/32/df7f18bd0e724e0d9748829765455d6643ec847b3f87e77456fc99d0edab/black-24.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:fe4d6476887de70546212c99ac9bd803d90b42fc4767f058a0baa895013fbb3e", size = 1414209, upload-time = "2024-10-07T19:24:42.54Z" }, - { url = "https://files.pythonhosted.org/packages/c2/cc/7496bb63a9b06a954d3d0ac9fe7a73f3bf1cd92d7a58877c27f4ad1e9d41/black-24.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5a2221696a8224e335c28816a9d331a6c2ae15a2ee34ec857dcf3e45dbfa99ad", size = 1607468, upload-time = "2024-10-07T19:26:14.966Z" }, - { url = "https://files.pythonhosted.org/packages/2b/e3/69a738fb5ba18b5422f50b4f143544c664d7da40f09c13969b2fd52900e0/black-24.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f9da3333530dbcecc1be13e69c250ed8dfa67f43c4005fb537bb426e19200d50", size = 1437270, upload-time = "2024-10-07T19:25:24.291Z" }, - { url = "https://files.pythonhosted.org/packages/c9/9b/2db8045b45844665c720dcfe292fdaf2e49825810c0103e1191515fc101a/black-24.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4007b1393d902b48b36958a216c20c4482f601569d19ed1df294a496eb366392", size = 1737061, upload-time = "2024-10-07T19:23:52.18Z" }, - { url = "https://files.pythonhosted.org/packages/a3/95/17d4a09a5be5f8c65aa4a361444d95edc45def0de887810f508d3f65db7a/black-24.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:394d4ddc64782e51153eadcaaca95144ac4c35e27ef9b0a42e121ae7e57a9175", size = 1423293, upload-time = "2024-10-07T19:24:41.7Z" }, - { url = "https://files.pythonhosted.org/packages/90/04/bf74c71f592bcd761610bbf67e23e6a3cff824780761f536512437f1e655/black-24.10.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e39e0fae001df40f95bd8cc36b9165c5e2ea88900167bddf258bacef9bbdc3", size = 1644256, upload-time = "2024-10-07T19:27:53.355Z" }, - { url = "https://files.pythonhosted.org/packages/4c/ea/a77bab4cf1887f4b2e0bce5516ea0b3ff7d04ba96af21d65024629afedb6/black-24.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d37d422772111794b26757c5b55a3eade028aa3fde43121ab7b673d050949d65", size = 1448534, upload-time = "2024-10-07T19:26:44.953Z" }, - { url = "https://files.pythonhosted.org/packages/4e/3e/443ef8bc1fbda78e61f79157f303893f3fddf19ca3c8989b163eb3469a12/black-24.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14b3502784f09ce2443830e3133dacf2c0110d45191ed470ecb04d0f5f6fcb0f", size = 1761892, upload-time = "2024-10-07T19:24:10.264Z" }, - { url = "https://files.pythonhosted.org/packages/52/93/eac95ff229049a6901bc84fec6908a5124b8a0b7c26ea766b3b8a5debd22/black-24.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:30d2c30dc5139211dda799758559d1b049f7f14c580c409d6ad925b74a4208a8", size = 1434796, upload-time = "2024-10-07T19:25:06.239Z" }, - { url = "https://files.pythonhosted.org/packages/d0/a0/a993f58d4ecfba035e61fca4e9f64a2ecae838fc9f33ab798c62173ed75c/black-24.10.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cbacacb19e922a1d75ef2b6ccaefcd6e93a2c05ede32f06a21386a04cedb981", size = 1643986, upload-time = "2024-10-07T19:28:50.684Z" }, - { url = "https://files.pythonhosted.org/packages/37/d5/602d0ef5dfcace3fb4f79c436762f130abd9ee8d950fa2abdbf8bbc555e0/black-24.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f93102e0c5bb3907451063e08b9876dbeac810e7da5a8bfb7aeb5a9ef89066b", size = 1448085, upload-time = "2024-10-07T19:28:12.093Z" }, - { url = "https://files.pythonhosted.org/packages/47/6d/a3a239e938960df1a662b93d6230d4f3e9b4a22982d060fc38c42f45a56b/black-24.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ddacb691cdcdf77b96f549cf9591701d8db36b2f19519373d60d31746068dbf2", size = 1760928, upload-time = "2024-10-07T19:24:15.233Z" }, - { url = "https://files.pythonhosted.org/packages/dd/cf/af018e13b0eddfb434df4d9cd1b2b7892bab119f7a20123e93f6910982e8/black-24.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:680359d932801c76d2e9c9068d05c6b107f2584b2a5b88831c83962eb9984c1b", size = 1436875, upload-time = "2024-10-07T19:24:42.762Z" }, - { url = "https://files.pythonhosted.org/packages/8d/a7/4b27c50537ebca8bec139b872861f9d2bf501c5ec51fcf897cb924d9e264/black-24.10.0-py3-none-any.whl", hash = "sha256:3bb2b7a1f7b685f85b11fed1ef10f8a9148bceb49853e47a294a3dd963c1dd7d", size = 206898, upload-time = "2024-10-07T19:20:48.317Z" }, -] - [[package]] name = "cachetools" version = "5.5.2" @@ -353,11 +319,11 @@ wheels = [ [[package]] name = "cloudpickle" -version = "3.1.1" +version = "3.1.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/52/39/069100b84d7418bc358d81669d5748efb14b9cceacd2f9c75f550424132f/cloudpickle-3.1.1.tar.gz", hash = "sha256:b216fa8ae4019d5482a8ac3c95d8f6346115d8835911fd4aefd1a445e4242c64", size = 22113, upload-time = "2025-01-14T17:02:05.085Z" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7e/e8/64c37fadfc2816a7701fa8a6ed8d87327c7d54eacfbfb6edab14a2f2be75/cloudpickle-3.1.1-py3-none-any.whl", hash = "sha256:c8c5a44295039331ee9dad40ba100a9c7297b6f988e50e87ccdf3765a668350e", size = 20992, upload-time = "2025-01-14T17:02:02.417Z" }, + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, ] [[package]] @@ -439,7 +405,7 @@ wheels = [ [[package]] name = "dask" -version = "2025.5.1" +version = "2026.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -451,9 +417,9 @@ dependencies = [ { name = "pyyaml" }, { name = "toolz" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3d/29/05feb8e2531c46d763547c66b7f5deb39b53d99b3be1b4ddddbd1cec6567/dask-2025.5.1.tar.gz", hash = "sha256:979d9536549de0e463f4cab8a8c66c3a2ef55791cd740d07d9bf58fab1d1076a", size = 10969324, upload-time = "2025-05-20T19:54:30.688Z" } +sdist = { url = "https://files.pythonhosted.org/packages/60/ca/58434f10ebb45d2ddc6edd6e2988abcd38da1ab1897a5f6f402711a594e9/dask-2026.6.0.tar.gz", hash = "sha256:ae3436bd31ebce2be75edf952bd1fc687a1f11ec03fe8b1bec2903d222344a45", size = 11544529, upload-time = "2026-06-11T17:48:43.316Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/38/30/53b0844a7a4c6b041b111b24ca15cc9b8661a86fe1f6aaeb2d0d7f0fb1f2/dask-2025.5.1-py3-none-any.whl", hash = "sha256:3b85fdaa5f6f989dde49da6008415b1ae996985ebdfb1e40de2c997d9010371d", size = 1474226, upload-time = "2025-05-20T19:54:20.309Z" }, + { url = "https://files.pythonhosted.org/packages/02/17/b82b537a30be67ba178fc0b0fbda8fcbaeda188fb039d6ad8a84e89353a2/dask-2026.6.0-py3-none-any.whl", hash = "sha256:1539859071065dca379ca592ff76e911cd7965dc466da0040354ab466179189b", size = 1488995, upload-time = "2026-06-11T17:48:41.008Z" }, ] [[package]] @@ -881,14 +847,14 @@ wheels = [ [[package]] name = "importlib-metadata" -version = "8.7.0" +version = "9.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "zipp", marker = "python_full_version < '3.12'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/76/66/650a33bd90f786193e4de4b3ad86ea60b53c89b669a5c7be931fac31cdb0/importlib_metadata-8.7.0.tar.gz", hash = "sha256:d13b81ad223b890aa16c5471f2ac3056cf76c5f10f82d6f9292f0b415f389000", size = 56641, upload-time = "2025-04-27T15:29:01.736Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/b0/36bd937216ec521246249be3bf9855081de4c5e06a0c9b4219dbeda50373/importlib_metadata-8.7.0-py3-none-any.whl", hash = "sha256:e5dd1551894c77868a30651cef00984d50e1002d06942a7101d34870c5f02afd", size = 27656, upload-time = "2025-04-27T15:29:00.214Z" }, + { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, ] [[package]] @@ -1322,15 +1288,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d8/30/9aec301e9772b098c1f5c0ca0279237c9766d94b97802e9888010c64b0ed/multidict-6.6.3-py3-none-any.whl", hash = "sha256:8db10f29c7541fc5da4defd8cd697e1ca429db743fa716325f236079b96f775a", size = 12313, upload-time = "2025-06-30T15:53:45.437Z" }, ] -[[package]] -name = "mypy-extensions" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, -] - [[package]] name = "netcdf4" version = "1.7.2" @@ -1953,25 +1910,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, ] -[[package]] -name = "pyink" -version = "24.10.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "black" }, - { name = "click" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "pathspec" }, - { name = "platformdirs" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d1/a1/e5e28626fca4266a94c2e1c9264fbf915b9e83e94f52e965190e48fd0cbf/pyink-24.10.1.tar.gz", hash = "sha256:5ec4339aa4953f796e88d90bcac3e3472161e4c36dbde203d80f5f76721ac718", size = 267230, upload-time = "2025-01-10T11:28:09.907Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/12/2f271b3601ae25731879f160d6b3941d80eb6b4f3e24be90289e33fb1dc4/pyink-24.10.1-py3-none-any.whl", hash = "sha256:6349bf6ab75e2ea39a5f0bc3dee7ede7f4af8529291472638026de5fd4af80d2", size = 137118, upload-time = "2025-01-10T11:28:06.138Z" }, -] - [[package]] name = "pymdown-extensions" version = "10.21" @@ -2150,6 +2088,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696, upload-time = "2025-04-16T09:51:17.142Z" }, ] +[[package]] +name = "ruff" +version = "0.15.17" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/a9/3abdf488f1bf3d24c699415e454ed554a6350d5d89ce183be1ee0a3361ac/ruff-0.15.17.tar.gz", hash = "sha256:2ec446937fd16c8c4de2674a209cc5af64d9c6f17d21fbf1151054fa0bcf5219", size = 4743346, upload-time = "2026-06-11T17:54:47.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/4d/e11259f5da07cb6afb2d074c31bf09da9671993f7329d4f15d2fdc458301/ruff-0.15.17-py3-none-linux_armv6l.whl", hash = "sha256:d9feddb927fc68bd295f5eebc587a7e42cfaf9b65f60ca4a2386febff575da8f", size = 10856677, upload-time = "2026-06-11T17:54:49.533Z" }, + { url = "https://files.pythonhosted.org/packages/29/3e/772d679e1a0dc058e58875bd2c0cb713a0530877b4a76fee3c7966df0d49/ruff-0.15.17-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:25805a226d741c47d274a35ad5c10a7dde175fcddfa511d7cf3da0a21eb3eab7", size = 11223443, upload-time = "2026-06-11T17:55:00.573Z" }, + { url = "https://files.pythonhosted.org/packages/68/58/bd41f7688b2fd5623012605130ed70e60aa7f2244baa3d5066bdd61530c8/ruff-0.15.17-py3-none-macosx_11_0_arm64.whl", hash = "sha256:f6ad73b14c2d18a3bf8ad7cb6974294d7f613a7898604826058e6ac64918ef4d", size = 10566458, upload-time = "2026-06-11T17:55:07.52Z" }, + { url = "https://files.pythonhosted.org/packages/d8/5b/733371013fcf1ec339e477ece6ab42bfe10bdd9bba8ee88a9516aa56bfc0/ruff-0.15.17-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ba0c1e4f95bcb3869d0d30cbd5917071ef2e28665abfec970cdab0492c713ed", size = 10914483, upload-time = "2026-06-11T17:55:05.501Z" }, + { url = "https://files.pythonhosted.org/packages/bd/cc/6f24251cc0252f7239391ccb85833f320efad14ebe5b443943f37ced6332/ruff-0.15.17-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:81647960f10bff57d2e51cadd0c3950fe598400c852863a038720ef5b8cca91e", size = 10647497, upload-time = "2026-06-11T17:54:57.733Z" }, + { url = "https://files.pythonhosted.org/packages/68/dd/0d10c17ce1a1624d6fc3156309c3f834fdb5dfaad026ec90c85684f3990e/ruff-0.15.17-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e01a84ddbc8c16c23055ba3924476850f1bbc1917cebbb9376665a63e74260d", size = 11416967, upload-time = "2026-06-11T17:54:51.461Z" }, + { url = "https://files.pythonhosted.org/packages/2f/91/556bfb156f6144f355e831c23db00b2fc4120f86b3ce81cc5f7fd2df51f3/ruff-0.15.17-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:84fe9f653152f8f294f9f7e03bf3a453d8b4a27f7a59c78c8666167f2b17b96c", size = 12335770, upload-time = "2026-06-11T17:54:45.793Z" }, + { url = "https://files.pythonhosted.org/packages/88/82/8b5999aa13355e926f06d9f42a32dcca862f623bf0363785ff89d607dffd/ruff-0.15.17-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c0fe88a7676e7a05b73174d4d4a59cb2ac21ff8263583f87a81a6018475a978", size = 11575441, upload-time = "2026-06-11T17:54:32.661Z" }, + { url = "https://files.pythonhosted.org/packages/11/93/f10377bb04109ca0e8cbc483ff1982c54b6d418210041776f93e8cdc7fa9/ruff-0.15.17-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ecfc3c7878fff94633ab0348524e093f9ce3243080416dd7d14f8ba400174719", size = 11557614, upload-time = "2026-06-11T17:54:34.698Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a6/eeeae7f7d5493df41649ab3db92f086b2d0a30199e4efdf8e3dd7a033f24/ruff-0.15.17-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:b8461180b22420b1bdc289909410930761629fddf2a5aaf60fae1ab26cedc4c4", size = 11544450, upload-time = "2026-06-11T17:54:39.042Z" }, + { url = "https://files.pythonhosted.org/packages/32/88/5991ce565129a24dd4a00db1254b3b5db2e53018cbe4018ea5a89738e727/ruff-0.15.17-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6eccbe50a038b503e7140b441aa9c7fc8c1f36edf23ebef9f4165c2f28f568b7", size = 10892524, upload-time = "2026-06-11T17:55:09.432Z" }, + { url = "https://files.pythonhosted.org/packages/f5/1d/0fdd248313425f55223968af04b0a42125466a8d88d21c1d99c6af0a51e8/ruff-0.15.17-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:382fc0521025f5a8ad447d8bdd523545d0d7646adb718eb1c2dac5065ec27c0f", size = 10659573, upload-time = "2026-06-11T17:54:36.824Z" }, + { url = "https://files.pythonhosted.org/packages/9e/0e/072e8260deb9461062ce9311ced27a8e541229a6ffd483013dd37661e43e/ruff-0.15.17-py3-none-musllinux_1_2_i686.whl", hash = "sha256:456d41fcd1b2777ad63f09a6e7121d43f7b688bbc76a800c10f7f8fb1f912c3f", size = 11127818, upload-time = "2026-06-11T17:55:03.124Z" }, + { url = "https://files.pythonhosted.org/packages/ab/b4/55060a34163121498014696b5f656db5b8c6963768f227dbf0d76b311073/ruff-0.15.17-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b1a04bcc94ae6194e9db05d16ad31f298a7194bfbcb08258bbe589cee1d587b8", size = 11655901, upload-time = "2026-06-11T17:54:53.562Z" }, + { url = "https://files.pythonhosted.org/packages/49/71/9b29d6b87cef468d697f43c6a91e3fae4a80185779d7d5a4ef27d173439f/ruff-0.15.17-py3-none-win32.whl", hash = "sha256:596065960ab1ff593f744220c9fe6580eda00a95003cffa9f4048bb5b1bf0392", size = 10925574, upload-time = "2026-06-11T17:54:55.723Z" }, + { url = "https://files.pythonhosted.org/packages/3d/b2/8fc77f3723228836fa5d12497eb71c808f83782e10d058d2b15cfa14640b/ruff-0.15.17-py3-none-win_amd64.whl", hash = "sha256:6769e5fa1710b179b92e0bfa5a51735b35baea9013dadb06d5f44cbcf9547084", size = 12058788, upload-time = "2026-06-11T17:54:41.042Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" }, +] + [[package]] name = "scipy" version = "1.15.3" @@ -2321,11 +2284,11 @@ wheels = [ [[package]] name = "toolz" -version = "1.0.0" +version = "1.1.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8a/0b/d80dfa675bf592f636d1ea0b835eab4ec8df6e9415d8cfd766df54456123/toolz-1.0.0.tar.gz", hash = "sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02", size = 66790, upload-time = "2024-10-04T16:17:04.001Z" } +sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/98/eb27cc78ad3af8e302c9d8ff4977f5026676e130d28dd7578132a457170c/toolz-1.0.0-py3-none-any.whl", hash = "sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236", size = 56383, upload-time = "2024-10-04T16:17:01.533Z" }, + { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" }, ] [[package]] @@ -2594,7 +2557,6 @@ io = [ name = "xarray-sql" source = { editable = "." } dependencies = [ - { name = "dask" }, { name = "datafusion" }, { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "xarray", version = "2025.7.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -2614,6 +2576,7 @@ docs = [ ] test = [ { name = "cftime" }, + { name = "dask" }, { name = "gcsfs" }, { name = "pytest" }, { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, extra = ["io"], marker = "python_full_version < '3.11'" }, @@ -2624,14 +2587,14 @@ test = [ dev = [ { name = "maturin" }, { name = "py-spy" }, - { name = "pyink" }, + { name = "ruff" }, { name = "xarray-sql", extra = ["docs", "test"] }, ] [package.metadata] requires-dist = [ { name = "cftime", marker = "extra == 'test'" }, - { name = "dask", specifier = ">=2024.8.0" }, + { name = "dask", marker = "extra == 'test'" }, { name = "datafusion", specifier = "==52.0.0" }, { name = "gcsfs", marker = "extra == 'test'" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" }, @@ -2650,7 +2613,7 @@ provides-extras = ["dev", "docs", "test"] dev = [ { name = "maturin", specifier = ">=1.9.1" }, { name = "py-spy", specifier = ">=0.4.0" }, - { name = "pyink", specifier = ">=24.10.1" }, + { name = "ruff", specifier = ">=0.15.10" }, { name = "xarray-sql", extras = ["docs"] }, { name = "xarray-sql", extras = ["test"] }, ] @@ -2825,9 +2788,9 @@ wheels = [ [[package]] name = "zipp" -version = "3.23.0" +version = "4.1.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, + { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, ] diff --git a/xarray_sql/df.py b/xarray_sql/df.py index f55b6ad..c592fb2 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -26,6 +26,32 @@ def _get_chunk_slicer( return slice(None) +def _chunk_tuples( + sizes: Mapping[Hashable, int], spec: dict[str, int] +) -> dict[Hashable, tuple[int, ...]]: + """Per-dimension chunk-size tuples for ``spec`` over ``sizes``. + + The dask-free equivalent of ``Dataset.chunk(spec).chunks``, computed + arithmetically. A dimension absent from ``spec`` (or whose requested size is + non-positive or at least its length) becomes a single full-length chunk; + otherwise it splits into ``size``-length chunks with a short final chunk for + any remainder. + + Why not just call ``.chunk()``? -- that constructs a transient dask graph (one + task per chunk, per variable) purely to read these tuples back. This keeps the + forward path dask-free. + """ + out: dict[Hashable, tuple[int, ...]] = {} + for dim, length in sizes.items(): + size = spec.get(str(dim)) + if size is None or size <= 0 or size >= length: + out[dim] = (length,) + else: + n_full, rem = divmod(length, size) + out[dim] = (size,) * n_full + ((rem,) if rem else ()) + return out + + # Adapted from Xarray `map_blocks` implementation. def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]: """Compute block slices for a chunked Dataset.""" @@ -35,9 +61,9 @@ def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]: # contain every dimension named in the chunks spec. chunks = {dim: size for dim, size in chunks.items() if dim in ds.sizes} if chunks: - for_chunking = ds.copy(data=None, deep=False).chunk(chunks) - chunks = for_chunking.chunks - del for_chunking + # Arithmetic, dask-free. A dask-backed input is not required, and no + # transient dask graph is built just to recover chunk boundaries. + chunks = _chunk_tuples(ds.sizes, chunks) else: chunks = ds.chunks diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py new file mode 100644 index 0000000..209a14c --- /dev/null +++ b/xarray_sql/ds.py @@ -0,0 +1,818 @@ +"""Reconstruct xarray Datasets from SQL query results. + +The inverse of the forward Dataset-to-table pivot done by +:func:`xarray_sql.df.pivot`. Internally defines an :class:`XarrayDataFrame` +wrapper around the DataFusion ``DataFrame`` returned by +:meth:`XarrayContext.sql`, with a :meth:`XarrayDataFrame.to_dataset` +method that round-trips a query result back to ``xr.Dataset``. + +Reconstruction is controlled by the ``chunks`` argument to +:meth:`XarrayDataFrame.to_dataset` -- the xarray idiom for tuning how a +result is partitioned: + +* **Eager** (``chunks=None``, or the default ``"inherit"`` when the + result keeps no multi-chunk source dimension): the plan executes + exactly once via ``execute_stream`` and the result is scattered into a + dense in-memory Dataset. This is the right default for reductions + (aggregations), whose results are small, and it never re-executes. +* **Lazy / chunked** (``chunks`` is a mapping, ``"auto"``, or + ``"inherit"`` over a multi-chunk source dimension): data variables are + backed by :class:`SQLBackendArray` wrapped in + ``xarray.core.indexing.LazilyIndexedArray`` and chunked via xarray's + configured chunk manager (dask, cubed, ...). Each chunk maps onto the + source partitions and reads its coordinate range on access by + translating the indexer into a DataFusion ``filter`` expression, so only + the requested partitions are materialized as Arrow ``RecordBatch`` es + and scattered into numpy. + +``.compute()`` materializes the whole Dataset in memory. +""" + +from __future__ import annotations + +import warnings +from collections.abc import Mapping +from typing import Any, cast + +import numpy as np +import pandas as pd +import pyarrow as pa +import xarray as xr +from datafusion import DataFrame, col, literal + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + + +def _ds_var_dims(ds: xr.Dataset) -> list[str]: + """Return a Dataset's data-variable dim order. + + The forward path validates that all data variables share the same dims + tuple, so the first var's dim order is canonical. Falls back to + ``ds.dims`` keys for empty Datasets. Always use this rather than + ``list(ds.dims)`` when round-tripping, since the latter is in + canonical name order and may not match the variable's axis order. + """ + if ds.data_vars: + return list(next(iter(ds.data_vars.values())).dims) + return list(ds.dims) + + +def _apply_template(ds: xr.Dataset, template: xr.Dataset) -> xr.Dataset: + """Recover metadata that the forward SQL pivot strips. + + Adds back, where unambiguous: + + * Data-variable ``attrs`` and ``encoding`` for vars present in + ``template`` (aggregation aliases like ``air_avg`` get nothing). + Dtype-bound encoding keys (``dtype``, ``_FillValue``, + ``missing_value``) are intentionally dropped: SQL may have + changed the column's dtype (e.g. ``int16`` -> ``float64`` after + ``AVG`` or a null-introducing filter), and reattaching the + source's packing would make a later ``ds.to_netcdf()`` write + corrupt values. + * Dim-coordinate dtype, where SQL upcasted (datetime is the + canonical case). + * Non-dim coordinates whose dims are all present in ``ds`` (scalar + coords attach as-is; vector coords use ``.sel``). + * Dataset-level ``attrs``. + + Skipped coords are warned about once per call. + """ + out = ds.copy() + + # 1. Data-var attrs / encoding for vars present in the template. + # Aggregation aliases absent from template intentionally inherit nothing. + for name in list(out.data_vars): + if name in template.data_vars: + out[name].attrs = dict(template[name].attrs) + # Drop dtype-bound encoding keys; SQL may have changed dtype. + enc = { + k: v + for k, v in template[name].encoding.items() + if k not in {"dtype", "_FillValue", "missing_value"} + } + out[name].encoding = enc + + # 2. Restore dim-coordinate dtype when SQL changed it (e.g. datetime + # upcast through pyarrow / pandas) and copy the source's dim-coord + # attrs (``standard_name``, ``long_name``, ``units``, etc.). + for d in list(out.dims): + if d in template.coords: + tdt = template.coords[d].dtype + if out.coords[d].dtype != tdt: + try: + out = out.assign_coords({d: out.coords[d].astype(tdt)}) + except (ValueError, TypeError): + pass # incompatible cast; leave as-is + out[d].attrs = dict(template.coords[d].attrs) + + # 3. Non-dim coordinates whose dims are all present in the result. + out_dims = set(out.dims) + skipped: list[str] = [] + for cname, coord in template.coords.items(): + if cname in template.dims: + continue # dim coord; already in out + if not set(coord.dims) <= out_dims: + continue # spans dims the result lacks + try: + if not coord.dims: + # Scalar coord (e.g. weather_dataset.reference_time). + out = out.assign_coords({cname: coord}) + else: + sel = {d: out.coords[d] for d in coord.dims} + out = out.assign_coords({cname: coord.sel(sel)}) + except (KeyError, ValueError, TypeError): + skipped.append(cname) + + # 4. Dataset-level attrs. + out.attrs = dict(template.attrs) + + if skipped: + warnings.warn( + f"Could not re-attach non-dim coordinates from template: {skipped}", + stacklevel=3, + ) + return out + + +def _scatter_batches_to_ndarray( + batches: list[pa.RecordBatch], + dimension_columns: list[str], + requested: dict[str, np.ndarray], + var_name: str, + out_shape: tuple[int, ...], + dtype: np.dtype, + drop_axes: list[int], +) -> np.ndarray: + """Convert filtered Arrow ``RecordBatch`` rows into a dense N-D numpy array. + + SQL query results arrive as flat rows; xarray expects N-D arrays. + This bridges the two: each row carries the dim-coord values that + identify its cell in the output cube plus the value to write there. + We look up the row's N-D position by binary-searching its coord + values within the caller's requested coord arrays + (``np.searchsorted``), then scatter-write the value at that index. + + The result must be a clean grid: exactly one row per cell. We guard both + failure modes rather than corrupt silently: + + * **Missing cells** (a filter or non-rectangular result leaves holes). + Float outputs pre-fill with ``NaN``; integer outputs have no missing + sentinel, so a hole would be silent garbage -- we raise instead. + * **Duplicate cells** (more rows than the grid holds, e.g. a query that + drops a dimension) -- we raise rather than last-write-wins. + + Raises: + ValueError: the result does not fill an integer grid completely, or + contains more rows than cells (duplicate dimension tuples), or + carries a coordinate value absent from ``requested``. + """ + n_cells = int(np.prod(out_shape)) if out_shape else 1 + n_rows = sum(b.num_rows for b in batches) + if n_rows > n_cells: + raise ValueError( + f"the result has {n_rows} rows for a {n_cells}-cell grid while " + f"reconstructing {var_name!r}: the dimension columns do not " + f"uniquely identify a cell (duplicate dimension tuples). Aggregate " + f"or add dimension columns so each row maps to one cell." + ) + if n_rows < n_cells and not np.issubdtype(dtype, np.floating): + raise ValueError( + f"data variable {var_name!r} has integer dtype {dtype} but the " + f"result fills only {n_rows} of {n_cells} cells; integers cannot " + f"represent the missing cells. Cast it to a floating type in SQL " + f"(e.g. CAST({var_name} AS DOUBLE)), or return a complete grid." + ) + + # NaN fill for float outputs (sparse cells stay NaN); integer outputs are + # guaranteed complete by the check above, so every cell is written below. + out = ( + np.full(out_shape, np.nan, dtype=dtype) + if np.issubdtype(dtype, np.floating) + else np.empty(out_shape, dtype=dtype) + ) + + # ``requested[d]`` may be in any order (callers can iselect arbitrary + # positions, and template coords like air_temperature.lat are descending). + # ``np.searchsorted`` requires ascending input, so we sort each requested + # array once, search there, and remap back to the original positions. + sorted_idx = {d: np.argsort(requested[d]) for d in dimension_columns} + sorted_req = {d: requested[d][sorted_idx[d]] for d in dimension_columns} + + for batch in batches: + if batch.num_rows == 0: + continue + schema_names = batch.schema.names + # Build per-dim position arrays for this batch (positions within + # the caller's requested coord order). + positions = [] + for d in dimension_columns: + col_arr = batch.column(schema_names.index(d)) + vals = col_arr.to_numpy(zero_copy_only=False) + hi = len(sorted_req[d]) + pos_in_sorted = np.searchsorted(sorted_req[d], vals) + # ``searchsorted`` returns an insertion point, not a match. Verify + # the row's coordinate actually exists in ``requested`` so a stray + # value can't scatter into the wrong cell or index out of bounds. + if hi == 0 or not np.array_equal( + sorted_req[d][np.clip(pos_in_sorted, 0, hi - 1)], vals + ): + raise ValueError( + f"result contains {d!r} values absent from the discovered " + f"coordinates while reading {var_name!r}; the coordinate " + f"set and the data are inconsistent." + ) + positions.append(sorted_idx[d][pos_in_sorted]) + value_arr = batch.column(schema_names.index(var_name)).to_numpy( + zero_copy_only=False + ) + out[tuple(positions)] = value_arr.astype(dtype, copy=False) + + if drop_axes: + out = np.squeeze(out, axis=tuple(drop_axes)) + return cast(np.ndarray, out) + + +class SQLBackendArray(xr.backends.BackendArray): + """Read-only lazy N-D array view over a DataFusion DataFrame. + + Bridges xarray's lazy-indexing interface + (:class:`xarray.backends.BackendArray`) to a DataFusion query result, + so an xarray ``Dataset`` can present a SQL query as if it were a + materialized N-D array without actually loading any data until the + caller asks for it. This is the workhorse that lets + :meth:`XarrayDataFrame.to_dataset` return a Dataset cheaply. + + On each ``__getitem__`` call, the requested xarray indexer is + translated into a DataFusion filter expression (``df.filter(expr)``) + and a column projection (``df.select(*cols)``). The filtered + DataFrame is consumed via ``execute_stream`` as a sequence of Arrow + ``RecordBatch`` es and scattered into a preallocated numpy buffer, + so only the requested data is materialized. + + Constraints and caveats: + + - Read-only: there is no write path; the backend exists to surface + query results, not to round-trip writes into a SQL store. + - The underlying DataFusion ``DataFrame`` holds a reference to its + originating ``SessionContext``, which is not picklable. The class + therefore overrides ``__copy__`` and ``__deepcopy__`` to return + ``self`` -- this is safe because the backend is read-only. + - ``IndexingSupport.OUTER``: ``BasicIndexer`` and ``OuterIndexer`` + are translated to filter predicates directly; ``VectorizedIndexer`` + paths through xarray's adapter to outer-then-gather and so still + works, just less efficiently. + + Raises: + ValueError, datafusion exceptions: propagated from the + underlying ``df.filter().select().execute_stream()`` chain + if a predicate refers to a missing column, the dtype of a + literal is incompatible, or the execution itself fails. + AssertionError: from ``np.searchsorted`` mis-alignment, which + indicates the result contains coordinate values not present + in the wrapper's pre-computed coord arrays -- usually a + symptom of a filtered query whose coord discovery missed a + value. + + Constructed by :func:`_build_lazy_scan`; users should not instantiate + this class directly. + """ + + def __init__( + self, + inner_df: Any, + var_name: str, + dimension_columns: list[str], + coord_arrays: dict[str, np.ndarray], + shape: tuple[int, ...], + dtype: np.dtype, + ) -> None: + self._inner_df = inner_df + self._var_name = var_name + self._dimension_columns = list(dimension_columns) + self._coord_arrays = coord_arrays + self.shape = tuple(shape) + self.dtype = np.dtype(dtype) + + def __getitem__(self, key: Any) -> np.ndarray: + return cast( + np.ndarray, + xr.core.indexing.explicit_indexing_adapter( + key, + self.shape, + xr.core.indexing.IndexingSupport.OUTER, + self._raw_getitem, + ), + ) + + def __copy__(self) -> "SQLBackendArray": + # The backend is read-only; the underlying DataFusion DataFrame + # holds a non-picklable SessionContext reference, so sharing the + # same backend across a copy is both safe and necessary. + return self + + def __deepcopy__(self, memo: dict) -> "SQLBackendArray": + return self + + # ------------------------------------------------------------------ + + def _raw_getitem(self, key: tuple) -> np.ndarray: + """Materialize the indexed region described by *key* via DataFusion + Arrow. + + ``key`` is a tuple of ``int``/``slice``/1-D integer-array, one per + dim, in :attr:`_dimension_columns` order. + """ + requested: dict[str, np.ndarray] = {} + # Dims whose indexer covers the full extent (slice(None) or + # equivalent). For these we omit the filter predicate entirely + # so DataFusion doesn't have to evaluate a tautology. + full_dims: set[str] = set() + drop_axes: list[int] = [] + for axis, (dim, k) in enumerate( + zip(self._dimension_columns, key, strict=True) + ): + coord = self._coord_arrays[dim] + if isinstance(k, slice): + start = 0 if k.start is None else k.start + stop = len(coord) if k.stop is None else k.stop + step = 1 if k.step is None else k.step + requested[dim] = np.asarray(coord[start:stop:step]) + if start == 0 and stop >= len(coord) and step == 1: + full_dims.add(dim) + elif isinstance(k, (int, np.integer)): + requested[dim] = np.asarray([coord[int(k)]]) + drop_axes.append(axis) + else: + arr = np.asarray(k) + requested[dim] = np.asarray(coord[arr]) + if ( + len(arr) == len(coord) + and (arr == np.arange(len(coord))).all() + ): + full_dims.add(dim) + + out_shape = tuple(len(requested[d]) for d in self._dimension_columns) + if any(n == 0 for n in out_shape): + empty = np.empty(out_shape, dtype=self.dtype) + squeezed = ( + np.squeeze(empty, axis=tuple(drop_axes)) if drop_axes else empty + ) + return cast(np.ndarray, squeezed) + + # Build a single DataFusion filter expression as the AND of per-dim + # predicates. For a single requested value: equality. For multiple: + # OR-chain of equalities (DataFusion 52.0.0 does not expose a clean + # ``Expr.in_list`` from Python; OR-chained equalities constant-fold + # equivalently and stay typed). + predicates = [] + for dim in self._dimension_columns: + if dim in full_dims: + continue + vals = requested[dim] + if len(vals) == 1: + predicates.append(col(f'"{dim}"') == literal(vals[0])) + else: + eq = col(f'"{dim}"') == literal(vals[0]) + for v in vals[1:]: + eq = eq | (col(f'"{dim}"') == literal(v)) + predicates.append(eq) + + filtered = self._inner_df + if predicates: + combined = predicates[0] + for p in predicates[1:]: + combined = combined & p + filtered = filtered.filter(combined) + projected = filtered.select( + *(col(f'"{c}"') for c in self._dimension_columns + [self._var_name]) + ) + + # Consume the projected DataFrame as Arrow RecordBatches. The + # DataFusion wrapper exposes ``.to_pyarrow()`` to convert each + # batch into a true ``pyarrow.RecordBatch``. + batches = [b.to_pyarrow() for b in projected.execute_stream()] + return _scatter_batches_to_ndarray( + batches=batches, + dimension_columns=self._dimension_columns, + requested=requested, + var_name=self._var_name, + out_shape=out_shape, + dtype=self.dtype, + drop_axes=drop_axes, + ) + + +def _coords_from_batches( + batches: list[pa.RecordBatch], dimension_columns: list[str] +) -> dict[str, np.ndarray]: + """Distinct coordinate values per dim, in first-appearance order. + + This preserves stability in result orders between chunked and non-chunked + versions of a sorted query (i.e. queries that use ORDER BY). + """ + coord_arrays: dict[str, np.ndarray] = {} + for d in dimension_columns: + if not batches: + coord_arrays[d] = np.asarray([]) + continue + vals = np.concatenate( + [ + b.column(b.schema.names.index(d)).to_numpy(zero_copy_only=False) + for b in batches + ] + ) + coord_arrays[d] = np.asarray(pd.unique(vals)) + return coord_arrays + + +def _materialize( + inner_df: Any, + dimension_columns: list[str], + field_names: list[str], + field_types: dict[str, Any], +) -> xr.Dataset: + """Execute the query once and build a dense in-memory Dataset. + + Runs the plan exactly once via ``execute_stream()`` -- streaming the result + as Arrow ``RecordBatch`` es (``datafusion.RecordBatch.to_pyarrow()``) -- then + derives both the coordinates and every data variable from that single pass. + This is the eager path, used when no output chunking is requested. It never + re-executes, so an aggregation over a remote Zarr scan costs exactly one + scan, regardless of how many dimensions or variables the result has. + """ + batches = [b.to_pyarrow() for b in inner_df.execute_stream()] + coord_arrays = _coords_from_batches(batches, dimension_columns) + shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + + data_vars: dict[str, xr.Variable] = {} + for name in field_names: + if name in dimension_columns: + continue + np_dtype = np.dtype(field_types[name].to_pandas_dtype()) + dense = _scatter_batches_to_ndarray( + batches=batches, + dimension_columns=dimension_columns, + requested=coord_arrays, + var_name=name, + out_shape=shape, + dtype=np_dtype, + drop_axes=[], + ) + data_vars[name] = xr.Variable(dimension_columns, dense) + + coords_arg = {d: coord_arrays[d] for d in dimension_columns} + return xr.Dataset(data_vars=data_vars, coords=coords_arg) + + +def _build_lazy_scan( + inner_df: Any, + dimension_columns: list[str], + field_names: list[str], + field_types: dict[str, Any], +) -> xr.Dataset: + """Build a lazy Dataset whose data vars are :class:`SQLBackendArray`. + + Used when output chunking is requested: each data variable stays lazy and, + once wrapped by ``Dataset.chunk``, every chunk reads its coordinate range via + a pushdown filter on first access. Coordinates are discovered by projecting + the result to its dimension columns and streaming them once + (:func:`_coords_from_batches`), so discovery reads coordinate values only + (the provider skips data variables) and uses the same first-appearance + ordering as the eager path -- keeping ``to_dataset`` order-stable. + """ + dim_proj = inner_df.select(*(col(f'"{d}"') for d in dimension_columns)) + coord_batches = [b.to_pyarrow() for b in dim_proj.execute_stream()] + coord_arrays = _coords_from_batches(coord_batches, dimension_columns) + shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + + data_vars: dict[str, xr.Variable] = {} + for name in field_names: + if name in dimension_columns: + continue + np_dtype = field_types[name].to_pandas_dtype() + backend = SQLBackendArray( + inner_df=inner_df, + var_name=name, + dimension_columns=dimension_columns, + coord_arrays=coord_arrays, + shape=shape, + dtype=np_dtype, + ) + lazy = xr.core.indexing.LazilyIndexedArray(backend) + data_vars[name] = xr.Variable(dimension_columns, lazy) + + coords_arg = {d: coord_arrays[d] for d in dimension_columns} + return xr.Dataset(data_vars=data_vars, coords=coords_arg) + + +def _auto_chunk_target_bytes() -> int: + """Byte target for ``chunks="auto"`` (the chunk manager's, else 128 MiB).""" + try: + import dask + from dask.utils import parse_bytes + + return int(parse_bytes(dask.config.get("array.chunk-size"))) + except Exception: + return 128 * 1024 * 1024 + + +def _auto_chunks( + template: xr.Dataset | None, + dimension_columns: list[str], + field_types: dict[str, Any], +) -> dict[str, int] | None: + """Resolve ``chunks="auto"`` to a source-partition-aligned chunk spec. + + Sizes chunks to roughly the chunk manager's byte target (dask's + ``array.chunk-size``, default 128 MiB) but snaps boundaries to whole source + partitions, so every chunk is a union of source partitions -- no chunk splits + a partition (which would make adjacent chunks re-read it). This is what makes + ``"auto"`` useful for finely partitioned sources (e.g. ERA5 + ``chunks={"time": 1}``): it coarsens many tiny partitions into memory-sized, + aligned chunks. Returns ``None`` when there is no resolvable source grid to + align to, so the caller falls back to the chunk manager's own ``"auto"``. + """ + if template is None: + return None + part = template.chunksizes # dim -> tuple of source chunk lengths + chunked_dims = [ + d for d in dimension_columns if d in part and len(part[d]) > 1 + ] + if not chunked_dims: + return None + + itemsizes = [ + np.dtype(t.to_pandas_dtype()).itemsize + for name, t in field_types.items() + if name not in dimension_columns + ] + itemsize = max(itemsizes) if itemsizes else 8 + + # Bytes in one source-partition block: the nominal source chunk length per + # dimension (``part[d][0]``) multiplied across all dims, times itemsize. + block_bytes = itemsize + for d in dimension_columns: + if d in part: + block_bytes *= int(part[d][0]) + # Number of source partitions to merge per chunk to approach the target. + merge = max(1, _auto_chunk_target_bytes() // max(block_bytes, 1)) + + # Absorb the coarsening into the most finely partitioned dimension; the rest + # keep their source chunk length. xarray caps an oversize chunk at the dim + # length, so an over-large merge simply yields a single chunk on that dim. + primary = max(chunked_dims, key=lambda d: len(part[d])) + return { + d: int(part[d][0]) * (merge if d == primary else 1) + for d in chunked_dims + } + + +def _to_dataset( + inner_df: "XarrayDataFrame", + dims: list[str], + template: xr.Dataset | None, + chunks: Mapping[str, int] | str | None, +) -> xr.Dataset: + """Reconstruct an ``xr.Dataset`` from a SQL result. + + ``chunks`` (resolved by :func:`_resolve_chunks`) selects the execution + strategy: + + * ``None`` -> eager: execute once and materialize a dense Dataset + (:func:`_materialize`). Correct for any query and the right default for + reductions, whose results are small. + * a mapping (or ``"auto"``) -> lazy/chunked: build :class:`SQLBackendArray` + data variables (:func:`_build_lazy_scan`) and wrap them with + ``Dataset.chunk`` so each chunk reads its coordinate range via filter + pushdown. The chunk grid maps onto the source partitions. Chunking goes + through xarray's configured chunk manager (dask, cubed, ...), so no + chunked-array backend is imported directly here. + """ + chunks = _resolve_chunks(chunks, template, dims) + + schema = inner_df.schema() + field_names = [f.name for f in schema] + field_types = {f.name: f.type for f in schema} + + if chunks is None: + ds = _materialize( + inner_df, dims, field_names, field_types + ) + else: + ds = _build_lazy_scan( + inner_df, dims, field_names, field_types + ) + + if template is not None: + ds = _apply_template(ds, template) + + if chunks is not None: + if chunks == "auto": + # Snap the byte-budgeted "auto" sizing to source partition + # boundaries; fall back to the chunk manager's own "auto" when there + # is no source grid to align to. + chunks = ( + _auto_chunks(template, dims, field_types) or "auto" + ) + # Wrap the lazy data variables in the configured chunk manager (dask by + # default). Each chunk reads its coordinate range via pushdown on access. + ds = ds.chunk(chunks) + return ds + + +def _resolve_chunks( + chunks: Mapping[str, int] | str | None, + template: xr.Dataset | None, + dimension_columns: list[str], +) -> Mapping[str, int] | str | None: + """Resolve the ``chunks`` argument to a concrete spec or ``None``. + + ``None`` selects the eager path; anything else selects the lazy/chunked + path. ``"inherit"`` reuses the source Dataset's chunk sizes -- but only + for dimensions actually split into more than one chunk in the input + (a single full chunk is not "chunked"), so reductions that drop the + chunked dimension resolve to ``None`` (eager) automatically. Mappings + pass through unchanged; ``"auto"`` passes through here and is snapped to + source partition boundaries later (see :func:`_auto_chunks`). + """ + if chunks is None: + return None + if chunks == "inherit": + if template is None: + return None + sizes = template.chunksizes # dim -> tuple of chunk lengths + inherited = { + d: sizes[d][0] + for d in dimension_columns + if d in sizes and len(sizes[d]) > 1 + } + return inherited or None + return chunks + + +# --------------------------------------------------------------------------- +# Public wrapper +# --------------------------------------------------------------------------- + + +class XarrayDataFrame(DataFrame): + """A DataFusion ``DataFrame`` with xarray-aware helpers. + + Carries a private snapshot of the context's registered Datasets so + :meth:`to_dataset` can default ``dims`` and recover metadata + dropped by the forward pivot. + + Users should not construct this class directly; let + :meth:`XarrayContext.sql` produce it. + """ + + def __init__( + self, + inner: Any, + templates: dict[str, xr.Dataset] | None = None, + ) -> None: + """Construct a wrapper. + + Args: + inner: The underlying ``datafusion.DataFrame`` returned by + :meth:`XarrayContext.sql`. + templates: Snapshot of the registered Datasets on the producing + context, keyed by the SQL identifier each was registered + under. Used by :meth:`to_dataset` to recover metadata that + the forward pivot strips. ``None`` means no metadata + recovery is possible from registrations alone; callers may + still pass ``template=`` to :meth:`to_dataset` explicitly. + """ + super().__init__(inner.df) + self._templates = dict(templates or {}) + + def to_dataset( + self, + dims: list[str] | None = None, + template: xr.Dataset | str | None = None, + chunks: Mapping[str, int] | str | None = "inherit", + ) -> xr.Dataset: + """Convert the result to an ``xr.Dataset``. + + Args: + dims: Result columns to use as Dataset dimensions. When + ``None``, defaults to the dims of the registered Dataset + referenced by the SQL ``FROM`` clause (if exactly one + matches), or any single registered Dataset whose dims are + all present in the result columns. + template: Source to recover metadata (attrs, encoding, non-dim + coordinates, dim-coord dtype) from. Either an ``xr.Dataset`` + used directly, or the name of a registered table (e.g. + ``"era5.surface"``) whose Dataset is looked up. When ``None`` + and exactly one Dataset is registered, that one is used. + chunks: Output chunking, controlling laziness (an xarray idiom). + + * ``"inherit"`` (default): reuse the source Dataset's chunk + sizes, but only for dimensions that were genuinely split into + multiple chunks in the input -- so the output chunk grid maps + onto the source partitions. A reduction that drops the chunked + dimension (e.g. a global aggregation) inherits nothing and so + is materialized eagerly. Falls back to eager when no source + Dataset is resolvable. + * ``None``: eager. Execute the query once and return a dense + in-memory Dataset. Best for reductions (small results). + * a mapping (e.g. ``{"time": 100}``): chunk explicitly. Each + chunk reads its coordinate range lazily via filter pushdown on + access, through xarray's configured chunk manager (dask, + cubed, ...). + * ``"auto"``: size chunks to the chunk manager's byte target but + snap boundaries to whole source partitions, so each chunk is a + union of source partitions. Useful for finely partitioned + sources (e.g. ERA5 ``chunks={"time": 1}``), coarsening many + tiny partitions into memory-sized, aligned chunks. + + Returns: + An ``xr.Dataset`` with ``dims`` as dimensions and the + remaining result columns as data variables. + + Raises: + ValueError: ``dims`` cannot be inferred; ``template`` names an + unknown registered table; the result has duplicate dimension + tuples (more rows than the coordinate grid holds); or an + integer data variable does not fill the grid completely + (integers have no missing-value sentinel -- cast it to a + floating type in SQL). These are raised lazily on the chunked + path, when a chunk is first read. + """ + template = self._resolve_template(template) + dims = dims or self._infer_dims(preferred_template=template) + return _to_dataset( + inner_df=self, + dims=dims, + template=template, + chunks=chunks, + ) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _resolve_template( + self, candidate: xr.Dataset | str | None + ) -> xr.Dataset | None: + """Pick a template Dataset for metadata recovery (e.g. by name).""" + if isinstance(candidate, xr.Dataset): + return candidate + + # ``template`` is a registered-table name or None; look it up. + templates = self._templates + if candidate is not None: + if candidate not in templates: + raise ValueError( + f"template={candidate!r} is not a registered table on this " + f"context. Registered: {list(templates)}" + ) + return templates[candidate] + + if len(templates) == 1: + return next(iter(templates.values())) + + return None + + def _infer_dims( + self, preferred_template: xr.Dataset | None = None + ) -> list[str]: + """Pick a default ``dims`` from the registry, or raise. + + Uses the data variable's dim order (via :func:`_ds_var_dims`) so + the round-trip preserves the original axis order. + """ + result_cols = {field.name for field in self.schema()} + if ( + preferred_template is not None + and set(preferred_template.dims) <= result_cols + ): + return _ds_var_dims(preferred_template) + if not self._templates: + raise ValueError( + "dims cannot be inferred (no registered " + "Dataset on this result); pass dims=[...] " + "explicitly." + ) + candidates = [ + _ds_var_dims(t) + for t in self._templates.values() + if set(t.dims) <= result_cols + ] + if len(candidates) == 1: + return candidates[0] + if not candidates: + raise ValueError( + "dims cannot be inferred: no registered " + "Dataset has all of its dims present in the result " + "columns. Pass dims=[...] explicitly." + ) + raise ValueError( + "dims cannot be inferred unambiguously: multiple " + "registered Datasets are compatible with the result. Pass " + "dims=[...] explicitly." + ) + diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 403eaee..b7b2909 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -5,12 +5,23 @@ from . import cftime as cft from .df import Chunks +from .ds import XarrayDataFrame from .reader import read_xarray_table class XarrayContext(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Track registered xarray Datasets so XarrayDataFrame can recover + # defaults (dimension_columns) and metadata (var/dataset attrs, + # non-dim coords, dim-coord dtype) that the forward pivot drops. + # Keys are the fully-qualified table names users will reference + # in SQL (e.g. ``"air"`` for a uniform-dim Dataset, or + # ``"era5.surface"`` for one entry from a multi-dim-group split). + self._registered_datasets: dict[str, xr.Dataset] = {} + def from_dataset( self, name: str, @@ -89,8 +100,16 @@ def from_dataset( # 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") + # The SQL table is ``.``; key the template registry + # by that same fully-qualified name so ``template="era5.surface"`` + # matches the table the user queries (and so sub-tables of different + # datasets that share a sub-name don't collide). self._from_dataset( - sub_name, input_table[var_names], chunks, schema=schema + sub_name, + input_table[var_names], + chunks, + schema=schema, + registry_name=f"{name}.{sub_name}", ) return self @@ -101,12 +120,16 @@ def _from_dataset( input_table: xr.Dataset, chunks: Chunks = None, schema: Schema | None = None, + registry_name: str | None = None, ): """Register a Dataset as a single SQL table. Registers a top-level table by default, or a table inside ``schema`` - (a SQL namespace) when one is given. + (a SQL namespace) when one is given. ``registry_name`` is the key the + template registry uses (the fully-qualified ``.`` for + namespaced tables); it defaults to ``table_name`` for top-level tables. """ + self._registered_datasets[registry_name or table_name] = input_table register = ( self.register_table if schema is None else schema.register_table ) @@ -123,6 +146,27 @@ def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None: self.register_udf(cft.make_cftime_udf(units, cal)) break # One UDF per context is enough. + def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: + """Run a SQL query, returning an :class:`XarrayDataFrame` wrapper. + + Identical to ``datafusion.SessionContext.sql`` except the returned + object wraps the DataFusion DataFrame. The wrapper exposes + ``.to_pandas()`` (unchanged), forwards every other DataFusion + method via ``__getattr__``, and adds + ``.to_dataset(dimension_columns=[...])`` for round-tripping the + result back to an ``xr.Dataset``. + + Args: + query: A SQL query string. + *args: Forwarded to ``SessionContext.sql``. + **kwargs: Forwarded to ``SessionContext.sql``. + + Returns: + An :class:`XarrayDataFrame` wrapping the DataFusion DataFrame. + """ + inner = super().sql(query, *args, **kwargs) + return XarrayDataFrame(inner, templates=self._registered_datasets) + def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: """Group variables in the dataset based on shared dims.