From 42b4bfe22fa3331ba45ef5b57011e6bcc4d52523 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Thu, 8 Oct 2026 18:36:48 -0700 Subject: [PATCH 1/9] wip, claude adds wasm support. datafusion is now optional --- .github/workflows/ci-wasm.yml | 96 +++++++++++++++++++++++++++++++ CONTRIBUTING.md | 23 ++++++++ Cargo.lock | 2 +- Cargo.toml | 14 +++-- README.md | 4 +- docs/engines.md | 8 +-- docs/examples.md | 10 ++-- docs/geospatial.md | 2 +- docs/performance.md | 2 +- flake.nix | 40 +++++++++++++ pyproject.toml | 5 +- src/lib.rs | 2 + tests/test_sql.py | 5 ++ uv.lock | 6 +- xarray_sql/__init__.py | 20 ++++++- xarray_sql/backends/base.py | 2 +- xarray_sql/backends/datafusion.py | 29 +++++++--- xarray_sql/backends/flight.py | 15 ++++- xarray_sql/backends/pyarrow.py | 7 ++- xarray_sql/cftime.py | 2 +- xarray_sql/df.py | 4 ++ xarray_sql/ds.py | 8 +-- xarray_sql/lazyscan.py | 16 +++++- xarray_sql/proj.py | 4 +- xarray_sql/sql.py | 13 ++++- 25 files changed, 292 insertions(+), 47 deletions(-) create mode 100644 .github/workflows/ci-wasm.yml diff --git a/.github/workflows/ci-wasm.yml b/.github/workflows/ci-wasm.yml new file mode 100644 index 00000000..a81c59dd --- /dev/null +++ b/.github/workflows/ci-wasm.yml @@ -0,0 +1,96 @@ +name: ci-wasm + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +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@v4 + + # 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. + MATURIN_PEP517_ARGS: --no-default-features + 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 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], + }, + ).chunk({"time": 2}) + con = duckdb.connect() + xql.register(con, "air", ds) + rel = con.sql("SELECT time, lat, air FROM air ORDER BY time, lat") + xr.testing.assert_identical(xql.to_dataset(rel, template=ds), ds.compute()) + assert not hasattr(xql, "Context") + `); + console.log(`${wheel} round-trips through DuckDB in Pyodide ${py.version}`); + EOF + node smoke.mjs + + - uses: actions/upload-artifact@v4 + with: + name: wheels-pyodide + path: dist/*.whl diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b7f1722b..ba951dd6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -59,6 +59,29 @@ 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`. You need a Python interpreter with the same +version as that Pyodide release (for example, Python 3.14 for Pyodide 314). + +```shell +pip install 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 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 6a70b0c9..2e38e61b 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 c8bea272..6b047af3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,21 +19,27 @@ 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-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. +flight = ["dep:arrow-flight", "dep:prost", "dep:tonic", "tokio/net", "tokio/rt-multi-thread"] [build-dependencies] diff --git a/README.md b/README.md index 0fa10528..42fe3ca0 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,7 +110,7 @@ ds = xr.open_zarr( storage_options={'token': 'anon'} # Anonymous read from the public GCS bucket — no auth required. ) -ctx = xql.XarrayContext() +ctx = xql.Context() # Make sure to pass `chunks`! ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={ ('time', 'latitude', 'longitude'): 'surface', diff --git a/docs/engines.md b/docs/engines.md index cc6e9b26..8a17d1a9 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() ``` @@ -349,7 +349,7 @@ round-trips eagerly, and `chunks=` needs `spill=True`. 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) | diff --git a/docs/examples.md b/docs/examples.md index f2b925b7..7fb02c4a 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(''' @@ -76,7 +76,7 @@ full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) # across the year (terabytes from GCS). Always SELECT specific columns. ds = full.sel(time='2020').chunk({'time': 1}) -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 2ae86258..ca9b4663 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/performance.md b/docs/performance.md index 349589d7..c461de93 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 1bbdf65b..c393dc34 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 75365869..9d60a771 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,10 @@ classifiers = [ ] 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", ] diff --git a/src/lib.rs b/src/lib.rs index 61ab2c03..82458a62 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_sql.py b/tests/test_sql.py index 78042cfa..57d69b9b 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/uv.lock b/uv.lock index 614b82c0..e8dbe11e 100644 --- a/uv.lock +++ b/uv.lock @@ -2835,7 +2835,8 @@ 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'" }, ] @@ -2895,12 +2896,13 @@ requires-dist = [ { 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 = "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'" }, diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py index 28672586..af1a85dd 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 b86dd8a0..227228bd 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. diff --git a/xarray_sql/backends/datafusion.py b/xarray_sql/backends/datafusion.py index 4fc652f0..5f9660d8 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 294cae6b..98670d0f 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 diff --git a/xarray_sql/backends/pyarrow.py b/xarray_sql/backends/pyarrow.py index 9649dffd..e4eaa715 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. diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 27a6c94c..96591303 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 f875c98b..f256db2b 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,9 @@ from . import cftime as cft +_CAN_START_THREADS = sys.platform != "emscripten" +"""Whether this interpreter can start OS threads; Pyodide cannot.""" + Block = dict[Hashable, slice] Chunks = dict[str, int] | None TableNames = Mapping[tuple[str, ...], str] | None diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index e45b4d23..1c699280 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 @@ -1012,7 +1012,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 +1022,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 +1034,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 diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py index e8f36fc2..62589a61 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 diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py index 2470d1b8..ea63703b 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/sql.py b/xarray_sql/sql.py index c9911c12..cc07e055 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: @@ -190,3 +193,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.""" From e1948c1f12b88ed71a3be797e7e4a14fd48ba863 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 9 Oct 2026 17:20:44 -0700 Subject: [PATCH 2/9] Publish the Pyodide wheel to PyPI on release PyPI accepts pyemscripten wheels under PEP 783, so releases can ship the wasm wheel alongside the native ones and Pyodide users can install with `micropip.install("xarray-sql")`. publish.yml calls ci-wasm.yml as a reusable workflow, so every released wasm wheel has passed the Pyodide DuckDB smoke test. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- .github/workflows/ci-wasm.yml | 8 +++++--- .github/workflows/publish.yml | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci-wasm.yml b/.github/workflows/ci-wasm.yml index a81c59dd..541ce4c3 100644 --- a/.github/workflows/ci-wasm.yml +++ b/.github/workflows/ci-wasm.yml @@ -6,6 +6,8 @@ on: 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 }} @@ -21,7 +23,7 @@ jobs: name: Pyodide wheel runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 # pyodide-build requires the host Python to match Pyodide's Python. - uses: actions/setup-python@v5 @@ -90,7 +92,7 @@ jobs: EOF node smoke.mjs - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v6 with: - name: wheels-pyodide + name: dist-pyodide path: dist/*.whl diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index ab567ebd..4a07fc68 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 @@ -161,6 +167,7 @@ jobs: - build-artifacts-mac-win - build-artifacts-manylinux-x86_64 - build-artifacts-manylinux-arm64 + - build-artifacts-pyodide - build-sdist steps: - name: Merge Build Artifacts From a3aac49577b6a0cb48a1f4636f91bf9cce869939 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 9 Oct 2026 17:48:12 -0700 Subject: [PATCH 3/9] Upload the Pyodide wheel after the native release files PyPI's wasm wheels so far are all cp314-cp314; an abi3 one has not been tried. Publishing it in its own step, after the native wheels and sdist, keeps a rejection from leaving a partial release. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- .github/workflows/ci-wasm.yml | 2 +- .github/workflows/publish.yml | 13 +++++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci-wasm.yml b/.github/workflows/ci-wasm.yml index 541ce4c3..19773deb 100644 --- a/.github/workflows/ci-wasm.yml +++ b/.github/workflows/ci-wasm.yml @@ -94,5 +94,5 @@ jobs: - uses: actions/upload-artifact@v6 with: - name: dist-pyodide + name: pyodide-dist path: dist/*.whl diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 4a07fc68..20a8158b 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -167,7 +167,6 @@ jobs: - build-artifacts-mac-win - build-artifacts-manylinux-x86_64 - build-artifacts-manylinux-arm64 - - build-artifacts-pyodide - build-sdist steps: - name: Merge Build Artifacts @@ -200,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 @@ -218,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/* From 2dd827f4a3d80d27dea028ae2a1a5102d772fad5 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 9 Oct 2026 19:25:23 -0700 Subject: [PATCH 4/9] Shrink the Pyodide wheel from 10.8 MB to 4.3 MB Pyodide downloads the wheel on every page load. Without the Flight SQL server, the extension needs only DataFusion's core table-provider APIs, so DataFusion's default features (SQL, function libraries, Parquet, compression) now come with the `flight` feature, which native builds keep on by default. A `wasm` cargo profile optimizes for size (opt-level z, fat LTO, one codegen unit, stripped), taking the module from 44.3 MB to 15.1 MB uncompressed. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- .github/workflows/ci-wasm.yml | 5 +++-- CONTRIBUTING.md | 2 +- Cargo.toml | 15 +++++++++++++-- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/.github/workflows/ci-wasm.yml b/.github/workflows/ci-wasm.yml index 19773deb..32f21519 100644 --- a/.github/workflows/ci-wasm.yml +++ b/.github/workflows/ci-wasm.yml @@ -47,8 +47,9 @@ jobs: - name: Build env: - # The Flight SQL server needs sockets and threads. - MATURIN_PEP517_ARGS: --no-default-features + # 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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ba951dd6..8ddf9bcf 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -76,7 +76,7 @@ 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 pyodide build -o dist +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 diff --git a/Cargo.toml b/Cargo.toml index 6b047af3..231d2b64 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,7 +22,7 @@ arrow = { version = "58", features = ["pyarrow"] } 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" @@ -39,9 +39,20 @@ tonic = { version = "0.14", features = ["transport"], optional = true } default = ["flight"] # The Arrow Flight SQL server needs TCP sockets and OS threads, neither of # which exist under WebAssembly. Pyodide builds disable default features. -flight = ["dep:arrow-flight", "dep:prost", "dep:tonic", "tokio/net", "tokio/rt-multi-thread"] +# 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" From 6811d953f8a3069ce45fb72c085294d7dc871229 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 9 Oct 2026 19:58:02 -0700 Subject: [PATCH 5/9] Make dask an optional dependency Registration resolves chunk boundaries arithmetically, so dask was only needed to call .chunk() on a Dataset before registering it. An unchunked Dataset (in memory, or opened lazily with chunks=None) now registers along the storage chunks its backend recorded in encoding["preferred_chunks"], and as one chunk along other dimensions, instead of raising. dask moves to a `dask` extra (included in `test`), which lazy, chunked results from to_dataset(chunks=...) still need. Pyodide installs no longer download and import it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- .github/workflows/ci-wasm.yml | 9 ++++++--- README.md | 3 ++- docs/examples.md | 8 ++++---- pyproject.toml | 7 +++++-- tests/test_arrow_dataset.py | 10 ++++++++++ tests/test_df.py | 6 ++---- tests/test_duckdb_backend.py | 11 ++++++++--- uv.lock | 11 +++++++---- xarray_sql/backends/base.py | 6 +++++- xarray_sql/backends/flight.py | 3 ++- xarray_sql/backends/pyarrow.py | 9 +++------ xarray_sql/df.py | 31 ++++++++++++++++++++++--------- xarray_sql/reader.py | 9 ++++++--- xarray_sql/sql.py | 3 ++- 14 files changed, 84 insertions(+), 42 deletions(-) diff --git a/.github/workflows/ci-wasm.yml b/.github/workflows/ci-wasm.yml index 32f21519..d004b1b9 100644 --- a/.github/workflows/ci-wasm.yml +++ b/.github/workflows/ci-wasm.yml @@ -70,6 +70,8 @@ jobs: import micropip await micropip.install(["emfs:/tmp/${wheel}", "duckdb"]) + import importlib.util + import duckdb import numpy as np import pandas as pd @@ -82,12 +84,13 @@ jobs: "time": pd.date_range("2020-01-01", periods=4), "lat": [10.0, 20.0, 30.0], }, - ).chunk({"time": 2}) + ) con = duckdb.connect() - xql.register(con, "air", ds) + 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.compute()) + 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 diff --git a/README.md b/README.md index 42fe3ca0..1bc7b20c 100644 --- a/README.md +++ b/README.md @@ -111,7 +111,8 @@ ds = xr.open_zarr( ) ctx = xql.Context() -# Make sure to pass `chunks`! +# `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/examples.md b/docs/examples.md index 7fb02c4a..46876689 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -66,15 +66,15 @@ 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.Context() ctx.from_dataset('era5', ds, table_names={ diff --git a/pyproject.toml b/pyproject.toml index 9d60a771..42d64660 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,6 @@ classifiers = [ "Topic :: Database :: Front-Ends", ] dependencies = [ - "dask>=2024.8.0", # 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'", @@ -39,6 +38,10 @@ dependencies = [ ] [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", @@ -59,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/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 34aebc72..cc70d6fc 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -125,6 +125,16 @@ 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_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 099e3f73..555b8e62 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_duckdb_backend.py b/tests/test_duckdb_backend.py index e592cfc5..2723a6ec 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/uv.lock b/uv.lock index e8dbe11e..c62d5b07 100644 --- a/uv.lock +++ b/uv.lock @@ -2834,7 +2834,6 @@ io = [ name = "xarray-sql" source = { editable = "." } dependencies = [ - { name = "dask" }, { 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'" }, @@ -2845,6 +2844,9 @@ dependencies = [ adbc = [ { name = "adbc-driver-manager" }, ] +dask = [ + { name = "dask" }, +] dev = [ { name = "mkdocstrings", extra = ["python"] }, { name = "pre-commit" }, @@ -2871,6 +2873,7 @@ test = [ { name = "adbc-driver-manager" }, { name = "adbc-driver-sqlite" }, { name = "cftime" }, + { name = "dask" }, { name = "duckdb" }, { name = "gcsfs" }, { name = "polars" }, @@ -2895,7 +2898,7 @@ 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 = "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'" }, @@ -2909,11 +2912,11 @@ requires-dist = [ { 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/backends/base.py b/xarray_sql/backends/base.py index 227228bd..a733db66 100644 --- a/xarray_sql/backends/base.py +++ b/xarray_sql/backends/base.py @@ -120,7 +120,11 @@ 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 (in memory, or 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; 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/flight.py b/xarray_sql/backends/flight.py index 98670d0f..a1adbac2 100644 --- a/xarray_sql/backends/flight.py +++ b/xarray_sql/backends/flight.py @@ -120,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 e4eaa715..7b9c11d1 100644 --- a/xarray_sql/backends/pyarrow.py +++ b/xarray_sql/backends/pyarrow.py @@ -343,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() } @@ -1105,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). @@ -1189,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/df.py b/xarray_sql/df.py index f256db2b..f5472791 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -67,8 +67,10 @@ 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 (in memory, + or opened lazily without a chunk manager) is split along the storage + chunks its backend recorded, and is one chunk along other dimensions. Returns an empty mapping for scalar datasets; callers should treat that as "one block covering everything". @@ -77,7 +79,22 @@ 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)) + + +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 _ensure_default_indexes(ds: xr.Dataset) -> xr.Dataset: @@ -103,12 +120,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/reader.py b/xarray_sql/reader.py index 153ed8c3..d2c78c4a 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/sql.py b/xarray_sql/sql.py index cc07e055..9bdf1eb7 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -103,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. From 1c651cea7f5f4c3ef87ce0ddf0715abd78d6ccd9 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 9 Oct 2026 20:28:39 -0700 Subject: [PATCH 6/9] Add lazy=True to to_dataset for laziness without a chunk manager to_dataset(lazy=True) without chunks= returns the lazily indexed SQLBackendArray variables that the chunked path builds, without wrapping them in a chunk manager. Each access re-executes the query narrowed to the selection, so it reads only the source chunks under it, and dask is not needed. Because every access runs on the calling thread, DuckDB relations (which refuse chunked reconstruction) work too. lazy defaults to None, which keeps today's behavior: lazy exactly when chunks are given. lazy=False with an explicit chunking raises. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- docs/engines.md | 13 +++++++++ tests/test_ds.py | 13 +++++++++ tests/test_lazy_roundtrip.py | 23 +++++++++++++++ xarray_sql/ds.py | 30 ++++++++++++++++++-- xarray_sql/roundtrip.py | 55 ++++++++++++++++++++++++++++++------ 5 files changed, 123 insertions(+), 11 deletions(-) diff --git a/docs/engines.md b/docs/engines.md index 8a17d1a9..288e5649 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -515,6 +515,19 @@ 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. +**Lazy without a chunk manager:** `lazy=True` without `chunks=` returns +lazily indexed variables, the way `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, so it reads only the source chunks +under it. 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, lazy=True) +out.t2m.sel(time="2020-01-01").mean().item() # reads one day's chunks +``` + Two knobs matter at scale: - `coords="template"` trusts the template's coordinate arrays instead of diff --git a/tests/test_ds.py b/tests/test_ds.py index 3682dacf..c9ca08bb 100644 --- a/tests/test_ds.py +++ b/tests/test_ds.py @@ -350,6 +350,19 @@ def test_chunks_argument_controls_partitioning(synthetic_dataset): ) +def test_lazy_without_chunks_matches_eager(synthetic_dataset): + ctx = XarrayContext() + ctx.from_dataset("t", synthetic_dataset) + out = ctx.sql("SELECT * FROM t").to_dataset(chunks=None, lazy=True) + assert not out.chunks + xr.testing.assert_allclose( + out.load().sortby(["time", "lat", "lon"]), + synthetic_dataset.compute().sortby(["time", "lat", "lon"]), + ) + with pytest.raises(ValueError, match="lazy=False"): + ctx.sql("SELECT * FROM t").to_dataset(chunks={"time": 2}, lazy=False) + + 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_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py index e2a7db54..b511e3d5 100644 --- a/tests/test_lazy_roundtrip.py +++ b/tests/test_lazy_roundtrip.py @@ -61,6 +61,29 @@ 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") + out = xql.to_dataset(rel, template=source, lazy=True) + assert not out.chunks # no chunk manager involved + reads.clear() + 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_lazy_false_with_chunks_raises(source, registered): + con, _ = registered + with pytest.raises(ValueError, match="lazy=False"): + xql.to_dataset( + con.sql("SELECT * FROM t"), + template=source, + chunks={"time": 10}, + lazy=False, + ) + + def test_duckdb_eager_round_trip_through_handle(source, registered): con, _ = registered rel = con.sql("SELECT * FROM t") diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index 1c699280..9f23f3a1 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -915,15 +915,20 @@ def _result_to_xarray( fill_value: Any, chunks: Mapping[str, int] | str | None, templates: dict[str, xr.Dataset] | None = None, + lazy: bool = False, ) -> xr.Dataset: """Reconstruct an ``xr.Dataset`` from a SQL result. ``chunks`` (already resolved by ``XarrayDataFrame._resolve_chunks``) - selects the execution strategy: + and ``lazy`` select 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`` with ``lazy=True`` -> 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. No + chunk manager is involved. * 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,7 +949,7 @@ 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: + if chunks is None and not lazy: ds = _materialize(inner_df, dimension_columns, field_names, field_types) else: ds = _build_lazy_scan( @@ -1056,6 +1061,7 @@ def to_dataset( sparsity: Sparsity = "result", fill_value: Any = np.nan, chunks: Mapping[str, int] | str | None = "inherit", + lazy: bool | None = None, ) -> xr.Dataset: """Convert the result to an ``xr.Dataset``. @@ -1098,6 +1104,15 @@ def to_dataset( union of source partitions. Useful for finely partitioned sources (e.g. ERA5 ``chunks={"time": 1}``), coarsening many tiny partitions into memory-sized, aligned chunks. + lazy: Whether to defer execution until the data is accessed. + ``None`` (default) is lazy exactly when ``chunks`` resolves + to a chunking. ``True`` without chunking returns lazily + indexed variables: each access, such as + ``out.t2m.sel(time="2020-01").values``, re-executes the query + narrowed to that selection. This needs no chunk manager + (dask is not required). ``sparsity="template"`` loads the + data to reindex it. ``False`` is always eager and cannot be + combined with an explicit chunking. Returns: An ``xr.Dataset`` with ``dims`` as dimensions and the @@ -1115,7 +1130,15 @@ def to_dataset( template = self._resolve_template(template) if dims is None: dims = self._infer_dimension_columns(preferred_template=template) - resolved_chunks = self._resolve_chunks(chunks, template, dims) + if lazy is False: + if chunks not in (None, "inherit"): + raise ValueError( + "lazy=False cannot be combined with chunks=" + f"{chunks!r}: chunked output is always lazy." + ) + resolved_chunks = None + else: + resolved_chunks = self._resolve_chunks(chunks, template, dims) return _result_to_xarray( inner_df=self._inner, dimension_columns=dims, @@ -1124,6 +1147,7 @@ def to_dataset( fill_value=fill_value, chunks=resolved_chunks, templates=self._templates, + lazy=bool(lazy), ) # ------------------------------------------------------------------ diff --git a/xarray_sql/roundtrip.py b/xarray_sql/roundtrip.py index a592ccc2..d3772120 100644 --- a/xarray_sql/roundtrip.py +++ b/xarray_sql/roundtrip.py @@ -166,6 +166,7 @@ def to_dataset( coords: Literal["discover", "template"] = "discover", max_result_bytes: int | None = None, spill: bool | str | os.PathLike = False, + lazy: bool | None = None, ) -> xr.Dataset: """Convert an engine's Arrow result into a labeled ``xr.Dataset``. @@ -238,11 +239,23 @@ 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=``. + lazy: Whether to defer execution until the data is accessed. + ``None`` (default) is lazy exactly when ``chunks`` is given. + ``True`` without ``chunks`` returns lazily indexed variables: + each access, such as ``out.t2m.sel(time="2020-01").values``, + re-executes the query narrowed to that selection. This needs + no chunk manager (dask is not required) and, unlike + ``chunks=``, works with DuckDB relations, since every access + runs on the calling thread. Requires a re-executable + ``result``, as ``chunks=`` does. ``sparsity="template"`` + loads the data to reindex it. ``False`` is always eager and + cannot be combined 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. + default, lazily indexed with ``lazy=True``, and lazily chunked + when ``chunks`` is given. Raises: ValueError: When neither ``dims`` nor ``template`` resolves the @@ -268,6 +281,11 @@ def to_dataset( raise ValueError( "spill= only applies to chunked reconstruction; pass chunks=." ) + if lazy is False and chunks is not None: + raise ValueError( + "lazy=False cannot be combined with chunks=: chunked output is " + "always lazy." + ) if chunks is not None: if spill: @@ -284,6 +302,10 @@ def to_dataset( return _to_dataset_lazy( result, dims, template, sparsity, fill_value, chunks, coords ) + if lazy: + return _to_dataset_lazy( + result, dims, template, sparsity, fill_value, None, coords + ) schema, batches = _result_to_batches(result, max_result_bytes) field_names = [f.name for f in schema] @@ -367,20 +389,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=...)``. + + ``chunks=None`` (from ``lazy=True``) 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 +425,16 @@ 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. + 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 +451,10 @@ 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 lazy=True without chunks= (each access runs on the " + "calling thread), add spill=True to reconstruct from a " + "one-pass on-disk spill, use chunks=None (eager), or run the " + "query through " "Polars (pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))) " "or a DataFusion context." ) From 16c6020e62b99806888b577b12ae9bac107f6a55 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 9 Oct 2026 21:01:04 -0700 Subject: [PATCH 7/9] Make to_dataset(chunks=None) lazy without a chunk manager Follow xarray's convention, where chunks=None turns dask off rather than reading eagerly: to_dataset(chunks=None), the default, returns lazily indexed variables, and .load() reads them. This replaces the lazy= argument. Context.sql(...).to_dataset() falls back to the same when "inherit" finds nothing to inherit. Two changes keep the default from re-running queries: - When coordinates must be discovered, the result is first streamed once and kept in memory if it is small (up to 64 MiB), so an aggregation still executes once. DuckDBHandle gains stream() for this. - A full-extent read (.load()) fetches every result column once and shares it among the variables; windowed reads still fetch only their own column. One-shot Arrow streams, and calls with max_result_bytes=, are still read into memory directly. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- docs/engines.md | 30 ++++-- docs/limitations.md | 13 +-- tests/test_ds.py | 6 +- tests/test_lazy_roundtrip.py | 26 +++-- tests/test_to_dataset_perf.py | 6 +- xarray_sql/ds.py | 190 +++++++++++++++++++++------------- xarray_sql/lazyscan.py | 9 ++ xarray_sql/roundtrip.py | 73 +++++++------ 8 files changed, 217 insertions(+), 136 deletions(-) diff --git a/docs/engines.md b/docs/engines.md index 288e5649..673f4661 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -342,7 +342,7 @@ 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) @@ -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
(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,19 +517,27 @@ 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. -**Lazy without a chunk manager:** `lazy=True` without `chunks=` returns -lazily indexed variables, the way `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, so it reads only the source chunks -under it. 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. +**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, so it reads only the source chunks under it, and +`.load()` reads the whole result in one execution. 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, lazy=True) +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 ``` +Discovering a lazy result's coordinates executes the query once per +dimension, so a result whose coordinates must be discovered is first +streamed in one pass and kept in memory if it is small (up to 64 MiB). +An aggregation, which no window filter can push below, therefore +executes once. One-shot Arrow streams, which cannot re-execute, and +calls with `max_result_bytes=` are read into memory directly. + Two knobs matter at scale: - `coords="template"` trusts the template's coordinate arrays instead of diff --git a/docs/limitations.md b/docs/limitations.md index 54473e22..50dc16ba 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/tests/test_ds.py b/tests/test_ds.py index c9ca08bb..72dba124 100644 --- a/tests/test_ds.py +++ b/tests/test_ds.py @@ -350,17 +350,15 @@ def test_chunks_argument_controls_partitioning(synthetic_dataset): ) -def test_lazy_without_chunks_matches_eager(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").to_dataset(chunks=None, lazy=True) + out = ctx.sql("SELECT * FROM t").to_dataset(chunks=None) assert not out.chunks xr.testing.assert_allclose( out.load().sortby(["time", "lat", "lon"]), synthetic_dataset.compute().sortby(["time", "lat", "lon"]), ) - with pytest.raises(ValueError, match="lazy=False"): - ctx.sql("SELECT * FROM t").to_dataset(chunks={"time": 2}, lazy=False) def test_chunks_auto_snaps_to_source_partitions(): diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py index b511e3d5..13721cee 100644 --- a/tests/test_lazy_roundtrip.py +++ b/tests/test_lazy_roundtrip.py @@ -64,24 +64,28 @@ def test_duckdb_chunked_fails_fast_with_guidance(source, registered): def test_duckdb_lazy_reads_only_the_accessed_window(source, registered): con, reads = registered rel = con.sql("SELECT * FROM t") - out = xql.to_dataset(rel, template=source, lazy=True) + # 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 - reads.clear() + 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_lazy_false_with_chunks_raises(source, registered): - con, _ = registered - with pytest.raises(ValueError, match="lazy=False"): - xql.to_dataset( - con.sql("SELECT * FROM t"), - template=source, - chunks={"time": 10}, - lazy=False, - ) +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_eager_round_trip_through_handle(source, registered): diff --git a/tests/test_to_dataset_perf.py b/tests/test_to_dataset_perf.py index 9124e9c4..2878d1a2 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/xarray_sql/ds.py b/xarray_sql/ds.py index 9f23f3a1..bdfbfbff 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -30,6 +30,7 @@ from __future__ import annotations +import threading import warnings from collections.abc import Mapping from typing import Any, Literal, cast @@ -705,27 +706,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 +779,83 @@ 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_small_result( + handle: LazyResultHandle, columns: list[str] +) -> list[pa.RecordBatch] | None: + """The whole result if it fits in ``_ONE_PASS_BYTES``, else ``None``. + + Discovering a lazy result's coordinates executes the query once per + dimension, and loading it executes the query again. For a small + result -- an aggregation, which no window filter can push below, or a + narrow selection -- one streamed pass is cheaper than both. A result + that outgrows the budget stops streaming and stays lazy. + """ + stream = getattr(handle, "stream", None) + if stream is None: + return None + batches: list[pa.RecordBatch] = [] + total = 0 + it = stream(columns) + try: + for batch in it: + total += batch.nbytes + if total > _ONE_PASS_BYTES: + return None + batches.append(batch) + finally: + close = getattr(it, "close", None) + if close is not None: + close() + return batches + + +class _SharedFullFetch: + """Shares one full-extent fetch among a lazy Dataset's variables. + + Loading a whole lazy Dataset (``.load()``, or ``.values`` on every + variable) reads each variable at full extent. Fetched separately, that + executes the query once per variable -- for an aggregation, once per + variable over the whole source. Instead, the first full-extent read + fetches every result column, and the other variables take theirs from + it; the batches are released once every variable has. Windowed reads + still fetch only their own column, so projection pushdown is kept. + """ + + def __init__( + self, handle: LazyResultHandle, columns: list[str], variables: list[str] + ) -> None: + self._handle = handle + self._columns = columns + self._variables = frozenset(variables) + self._pending: set[str] = set() + self._batches: list[pa.RecordBatch] | None = None + self._lock = threading.Lock() + + def __getattr__(self, name: str) -> Any: + return getattr(self._handle, name) + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + if specs: + return self._handle.fetch(specs, columns) + (variable,) = (c for c in columns if c in self._variables) + with self._lock: + if self._batches is None: + self._batches = self._handle.fetch({}, self._columns) + self._pending = set(self._variables) + batches = self._batches + self._pending.discard(variable) + if not self._pending: + self._batches = None + return batches + + def _build_lazy_scan( handle: LazyResultHandle, dimension_columns: list[str], @@ -808,9 +865,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 @@ -824,6 +881,8 @@ def _build_lazy_scan( # positional slices map onto contiguous value ranges. coord_arrays[d] = np.sort(handle.distinct(d)) shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + variables = [n for n in field_names if n not in dimension_columns] + shared = _SharedFullFetch(handle, dimension_columns + variables, variables) data_vars: dict[str, xr.Variable] = {} for name in field_names: @@ -831,7 +890,7 @@ def _build_lazy_scan( continue np_dtype = field_types[name].to_pandas_dtype() backend = SQLBackendArray( - handle=handle, + handle=shared, var_name=name, dimension_columns=dimension_columns, coord_arrays=coord_arrays, @@ -915,20 +974,19 @@ def _result_to_xarray( fill_value: Any, chunks: Mapping[str, int] | str | None, templates: dict[str, xr.Dataset] | None = None, - lazy: bool = False, ) -> xr.Dataset: """Reconstruct an ``xr.Dataset`` from a SQL result. ``chunks`` (already resolved by ``XarrayDataFrame._resolve_chunks``) - and ``lazy`` select 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`` with ``lazy=True`` -> 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. No - chunk manager is involved. + selects the execution strategy: + + * ``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 ``.load()`` + executes it once for every variable. No chunk manager is involved. + A result whose coordinates must be discovered is first streamed in + one pass and kept in memory if small (``_read_small_result``), 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 @@ -949,17 +1007,24 @@ 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 and not lazy: - 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 = _read_small_result(handle, field_names) + 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, @@ -1061,7 +1126,6 @@ def to_dataset( sparsity: Sparsity = "result", fill_value: Any = np.nan, chunks: Mapping[str, int] | str | None = "inherit", - lazy: bool | None = None, ) -> xr.Dataset: """Convert the result to an ``xr.Dataset``. @@ -1090,11 +1154,13 @@ 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; ``.load()`` executes + it once for all variables. Dask is not required. * 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, @@ -1104,15 +1170,6 @@ def to_dataset( union of source partitions. Useful for finely partitioned sources (e.g. ERA5 ``chunks={"time": 1}``), coarsening many tiny partitions into memory-sized, aligned chunks. - lazy: Whether to defer execution until the data is accessed. - ``None`` (default) is lazy exactly when ``chunks`` resolves - to a chunking. ``True`` without chunking returns lazily - indexed variables: each access, such as - ``out.t2m.sel(time="2020-01").values``, re-executes the query - narrowed to that selection. This needs no chunk manager - (dask is not required). ``sparsity="template"`` loads the - data to reindex it. ``False`` is always eager and cannot be - combined with an explicit chunking. Returns: An ``xr.Dataset`` with ``dims`` as dimensions and the @@ -1130,15 +1187,7 @@ def to_dataset( template = self._resolve_template(template) if dims is None: dims = self._infer_dimension_columns(preferred_template=template) - if lazy is False: - if chunks not in (None, "inherit"): - raise ValueError( - "lazy=False cannot be combined with chunks=" - f"{chunks!r}: chunked output is always lazy." - ) - resolved_chunks = None - else: - resolved_chunks = self._resolve_chunks(chunks, template, dims) + resolved_chunks = self._resolve_chunks(chunks, template, dims) return _result_to_xarray( inner_df=self._inner, dimension_columns=dims, @@ -1147,7 +1196,6 @@ def to_dataset( fill_value=fill_value, chunks=resolved_chunks, templates=self._templates, - lazy=bool(lazy), ) # ------------------------------------------------------------------ @@ -1162,11 +1210,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 62589a61..77cf9d2c 100644 --- a/xarray_sql/lazyscan.py +++ b/xarray_sql/lazyscan.py @@ -275,6 +275,15 @@ 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.""" + import duckdb + + rel = self._rel.project(*(duckdb.ColumnExpression(n) for n in columns)) + batches = self._run(lambda: iter(self._to_arrow_reader(rel))) + while (batch := self._run(lambda: next(batches, None))) is not None: + yield batch + def spill_parquet(self, path: str) -> None: def run() -> None: reader = self._to_arrow_reader(self._rel) diff --git a/xarray_sql/roundtrip.py b/xarray_sql/roundtrip.py index d3772120..bdd4469e 100644 --- a/xarray_sql/roundtrip.py +++ b/xarray_sql/roundtrip.py @@ -41,6 +41,7 @@ _build_lazy_scan, _dataset_from_batches, _finish_dataset, + _read_small_result, _result_dims, ) from .lazyscan import LazyResultHandle, PolarsHandle, resolve_lazy_handle @@ -166,7 +167,6 @@ def to_dataset( coords: Literal["discover", "template"] = "discover", max_result_bytes: int | None = None, spill: bool | str | os.PathLike = False, - lazy: bool | None = None, ) -> xr.Dataset: """Convert an engine's Arrow result into a labeled ``xr.Dataset``. @@ -203,7 +203,14 @@ 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 + ``.load()`` executes it once for all variables. Dask is not + required, and DuckDB relations work, since every access runs + on the calling thread. 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 @@ -219,7 +226,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 — @@ -239,23 +247,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=``. - lazy: Whether to defer execution until the data is accessed. - ``None`` (default) is lazy exactly when ``chunks`` is given. - ``True`` without ``chunks`` returns lazily indexed variables: - each access, such as ``out.t2m.sel(time="2020-01").values``, - re-executes the query narrowed to that selection. This needs - no chunk manager (dask is not required) and, unlike - ``chunks=``, works with DuckDB relations, since every access - runs on the calling thread. Requires a re-executable - ``result``, as ``chunks=`` does. ``sparsity="template"`` - loads the data to reindex it. ``False`` is always eager and - cannot be combined 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 indexed with ``lazy=True``, and 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 @@ -281,11 +277,6 @@ def to_dataset( raise ValueError( "spill= only applies to chunked reconstruction; pass chunks=." ) - if lazy is False and chunks is not None: - raise ValueError( - "lazy=False cannot be combined with chunks=: chunked output is " - "always lazy." - ) if chunks is not None: if spill: @@ -302,10 +293,19 @@ def to_dataset( return _to_dataset_lazy( result, dims, template, sparsity, fill_value, chunks, coords ) - if lazy: - return _to_dataset_lazy( - result, dims, template, sparsity, fill_value, None, 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] @@ -427,10 +427,21 @@ def _to_dataset_lazy( if chunks is None: # Lazily indexed and unchunked: every access runs on the calling - # thread, so even handles that refuse chunking can serve it. - ds = _build_lazy_scan( - handle, dims, field_names, field_types, coord_arrays=coord_arrays - ) + # 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 = _read_small_result(handle, field_names) + 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 ) From 0a05aa1123901b60a91480b21c5783fec05c3b04 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 10 Oct 2026 09:36:43 -0700 Subject: [PATCH 8/9] Address review of lazy chunks=None reconstruction - Drop the shared full-extent fetch: it fetched every column for a one-variable read, cached the rest indefinitely, and held a lock that broke pickling of lazy and chunked results. Each variable read is now projected to that variable, as on the chunked path. - Narrow widened dtypes only for in-memory variables, without loading a lazy one through Variable.data. - Discover a large result's coordinates in the same streamed pass that checks the one-pass budget, instead of discarding it and running one DISTINCT query per dimension. - Close DuckDB's streaming reader on the engine thread. - Split in-memory Datasets with no storage chunks into partitions of about a million rows rather than one partition. - Fix stale lazy=True references and document that lazy results need the connection open and see later changes to the source. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- docs/engines.md | 36 ++++++---- tests/test_arrow_dataset.py | 14 ++++ tests/test_ds.py | 9 ++- tests/test_lazy_roundtrip.py | 47 +++++++++++++ xarray_sql/backends/base.py | 9 +-- xarray_sql/df.py | 30 +++++++-- xarray_sql/ds.py | 126 +++++++++++++++-------------------- xarray_sql/lazyscan.py | 16 +++-- xarray_sql/roundtrip.py | 51 +++++++------- 9 files changed, 211 insertions(+), 127 deletions(-) diff --git a/docs/engines.md b/docs/engines.md index 673f4661..ab1e9658 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -502,7 +502,7 @@ only the source chunks it maps onto. flowchart TB R["xql.to_dataset(result, ...)"] --> K{"chunks=?"} 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
(small results: read in one pass)"] + 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"} @@ -520,23 +520,35 @@ 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, so it reads only the source chunks under it, and -`.load()` reads the whole result in one execution. 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. +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 ``` -Discovering a lazy result's coordinates executes the query once per -dimension, so a result whose coordinates must be discovered is first -streamed in one pass and kept in memory if it is small (up to 64 MiB). -An aggregation, which no window filter can push below, therefore -executes once. One-shot Arrow streams, which cannot re-execute, and -calls with `max_result_bytes=` are read into memory directly. +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: diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index cc70d6fc..cf2fb4a7 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -135,6 +135,20 @@ def test_unchunked_zarr_scans_by_storage_chunks(ds, tmp_path): ) +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_ds.py b/tests/test_ds.py index 72dba124..57ce980f 100644 --- a/tests/test_ds.py +++ b/tests/test_ds.py @@ -353,12 +353,11 @@ 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").to_dataset(chunks=None) - assert not out.chunks - xr.testing.assert_allclose( - out.load().sortby(["time", "lat", "lon"]), - synthetic_dataset.compute().sortby(["time", "lat", "lon"]), + 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(): diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py index 13721cee..1815a7f7 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 @@ -88,6 +90,37 @@ def test_duckdb_small_result_scans_source_once(source, registered): 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") @@ -136,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/xarray_sql/backends/base.py b/xarray_sql/backends/base.py index a733db66..8c993228 100644 --- a/xarray_sql/backends/base.py +++ b/xarray_sql/backends/base.py @@ -121,10 +121,11 @@ def register( ds: An xarray Dataset. chunks: Xarray-like chunks specification controlling partition granularity. Defaults to the Dataset's existing chunks. An - unchunked Dataset (in memory, or 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; dask is not required. + 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/df.py b/xarray_sql/df.py index f5472791..d26da7f0 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -15,6 +15,9 @@ _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 @@ -68,9 +71,10 @@ 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 rechunks arithmetically via ``compute_chunks``. Without a - spec, a chunked dataset keeps its chunks; an unchunked one (in memory, - or opened lazily without a chunk manager) is split along the storage - chunks its backend recorded, and is one chunk along other dimensions. + 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". @@ -81,7 +85,7 @@ def resolve_chunks( return compute_chunks(ds, chunks) if ds.chunks: return {d: tuple(c) for d, c in ds.chunks.items()} - return compute_chunks(ds, _storage_chunks(ds)) + return compute_chunks(ds, _storage_chunks(ds) or _default_chunks(ds)) def _storage_chunks(ds: xr.Dataset) -> dict[str, int]: @@ -97,6 +101,24 @@ def _storage_chunks(ds: xr.Dataset) -> dict[str, int]: 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: """Attach a default integer index coordinate to every dimension lacking one. diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index bdfbfbff..6532d7da 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -30,7 +30,6 @@ from __future__ import annotations -import threading import warnings from collections.abc import Mapping from typing import Any, Literal, cast @@ -38,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 @@ -145,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 @@ -784,76 +785,51 @@ def _maybe_template_coords( read in a single pass instead.""" -def _read_small_result( - handle: LazyResultHandle, columns: list[str] -) -> list[pa.RecordBatch] | None: - """The whole result if it fits in ``_ONE_PASS_BYTES``, else ``None``. +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. - Discovering a lazy result's coordinates executes the query once per - dimension, and loading it executes the query again. For a small - result -- an aggregation, which no window filter can push below, or a - narrow selection -- one streamed pass is cheaper than both. A result - that outgrows the budget stops streaming and stays lazy. + 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 + return None, None batches: list[pa.RecordBatch] = [] total = 0 - it = stream(columns) - try: - for batch in it: + distinct: dict[str, pa.Array] | None = None + for batch in stream(columns): + if distinct is None: total += batch.nbytes - if total > _ONE_PASS_BYTES: - return None batches.append(batch) - finally: - close = getattr(it, "close", None) - if close is not None: - close() - return batches - - -class _SharedFullFetch: - """Shares one full-extent fetch among a lazy Dataset's variables. - - Loading a whole lazy Dataset (``.load()``, or ``.values`` on every - variable) reads each variable at full extent. Fetched separately, that - executes the query once per variable -- for an aggregation, once per - variable over the whole source. Instead, the first full-extent read - fetches every result column, and the other variables take theirs from - it; the batches are released once every variable has. Windowed reads - still fetch only their own column, so projection pushdown is kept. - """ - - def __init__( - self, handle: LazyResultHandle, columns: list[str], variables: list[str] - ) -> None: - self._handle = handle - self._columns = columns - self._variables = frozenset(variables) - self._pending: set[str] = set() - self._batches: list[pa.RecordBatch] | None = None - self._lock = threading.Lock() - - def __getattr__(self, name: str) -> Any: - return getattr(self._handle, name) - - def fetch( - self, specs: dict[str, DimSpec], columns: list[str] - ) -> list[pa.RecordBatch]: - if specs: - return self._handle.fetch(specs, columns) - (variable,) = (c for c in columns if c in self._variables) - with self._lock: - if self._batches is None: - self._batches = self._handle.fetch({}, self._columns) - self._pending = set(self._variables) - batches = self._batches - self._pending.discard(variable) - if not self._pending: - self._batches = None - return batches + 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( @@ -881,8 +857,6 @@ def _build_lazy_scan( # positional slices map onto contiguous value ranges. coord_arrays[d] = np.sort(handle.distinct(d)) shape = tuple(len(coord_arrays[d]) for d in dimension_columns) - variables = [n for n in field_names if n not in dimension_columns] - shared = _SharedFullFetch(handle, dimension_columns + variables, variables) data_vars: dict[str, xr.Variable] = {} for name in field_names: @@ -890,7 +864,7 @@ def _build_lazy_scan( continue np_dtype = field_types[name].to_pandas_dtype() backend = SQLBackendArray( - handle=shared, + handle=handle, var_name=name, dimension_columns=dimension_columns, coord_arrays=coord_arrays, @@ -982,10 +956,10 @@ def _result_to_xarray( * ``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 ``.load()`` - executes it once for every variable. No chunk manager is involved. - A result whose coordinates must be discovered is first streamed in - one pass and kept in memory if small (``_read_small_result``), so an + 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 @@ -1013,7 +987,9 @@ def _result_to_xarray( ) batches = None if chunks is None and coord_arrays is None: - batches = _read_small_result(handle, field_names) + 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 @@ -1159,8 +1135,10 @@ def to_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; ``.load()`` executes - it once for all variables. Dask is not required. + 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, diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py index 77cf9d2c..a967c16e 100644 --- a/xarray_sql/lazyscan.py +++ b/xarray_sql/lazyscan.py @@ -276,13 +276,21 @@ def fetch( ) def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: - """Execute once, yielding Arrow batches as DuckDB produces them.""" + """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)) - batches = self._run(lambda: iter(self._to_arrow_reader(rel))) - while (batch := self._run(lambda: next(batches, None))) is not None: - yield batch + 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: diff --git a/xarray_sql/roundtrip.py b/xarray_sql/roundtrip.py index bdd4469e..4f567580 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,7 +40,7 @@ _build_lazy_scan, _dataset_from_batches, _finish_dataset, - _read_small_result, + _read_or_discover, _result_dims, ) from .lazyscan import LazyResultHandle, PolarsHandle, resolve_lazy_handle @@ -205,12 +204,17 @@ def to_dataset( fill_value: Fill for ``sparsity="template"``. Defaults to NaN. 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 - ``.load()`` executes it once for all variables. Dask is not - required, and DuckDB relations work, since every access runs - on the calling thread. One-shot streams (Arrow tables, - readers, ADBC cursors) have nothing to re-execute and are - read eagerly. A mapping + 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 @@ -393,10 +397,10 @@ def _to_dataset_lazy( coords: Literal["discover", "template"], _handle: LazyResultHandle | None = None, ) -> xr.Dataset: - """The lazy reconstruction behind ``to_dataset(chunks=...)``. + """The lazy reconstruction behind ``to_dataset``. - ``chunks=None`` (from ``lazy=True``) returns lazily indexed variables - without chunking them. + ``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: @@ -431,7 +435,7 @@ def _to_dataset_lazy( # small result is cheaper to read once than to discover lazily. batches = None if coord_arrays is None: - batches = _read_small_result(handle, field_names) + 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: @@ -462,10 +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). " - "Use lazy=True without chunks= (each access runs on the " + "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, use chunks=None (eager), or run the " - "query through " + "one-pass on-disk spill, or run the query through " "Polars (pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))) " "or a DataFusion context." ) From 739137e87ff32a005feb174b6aa04cf35edb974c Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 10 Oct 2026 14:45:16 -0700 Subject: [PATCH 9/9] Install pyodide-build with uv in the wasm build instructions The pyodide command comes from pyodide-cli, so install it as a uv tool with pyodide-build alongside, on the Python version Pyodide targets. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- CONTRIBUTING.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8ddf9bcf..e5220821 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -67,11 +67,13 @@ 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`. You need a Python interpreter with the same -version as that Pyodide release (for example, Python 3.14 for Pyodide 314). +`.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 -pip install pyodide-build +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