diff --git a/.github/workflows/ci-wasm.yml b/.github/workflows/ci-wasm.yml new file mode 100644 index 0000000..d004b1b --- /dev/null +++ b/.github/workflows/ci-wasm.yml @@ -0,0 +1,102 @@ +name: ci-wasm + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + # publish.yml calls this to add the Pyodide wheel to each release. + workflow_call: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + # The Pyodide release to build for. Its Python version, Emscripten version, + # and Rust toolchain all come from this release's cross-build environment. + PYODIDE_VERSION: "314.0.7" + +jobs: + build: + name: Pyodide wheel + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + # pyodide-build requires the host Python to match Pyodide's Python. + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + + - name: Install pyodide-build and the cross-build environment + run: | + pip install pyodide-build + pyodide xbuildenv install "$PYODIDE_VERSION" + echo "RUST_TOOLCHAIN=$(pyodide config get rust_toolchain)" >> "$GITHUB_ENV" + + - name: Install Rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: ${{ env.RUST_TOOLCHAIN }} + targets: wasm32-unknown-emscripten + + - name: Install Emscripten + run: pyodide xbuildenv install-emscripten + + - name: Build + env: + # The Flight SQL server needs sockets and threads; the wasm + # profile optimizes for download size (see Cargo.toml). + MATURIN_PEP517_ARGS: --no-default-features --profile wasm + run: pyodide build -o dist + + - uses: actions/setup-node@v4 + with: + node-version: "24" + + - name: Query a Dataset with DuckDB in Pyodide + run: | + npm install --no-save "pyodide@$PYODIDE_VERSION" + cat > smoke.mjs <<'EOF' + import { loadPyodide } from "pyodide"; + import { readdirSync, readFileSync } from "node:fs"; + const wheel = readdirSync("dist").find((f) => f.endsWith(".whl")); + const py = await loadPyodide(); + await py.loadPackage("micropip"); + py.FS.writeFile(`/tmp/${wheel}`, readFileSync(`dist/${wheel}`)); + await py.runPythonAsync(` + import micropip + await micropip.install(["emfs:/tmp/${wheel}", "duckdb"]) + + import importlib.util + + import duckdb + import numpy as np + import pandas as pd + import xarray as xr + import xarray_sql as xql + + ds = xr.Dataset( + {"air": (("time", "lat"), np.arange(12.0).reshape(4, 3))}, + coords={ + "time": pd.date_range("2020-01-01", periods=4), + "lat": [10.0, 20.0, 30.0], + }, + ) + con = duckdb.connect() + xql.register(con, "air", ds, chunks={"time": 2}) + rel = con.sql("SELECT time, lat, air FROM air ORDER BY time, lat") + xr.testing.assert_identical(xql.to_dataset(rel, template=ds), ds) + assert not hasattr(xql, "Context") + assert importlib.util.find_spec("dask") is None + `); + console.log(`${wheel} round-trips through DuckDB in Pyodide ${py.version}`); + EOF + node smoke.mjs + + - uses: actions/upload-artifact@v6 + with: + name: pyodide-dist + path: dist/*.whl diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index ab567eb..20a8158 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -124,6 +124,12 @@ jobs: name: dist-manylinux-aarch64 path: dist/* + # Runs the Pyodide smoke test too, so a release never ships an untested + # wasm wheel. PyPI accepts these wheels under PEP 783, so Pyodide users + # can `await micropip.install("xarray-sql")`. + build-artifacts-pyodide: + uses: ./.github/workflows/ci-wasm.yml + build-sdist: name: Source distribution runs-on: ubuntu-latest @@ -193,7 +199,7 @@ jobs: uv pip install --extra-index-url https://test.pypi.org/simple --upgrade de python -c "import xarray_sql; print(xarray_sql.__version__)" upload-to-pypi: - needs: verify-built-dist + needs: [verify-built-dist, build-artifacts-pyodide] runs-on: ubuntu-latest # Only publish for real releases. `workflow_dispatch` runs build and # verify the wheels (uploaded as artifacts) without pushing to PyPI, so @@ -211,3 +217,13 @@ jobs: - name: Publish package to PyPI run: uv publish --token ${{ secrets.PYPI_TOKEN }} + + # Uploaded separately and after the native wheels, so a rejected wasm + # wheel cannot leave a partial release. + - uses: actions/download-artifact@v6 + with: + name: pyodide-dist + path: pyodide-dist + + - name: Publish Pyodide wheel to PyPI + run: uv publish --token ${{ secrets.PYPI_TOKEN }} pyodide-dist/* diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b7f1722..e522082 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -59,6 +59,31 @@ already on `PATH`) and start at step 2. You can also run the hooks manually with: `uvx pre-commit run --all-files` 7. Build and serve docs locally: `uvx zensical serve` +### Building for WebAssembly (Pyodide) + +The `ci-wasm` workflow builds a Pyodide wheel. The target Pyodide release +fixes the Python version, the exact Emscripten version, and the Rust +toolchain. The Flight SQL server needs sockets and threads, so wasm builds +turn off the `flight` cargo feature. + +To build the wheel locally, set `PYODIDE_VERSION` to the value in +`.github/workflows/ci-wasm.yml`. pyodide-build must run on the same Python +version as that Pyodide release (for example, Python 3.14 for Pyodide 314), +which `uv tool install --python` sets up. The `pyodide` command comes from +`pyodide-cli`, so install it with `pyodide-build`: + +```shell +uv tool install --python 3.14 pyodide-cli --with pyodide-build +pyodide xbuildenv install "$PYODIDE_VERSION" +rustup toolchain install "$(pyodide config get rust_toolchain)" \ + --target wasm32-unknown-emscripten +pyodide xbuildenv install-emscripten +MATURIN_PEP517_ARGS="--no-default-features --profile wasm" pyodide build -o dist +``` + +On NixOS, run these commands in `nix run .#wasm-shell`. This shell is an +FHS environment, which lets the prebuilt Emscripten and rustup binaries run. + ## Before submitting a pull request... diff --git a/Cargo.lock b/Cargo.lock index 6a70b0c..2e38e61 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3703,7 +3703,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] name = "xarray_sql" -version = "0.4.0" +version = "0.5.0" dependencies = [ "arrow", "arrow-flight", diff --git a/Cargo.toml b/Cargo.toml index c8bea27..231d2b6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,22 +19,39 @@ exclude = [ [dependencies] arrow = { version = "58", features = ["pyarrow"] } -arrow-flight = { version = "58", features = ["flight-sql"] } +arrow-flight = { version = "58", features = ["flight-sql"], optional = true } async-stream = "0.3" async-trait = "0.1" -datafusion = { version = "54.0.0" } +datafusion = { version = "54.0.0", default-features = false } datafusion-ffi = { version = "54.0.0" } futures = { version = "0.3" } half = "2.7" -prost = "0.14" +prost = { version = "0.14", optional = true } # `abi3-py310` builds against CPython's stable ABI, so a single wheel per # platform works on all CPython >= 3.10 (matching `requires-python`). Maturin # enables `pyo3/extension-module` through pyproject.toml for wheel builds; it # must stay disabled for ordinary Cargo test binaries so they link libpython. pyo3 = { version = "0.28.0", features = ["abi3-py310"] } -tokio = { version = "1.46.1", features = ["macros", "net", "rt", "rt-multi-thread", "sync", "time"] } -tonic = { version = "0.14", features = ["transport"] } +tokio = { version = "1.46.1", features = ["macros", "rt", "sync", "time"] } +tonic = { version = "0.14", features = ["transport"], optional = true } +[features] +default = ["flight"] +# The Arrow Flight SQL server needs TCP sockets and OS threads, neither of +# which exist under WebAssembly. Pyodide builds disable default features. +# The server plans SQL, so it also needs DataFusion's default features +# (SQL, function libraries, Parquet, compression). +flight = ["dep:arrow-flight", "dep:prost", "dep:tonic", "datafusion/default", "tokio/net", "tokio/rt-multi-thread"] + + +# Pyodide downloads the wheel on every page load, so wasm builds trade +# speed for size: `MATURIN_PEP517_ARGS="--no-default-features --profile wasm"`. +[profile.wasm] +inherits = "release" +opt-level = "z" +lto = true +codegen-units = 1 +strip = true [build-dependencies] pyo3-build-config = "0.28" diff --git a/README.md b/README.md index 0fa1052..1bc7b20 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ import xarray_sql as xql # 4x-daily surface air temperature on a lat/lon grid, 2013-2014. ds = xr.tutorial.open_dataset('air_temperature') -ctx = xql.XarrayContext() +ctx = xql.Context() ctx.from_dataset('air', ds, chunks=dict(time=100)) # A climatology — the mean annual cycle — computed in SQL: average air @@ -110,8 +110,9 @@ ds = xr.open_zarr( storage_options={'token': 'anon'} # Anonymous read from the public GCS bucket — no auth required. ) -ctx = xql.XarrayContext() -# Make sure to pass `chunks`! +ctx = xql.Context() +# `chunks` sets the partition size; without it, partitions follow the +# store's own chunks (one hour each here). ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', diff --git a/docs/engines.md b/docs/engines.md index cc6e9b2..ab1e965 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -23,7 +23,7 @@ DataFusion is the built-in engine, wrapped in a session: ```python import xarray_sql as xql -ctx = xql.XarrayContext() +ctx = xql.Context() ctx.from_dataset("era5", ds, chunks={"time": 24}) result = ctx.sql("SELECT ... FROM era5").to_dataset() ``` @@ -342,14 +342,14 @@ stores `float32` as `float64` and `bool` as an integer, MySQL `bool` as keep the result's type. The cursor is a one-shot Arrow stream: `xql.to_dataset(cur, ...)` -round-trips eagerly, and `chunks=` needs `spill=True`. +reads it into memory, and `chunks=` needs `spill=True`. ## Serving over Flight SQL (no copy) The adapters above bring an engine to the data. `xql.serve` does the reverse and brings the data to remote clients: it starts an [Arrow Flight SQL](https://arrow.apache.org/docs/format/FlightSql.html) -server over the same lazy tables `XarrayContext` uses. +server over the same lazy tables `Context` uses. ```python import xarray_sql as xql @@ -453,7 +453,7 @@ to queries that filter on time. TABLE`, `COPY`, `SET`, ...) are rejected, so clients cannot read or write the server's filesystem. - **DataFusion's SQL, without xarray-sql's Python UDFs.** The - `cftime()` and `reproject()` functions `XarrayContext` registers are + `cftime()` and `reproject()` functions `Context` registers are not available on the server. - **Bound its memory.** Any client can send an expensive `ORDER BY`, join, or aggregation. `memory_limit=` (bytes) caps what those hold at @@ -475,7 +475,7 @@ What each integration provides. Known issues and constraints live on | | DataFusion | DuckDB | Polars | ADBC | |---|---|---|---|---| -| Register | `XarrayContext` / any `SessionContext` | `xql.register(con, name, ds)` | `pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` | `xql.register(con, name, ds)` (copies into the database) | +| Register | `Context` / any `SessionContext` | `xql.register(con, name, ds)` | `pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` | `xql.register(con, name, ds)` (copies into the database) | | Projection pushdown | yes | yes | yes | n/a (the database's own tables) | | Chunk pruning on dim predicates | yes | yes | yes | n/a (the database's own indexes) | | Eager round-trip (`xql.to_dataset`) | yes | yes | yes | yes (pass the cursor) | @@ -501,7 +501,9 @@ only the source chunks it maps onto. ```mermaid flowchart TB R["xql.to_dataset(result, ...)"] --> K{"chunks=?"} - K -- "None (default)" --> E["eager: materialize once
max_result_bytes= guards both the
Arrow stream and the dense grid"] + K -- "None (default)" --> LZ{"result type"} + LZ -- "re-executable
(DuckDB, Polars, DataFusion)" --> LI["lazily indexed, no dask:
each access re-runs the query
for its selection; .load() runs it
once per variable (small results:
read in one pass)"] + LZ -- "one-shot Arrow stream,
or max_result_bytes=" --> E["in memory: read once
max_result_bytes= guards both the
Arrow stream and the dense grid"] K -- "mapping / auto / inherit" --> SP{"spill=?"} SP -- "False (default)" --> HD{"result type"} HD -- "Polars LazyFrame/DataFrame
DataFusion DataFrame" --> RX["re-execution: each window
re-runs the query narrowed to its
coordinate range (flows back into
chunk pruning at the source)"] @@ -515,6 +517,39 @@ few windows of a huge result. Spill pays one full pass plus temporary disk — right when you'll touch most of the result, when the producer is a DuckDB relation, or when all you have is a one-shot stream. +**Without chunks:** `chunks=None`, the default, returns lazily indexed +variables, as `xr.open_dataset(chunks=None)` does. Each access, such as +`out.t2m.sel(time="2020-01").values`, re-executes the query narrowed to +that selection and projected to the variable it reads, so it reads only +the source chunks under it, and `.load()` executes the query once per +variable. Nothing is chunked, so dask is not needed, and DuckDB +relations work too: every access runs on the calling thread, unlike +chunked windows, which dask computes on worker threads. + +```python +out = xql.to_dataset(con.sql("SELECT * FROM era5"), template=ds) +out.t2m.sel(time="2020-01-01").mean().item() # reads one day's chunks +``` + +A lazy result must know its coordinates before any data is read, so +unless `coords="template"` supplies them, the query is first streamed +once. A small result (up to 64 MiB) is kept in memory from that pass, +so an aggregation, which no window filter can push below, executes +once. A larger one keeps only each dimension's distinct values and +stays lazy. One-shot Arrow streams, which cannot re-execute, and calls +with `max_result_bytes=` are read into memory directly. + +Because a lazy result re-executes its query on access: + +- the connection must stay open while the result is used: after + `con.close()`, reading a variable raises; +- changes to the source tables after `to_dataset` show through in + later reads; +- a query whose rows vary between runs, such as `LIMIT` without + `ORDER BY` or random sampling, can disagree with the coordinates + discovered for it. Read such a result once, in memory, with + `max_result_bytes=`. + Two knobs matter at scale: - `coords="template"` trusts the template's coordinate arrays instead of diff --git a/docs/examples.md b/docs/examples.md index f2b925b..4687668 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -15,7 +15,7 @@ import xarray_sql as xql ds = xr.tutorial.open_dataset('air_temperature') -ctx = xql.XarrayContext() +ctx = xql.Context() ctx.from_dataset('air', ds, chunks=dict(time=100)) clim = ctx.sql(''' @@ -66,17 +66,17 @@ url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) # A full year of hourly ERA5 — all 273 variables. No spatial slicing on the -# xarray side; SQL WHERE clauses below express the filters. `chunks={'time': 1}` -# aligns Dask chunks to native Zarr chunks of shape (1, 37, 721, 1440) so -# chunk reads from GCS happen concurrently. +# xarray side; SQL WHERE clauses below express the filters. Each table +# partition is one native Zarr chunk of shape (1, 37, 721, 1440); pass +# `chunks=` to `from_dataset` to group several into one partition. # # Heads up: 262 of those variables are surface and 11 are atmospheric. The # library pushes column projection down, so SELECT only fetches what you ask # for — but `SELECT * FROM era5.surface` would try to pull every variable # across the year (terabytes from GCS). Always SELECT specific columns. -ds = full.sel(time='2020').chunk({'time': 1}) +ds = full.sel(time='2020') -ctx = xql.XarrayContext() +ctx = xql.Context() ctx.from_dataset('era5', ds, table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', @@ -123,7 +123,7 @@ scalars into a single one-row table named `scalar`: ```python import fsspec import xarray as xr -from xarray_sql import XarrayContext +from xarray_sql import Context # A real GOES-16 ABI cloud-and-moisture file from NOAA's public bucket: # (y, x) image bands alongside dozens of scalar metadata variables. @@ -135,7 +135,7 @@ ds = xr.open_dataset(fsspec.open_local(f'simplecache::{url}')).chunk( {'y': 250, 'x': 250} ) -ctx = XarrayContext() +ctx = Context() ctx.from_dataset('goes', ds) # The gridded bands and the scalar metadata are separate tables. @@ -153,7 +153,7 @@ A runnable version of the ERA5 example lives at ## The same tables on DuckDB and Polars -Every example above registers through an `XarrayContext`, but the tables are +Every example above registers through a `Context`, but the tables are not DataFusion-specific: `xql.register(con, name, ds)` attaches the same lazy, pushdown-scanned table to a DuckDB connection, and `pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` serves Polars — same diff --git a/docs/geospatial.md b/docs/geospatial.md index 2ae8625..ca9b466 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -215,7 +215,7 @@ paradigm. They split cleanly along one line: **is the operation row-independent? that coordinate, so it is a *scalar function* — exactly what PostGIS and DuckDB-spatial already ship as `ST_Transform`. xarray-sql ships it as an optional geo extension (`pip install xarray-sql[geo]`): with pyproj -installed, every `XarrayContext` registers a PROJ-backed +installed, every `Context` registers a PROJ-backed `reproject(x, y, src_crs, dst_crs)` scalar UDF, so the CRS pair — any CRS pyproj understands — is part of the query rather than baked into the function: diff --git a/docs/limitations.md b/docs/limitations.md index 54473e2..50dc16b 100644 --- a/docs/limitations.md +++ b/docs/limitations.md @@ -33,8 +33,9 @@ Pick your engine: `spill=True` provides the chunked path without ever re-executing: the result is streamed once (bounded memory, on the handle's dedicated engine thread) into a temporary Parquet - file that windows re-execute against. The eager round-trip is - unaffected. + file that windows re-execute against. The unchunked round-trip + (`chunks=None`) is unaffected: every access runs on the calling + thread. **Derived relations break under concurrent materialization.** @@ -183,11 +184,11 @@ scans accept the extra reads. ### Sparse results can explode the dense grid -The eager round-trip reconstructs the coordinate-product grid: a +Loading a round-trip reconstructs the coordinate-product grid: a diagonal of n rows becomes an n×n array that can dwarf its Arrow -payload. `max_result_bytes=` raises cleanly at both danger points -(stream collection and dense allocation); it is opt-in and unlimited -by default. +payload. `max_result_bytes=` reads the result into memory and raises +cleanly at both danger points (stream collection and dense +allocation); it is opt-in and unlimited by default. ### One-shot Arrow streams cannot re-execute diff --git a/docs/performance.md b/docs/performance.md index 349589d..c461de9 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -194,7 +194,7 @@ adds value when the question is relational. === "DataFusion" - **Two registration paths.** `XarrayContext.from_dataset` uses the + **Two registration paths.** `Context.from_dataset` uses the native Rust table provider — partition-parallel, with `chunks=` controlling partition granularity; the `prefetch`/`coalesce_rows` scanner knobs on this page apply to the *pyarrow-dataset* path diff --git a/flake.nix b/flake.nix index 1bbdf65..c393dc3 100644 --- a/flake.nix +++ b/flake.nix @@ -20,8 +20,48 @@ zlib openssl ]; + + # Pyodide pins an exact Emscripten release and Rust toolchain, both of + # which ship as prebuilt generic-Linux binaries (emsdk's clang/node, + # rustup toolchains). An FHS sandbox lets those binaries run unmodified + # on NixOS. See "Building for WebAssembly" in CONTRIBUTING.md. + # + # nix run .#wasm-shell # interactive + # nix run .#wasm-shell -- -c 'pyodide build' # one-shot + wasmShell = pkgs.buildFHSEnv { + name = "xarray-sql-wasm"; + targetPkgs = p: with p; [ + bash + bzip2 + cacert + coreutils + curl + git + gnumake + gnutar + nodejs + python314 + rustup + stdenv.cc + stdenv.cc.cc.lib + uv + which + xz + zlib + ]; + runScript = "bash"; + profile = '' + export UV_PYTHON_PREFERENCE=only-system + export SSL_CERT_FILE=/etc/ssl/certs/ca-bundle.crt + ''; + }; in { + # buildFHSEnv relies on Linux namespaces. + packages = pkgs.lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux { + wasm-shell = wasmShell; + }; + devShells.default = pkgs.mkShell { packages = with pkgs; [ uv diff --git a/pyproject.toml b/pyproject.toml index 7536586..42d6466 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,12 +30,18 @@ classifiers = [ "Topic :: Database :: Front-Ends", ] dependencies = [ - "dask>=2024.8.0", - "datafusion==54.0.0", # This needs to match the cargo datafusion version!! + # This needs to match the cargo datafusion version!! DataFusion's Python + # bindings have no Pyodide wheels, so they are optional there. + "datafusion==54.0.0; sys_platform != 'emscripten'", + "pyarrow>=16.0.0", "xarray>=2024.7.0", ] [project.optional-dependencies] +# Lazy, chunked results from `to_dataset(chunks=...)` need a chunk manager. +dask = [ + "dask>=2024.8.0", +] adbc = [ # Plus a driver package for your database, e.g. adbc-driver-postgresql. "adbc-driver-manager>=1.12", @@ -56,7 +62,7 @@ test = [ "adbc-driver-flightsql>=1.12", "adbc-driver-sqlite>=1.12", "cftime", - "xarray-sql[adbc,duckdb,polars,geo]", + "xarray-sql[adbc,dask,duckdb,polars,geo]", "pytest", "xarray[io]", "gcsfs", diff --git a/src/lib.rs b/src/lib.rs index 61ab2c0..82458a6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,6 +41,7 @@ //! Will skip loading partitions whose time ranges are entirely before 2020-02-01. //! Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`. +#[cfg(feature = "flight")] mod flight; use std::collections::{HashMap, HashSet}; @@ -1645,6 +1646,7 @@ impl LazyArrowStreamTable { #[pymodule] fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + #[cfg(feature = "flight")] m.add_class::()?; Ok(()) } diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 34aebc7..cf2fb4a 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -125,6 +125,30 @@ def test_get_fragments_prunes_and_scans(ds): assert dataset.get_fragments(filter=pc.field("lat") > 100) == [] +def test_unchunked_zarr_scans_by_storage_chunks(ds, tmp_path): + ds.to_zarr(tmp_path / "store.zarr") + opened = xr.open_zarr(tmp_path / "store.zarr", chunks=None) + dataset = xql.arrow_dataset(opened) + assert len(dataset.get_fragments()) == 4 # stored with time chunks of 5 + xr.testing.assert_identical( + xql.to_dataset(dataset.to_table(), template=ds), ds.compute() + ) + + +def test_unchunked_in_memory_dataset_is_partitioned(): + big = xr.Dataset( + {"v": (["time", "y", "x"], np.zeros((40, 256, 256), np.float32))}, + coords={ + "time": np.arange(40), + "y": np.arange(256), + "x": np.arange(256), + }, + ) + dataset = xql.arrow_dataset(big) + assert len(dataset.get_fragments()) > 1 + assert dataset.count_rows() == big.v.size + + def test_scanner_honors_batch_size(ds): dataset = xql.arrow_dataset(ds) batches = list(dataset.scanner(batch_size=7).to_batches()) diff --git a/tests/test_df.py b/tests/test_df.py index 099e3f7..555b8e6 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -88,11 +88,9 @@ def test_block_slices_filters_chunk_keys_to_dataset_dims(air_small): assert len(extra) == len(base) -def test_block_slices_dimensional_unchunked_raises(): - # A dataset with dimensions but no chunking is still a user error. +def test_block_slices_unchunked_is_one_block(): ds = xr.Dataset({"v": (["x"], np.arange(3))}, coords={"x": np.arange(3)}) - with pytest.raises(AssertionError): - list(block_slices(ds)) + assert list(block_slices(ds)) == [{"x": slice(0, 3)}] def test_from_map_basic(): diff --git a/tests/test_ds.py b/tests/test_ds.py index 3682dac..57ce980 100644 --- a/tests/test_ds.py +++ b/tests/test_ds.py @@ -350,6 +350,16 @@ def test_chunks_argument_controls_partitioning(synthetic_dataset): ) +def test_chunks_none_is_lazy_without_a_chunk_manager(synthetic_dataset): + ctx = XarrayContext() + ctx.from_dataset("t", synthetic_dataset) + out = ctx.sql("SELECT * FROM t ORDER BY time, lat, lon").to_dataset( + chunks=None + ) + assert not out.chunks + xr.testing.assert_allclose(out.load(), synthetic_dataset.compute()) + + 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).""" diff --git a/tests/test_duckdb_backend.py b/tests/test_duckdb_backend.py index e592cfc..2723a6e 100644 --- a/tests/test_duckdb_backend.py +++ b/tests/test_duckdb_backend.py @@ -271,9 +271,14 @@ def test_or_and_in_filters_round_trip(ds): xr.testing.assert_allclose(out, expected) -def test_pushdown_dataset_rejects_unchunked_dataset(ds): - with pytest.raises(ValueError, match="must be chunked"): - XarrayPushdownDataset(ds.compute()) +def test_unchunked_dataset_round_trips(ds): + con = duckdb.connect() + xql.register(con, "weather", ds.compute()) + rel = con.sql( + "SELECT time, lat, lon, temperature, precipitation FROM weather " + "ORDER BY time, lat, lon" + ) + xr.testing.assert_identical(xql.to_dataset(rel, template=ds), ds.compute()) def test_fully_pruned_scan_returns_empty(con, ds): diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py index e2a7db5..1815a7f 100644 --- a/tests/test_lazy_roundtrip.py +++ b/tests/test_lazy_roundtrip.py @@ -8,6 +8,8 @@ streams are rejected with a clear error. """ +import pickle + import numpy as np import pandas as pd import pyarrow as pa @@ -61,6 +63,64 @@ def test_duckdb_chunked_fails_fast_with_guidance(source, registered): xql.to_dataset(rel, template=source, chunks={"time": 10}) +def test_duckdb_lazy_reads_only_the_accessed_window(source, registered): + con, reads = registered + rel = con.sql("SELECT * FROM t") + # Trusting the template's coords skips discovery, so nothing is read + # until the data is accessed. + out = xql.to_dataset(rel, template=source, coords="template") + assert not out.chunks # no chunk manager involved + assert reads == [] + window = out.t2m.isel(time=slice(0, 5)) + xr.testing.assert_allclose(window.load(), source.t2m.isel(time=slice(0, 5))) + assert len(reads) == 1 # the first of ten time chunks + xr.testing.assert_allclose(out.load(), source) + + +def test_duckdb_small_result_scans_source_once(source, registered): + con, reads = registered + rel = con.sql( + "SELECT lat, AVG(t2m) AS t2m FROM t GROUP BY lat ORDER BY lat" + ) + out = xql.to_dataset(rel, template=source) + assert len(reads) == 10 # each source chunk once + xr.testing.assert_allclose( + out.t2m.load(), source.t2m.mean("time"), check_dim_order=False + ) + assert len(reads) == 10 + + +def test_duckdb_large_result_discovers_coords_in_one_pass( + source, registered, monkeypatch +): + con, reads = registered + # A result over the one-pass budget stays lazy. + monkeypatch.setattr("xarray_sql.ds._ONE_PASS_BYTES", 0) + rel = con.sql("SELECT * FROM t WHERE lat > 0") + out = xql.to_dataset(rel, template=source) + assert len(reads) == 10 # one pass over the source finds every dim + xr.testing.assert_allclose(out.load(), source.sel(lat=source.lat > 0)) + + +def test_lazy_result_keeps_widened_dtype_without_reading(source, registered): + con, reads = registered + rel = con.sql("SELECT time, lat, CAST(t2m AS FLOAT) AS t2m FROM t") + out = xql.to_dataset(rel, template=source, coords="template") + assert reads == [] + assert out.t2m.dtype == np.float32 + np.testing.assert_allclose(out.t2m.values, source.t2m.values, rtol=1e-6) + + +def test_duckdb_lazy_result_needs_an_open_connection(source, registered): + con, _ = registered + out = xql.to_dataset( + con.sql("SELECT * FROM t"), template=source, coords="template" + ) + con.close() + with pytest.raises(Exception, match="closed"): + out.t2m.values + + def test_duckdb_eager_round_trip_through_handle(source, registered): con, _ = registered rel = con.sql("SELECT * FROM t") @@ -109,6 +169,20 @@ def test_polars_eager_frame_is_reexecutable(source): xr.testing.assert_allclose(out.compute(), source) +def test_chunked_polars_result_pickles(source): + pl = pytest.importorskip("polars") + + frame = pl.DataFrame( + { + "time": np.repeat(source.time.values, 6), + "lat": np.tile(source.lat.values, 100), + "t2m": source.t2m.values.ravel(), + } + ) + out = xql.to_dataset(frame.lazy(), template=source, chunks={"time": 50}) + xr.testing.assert_allclose(pickle.loads(pickle.dumps(out)).load(), source) + + def test_one_shot_stream_with_chunks_raises(source, registered): con, _ = registered table = con.sql("SELECT * FROM t").to_arrow_table() diff --git a/tests/test_sql.py b/tests/test_sql.py index 78042cf..57d69b9 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -5,9 +5,14 @@ import pytest import xarray as xr +import xarray_sql from xarray_sql import XarrayContext +def test_context_is_the_xarray_context(): + assert xarray_sql.Context is XarrayContext + + def test_sanity(air_dataset_small): ctx = XarrayContext() ctx.from_dataset("air", air_dataset_small) diff --git a/tests/test_to_dataset_perf.py b/tests/test_to_dataset_perf.py index 9124e9c..2878d1a 100644 --- a/tests/test_to_dataset_perf.py +++ b/tests/test_to_dataset_perf.py @@ -44,8 +44,8 @@ def test_lazy_chunk_peak_memory_is_bounded(air_source): Reference observation: lazy chunk peak is ~1.8 MB on a 31 MB dense source. A regression that quietly buffers the whole result - would push past 10 MB. Eager ``to_dataset(chunks=None)`` is - measured too as a sanity floor for the gap: the eager path should + would push past 10 MB. Loading ``to_dataset(chunks=None)`` whole is + measured too as a sanity floor for the gap: loading the whole grid should be at least an order of magnitude heavier than the lazy chunk, otherwise the lazy path isn't actually lazy. """ @@ -62,7 +62,7 @@ def test_lazy_chunk_peak_memory_is_bounded(air_source): ) _, eager_peak = _peak_mb( - lambda: ctx.sql('SELECT * FROM "air"').to_dataset(chunks=None) + lambda: ctx.sql('SELECT * FROM "air"').to_dataset(chunks=None).load() ) assert eager_peak > 50.0, ( f"eager whole-grid materialization should peak above 50 MB; " diff --git a/uv.lock b/uv.lock index 614b82c..c62d5b0 100644 --- a/uv.lock +++ b/uv.lock @@ -2834,8 +2834,8 @@ io = [ name = "xarray-sql" source = { editable = "." } dependencies = [ - { name = "dask" }, - { name = "datafusion" }, + { name = "datafusion", marker = "sys_platform != 'emscripten'" }, + { name = "pyarrow" }, { 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'" }, ] @@ -2844,6 +2844,9 @@ dependencies = [ adbc = [ { name = "adbc-driver-manager" }, ] +dask = [ + { name = "dask" }, +] dev = [ { name = "mkdocstrings", extra = ["python"] }, { name = "pre-commit" }, @@ -2870,6 +2873,7 @@ test = [ { name = "adbc-driver-manager" }, { name = "adbc-driver-sqlite" }, { name = "cftime" }, + { name = "dask" }, { name = "duckdb" }, { name = "gcsfs" }, { name = "polars" }, @@ -2894,24 +2898,25 @@ requires-dist = [ { name = "adbc-driver-manager", marker = "extra == 'adbc'", specifier = ">=1.12" }, { name = "adbc-driver-sqlite", marker = "extra == 'test'", specifier = ">=1.12" }, { name = "cftime", marker = "extra == 'test'" }, - { name = "dask", specifier = ">=2024.8.0" }, - { name = "datafusion", specifier = "==54.0.0" }, + { name = "dask", marker = "extra == 'dask'", specifier = ">=2024.8.0" }, + { name = "datafusion", marker = "sys_platform != 'emscripten'", specifier = "==54.0.0" }, { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.4.0" }, { name = "gcsfs", marker = "extra == 'test'" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" }, { name = "polars", marker = "extra == 'polars'", specifier = ">=1.33" }, { name = "pre-commit", marker = "extra == 'dev'" }, + { name = "pyarrow", specifier = ">=16.0.0" }, { name = "pyproj", marker = "extra == 'geo'" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "pytest", marker = "extra == 'test'" }, { name = "watchfiles", marker = "extra == 'dev'" }, { name = "xarray", specifier = ">=2024.7.0" }, { name = "xarray", extras = ["io"], marker = "extra == 'test'" }, - { name = "xarray-sql", extras = ["adbc", "duckdb", "polars", "geo"], marker = "extra == 'test'" }, + { name = "xarray-sql", extras = ["adbc", "dask", "duckdb", "polars", "geo"], marker = "extra == 'test'" }, { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" }, { name = "zensical", marker = "extra == 'docs'" }, ] -provides-extras = ["adbc", "dev", "docs", "duckdb", "geo", "polars", "test"] +provides-extras = ["adbc", "dask", "dev", "docs", "duckdb", "geo", "polars", "test"] [package.metadata.requires-dev] dev = [ diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py index 2867258..af1a85d 100644 --- a/xarray_sql/__init__.py +++ b/xarray_sql/__init__.py @@ -10,10 +10,10 @@ from .df import from_map from .reader import read_xarray, read_xarray_table from .roundtrip import to_dataset -from .sql import XarrayContext __all__ = [ "cftime", + "Context", "XarrayContext", "FlightSQLServer", "read_xarray_table", @@ -26,3 +26,21 @@ "to_dataset", "from_map", # deprecated ] + +try: + from .sql import Context, XarrayContext +except ModuleNotFoundError as e: + # DataFusion's Python bindings have no wheels for some platforms + # (e.g. Pyodide); every other engine adapter still works there. + if e.name != "datafusion": + raise + __all__.remove("Context") + __all__.remove("XarrayContext") + + def __getattr__(name: str): + if name in ("Context", "XarrayContext"): + raise AttributeError( + f"{name} requires the datafusion package, which is " + "not installed." + ) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/xarray_sql/backends/base.py b/xarray_sql/backends/base.py index b86dd8a..8c99322 100644 --- a/xarray_sql/backends/base.py +++ b/xarray_sql/backends/base.py @@ -109,7 +109,7 @@ def register( Args: con: An engine connection: a ``datafusion.SessionContext`` (or - [xarray_sql.XarrayContext][]), a + [xarray_sql.Context][]), a ``duckdb.DuckDBPyConnection``, or an ADBC DBAPI connection (``adbc_driver_manager.dbapi.Connection``), into which the Dataset is ingested as a table. @@ -120,7 +120,12 @@ def register( spelling, since its registration namespace is flat). ds: An xarray Dataset. chunks: Xarray-like chunks specification controlling partition - granularity. Defaults to the Dataset's existing chunks. + granularity. Defaults to the Dataset's existing chunks. An + unchunked Dataset opened lazily with ``chunks=None`` is split + along the storage chunks its backend recorded in + ``encoding["preferred_chunks"]``, and is one chunk along other + dimensions; one built in memory, which records none, is split + into partitions of about a million rows. Dask is not required. table_names: Maps a dimension group's exact dim tuple to the name its table takes. Groups left unnamed take their dimensions joined by underscores (``time_latitude_longitude``); the diff --git a/xarray_sql/backends/datafusion.py b/xarray_sql/backends/datafusion.py index 4fc652f..5f9660d 100644 --- a/xarray_sql/backends/datafusion.py +++ b/xarray_sql/backends/datafusion.py @@ -5,15 +5,17 @@ dimension predicates, projection pushdown, and exact per-partition statistics for the optimizer. This module only routes the generic [xarray_sql.register][] seam onto that existing machinery. + +Like the other adapters, it never imports ``datafusion`` until a +DataFusion connection is registered, so the rest of xarray-sql works +where DataFusion's Python bindings are unavailable (e.g. Pyodide). """ from __future__ import annotations -from typing import Any, TypeGuard +from typing import TYPE_CHECKING, Any, TypeGuard import xarray as xr -from datafusion import SessionContext -from datafusion.catalog import Schema from ..df import ( Chunks, @@ -23,9 +25,11 @@ shared_coord_arrays, ) from ..reader import read_xarray_table -from ..sql import XarrayContext from .base import register_adapter +if TYPE_CHECKING: + from datafusion import SessionContext + @register_adapter class DataFusionAdapter: @@ -33,7 +37,12 @@ class DataFusionAdapter: @staticmethod def matches(con: object) -> TypeGuard[SessionContext]: - return isinstance(con, SessionContext) + # Context and any other subclass carry this in their MRO. + return any( + cls.__module__ == "datafusion.context" + and cls.__qualname__ == "SessionContext" + for cls in type(con).__mro__ + ) @staticmethod def register( @@ -45,10 +54,14 @@ def register( table_names: TableNames = None, **kwargs: Any, ) -> SessionContext: - # XarrayContext.from_dataset adds cftime UDF registration and + from datafusion.catalog import Schema + + from ..sql import Context + + # Context.from_dataset adds cftime UDF registration and # round-trip metadata tracking on top of the split below; use it # when available. - if isinstance(con, XarrayContext): + if isinstance(con, Context): return con.from_dataset( name, ds, chunks=chunks, table_names=table_names, **kwargs ) @@ -60,7 +73,7 @@ def register( return con # A mixed-dimension Dataset becomes one table per dimension group - # in a schema named after the Dataset, exactly as XarrayContext + # in a schema named after the Dataset, exactly as Context # registers it — so ``name.group`` is the same SQL either way. coord_arrays = shared_coord_arrays(ds) schema = Schema.memory_schema(con) diff --git a/xarray_sql/backends/flight.py b/xarray_sql/backends/flight.py index 294cae6..a1adbac 100644 --- a/xarray_sql/backends/flight.py +++ b/xarray_sql/backends/flight.py @@ -5,7 +5,7 @@ [FlightSQLServer][xarray_sql.backends.flight.FlightSQLServer] turns that around and brings the data to remote clients. It hosts a DataFusion session over the same lazy tables -[XarrayContext][xarray_sql.XarrayContext] uses and speaks +[Context][xarray_sql.Context] uses and speaks [Flight SQL](https://arrow.apache.org/docs/format/FlightSql.html), so any Flight SQL client — ADBC's Flight SQL driver in Python, R, Go, or Java, the Flight SQL JDBC and ODBC drivers, and the SQL tools built on them — @@ -37,7 +37,6 @@ import xarray as xr -from .._native import FlightSqlServer as _NativeFlightSqlServer from ..df import ( Chunks, TableNames, @@ -50,6 +49,11 @@ from .adbc import _warn_on_folded_names from .base import register_adapter +try: + from .._native import FlightSqlServer as _NativeFlightSqlServer +except ImportError: # WebAssembly builds have no sockets or threads. + _NativeFlightSqlServer = None + __all__ = ["FlightSQLServer", "serve"] _WILDCARD_HOSTS = ("", "0.0.0.0", "::") @@ -82,6 +86,11 @@ def __init__(self, memory_limit: int | None = None) -> None: externally too (e.g. a container memory limit) before serving untrusted clients. """ + if _NativeFlightSqlServer is None: + raise NotImplementedError( + "This build of xarray-sql has no Flight SQL server " + "(it is not available under WebAssembly)." + ) self._native = _NativeFlightSqlServer(memory_limit) self._host: str | None = None self._port: int | None = None @@ -99,7 +108,7 @@ def register( A Dataset whose variables sit on different dimensions is split into one table per dimension group, in a SQL schema named ``name`` (``name.group``), exactly as - [XarrayContext.from_dataset][xarray_sql.XarrayContext.from_dataset] + [Context.from_dataset][xarray_sql.Context.from_dataset] registers it. Names are registered exactly as given. The server's SQL folds @@ -111,7 +120,8 @@ def register( mixed-dimension Dataset. ds: An xarray Dataset. chunks: Xarray-like chunks specification controlling partition - granularity. Defaults to the Dataset's existing chunks. + granularity. Defaults to the Dataset's existing chunks, or + for an unchunked Dataset its storage chunks. table_names: Maps a dimension group's exact dim tuple to the name its table takes. diff --git a/xarray_sql/backends/pyarrow.py b/xarray_sql/backends/pyarrow.py index 9649dff..7b9c11d 100644 --- a/xarray_sql/backends/pyarrow.py +++ b/xarray_sql/backends/pyarrow.py @@ -42,6 +42,7 @@ import xarray as xr from ..df import ( + _CAN_START_THREADS, Block, Chunks, DEFAULT_BATCH_SIZE, @@ -57,8 +58,10 @@ from ..geometry import GEOMETRY_COLUMN, build_geometry, geometry_field from ..reader import XarrayRecordBatchReader -DEFAULT_PREFETCH = 4 -"""Chunk loads kept in flight ahead of the consumer during a scan.""" +DEFAULT_PREFETCH = 4 if _CAN_START_THREADS else 1 +"""Chunk loads kept in flight ahead of the consumer during a scan. + +Where threads are unavailable (Pyodide), chunks load one at a time.""" _SHADOW_FANOUT = 1024 """Maximum fragments per shadow level. @@ -340,10 +343,6 @@ def __init__( geometry_field(geometry_encoding, geometry_crs) ) self._resolved = resolve_chunks(ds, chunks) - if not self._resolved and ds.sizes: - raise ValueError( - "Dataset `ds` must be chunked or `chunks` must be provided." - ) self._chunk_bounds = { d: np.cumsum((0, *sizes)) for d, sizes in self._resolved.items() } @@ -1102,7 +1101,8 @@ def arrow_dataset( ds: An xarray Dataset. All data variables must share the same dimensions (select a variable subset first otherwise). chunks: Xarray-like chunks specification controlling partition - granularity. Defaults to the Dataset's existing chunks. + granularity. Defaults to the Dataset's existing chunks, or for + an unchunked Dataset its storage chunks. batch_size: Maximum rows per emitted Arrow RecordBatch. prefetch: Chunk loads kept in flight ahead of the consumer (memory scales with ``prefetch`` x pivoted chunk size). @@ -1186,7 +1186,7 @@ def arrow_datasets( chunks: Xarray-like chunks specification controlling partition granularity. Keys naming dimensions a group does not have are ignored for that group. Defaults to the Dataset's existing - chunks. + chunks, or for an unchunked Dataset its storage chunks. table_names: Maps a dimension group's exact dim tuple to the name its table takes, e.g. ``{('time', 'latitude', 'longitude'): 'surface'}``. Groups diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 27a6c94..9659130 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -15,7 +15,7 @@ with ``xarray:units`` and ``xarray:calendar`` metadata on the Arrow field. This preserves the original CF-convention encoding losslessly. A ``cftime()`` DataFusion UDF (registered automatically by - ``XarrayContext.from_dataset``) provides ergonomic SQL filtering. + ``Context.from_dataset``) provides ergonomic SQL filtering. """ from __future__ import annotations diff --git a/xarray_sql/df.py b/xarray_sql/df.py index f875c98..d26da7f 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,4 +1,5 @@ import itertools +import sys from collections import defaultdict from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping from typing import Any @@ -11,6 +12,12 @@ from . import cftime as cft +_CAN_START_THREADS = sys.platform != "emscripten" +"""Whether this interpreter can start OS threads; Pyodide cannot.""" + +_DEFAULT_BLOCK_ROWS = 2**20 +"""Rows per partition for a Dataset that records no storage chunks.""" + Block = dict[Hashable, slice] Chunks = dict[str, int] | None TableNames = Mapping[tuple[str, ...], str] | None @@ -63,8 +70,11 @@ def resolve_chunks( Filters out keys for dims this dataset doesn't have (sub-datasets in a heterogeneous group need not contain every dimension named in the - spec), then either rechunks arithmetically via ``compute_chunks`` or - falls back to the dataset's existing dask chunks. + spec), then rechunks arithmetically via ``compute_chunks``. Without a + spec, a chunked dataset keeps its chunks; an unchunked one opened from + storage is split along the storage chunks its backend recorded, and is + one chunk along other dimensions. A dataset that records none (built in + memory) is split into blocks of about ``_DEFAULT_BLOCK_ROWS`` rows. Returns an empty mapping for scalar datasets; callers should treat that as "one block covering everything". @@ -73,7 +83,40 @@ def resolve_chunks( chunks = {dim: size for dim, size in chunks.items() if dim in ds.sizes} if chunks: return compute_chunks(ds, chunks) - return {d: tuple(c) for d, c in ds.chunks.items()} + if ds.chunks: + return {d: tuple(c) for d, c in ds.chunks.items()} + return compute_chunks(ds, _storage_chunks(ds) or _default_chunks(ds)) + + +def _storage_chunks(ds: xr.Dataset) -> dict[str, int]: + """Chunk sizes from the backend's ``preferred_chunks`` encoding. + + Zarr and NetCDF backends record each variable's on-disk chunking there. + Where variables disagree on a dimension, the first one wins. + """ + chunks: dict[str, int] = {} + for var in ds.data_vars.values(): + for dim, size in var.encoding.get("preferred_chunks", {}).items(): + chunks.setdefault(dim, size) + return chunks + + +def _default_chunks(ds: xr.Dataset) -> dict[str, int]: + """Chunk sizes giving blocks of at most about ``_DEFAULT_BLOCK_ROWS`` rows. + + Inner dimensions stay whole while they fit, so each block is a + contiguous run of the row-major grid; the outermost dimension that + does not fit is split, and the dimensions outside it get size one. + """ + chunks: dict[str, int] = {} + inner = 1 + for dim in reversed(list(ds.sizes)): + size = ds.sizes[dim] + step = max(1, min(size, _DEFAULT_BLOCK_ROWS // inner)) + if step < size: + chunks[str(dim)] = step + inner *= step + return chunks def _ensure_default_indexes(ds: xr.Dataset) -> xr.Dataset: @@ -99,12 +142,8 @@ def _block_slices_from_resolved( ) -> Iterator[Block]: """Emit blocks given pre-resolved per-dim chunk tuples.""" if not resolved: - # No chunkable dimensions. A dimensionless dataset (e.g. scalar - # metadata variables) is a single block; a dataset that has - # dimensions but no chunking is a user error. - assert not ds.sizes, ( - "Dataset `ds` must be chunked or `chunks` must be provided." - ) + # A dimensionless dataset (e.g. scalar metadata variables) is a + # single block. yield {} return diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index e45b4d2..6532d7d 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -3,7 +3,7 @@ The inverse of the forward Dataset-to-table pivot done by [xarray_sql.df.pivot][]. Internally defines an [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapper around the DataFusion ``DataFrame`` returned by -[XarrayContext.sql][xarray_sql.sql.XarrayContext.sql], with a [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] +[Context.sql][xarray_sql.sql.Context.sql], with a [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] method that round-trips a query result back to ``xr.Dataset``. Reconstruction is controlled by the ``chunks`` argument to @@ -37,6 +37,7 @@ import numpy as np import pandas as pd import pyarrow as pa +import pyarrow.compute as pc import xarray as xr from .lazyscan import DataFusionHandle, DimSpec, LazyResultHandle @@ -144,7 +145,8 @@ def _restore_dtype(var: xr.DataArray, template: xr.DataArray) -> xr.DataArray: keeps the result's dtype too. """ dtype = template.dtype - if var.dtype == dtype or not isinstance(var.data, np.ndarray): + # ``var.data`` would load a lazy variable; check without reading it. + if var.dtype == dtype or not var.variable._in_memory: return var if set(var.dims) != set(template.dims): return var @@ -705,27 +707,6 @@ def _dataset_from_batches( return xr.Dataset(data_vars=data_vars, coords=coords_arg) -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()] - return _dataset_from_batches( - batches, dimension_columns, field_names, field_types - ) - - _PURE_SCAN_NODES = {"Projection", "Sort", "TableScan", "SubqueryAlias"} @@ -799,6 +780,58 @@ def _maybe_template_coords( return {d: np.asarray(source.coords[d].values) for d in dimension_columns} +_ONE_PASS_BYTES = 64 * 2**20 +"""Size up to which a lazy result that must discover its coordinates is +read in a single pass instead.""" + + +def _read_or_discover( + handle: LazyResultHandle, columns: list[str], dimension_columns: list[str] +) -> tuple[list[pa.RecordBatch] | None, dict[str, np.ndarray] | None]: + """Execute the query once: the whole result if small, else its coordinates. + + A lazy result must know each dimension's coordinates before any data + is read. One streamed pass serves both cases. A result that fits in + ``_ONE_PASS_BYTES`` -- an aggregation, which no window filter can push + below, or a narrow selection -- is returned whole, to be built in + memory. Past the budget, the pass keeps going but keeps only each + dimension's distinct values, so discovery executes the query once + rather than once per dimension. Returns ``(None, None)`` for handles + that cannot stream, which then discover one dimension at a time. + """ + stream = getattr(handle, "stream", None) + if stream is None: + return None, None + batches: list[pa.RecordBatch] = [] + total = 0 + distinct: dict[str, pa.Array] | None = None + for batch in stream(columns): + if distinct is None: + total += batch.nbytes + batches.append(batch) + if total <= _ONE_PASS_BYTES: + continue + # Over budget: from here on, keep only distinct coordinates. + distinct = {} + pending, batches = batches, [] + else: + pending = [batch] + for b in pending: + for d in dimension_columns: + values = pc.unique(b.column(b.schema.names.index(d))) + if d in distinct: + values = pc.unique(pa.concat_arrays([distinct[d], values])) + distinct[d] = values + if distinct is None: + return batches, None + return None, { + # Sorted, as from ``distinct``, so positional slices map onto + # contiguous value ranges. + d: np.sort(distinct[d].to_numpy(zero_copy_only=False)) + for d in dimension_columns + } + + def _build_lazy_scan( handle: LazyResultHandle, dimension_columns: list[str], @@ -808,9 +841,9 @@ def _build_lazy_scan( ) -> xr.Dataset: """Build a lazy Dataset whose data vars are [SQLBackendArray][xarray_sql.ds.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 come either from the + Each data variable stays lazy: every access reads its coordinate range + via a pushdown filter, either directly or, once wrapped by + ``Dataset.chunk``, one chunk at a time. Coordinates come either from the caller (the scanned table's registered Dataset for unfiltered DataFusion scans -- see ``_maybe_template_coords`` -- or an explicitly trusted template) or from per-dim distinct queries through the handle; over a @@ -921,9 +954,13 @@ def _result_to_xarray( ``chunks`` (already resolved by ``XarrayDataFrame._resolve_chunks``) selects the execution strategy: - * ``None`` -> eager: execute once and materialize a dense Dataset - (``_materialize``). Correct for any query and the right default for - reductions, whose results are small. + * ``None`` -> lazy, unchunked: the data variables are lazily indexed + [SQLBackendArray][xarray_sql.ds.SQLBackendArray]s, so each access + re-executes the query narrowed to the selection and projected to + the variable read. No chunk manager is involved. A result whose + coordinates must be discovered is first streamed in one pass, which + keeps it in memory if small (``_read_or_discover``), so an + aggregation executes once. * a mapping (or ``"auto"``) -> lazy/chunked: build [SQLBackendArray][xarray_sql.ds.SQLBackendArray] data variables (``_build_lazy_scan``) and wrap them with ``Dataset.chunk`` so each chunk reads its coordinate range via filter @@ -944,17 +981,26 @@ def _result_to_xarray( 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, dimension_columns, field_names, field_types) + handle = DataFusionHandle(inner_df) + coord_arrays = _maybe_template_coords( + templates, dimension_columns, inner_df + ) + batches = None + if chunks is None and coord_arrays is None: + batches, coord_arrays = _read_or_discover( + handle, field_names, dimension_columns + ) + if batches is not None: + ds = _dataset_from_batches( + batches, dimension_columns, field_names, field_types + ) else: ds = _build_lazy_scan( - DataFusionHandle(inner_df), + handle, dimension_columns, field_names, field_types, - coord_arrays=_maybe_template_coords( - templates, dimension_columns, inner_df - ), + coord_arrays=coord_arrays, ) return _finish_dataset( ds, @@ -1012,7 +1058,7 @@ def _finish_dataset( class XarrayDataFrame: """Wrapper around a DataFusion ``DataFrame`` with xarray-aware helpers. - Returned by [xarray_sql.XarrayContext.sql][]. Forwards every + Returned by [xarray_sql.Context.sql][]. Forwards every attribute it does not define itself to the wrapped DataFrame, so ``.collect()``, ``.schema()``, ``.show()``, ``.count()`` all work unchanged. @@ -1022,7 +1068,7 @@ class XarrayDataFrame: dropped by the forward pivot. Users should not construct this class directly; let - [XarrayContext.sql][xarray_sql.sql.XarrayContext.sql] produce it. + [Context.sql][xarray_sql.sql.Context.sql] produce it. """ def __init__( @@ -1034,7 +1080,7 @@ def __init__( Args: inner: The underlying ``datafusion.DataFrame`` returned by - [XarrayContext.sql][xarray_sql.sql.XarrayContext.sql]. + [Context.sql][xarray_sql.sql.Context.sql]. templates: Snapshot of the registered Datasets on the producing context, keyed by the SQL identifier each was registered under. Used by ``to_dataset`` to recover metadata that @@ -1084,11 +1130,15 @@ def to_dataset( 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). + dimension (e.g. a global aggregation) inherits nothing, and + neither does a result with no resolvable source Dataset; + both are returned as with ``None``. + * ``None``: lazy without a chunk manager, as with + ``xr.open_dataset(chunks=None)``. Each access re-executes + the query narrowed to the selection and projected to the + variable read, so ``.load()`` executes it once per + variable. Dask is not required. The context must stay + usable while the result is read. * 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, @@ -1138,11 +1188,11 @@ def _resolve_chunks( ) -> 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 + ``None`` selects the lazy, unchunked path; anything else selects the + 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`` automatically. Mappings pass through unchanged; ``"auto"`` passes through here and is snapped to source partition boundaries later (see ``_auto_chunks``). """ diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py index e8f36fc..a967c16 100644 --- a/xarray_sql/lazyscan.py +++ b/xarray_sql/lazyscan.py @@ -37,7 +37,8 @@ import pandas as pd import pyarrow as pa import pyarrow.parquet as pq -from datafusion import col, literal + +from .df import _CAN_START_THREADS DimSpec = tuple[Literal["range", "values"], Any, Any] """One dimension's window: ``("range", lo, hi)`` (inclusive bounds; the @@ -113,6 +114,8 @@ def schema(self) -> pa.Schema: return self._df.schema() def distinct(self, column: str) -> np.ndarray: + from datafusion import col + dim_only = self._df.select(col(f'"{column}"')).distinct() batches = [b.to_pyarrow() for b in dim_only.execute_stream()] if not batches: @@ -127,6 +130,8 @@ def distinct(self, column: str) -> np.ndarray: def fetch( self, specs: dict[str, DimSpec], columns: list[str] ) -> list[pa.RecordBatch]: + from datafusion import col, literal + predicate = None for dim, (kind, a, b) in specs.items(): c = col(f'"{dim}"') @@ -146,6 +151,8 @@ def fetch( def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: """Execute once, yielding Arrow batches as the plan produces them.""" + from datafusion import col + out = self._df.select(*(col(f'"{n}"') for n in columns)) return (b.to_pyarrow() for b in out.execute_stream()) @@ -195,6 +202,11 @@ class DuckDBHandle: def __init__(self, rel: Any) -> None: self._rel = rel + self._runner: ThreadPoolExecutor | None = None + if not _CAN_START_THREADS: + # Single-threaded interpreters (Pyodide) have no other + # threads to deadlock against; run engine calls inline. + return self._runner = ThreadPoolExecutor(max_workers=1) self._runner.submit(lambda: None).result() # start the thread now # Stop the dedicated engine thread when the handle dies; it @@ -206,6 +218,8 @@ def __init__(self, rel: Any) -> None: ) def _run(self, fn: Any) -> Any: + if self._runner is None: + return fn() return self._runner.submit(fn).result() @staticmethod @@ -261,6 +275,23 @@ def fetch( self._run(lambda: list(self._to_arrow_reader(rel))), ) + def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: + """Execute once, yielding Arrow batches as DuckDB produces them. + + The reader is closed on the engine thread, even when the consumer + stops early, so the query is never torn down from another thread. + """ + import duckdb + + rel = self._rel.project(*(duckdb.ColumnExpression(n) for n in columns)) + reader = self._run(lambda: self._to_arrow_reader(rel)) + batches = iter(reader) + try: + while (batch := self._run(lambda: next(batches, None))) is not None: + yield batch + finally: + self._run(reader.close) + def spill_parquet(self, path: str) -> None: def run() -> None: reader = self._to_arrow_reader(self._rel) diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py index 2470d1b..ea63703 100644 --- a/xarray_sql/proj.py +++ b/xarray_sql/proj.py @@ -42,7 +42,7 @@ would return ``inf``); NULL CRS arguments yield NaN as well. Requires ``pyproj`` (``pip install xarray-sql[geo]``). When pyproj is -installed, [xarray_sql.XarrayContext][] registers ``reproject()`` +installed, [xarray_sql.Context][] registers ``reproject()`` automatically; [register][xarray_sql.proj.register] is the explicit hook for plain DataFusion ``SessionContext`` objects or custom UDF names. """ @@ -212,7 +212,7 @@ def _reproject( def register(ctx, name: str = "reproject") -> None: """Register the ``reproject(x, y, src_crs, dst_crs)`` scalar UDF. - Works on any DataFusion ``SessionContext`` (``XarrayContext`` + Works on any DataFusion ``SessionContext`` (``Context`` registers it automatically when pyproj is installed). The UDF returns a ``{x, y}`` struct of destination coordinates, so a query selects components with subscripts:: diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 153ed8c..d2c78c4 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -77,7 +77,8 @@ def __init__( Args: ds: An xarray Dataset. All data_vars must share the same dimensions. chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. + the Dataset's existing chunks, or for an unchunked Dataset + its storage chunks. batch_size: Maximum rows per emitted Arrow RecordBatch. Smaller values let DataFusion start processing earlier at the cost of more Python→Arrow conversion calls. @@ -179,7 +180,8 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: Args: ds: An Xarray Dataset. All `data_vars` must share the same dimensions. chunks: Xarray-like chunks. If not provided, will default to the - Dataset's chunks. The product of the chunk sizes becomes the + Dataset's chunks, or for an unchunked Dataset its storage chunks. + The product of the chunk sizes becomes the standard length of each dataframe partition. Returns: @@ -223,7 +225,8 @@ def read_xarray_table( Args: ds: An xarray Dataset. All data_vars must share the same dimensions. chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. + the Dataset's existing chunks, or for an unchunked Dataset its + storage chunks. batch_size: Maximum rows per Arrow RecordBatch emitted per partition. Smaller values let DataFusion start processing earlier; the default (65 536) works well for most datasets. diff --git a/xarray_sql/roundtrip.py b/xarray_sql/roundtrip.py index a592ccc..4f56758 100644 --- a/xarray_sql/roundtrip.py +++ b/xarray_sql/roundtrip.py @@ -7,17 +7,16 @@ Dataset. Nothing here is engine-specific: results arrive as Arrow record batches regardless of which engine executed the SQL. -Reconstruction is eager by default (the result is materialized once -into a dense in-memory Dataset). Passing ``chunks=`` selects the -lazy/chunked path instead: data variables are reconstructed on access, -window by window, by re-executing the engine's query narrowed to each -chunk's coordinate range. That requires the result to be -*re-executable* — a Polars LazyFrame (or eager DataFrame) or a -DataFusion DataFrame — not a one-shot Arrow stream; see -[xarray_sql.lazyscan][]. DuckDB relations are re-executable but -refuse the chunked path (a thread-safety limitation noted on +Reconstruction of a *re-executable* result — a DuckDB relation, a +Polars LazyFrame (or eager DataFrame), or a DataFusion DataFrame; see +[xarray_sql.lazyscan][] — is lazy by default: data variables are read +on access by re-executing the engine's query narrowed to the selection. +Passing ``chunks=`` additionally wraps them in a chunk manager, one +window per chunk. DuckDB relations refuse the chunked path (a +thread-safety limitation noted on [DuckDBHandle][xarray_sql.lazyscan.DuckDBHandle]); pair them with -``spill=True`` instead. +``spill=True`` instead. One-shot Arrow streams have nothing to +re-execute and are materialized once into a dense in-memory Dataset. """ from __future__ import annotations @@ -41,6 +40,7 @@ _build_lazy_scan, _dataset_from_batches, _finish_dataset, + _read_or_discover, _result_dims, ) from .lazyscan import LazyResultHandle, PolarsHandle, resolve_lazy_handle @@ -202,7 +202,19 @@ def to_dataset( the result. ``"template"`` reindexes to the template's full coord ranges, filling absent cells with ``fill_value``. fill_value: Fill for ``sparsity="template"``. Defaults to NaN. - chunks: ``None`` (default) materializes eagerly. A mapping + chunks: ``None`` (default) is lazy without a chunk manager, as + with ``xr.open_dataset(chunks=None)``: each access + re-executes the query narrowed to the selection and + projected to the variable read, so ``.load()`` executes it + once per variable. Dask is not required, and DuckDB + relations work, since every access runs on the calling + thread. Because reads re-execute, the connection must stay + open while the result is used, later changes to the source + show through, and a query whose rows vary between runs + (``LIMIT`` without ``ORDER BY``, random sampling) should be + read once with ``max_result_bytes=`` instead. One-shot + streams (Arrow tables, readers, ADBC cursors) have nothing + to re-execute and are read eagerly. A mapping (e.g. ``{"time": 100}``), ``"auto"``, or ``"inherit"`` selects the lazy/chunked path: data variables are reconstructed window by window on access, each window @@ -218,7 +230,8 @@ def to_dataset( template's coord arrays instead, skipping discovery; only valid when the result spans the template's full extent (an unfiltered scan), and requires ``template=``. - max_result_bytes: Optional budget for the eager path. Raises a + max_result_bytes: Optional budget for materializing the result. + Setting it reads the result eagerly, in memory. Raises a clean ``ValueError`` (with the running size) as soon as the materializing result exceeds it — both while collecting the Arrow stream and before allocating the dense arrays — @@ -238,11 +251,11 @@ def to_dataset( to the system temp dir; a path spills into that directory. The file is removed when the returned Dataset is garbage collected. Requires Polars; only valid with ``chunks=``. - Returns: An ``xr.Dataset`` with ``dims`` as dimensions and the remaining - result columns as data variables — dense and in-memory by - default, lazily chunked when ``chunks`` is given. + result columns as data variables: lazily indexed by default + (in memory for one-shot streams or with ``max_result_bytes``), and + lazily chunked when ``chunks`` is given. Raises: ValueError: When neither ``dims`` nor ``template`` resolves the @@ -284,6 +297,19 @@ def to_dataset( return _to_dataset_lazy( result, dims, template, sparsity, fill_value, chunks, coords ) + if max_result_bytes is None: + handle = resolve_lazy_handle(result) + if handle is not None: + return _to_dataset_lazy( + result, + dims, + template, + sparsity, + fill_value, + None, + coords, + _handle=handle, + ) schema, batches = _result_to_batches(result, max_result_bytes) field_names = [f.name for f in schema] @@ -367,20 +393,25 @@ def _to_dataset_lazy( template: xr.Dataset | None, sparsity: Sparsity, fill_value: Any, - chunks: Mapping[str, int] | str, + chunks: Mapping[str, int] | str | None, coords: Literal["discover", "template"], _handle: LazyResultHandle | None = None, ) -> xr.Dataset: - """The chunked reconstruction behind ``to_dataset(chunks=...)``.""" + """The lazy reconstruction behind ``to_dataset``. + + ``chunks=None`` returns lazily indexed variables without chunking + them. + """ handle = _handle if _handle is not None else resolve_lazy_handle(result) if handle is None: raise TypeError( - "chunks= requires a re-executable engine result (a Polars " - "LazyFrame/DataFrame or a DataFusion DataFrame); got " + "Lazy reconstruction requires a re-executable engine result (a " + "DuckDB relation, Polars LazyFrame/DataFrame, or DataFusion " + "DataFrame); got " f"{type(result).__qualname__}, which is a one-shot stream. " "Pass the engine's lazy handle instead of a materialized " "result, add spill=True to reconstruct from a one-pass " - "on-disk spill, or use chunks=None." + "on-disk spill, or reconstruct eagerly (the default)." ) schema = handle.schema() field_names = [f.name for f in schema] @@ -398,6 +429,27 @@ def _to_dataset_lazy( ) coord_arrays = {d: np.asarray(template.coords[d].values) for d in dims} + if chunks is None: + # Lazily indexed and unchunked: every access runs on the calling + # thread, so even handles that refuse chunking can serve it. A + # small result is cheaper to read once than to discover lazily. + batches = None + if coord_arrays is None: + batches, coord_arrays = _read_or_discover(handle, field_names, dims) + if batches is not None: + ds = _dataset_from_batches(batches, dims, field_names, field_types) + else: + ds = _build_lazy_scan( + handle, + dims, + field_names, + field_types, + coord_arrays=coord_arrays, + ) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, None, field_types + ) + resolved = XarrayDataFrame._resolve_chunks(chunks, template, dims) if resolved is None: # "inherit" with no chunked source dimension to inherit from: @@ -414,8 +466,9 @@ def _to_dataset_lazy( "relation from worker threads intermittently deadlocks in " "duckdb-python when the query scans a Python-backed table " "(see xarray_sql.lazyscan.DuckDBHandle.supports_chunked). " - "Add spill=True to reconstruct from a one-pass on-disk " - "spill, use chunks=None (eager), or run the query through " + "Use chunks=None (lazy and unchunked: each access runs on the " + "calling thread), add spill=True to reconstruct from a " + "one-pass on-disk spill, or run the query through " "Polars (pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))) " "or a DataFusion context." ) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index c9911c1..9bdf1eb 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -23,8 +23,11 @@ _proj = None -class XarrayContext(SessionContext): - """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" +class Context(SessionContext): + """A datafusion `SessionContext` that also supports `xarray.Dataset`s. + + Also available under its original name, ``XarrayContext``. + """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -86,7 +89,7 @@ def from_dataset( non-Gregorian calendars (e.g. one 360_day and one julian), the UDF from the first registration will be used for all subsequent ``cftime()`` calls and may produce incorrect offsets for the - other dataset. In that case, create a separate ``XarrayContext`` + other dataset. In that case, create a separate ``Context`` for each calendar. Args: @@ -100,7 +103,8 @@ def from_dataset( table names within the schema, used when the dataset has variables with differing dimensions. chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. + the Dataset's existing chunks, or for an unchunked Dataset + its storage chunks (see [register][xarray_sql.register]). Returns: self, to allow chaining. @@ -190,3 +194,7 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: """ inner = super().sql(query, *args, **kwargs) return XarrayDataFrame(inner, templates=self._registered_datasets) + + +XarrayContext = Context +"""The original name of [Context][xarray_sql.Context], kept for compatibility."""