diff --git a/.gitignore b/.gitignore index 92110e5..8999c64 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,5 @@ CHANGELOG.md *.ipynb /site .direnv +# Pyodide wheels copied in to test the browser demo locally +docs/demo/*.whl diff --git a/docs/browser.md b/docs/browser.md new file mode 100644 index 0000000..267fcf2 --- /dev/null +++ b/docs/browser.md @@ -0,0 +1,153 @@ +# Run in the browser + +xarray-sql runs in [Pyodide](https://pyodide.org), Python compiled to +WebAssembly. The demo below runs SQL over NOAA GFS forecasts read straight +from [dynamical.org](https://dynamical.org)'s public Zarr store. xarray-sql +and DuckDB run in a Web Worker in your browser, so there is no server. Each +query downloads only the chunks that its filters on `init_time`, `latitude`, +and `longitude` reach. + +The demo needs Chrome or Edge 137 or newer. It takes about 15 seconds to +start, most of it downloading and loading Pyodide, pyarrow, DuckDB, and +pandas. + + + +[Open the demo in its own page](demo/){ #demo-link target="_blank" } + + + +## Install in Pyodide + +Releases after 0.5.0 publish a Pyodide wheel to PyPI. The DuckDB engine +adapter works in Pyodide, including the round-trip back to a Dataset with +[to_dataset][xarray_sql.to_dataset]. DataFusion does not yet, because its +Python bindings have no Pyodide build, so [Context][xarray_sql.Context] is +unavailable there. + +Copy [`fetch_store.py`](https://github.com/xqlsystems/xarray-sql/tree/main/docs/demo/fetch_store.py) +from the demo to open Zarr stores over HTTP in Pyodide, since fsspec's HTTP +filesystem needs sockets. Then open a window of the dataset, read it into +memory, and query it: + +```python +import micropip +await micropip.install(["xarray-sql", "duckdb", "zarr"]) + +import duckdb +import xarray as xr +import xarray_sql as xql +from fetch_store import FetchStore + +store = FetchStore("https://data.dynamical.org/noaa/gfs/forecast/latest.zarr") +ds = ( + xr.open_zarr(store, chunks=None)[["temperature_2m"]] + .isel(init_time=-1, lead_time=slice(0, 24)) + .sel(latitude=slice(50, 25), longitude=slice(-100, -75)) +) +ds = await ds.load_async() + +con = duckdb.connect() +xql.register(con, "gfs", ds) +hottest = con.sql(""" + SELECT latitude, longitude, max(temperature_2m) AS t2m_max + FROM gfs + GROUP BY latitude, longitude +""") +out = xql.to_dataset(hottest, template=ds) # dims latitude, longitude +``` + +The window is read into memory first because a DuckDB scan cannot wait on +the network in Pyodide (see the next section). Pyodide also cannot start +threads, so xarray-sql loads chunks one at a time there. To query a larger +store, read it through a cache like the demo's. + +## How lazy reads work in the browser + +Browsers can only read over the network asynchronously. zarr's synchronous +API waits on those reads with Pyodide's JavaScript Promise Integration +(JSPI), but JSPI cannot suspend while DuckDB or pyarrow is calling back into +Python: their C++ frames sit between the suspension point and Python. So a +query cannot read chunks in the middle of a scan. + +The demo's `ChunkCache` runs each query against a cache of decoded chunks +instead: + +1. Run the query. A chunk that is not cached is recorded as missing and + scanned as fill values, so one run records every chunk that survives + xarray-sql's chunk pruning. +2. Fetch the missing chunks concurrently with `Dataset.load_async`, then run + the query again. +3. A run with no missing chunks used only real data, so its result is + exact. + +Most queries need two runs. The stats line under the editor shows the runs, +chunks, requests, and bytes for each query. + +The demo's files are in +[`docs/demo/`](https://github.com/xqlsystems/xarray-sql/tree/main/docs/demo): + +| File | Role | +|------|------| +| `index.html` | The page: example queries, SQL editor, results table, and an interactive map ([deck.gl](https://deck.gl) over a [MapLibre](https://maplibre.org) basemap) for `latitude, longitude, value` results. | +| `worker.js` | Loads Pyodide, installs xarray-sql, DuckDB, and zarr, and runs queries off the main thread. | +| `app.py` | Opens the dataset, registers it on DuckDB as `gfs`, and runs queries. | +| `chunk_cache.py` | Lazy reads that work in the browser, as described above. | +| `fetch_store.py` | A read-only Zarr store that uses `fetch` with `Range` requests. | +| `smoke.mjs` | Runs every example query headlessly in Node. | + +## Run it locally + +Serve `docs/demo` over HTTP and open it in Chrome or Edge: + +```shell +python -m http.server 8000 --directory docs/demo +``` + +By default the worker installs xarray-sql from PyPI. To test a local build, +copy the wheel from `pyodide build` (see "Building for WebAssembly" in the +[contributing guide](contributing.md)) into `docs/demo` and pass it in the +URL: + +``` +http://localhost:8000/?wheel=./xarray_sql-0.5.0-cp310-abi3-pyemscripten_2026_0_wasm32.whl +``` + +On this docs page, `?wheel=` is passed through to the embedded demo, so a +locally built docs site can use a local wheel too (for example, +`/browser/?wheel=../demo/xarray_sql-...whl` after copying the wheel into the +built `demo/` directory). + +To run the examples without a browser: + +```shell +cd docs/demo +npm install --no-save pyodide@314.0.7 +node smoke.mjs path/to/xarray_sql-*-pyemscripten_*.whl +``` + +## Known limits + +- **`lead_time` filters do not prune chunks.** DuckDB does not push + `INTERVAL` predicates down to Arrow scans. The demo registers only the + first lead-time chunk (forecast hours 0–104), so every query reads at most + that one. +- **Queries that touch the whole grid are slow.** For example, + `SELECT max(init_time) FROM gfs` builds every coordinate row (about 440 + million) even though it downloads nothing. Use the literal timestamps that + the example queries use. +- **The cache holds 768 MiB of decoded chunks.** A query that needs more + fails with a message to narrow its filters. diff --git a/docs/demo/app.py b/docs/demo/app.py new file mode 100644 index 0000000..8ca6bc6 --- /dev/null +++ b/docs/demo/app.py @@ -0,0 +1,217 @@ +"""The Python half of the browser demo: open a remote Zarr, query it with SQL. + +``worker.js`` loads this module into Pyodide and calls :func:`setup` once +and :func:`run` per query. Both return JSON strings so the worker can hand +results to the page without converting Python objects. +""" + +from __future__ import annotations + +import json +import math +from decimal import Decimal +from typing import Any + +import duckdb +import numpy as np +import pandas as pd +import pyarrow as pa +import xarray as xr +from chunk_cache import ChunkCache +from fetch_store import FetchStore + +import xarray_sql as xql + +GFS_URL = "https://data.dynamical.org/noaa/gfs/forecast/latest.zarr" +GFS_VARIABLES = [ + "temperature_2m", + "precipitation_surface", + "wind_u_10m", + "wind_v_10m", +] +# The store's inner chunks; one cached chunk is then one range request. +GFS_CHUNKS = { + "init_time": 1, + "lead_time": 105, + "latitude": 121, + "longitude": 121, +} + +_state: dict[str, Any] = {} + + +def _ts(value: Any) -> str: + return str(pd.Timestamp(value))[:19] + + +def _examples(ds: xr.Dataset) -> list[dict[str, str]]: + latest = _ts(ds.init_time.values[-1]) + earliest = _ts(ds.init_time.values[0]) + return [ + { + "title": "San Francisco: 2 m temperature for the next 3 days", + "sql": f"""SELECT init_time + lead_time AS valid_time, + round(avg(temperature_2m), 1) AS temperature_c +FROM gfs +WHERE init_time = TIMESTAMP '{latest}' + AND lead_time <= INTERVAL 72 HOUR + AND latitude BETWEEN 37.5 AND 38 AND longitude BETWEEN -122.75 AND -122.25 +GROUP BY ALL ORDER BY valid_time""", + }, + { + "title": "Map: Europe 2 m temperature, 24 h ahead", + "sql": f"""SELECT latitude, longitude, temperature_2m +FROM gfs +WHERE init_time = TIMESTAMP '{latest}' + AND lead_time = INTERVAL 24 HOUR + AND latitude BETWEEN 35 AND 70 AND longitude BETWEEN -10 AND 40""", + }, + { + "title": "Map: contiguous US wind speed, 12 h ahead", + "sql": f"""SELECT latitude, longitude, + round(sqrt(wind_u_10m ^ 2 + wind_v_10m ^ 2), 1) AS wind_speed +FROM gfs +WHERE init_time = TIMESTAMP '{latest}' + AND lead_time = INTERVAL 12 HOUR + AND latitude BETWEEN 24 AND 50 AND longitude BETWEEN -125 AND -66""", + }, + { + "title": "Wettest places in the contiguous US, next 24 h", + "sql": f"""SELECT latitude, longitude, + round(sum(precipitation_surface) * 3600, 1) AS approx_precip_mm +FROM gfs +WHERE init_time = TIMESTAMP '{latest}' + AND lead_time BETWEEN INTERVAL 1 HOUR AND INTERVAL 24 HOUR + AND latitude BETWEEN 24 AND 50 AND longitude BETWEEN -125 AND -66 +GROUP BY ALL ORDER BY approx_precip_mm DESC LIMIT 10""", + }, + { + "title": "How the forecast for San Francisco changed across runs", + "sql": f"""SELECT init_time, + round(max(temperature_2m), 1) AS max_temperature_c +FROM gfs +WHERE init_time BETWEEN TIMESTAMP '{earliest}' AND TIMESTAMP '{latest}' + AND init_time + lead_time BETWEEN TIMESTAMP '{latest}' + INTERVAL 24 HOUR + AND TIMESTAMP '{latest}' + INTERVAL 48 HOUR + AND lead_time <= INTERVAL 72 HOUR + AND latitude BETWEEN 37.5 AND 38 AND longitude BETWEEN -122.75 AND -122.25 +GROUP BY ALL ORDER BY init_time""", + }, + ] + + +async def setup( + url: str = GFS_URL, + variables: list[str] | None = None, + runs: int = 4, + max_cache_mib: int = 768, +) -> str: + """Open the dataset and register it on a DuckDB connection as ``gfs``.""" + store = FetchStore(url) + ds = xr.open_zarr(store, chunks=None) + # DuckDB does not push INTERVAL predicates down to Arrow scans, so a + # lead_time filter cannot prune chunks. Keeping only the first lead-time + # chunk (forecast hours 0-104) halves what every query downloads. + ds = ds[variables or GFS_VARIABLES].isel( + init_time=slice(-runs, None), + lead_time=slice(0, GFS_CHUNKS["lead_time"]), + ) + cache = ChunkCache(ds, GFS_CHUNKS, max_bytes=max_cache_mib * 2**20) + con = duckdb.connect() + con.sql("SET enable_progress_bar = false") + xql.register(con, "gfs", cache.dataset(), chunks=GFS_CHUNKS) + _state.update(store=store, cache=cache, con=con) + return json.dumps( + { + "url": url, + "variables": { + name: { + "dims": list(var.dims), + "units": var.attrs.get("units", ""), + "long_name": var.attrs.get("long_name", name), + } + for name, var in ds.data_vars.items() + }, + "sizes": dict(ds.sizes), + "init_times": [_ts(t) for t in ds.init_time.values], + "examples": _examples(ds), + } + ) + + +def _jsonable(value: Any) -> Any: + """A Python value from Arrow, made safe for JSON (which has no NaN).""" + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, Decimal): + return float(value) + # Exactly tuple: Arrow intervals are named tuples, shown as text below. + if isinstance(value, list) or type(value) is tuple: + return [_jsonable(v) for v in value] + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + if value is None or isinstance(value, (str, int)): + return value + return str(value) # Timestamps, intervals, bytes, ... + + +def _grid(table: pa.Table) -> dict[str, Any] | None: + """Columns for the page's map: latitude, longitude, and one number.""" + names = table.column_names + if "latitude" not in names or "longitude" not in names: + return None + values = [ + n + for n in names + if n not in ("latitude", "longitude") + and ( + pa.types.is_floating(table[n].type) + or pa.types.is_integer(table[n].type) + or pa.types.is_decimal(table[n].type) + ) + ] + if len(values) != 1 or table.num_rows < 4: + return None + lat = table["latitude"].to_numpy().astype("float64") + lon = table["longitude"].to_numpy().astype("float64") + # One pixel per point: with several rows per point (e.g. every lead + # time), the table is the honest view. + if len(np.unique(np.stack([lat, lon]), axis=1)[0]) != table.num_rows: + return None + value = table[values[0]].to_numpy().astype("float64") + finite = np.isfinite(value) + if not finite.any(): + return None + return { + "name": values[0], + "latitude": lat[finite].tolist(), + "longitude": lon[finite].tolist(), + "value": value[finite].tolist(), + } + + +async def run(sql: str, max_rows: int = 500) -> str: + """Run ``sql``, fetching whatever chunks it needs, and return the result.""" + store, cache, con = _state["store"], _state["cache"], _state["con"] + requests, nbytes = store.requests, store.bytes_read + table, stats = await cache.query(lambda: con.sql(sql).to_arrow_table()) + shown = table.slice(0, max_rows) + return json.dumps( + { + "columns": shown.column_names, + "rows": _jsonable( + list(zip(*(shown[c].to_pylist() for c in shown.column_names))) + ), + "num_rows": table.num_rows, + "grid": _grid(table), + "stats": { + "runs": stats.rounds, + "chunks_fetched": stats.chunks_fetched, + "requests": store.requests - requests, + "mb_downloaded": round((store.bytes_read - nbytes) / 1e6, 1), + "seconds_fetching": round(stats.seconds_fetching, 2), + "seconds_querying": round(stats.seconds_querying, 2), + "cached_chunks": len(cache), + }, + } + ) diff --git a/docs/demo/chunk_cache.py b/docs/demo/chunk_cache.py new file mode 100644 index 0000000..42daf1b --- /dev/null +++ b/docs/demo/chunk_cache.py @@ -0,0 +1,238 @@ +"""Lazy SQL over remote Zarr in Pyodide, by faulting chunks in between runs. + +In the browser, Zarr reads are asynchronous (``fetch``). zarr's +synchronous API blocks on them with Pyodide's JSPI support, but that +cannot suspend while the engine is calling back into Python: DuckDB and +pyarrow pull record batches through C++ frames that JSPI cannot unwind. +So a query cannot read chunks in the middle of a scan. + +Instead, every scan reads from a cache of decoded chunks: + +1. Run the query. A chunk that is not cached is recorded as missing and + scanned as fill values, so the scan keeps going and records every + chunk the query touches in one pass. xarray-sql's chunk pruning still + applies, so only chunks that survive the query's coordinate filters + are touched. +2. If anything was missing, fetch those chunks concurrently with + ``Dataset.load_async`` (plain top-level ``await``, no JSPI), then run + the query again. +3. A run with no misses saw only real data, so its result is exact. + +Placeholder data can change data-dependent filters in the discarded +runs, which may fetch a few extra chunks, or make a run fail outright +(casting NaN to an integer, say). A failed run that missed chunks is +treated like any other: fetch, then run again. Neither changes the final +answer. +""" + +from __future__ import annotations + +import asyncio +import itertools +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import xarray as xr + +ChunkKey = tuple[str, tuple[tuple[int, int], ...]] + + +def _fill_value(dtype: np.dtype) -> Any: + if dtype.kind in "fc": + return np.nan + if dtype.kind == "M": + return np.datetime64("NaT") + if dtype.kind == "m": + return np.timedelta64("NaT") + return 0 + + +class _CachedVariable(xr.backends.BackendArray): + """A lazily indexed variable whose reads are served from the cache.""" + + def __init__(self, cache: ChunkCache, name: str, source: xr.Variable): + self._cache = cache + self._name = name + self.shape = source.shape + self.dtype = source.dtype + self._chunks = tuple( + cache.chunks.get(str(d), n) + for d, n in zip(source.dims, source.shape) + ) + + def __getitem__(self, key: Any) -> np.ndarray: + return xr.core.indexing.explicit_indexing_adapter( + key, + self.shape, + xr.core.indexing.IndexingSupport.BASIC, + self._raw_getitem, + ) + + def _raw_getitem(self, key: tuple) -> np.ndarray: + """Assemble a basic selection from the cache chunks it overlaps.""" + ranges, steps, squeeze = [], [], [] + for axis, (k, n) in enumerate(zip(key, self.shape)): + if isinstance(k, (int, np.integer)): + k = slice(int(k), int(k) + 1) + squeeze.append(axis) + start, stop, step = k.indices(n) + ranges.append((start, max(start, stop))) + steps.append(step) + out = np.empty([hi - lo for lo, hi in ranges], dtype=self.dtype) + per_axis = [ + range(lo // c, (hi - 1) // c + 1) if hi > lo else range(0) + for (lo, hi), c in zip(ranges, self._chunks) + ] + for index in itertools.product(*per_axis): + bounds = tuple( + (i * c, min((i + 1) * c, n)) + for i, c, n in zip(index, self._chunks, self.shape) + ) + chunk = self._cache._get((self._name, bounds)) + src, dst = [], [] + for (lo, hi), (b0, b1) in zip(ranges, bounds): + a, b = max(lo, b0), min(hi, b1) + src.append(slice(a - b0, b - b0)) + dst.append(slice(a - lo, b - lo)) + out[tuple(dst)] = chunk[tuple(src)] + out = out[tuple(slice(None, None, s) for s in steps)] + return out.squeeze(axis=tuple(squeeze)) if squeeze else out + + +@dataclass +class QueryStats: + rounds: int = 0 + chunks_fetched: int = 0 + seconds_fetching: float = 0.0 + seconds_querying: float = 0.0 + misses_per_round: list[int] = field(default_factory=list) + + +class ChunkCache: + """Serves a remote Dataset's chunks to synchronous scans from memory. + + Args: + source: A lazily opened Dataset (e.g. ``xr.open_zarr(..., + chunks=None)``). Its coordinates should already be in memory; + only data variables are fetched through the cache. + chunks: Chunk sizes per dimension. Align them with the store's + (inner) chunks so one cached chunk is one fetch. + max_bytes: Memory budget for cached chunks; least recently used + chunks are evicted first. + """ + + def __init__( + self, + source: xr.Dataset, + chunks: dict[str, int], + max_bytes: int = 512 * 2**20, + ) -> None: + self.source = source + self.chunks = chunks + self.max_bytes = max_bytes + self._chunks: OrderedDict[ChunkKey, np.ndarray] = OrderedDict() + self._nbytes = 0 + self.missing: set[ChunkKey] = set() + + def __len__(self) -> int: + """Number of chunks currently cached.""" + return len(self._chunks) + + def dataset(self) -> xr.Dataset: + """The source Dataset with every data variable read through the cache. + + Register it with ``chunks=`` set to this cache's chunks, so each + scanned block is one cached chunk. + """ + variables = { + name: xr.Variable( + var.dims, + xr.core.indexing.LazilyIndexedArray( + _CachedVariable(self, str(name), var.variable) + ), + attrs=var.attrs, + ) + for name, var in self.source.data_vars.items() + } + return self.source.assign(variables) + + def _get(self, key: ChunkKey) -> np.ndarray: + if key in self._chunks: + self._chunks.move_to_end(key) + return self._chunks[key] + self.missing.add(key) + dtype = self.source[key[0]].dtype + shape = tuple(hi - lo for lo, hi in key[1]) + return np.full(shape, _fill_value(dtype), dtype=dtype) + + async def _fetch(self, key: ChunkKey) -> None: + name, bounds = key + var = self.source[name] + region = {d: slice(lo, hi) for d, (lo, hi) in zip(var.dims, bounds)} + loaded = await var.isel(region).load_async() + values = np.asarray(loaded.values) + self._chunks[key] = values + self._nbytes += values.nbytes + while self._nbytes > self.max_bytes and len(self._chunks) > 1: + _, evicted = self._chunks.popitem(last=False) + self._nbytes -= evicted.nbytes + + def _nbytes_of(self, keys: set[ChunkKey]) -> int: + return sum( + int(np.prod([hi - lo for lo, hi in bounds])) + * self.source[name].dtype.itemsize + for name, bounds in keys + ) + + async def fill(self) -> int: + """Fetch every chunk recorded as missing; returns how many.""" + keys = list(self.missing) + self.missing.clear() + await asyncio.gather(*(self._fetch(k) for k in keys)) + return len(keys) + + async def query(self, run: Any) -> tuple[Any, QueryStats]: + """Call ``run()`` until it completes without a cache miss. + + ``run`` executes the query and materializes its result, e.g. + ``lambda: con.sql(sql).to_arrow_table()``. + + Every retry fetches chunks this query has not fetched before, so + the loop ends. If a run misses a chunk this query already fetched, + the cache evicted it to make room for the others: the query needs + more than the cache holds. + """ + stats = QueryStats() + fetched: set[ChunkKey] = set() + self.missing.clear() + while True: + stats.rounds += 1 + start = time.perf_counter() + try: + result = run() + except Exception: + # Fill values can fail a query that real data would not. + if not self.missing: + raise + stats.seconds_querying += time.perf_counter() - start + stats.misses_per_round.append(len(self.missing)) + if not self.missing: + return result, stats + if ( + self._nbytes_of(self.missing) > self.max_bytes + or self.missing & fetched + ): + needed = self._nbytes_of(fetched | self.missing) + raise MemoryError( + f"The query needs about {needed / 2**20:.0f} MiB of " + f"chunks, more than the cache's " + f"{self.max_bytes / 2**20:.0f} MiB. " + "Narrow its filters on the dimension columns." + ) + fetched |= self.missing + start = time.perf_counter() + stats.chunks_fetched += await self.fill() + stats.seconds_fetching += time.perf_counter() - start diff --git a/docs/demo/fetch_store.py b/docs/demo/fetch_store.py new file mode 100644 index 0000000..d21a77c --- /dev/null +++ b/docs/demo/fetch_store.py @@ -0,0 +1,148 @@ +"""A read-only Zarr store over HTTP for Pyodide, built on the browser's fetch. + +zarr's own HTTP path goes through fsspec and aiohttp, which open raw +sockets that a browser does not have. This store issues ``fetch`` +requests instead (with ``Range`` headers for partial reads, which sharded +Zarr v3 arrays rely on). zarr drives it from synchronous code through +Pyodide's JSPI support, so ``xr.open_zarr(FetchStore(url))`` and lazy +reads work unchanged. + +The server must allow cross-origin requests (``Access-Control-Allow-Origin``) +and expose range requests. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Iterable + +from pyodide.http import pyfetch +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr.core.buffer import Buffer, BufferPrototype + +# S3-style object stores answer 403 rather than 404 for a missing key +# when listing is not public, and zarr probes for keys that may not exist. +# A key the store refuses for any other reason also reads as missing, so +# its chunks come back as fill values. +_MISSING = (403, 404) +_ATTEMPTS = 4 + + +def _range_header(byte_range: ByteRequest | None) -> dict[str, str]: + if byte_range is None: + return {} + if isinstance(byte_range, RangeByteRequest): + return {"Range": f"bytes={byte_range.start}-{byte_range.end - 1}"} + if isinstance(byte_range, OffsetByteRequest): + return {"Range": f"bytes={byte_range.offset}-"} + if isinstance(byte_range, SuffixByteRequest): + return {"Range": f"bytes=-{byte_range.suffix}"} + raise TypeError(f"Unexpected byte range: {byte_range!r}") + + +def _range_slice(byte_range: ByteRequest) -> slice: + if isinstance(byte_range, RangeByteRequest): + return slice(byte_range.start, byte_range.end) + if isinstance(byte_range, OffsetByteRequest): + return slice(byte_range.offset, None) + if isinstance(byte_range, SuffixByteRequest): + if byte_range.suffix == 0: + return slice(0, 0) # data[-0:] would be everything. + return slice(-byte_range.suffix, None) + raise TypeError(f"Unexpected byte range: {byte_range!r}") + + +class FetchStore(Store): + """Read-only Zarr store for an HTTP(S) URL, using ``fetch``.""" + + supports_writes = False + supports_deletes = False + supports_partial_writes = False + supports_listing = False + + def __init__(self, url: str) -> None: + super().__init__(read_only=True) + self.url = url.rstrip("/") + self.requests = 0 + """Number of HTTP requests issued, for observing chunk pruning.""" + self.bytes_read = 0 + """Bytes received, for observing chunk pruning.""" + + def __eq__(self, other: object) -> bool: + return isinstance(other, FetchStore) and other.url == self.url + + def __repr__(self) -> str: + return f"FetchStore({self.url!r})" + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + url = f"{self.url}/{key}" + for attempt in range(_ATTEMPTS): + if attempt: + await asyncio.sleep(0.25 * 2**attempt) + self.requests += 1 + try: + response = await pyfetch(url, headers=_range_header(byte_range)) + except OSError: # Network failure; retry. + if attempt == _ATTEMPTS - 1: + raise + continue + if response.status == 429 or response.status >= 500: + continue + break + if response.status in _MISSING: + return None + if not response.ok: + raise OSError(f"GET {url} failed: HTTP {response.status}") + data = await response.bytes() + self.bytes_read += len(data) + if byte_range is not None and response.status != 206: + # The server ignored the Range header and sent the whole object. + data = data[_range_slice(byte_range)] + return prototype.buffer.from_bytes(data) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + return list( + await asyncio.gather( + *(self.get(key, prototype, rng) for key, rng in key_ranges) + ) + ) + + async def exists(self, key: str) -> bool: + response = await pyfetch(f"{self.url}/{key}", method="HEAD") + return bool(response.ok) + + async def set(self, key: str, value: Buffer) -> None: + raise NotImplementedError("FetchStore is read-only") + + async def set_partial_values(self, key_start_values) -> None: + raise NotImplementedError("FetchStore is read-only") + + async def delete(self, key: str) -> None: + raise NotImplementedError("FetchStore is read-only") + + async def list(self) -> AsyncIterator[str]: + raise NotImplementedError("HTTP has no listing") + yield # pragma: no cover - makes this an async generator + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + raise NotImplementedError("HTTP has no listing") + yield # pragma: no cover + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + raise NotImplementedError("HTTP has no listing") + yield # pragma: no cover diff --git a/docs/demo/index.html b/docs/demo/index.html new file mode 100644 index 0000000..cb22e07 --- /dev/null +++ b/docs/demo/index.html @@ -0,0 +1,444 @@ + + + + + +xarray-sql in the browser + + + +
+

xarray-sql in the browser

+

+ SQL over NOAA GFS forecasts, read straight from + dynamical.org's Zarr store. xarray-sql and + DuckDB run in Pyodide in your browser; there is + no server. Each query downloads only the chunks its filters on + init_time, latitude, and longitude reach. +

+
+ +
+ +
+ + Starting… +
+ +
+ + + +
+
+
+ + + diff --git a/docs/demo/smoke.mjs b/docs/demo/smoke.mjs new file mode 100644 index 0000000..85200e1 --- /dev/null +++ b/docs/demo/smoke.mjs @@ -0,0 +1,52 @@ +// Run the demo's example queries headlessly in Node, against the live dataset. +// +// npm install pyodide@314.0.7 +// node smoke.mjs path/to/xarray_sql-*-pyemscripten_*.whl +// +// Without a wheel path, xarray-sql is installed from PyPI. +import { loadPyodide } from "pyodide"; +import { readFileSync } from "node:fs"; +import { basename, dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; + +const here = dirname(fileURLToPath(import.meta.url)); +const wheel = process.argv[2]; + +const py = await loadPyodide(); +await py.loadPackage("micropip"); +for (const module of ["app.py", "chunk_cache.py", "fetch_store.py"]) { + py.FS.writeFile(module, readFileSync(join(here, module))); +} +let requirement = "xarray-sql"; +if (wheel) { + requirement = `emfs:/tmp/${basename(wheel)}`; + py.FS.writeFile(`/tmp/${basename(wheel)}`, readFileSync(wheel)); +} +const micropip = py.pyimport("micropip"); +await micropip.install([requirement, "duckdb", "zarr"]); + +const app = py.pyimport("app"); +let start = performance.now(); +const info = JSON.parse(await app.setup()); +console.log( + `setup: ${((performance.now() - start) / 1000).toFixed(1)}s,`, + `runs ${info.init_times.join(", ")}`, +); + +let failed = false; +for (const { title, sql } of info.examples) { + start = performance.now(); + try { + const result = JSON.parse(await app.run(sql)); + const seconds = ((performance.now() - start) / 1000).toFixed(1); + console.log(`\n${title}\n ${seconds}s, ${result.num_rows} rows`, result.stats); + console.log(" first row:", result.columns, result.rows[0]); + if (result.num_rows === 0) failed = true; + } catch (error) { + failed = true; + // Python tracebacks end with the exception line. + const lines = String(error.message ?? error).trim().split("\n"); + console.error(`\n${title}\n FAILED: ${lines.at(-1)}\n${lines.slice(-12).join("\n")}`); + } +} +process.exit(failed ? 1 : 0); diff --git a/docs/demo/worker.js b/docs/demo/worker.js new file mode 100644 index 0000000..f35b42e --- /dev/null +++ b/docs/demo/worker.js @@ -0,0 +1,67 @@ +// Runs Pyodide, xarray-sql, and DuckDB off the page's main thread. +// +// Messages in: {id, sql} +// Messages out: {type: "status", message, step} +// {type: "ready", info} | {type: "failed", error} +// {type: "result", id, result} | {type: "error", id, error} +import { loadPyodide } from "https://cdn.jsdelivr.net/pyodide/v314.0.7/full/pyodide.mjs"; + +// xarray-sql comes from PyPI unless the page passes a wheel URL, e.g. a +// locally built one: index.html?wheel=./xarray_sql-...-pyemscripten_2026_0_wasm32.whl +const wheel = new URL(import.meta.url).searchParams.get("wheel"); +const status = (message, step) => postMessage({ type: "status", message, step }); + +async function start() { + // zarr's synchronous API, which xarray uses to open the store, blocks on + // fetch() through JavaScript Promise Integration. + if (typeof WebAssembly.Suspending !== "function") { + throw new Error( + "This demo needs WebAssembly JavaScript Promise Integration (JSPI), " + + "available in Chrome and Edge 137 or newer.", + ); + } + status("Loading Pyodide and packages…", 0); + // Everything Pyodide ships is fetched in parallel while the runtime + // boots; micropip then only resolves xarray-sql from PyPI. + const modules = Promise.all( + ["app.py", "chunk_cache.py", "fetch_store.py"].map(async (name) => { + const response = await fetch(new URL(name, import.meta.url)); + if (!response.ok) throw new Error(`Could not load ${name}: HTTP ${response.status}`); + return [name, await response.text()]; + }), + ); + const py = await loadPyodide({ + packages: ["micropip", "duckdb", "pyarrow", "xarray", "zarr"], + }); + for (const [name, source] of await modules) py.FS.writeFile(name, source); + status("Installing xarray-sql…", 1); + const micropip = py.pyimport("micropip"); + await micropip.install(wheel ?? "xarray-sql"); + status("Opening the dataset…", 2); + const app = py.pyimport("app"); + postMessage({ type: "ready", info: JSON.parse(await app.setup()) }); + return app; +} + +// Python errors carry the full traceback. Show the exception line (some +// messages add hints after it) and keep the rest for the console. +const describe = (error) => { + const detail = String(error.message ?? error).trim(); + const lines = detail.split("\n"); + const exception = lines.findLast((line) => /^[\w.]*(Error|Exception)\b/.test(line)); + return { error: exception ?? lines.at(-1), detail }; +}; + +const app = start().catch((error) => { + postMessage({ type: "failed", ...describe(error) }); +}); + +onmessage = async ({ data: { id, sql } }) => { + const loaded = await app; + if (!loaded) return; + try { + postMessage({ type: "result", id, result: JSON.parse(await loaded.run(sql)) }); + } catch (error) { + postMessage({ type: "error", id, ...describe(error) }); + } +}; diff --git a/docs/explainer/era5.json b/docs/explainer/era5.json new file mode 100644 index 0000000..c92a1e4 --- /dev/null +++ b/docs/explainer/era5.json @@ -0,0 +1 @@ +{"source":"gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3","variable":"t2m","long_name":"2 metre temperature","units":"K","time":["2024-07-01","2024-07-02","2024-07-03","2024-07-04","2024-07-05","2024-07-06","2024-07-07","2024-07-08"],"latitude":[75.0,45.0,15.0,-15.0,-45.0,-75.0],"longitude":[0.0,45.0,90.0,135.0,180.0,225.0,270.0,315.0],"values":[[[277.7,275.6,273.9,274.5,273.4,273.5,273.8,269.1],[295.5,305.7,302.5,288.4,283.5,285.1,284.5,292.7],[309.2,308.2,302.2,302.0,300.4,298.8,292.9,299.7],[294.8,298.2,298.5,295.3,299.4,298.0,293.0,296.4],[283.3,279.4,283.5,283.7,285.5,285.2,279.3,278.8],[224.3,227.9,220.9,221.3,244.4,248.6,247.4,250.5]],[[278.2,275.7,275.5,274.1,274.2,273.4,274.0,263.9],[293.2,302.4,299.0,291.6,285.1,286.5,287.0,293.6],[311.3,308.5,302.0,302.2,299.8,299.0,291.8,299.8],[293.2,298.1,298.3,291.9,300.0,298.5,292.9,295.8],[280.3,279.2,279.5,284.2,282.7,283.4,281.4,280.0],[223.4,227.2,216.2,214.8,246.6,256.1,250.8,251.1]],[[278.0,275.8,279.0,273.8,273.4,273.0,276.7,265.2],[295.3,304.6,297.2,292.9,286.6,285.4,291.2,294.6],[302.7,307.4,301.7,302.3,300.5,298.2,291.6,299.3],[293.3,298.0,299.0,292.6,300.8,299.1,293.0,295.3],[282.9,281.3,281.4,284.2,281.2,282.0,279.6,277.2],[221.3,225.4,212.3,213.1,250.3,255.9,255.9,253.4]],[[277.5,276.1,289.5,274.2,273.7,273.5,274.0,266.6],[295.5,310.2,299.3,291.3,287.4,286.5,289.7,294.1],[309.4,306.8,300.3,300.9,300.6,297.3,288.6,299.4],[293.5,298.2,299.6,290.0,300.7,298.7,292.5,295.0],[280.1,281.2,278.1,283.3,282.0,284.1,281.7,278.6],[223.1,224.8,212.8,215.9,246.3,247.8,251.2,246.6]],[[278.4,276.7,289.2,274.2,273.4,273.3,277.2,267.5],[297.7,312.5,301.3,289.7,286.5,287.4,289.7,294.9],[307.1,307.5,300.4,301.2,300.2,299.0,290.2,299.8],[294.2,299.0,298.1,293.6,300.7,297.8,293.3,294.9],[278.6,282.4,279.0,283.2,282.1,283.6,282.2,280.4],[230.9,226.3,218.6,215.3,248.1,255.6,249.3,250.6]],[[277.0,280.0,291.0,275.4,273.4,273.3,280.6,266.2],[291.9,312.8,302.0,290.9,286.3,288.3,286.5,294.0],[310.8,306.7,300.1,303.0,300.6,298.6,292.2,300.3],[292.9,298.6,297.6,295.9,298.4,298.1,293.6,295.7],[276.1,281.4,280.8,281.8,281.7,282.3,280.7,280.9],[232.0,224.1,211.9,217.9,249.0,256.4,254.1,245.4]],[[276.9,279.3,288.2,274.4,273.2,273.1,280.2,265.0],[294.0,312.2,297.7,290.2,285.8,287.8,290.1,291.9],[310.5,308.1,300.7,301.4,300.5,298.4,291.1,299.8],[293.8,298.9,298.2,296.0,299.0,297.8,292.8,296.0],[277.3,281.9,279.6,283.8,281.1,284.8,281.5,277.7],[242.3,238.9,211.9,217.8,251.6,258.6,255.1,242.2]],[[276.5,279.6,289.3,272.6,273.2,272.9,284.2,261.1],[300.5,313.5,304.3,290.7,286.5,287.7,290.9,289.9],[309.0,308.2,300.9,302.0,300.2,298.4,291.8,299.9],[293.1,298.6,299.1,295.3,299.2,297.5,292.7,296.5],[280.1,281.9,277.6,283.9,282.5,285.3,281.2,278.4],[234.3,226.2,211.6,212.9,253.5,251.4,256.7,245.3]]]} diff --git a/docs/explainer/index.html b/docs/explainer/index.html new file mode 100644 index 0000000..f2f22c2 --- /dev/null +++ b/docs/explainer/index.html @@ -0,0 +1,537 @@ + + + + + +How xarray-sql works + + + +
+
+

1A Dataset is a table

+

+ Registering a Dataset gives SQL a table with one row per grid cell: a column + for each dimension and one for each data variable. Nothing is copied. Each + chunk becomes a partition of the table, read only when a query needs it. + Hover over the grid or the table to match cells to rows. +

+ 
import duckdb, xarray as xr, xarray_sql as xql
+
+era5 = xr.open_zarr(
+    "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3",
+    chunks=None, storage_options={"token": "anon"},
+)
+# A grid small enough to see: 12:00 UTC on 8 days, 6 latitudes, 8 longitudes.
+ds = era5[["2m_temperature"]].rename({"2m_temperature": "t2m"}).sel(
+    time=xr.date_range("2024-07-01T12", periods=8, freq="D"),
+    latitude=[75, 45, 15, -15, -45, -75],
+    longitude=range(0, 360, 45),
+)
+con = duckdb.connect()
+xql.register(con, "era5", ds, chunks={"time": 4, "latitude": 3, "longitude": 4})
+
+
+
ds.t2m, one time step
+ +
· lines mark chunks
+ +
+
+
table era5 ( rows, 8 partitions)
+
+
timelatitudelongitudet2mpartition
+
+
+
+
+ +
+

2SQL reads only the chunks it needs

+

+ Filters on dimension columns are pushed down to the table, and xarray-sql + compares them with each partition's coordinate range. A partition that + cannot match is skipped, so its chunk is never read. Pick a query. +

+

+ Partitions follow chunks=. ARCO-ERA5 stores each hour of the + whole globe as one chunk, so registered with its storage chunks, a time + filter would still skip chunks but a bounding box could not. +

+
+ 

+    

+      

+        
partitions of era5

+        
+        

+          read
+          skipped
+          front: Jul 1–4, back: Jul 5–8
+        

+      

+      

+        

+        

+      

+    

+  

+
+  

+    
3Results come back as a Dataset

+    

+      to_dataset pivots the result rows back onto a grid. The
+      dimension columns that survive the query become the dimensions, and the
+      registered Dataset, passed as the template, restores attributes such as
+      units.
+    

+    
out = xql.to_dataset(con.sql(query), template=ds)

+    

+      

+      

+        
+        

+        
+      

+    

+  

+  

+    Data: ERA5 2 m temperature from
+    ARCO-ERA5.
+    Contains modified Copernicus Climate Change Service information 2024.
+  

+

+
+
+
diff --git a/docs/explainer/make_snapshot.py b/docs/explainer/make_snapshot.py
new file mode 100644
index 0000000..280ffc0
--- /dev/null
+++ b/docs/explainer/make_snapshot.py
@@ -0,0 +1,46 @@
+"""Writes the explainer's data: a tiny, real slice of ARCO-ERA5.
+
+The explainer shows every row of its table and every partition, so it
+needs a grid small enough to see: 2 m temperature at 12:00 UTC on eight
+days, at six latitudes and eight longitudes. The selection is the code the
+explainer shows. Run from the repository root:
+
+    python docs/explainer/make_snapshot.py
+"""
+
+import json
+import pathlib
+
+import numpy as np
+import xarray as xr
+
+ARCO_ERA5 = (
+    "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+)
+
+era5 = xr.open_zarr(ARCO_ERA5, chunks=None, storage_options={"token": "anon"})
+ds = (
+    era5[["2m_temperature"]]
+    .rename({"2m_temperature": "t2m"})
+    .sel(
+        time=xr.date_range("2024-07-01T12", periods=8, freq="D"),
+        latitude=[75, 45, 15, -15, -45, -75],
+        longitude=range(0, 360, 45),
+    )
+    .load()
+)
+
+snapshot = {
+    "source": ARCO_ERA5,
+    "variable": "t2m",
+    "long_name": ds.t2m.attrs.get("long_name", "2 metre temperature"),
+    "units": ds.t2m.attrs.get("units", "K"),
+    "time": [str(t)[:10] for t in ds.time.values.astype("datetime64[D]")],
+    "latitude": ds.latitude.values.tolist(),
+    "longitude": ds.longitude.values.tolist(),
+    # Rounded to 0.1 K, in (time, latitude, longitude) order.
+    "values": np.round(ds.t2m.values.astype(float), 1).tolist(),
+}
+out = pathlib.Path(__file__).with_name("era5.json")
+out.write_text(json.dumps(snapshot, separators=(",", ":")) + "\n")
+print(f"wrote {out} ({out.stat().st_size} bytes)")
diff --git a/docs/how-it-works.md b/docs/how-it-works.md
new file mode 100644
index 0000000..3af3c6b
--- /dev/null
+++ b/docs/how-it-works.md
@@ -0,0 +1,41 @@
+# How it works
+
+xarray-sql lets SQL engines query Xarray Datasets as if they were tables,
+and turns query results back into Datasets. The explainer below walks
+through the three steps with a tiny slice of
+[ARCO-ERA5](https://github.com/google-research/arco-era5) 2 m temperature:
+eight days at 48 points on the globe, small enough to show every row. The
+slice is saved with the page, so it loads instantly and needs no Python.
+
+
+
+
+
+## What happens in each step
+
+1. **Registering** a Dataset with [register][xarray_sql.register] (or
+   [Context.from_dataset][xarray_sql.Context.from_dataset]) creates a table
+   whose rows are the Dataset's grid cells. The table reads from the
+   Dataset as the query runs, one chunk at a time, so registering copies
+   nothing.
+2. **Querying** runs in the engine: DuckDB, DataFusion, or any engine
+   that reads Arrow. Filters on dimension columns are pushed down to the
+   table, which skips every partition whose coordinate range cannot match,
+   so a selective query reads few chunks. See [Engines](engines.md).
+3. **Round-tripping** with [to_dataset][xarray_sql.to_dataset] pivots the
+   result rows back onto a grid and restores metadata from the registered
+   Dataset. By default the result is lazy: each access reruns the query
+   for just the part you select.
+
+To run real queries over a public Zarr store, try
+[xarray-sql in the browser](browser.md).
diff --git a/flake.nix b/flake.nix
index c393dc3..52b755a 100644
--- a/flake.nix
+++ b/flake.nix
@@ -48,6 +48,31 @@
             which
             xz
             zlib
+            # Shared libraries for Playwright's headless Chromium, which
+            # drives the browser demo in docs/demo.
+            alsa-lib
+            at-spi2-atk
+            at-spi2-core
+            atk
+            cairo
+            cups
+            dbus
+            expat
+            glib
+            libdrm
+            libgbm
+            libxkbcommon
+            nspr
+            nss
+            pango
+            systemdLibs
+            xorg.libX11
+            xorg.libXcomposite
+            xorg.libXdamage
+            xorg.libXext
+            xorg.libXfixes
+            xorg.libXrandr
+            xorg.libxcb
           ];
           runScript = "bash";
           profile = ''
diff --git a/zensical.toml b/zensical.toml
index a820e55..511daeb 100644
--- a/zensical.toml
+++ b/zensical.toml
@@ -9,8 +9,10 @@ edit_uri = "edit/main/docs/"
 nav = [
   {"Home" = "index.md"},
   {"Getting started" = "examples.md"},
+  {"How it works" = "how-it-works.md"},
   {"Concepts" = [
     {"Engines" = "engines.md"},
+    {"Run in the browser" = "browser.md"},
   ]},
   {"Guides" = [
     {"Geospatial in SQL" = "geospatial.md"},