diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c2865b4..5e39cb7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,7 +2,7 @@ ## Where to start? -Please check out the [issues tab](https://github.com/alxmrs/xarray-sql/issues). +Please check out the [issues tab](https://github.com/xqlsystems/xarray-sql/issues). Let's have a discussion over there before proceeding with any changes. Great minds think alike -- someone may have already created an issue related to your inquiry. If there's a bug, please let us know. diff --git a/Cargo.lock b/Cargo.lock index f1c8957..0022dae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3367,7 +3367,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] name = "xarray_sql" -version = "0.3.2" +version = "0.3.3" dependencies = [ "arrow", "async-stream", diff --git a/README.md b/README.md index e8dfc06..794605e 100644 --- a/README.md +++ b/README.md @@ -3,10 +3,10 @@ _Query [Xarray](https://xarray.dev/) with SQL_ ![PyPI Version](https://img.shields.io/pypi/v/xarray-sql?color=green) -[![ci](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml) -[![lint](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml) -[![ci-build](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-build.yml) -[![ci-rust](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci-rust.yml) +[![ci](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci.yml) +[![lint](https://github.com/xqlsystems/xarray-sql/actions/workflows/lint.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/lint.yml) +[![ci-build](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-build.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-build.yml) +[![ci-rust](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-rust.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-rust.yml) [![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=total&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads)](https://pepy.tech/projects/xarray-sql) [![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=monthly&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads%2Fmonth)](https://pepy.tech/projects/xarray-sql) @@ -18,7 +18,11 @@ pip install xarray-sql This is an experiment to provide a SQL interface for array datasets. Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run -SQL queries against them. +SQL queries against them — on the query engine of your choice. xarray-sql +translates data, not queries: it registers a lazy Dataset as a table on +DataFusion (built in), DuckDB, or Polars, and turns any engine's Arrow result +back into a labeled Dataset. Dialects, geometry functions, and optimizers stay +with the engine. ## Quickstart @@ -61,6 +65,21 @@ clim_ds["air"].plot() # in a script, call matplotlib.pyplot.show() to display That's the round trip — Xarray in, SQL in the middle, Xarray (and a plot) back out. +The same Dataset registers on other engines with one call — DuckDB gets a +native lazy table with predicate pushdown, Polars scans the same object: + +```python +import duckdb + +con = duckdb.connect() +xql.register(con, 'air', ds, chunks=dict(time=100)) +rel = con.sql('SELECT time, AVG("air") AS air FROM air GROUP BY time ORDER BY time') +xql.to_dataset(rel, template=ds) # any engine's Arrow result round-trips +``` + +See [Engines](https://xqlsystems.github.io/xarray-sql/engines/) for the support matrix, DuckDB/Polars details, +and the lazy chunked round-trip. + ## A bigger example: ARCO-ERA5 The same interface scales to cloud-native datasets with hundreds of variables, @@ -133,6 +152,8 @@ result = ctx.sql(''' # | 775 | -2.3064649711534457 | # +-------+----------------------+ +# `latitude`/`longitude` are inferred from the registered table's surviving +# dims; `template` is kept only to recover metadata (attrs, encoding). ctx.sql(''' SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c FROM era5.surface @@ -140,8 +161,6 @@ ctx.sql(''' AND TIMESTAMP '2020-01-01 05:00:00' GROUP BY latitude, longitude ORDER BY latitude DESC, longitude -# `latitude`/`longitude` are inferred from the registered table's surviving -# dims; `template` is kept only to recover metadata (attrs, encoding). ''').to_dataset(template=ds) # Size: 8MB # Dimensions: (latitude: 721, longitude: 1440) @@ -158,7 +177,7 @@ ctx.sql(''' ``` _(A runnable version of this example lives at -[`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_ +[`perf_tests/era5_temp_profile.py`](https://github.com/xqlsystems/xarray-sql/blob/main/perf_tests/era5_temp_profile.py).)_ ## Why build this? @@ -191,6 +210,9 @@ pure DataFusion and PyArrow, but works with the same principle! _2026 update_: Instead of `from_map()`, we create a way to translate Xarray chunks into Arrow RecordBatches. We pass a Python callback into a DataFusion `TableProvider` that lets the DB engine translate the underlying Dataset arrays into DataFusion partitions. +The same chunks-to-batches translation is also exposed as a +`pyarrow.dataset.Dataset` with predicate and projection pushdown, which is how +DuckDB and Polars consume registered Datasets with no engine-specific code. Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be translated into a 2D table -- underlies this performant query mechanism. @@ -222,8 +244,8 @@ against an xarray/array reference** to floating-point tolerance: Every case matches its array reference. The headline finding: these operations are not really "array" operations at all — they are `GROUP BY`, `JOIN`, window functions, and `CASE` in disguise, and a query engine runs them at scale. See -[`benchmarks/geospatial/`](benchmarks/geospatial/) and the write-up, -[Geospatial operations are relational operations](docs/geospatial.md). +[`benchmarks/geospatial/`](https://github.com/xqlsystems/xarray-sql/tree/main/benchmarks/geospatial/) and the write-up, +[Geospatial operations are relational operations](https://xqlsystems.github.io/xarray-sql/geospatial/). ## Why does this work? @@ -231,15 +253,17 @@ Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in chunks and represented contiguously in memory. It is only a matter of metadata that breaks them up into ndarrays. `pivot()`, which uses `to_dataframe()`, just changes this metadata (via a `ravel()`/`reshape()`), back into a column -amenable to a DataFrame. We take advantage of this light weight metadata change to -make chunked information scannable by a DB engine (DataFusion). +amenable to a DataFrame. We take advantage of this lightweight metadata change to +make chunked information scannable by a DB engine (DataFusion, DuckDB, Polars — +anything that speaks Arrow). ## What are the current limitations? -TBD, DataFusion provides a whole new world! Currently, we're looking for +The sharp edges we know about — per engine and fundamental — are cataloged in +[Known issues & limitations](https://xqlsystems.github.io/xarray-sql/limitations/). Currently, we're looking for early users – "tire kickers", if you will. We'd love your input to shape the direction of this -project! Please, give this a try and [file issues](https://github.com/alxmrs/xarray-sql/issues) as -you see fit. Check out our [contributing guide](CONTRIBUTING.md), too 😉. +project! Please, give this a try and [file issues](https://github.com/xqlsystems/xarray-sql/issues) as +you see fit. Check out our [contributing guide](https://xqlsystems.github.io/xarray-sql/contributing/), too 😉. ## What would a deeper integration look like? @@ -252,7 +276,7 @@ a [virtual](https://fsspec.github.io/kerchunk/) filesystem for parquet that would internally map to Zarr. Raster-backed virtual parquet would open up integrations to numerous tools like dask, pyarrow, duckdb, and BigQuery. More thoughts on this -in [#4](https://github.com/alxmrs/xarray-sql/issues/4). +in [#4](https://github.com/xqlsystems/xarray-sql/issues/4). _2025 update_: Something like this is being built across a few projects! The ones I know about are: @@ -262,18 +286,18 @@ _2025 update_: Something like this is being built across a few projects! The one _2026 update_: A colleague and I are experimenting with native Zarr RDBMS engines. Check out: - [Zarr-Datafusion](https://lib.rs/crates/zarr-datafusion) -- [DuckDB-Zarr](https://github.com/alxmrs/duckdb-zarr) +- [DuckDB-Zarr](https://github.com/xqlsystems/duckdb-zarr) ## Roadmap -- [x] ~Lazy evaluation via the pyarrow Dataset interface [#93](https://github.com/alxmrs/xarray-sql/issues/93).~ _Implemented in [#100](https://github.com/alxmrs/xarray-sql/pull/100)_ -- [x] Support proper parallelism via proper partition handling on the rust/datafusion side. [#106](https://github.com/alxmrs/xarray-sql/issues/106) -- [x] Support core datafusion optimizations to scan less data, like [104](https://github.com/alxmrs/xarray-sql/issues/104), ... -- [x] Translate a single Zarr to a collection of tables [#85](https://github.com/alxmrs/xarray-sql/issues/85). -- [ ] Distributed beyond a single node through the DataFusion integration with Ray Datasets [#68](https://github.com/alxmrs/xarray-sql/issues/68) or Apache Ballista [#98](https://github.com/alxmrs/xarray-sql/issues/98). -- [ ] Demo: calculate Sea Surface Temperature from 1940 - Present in SQL [#36](https://github.com/alxmrs/xarray-sql/issues/36). -- [ ] Provide an option to integrate DataFusion directly to Zarr via Rust [#4](https://github.com/alxmrs/xarray-sql/issues/4). -- [ ] (To be formally announced eventually): The 100 Trillion Row Challenge [#34](https://github.com/alxmrs/xarray-sql/issues/34). +- [x] ~Lazy evaluation via the pyarrow Dataset interface [#93](https://github.com/xqlsystems/xarray-sql/issues/93).~ _Implemented in [#100](https://github.com/xqlsystems/xarray-sql/pull/100)_ +- [x] Support proper parallelism via proper partition handling on the rust/datafusion side. [#106](https://github.com/xqlsystems/xarray-sql/issues/106) +- [x] Support core datafusion optimizations to scan less data, like [#104](https://github.com/xqlsystems/xarray-sql/issues/104), ... +- [x] Translate a single Zarr to a collection of tables [#85](https://github.com/xqlsystems/xarray-sql/issues/85). +- [ ] Distributed beyond a single node through the DataFusion integration with Ray Datasets [#68](https://github.com/xqlsystems/xarray-sql/issues/68) or Apache Ballista [#98](https://github.com/xqlsystems/xarray-sql/issues/98). +- [ ] Demo: calculate Sea Surface Temperature from 1940 - Present in SQL [#36](https://github.com/xqlsystems/xarray-sql/issues/36). +- [ ] Provide an option to integrate DataFusion directly to Zarr via Rust [#4](https://github.com/xqlsystems/xarray-sql/issues/4). +- [ ] (To be formally announced eventually): The 100 Trillion Row Challenge [#34](https://github.com/xqlsystems/xarray-sql/issues/34). ## Sponsors & Contributors diff --git a/benchmarks/duckdb_pushdown.py b/benchmarks/duckdb_pushdown.py new file mode 100644 index 0000000..e1db700 --- /dev/null +++ b/benchmarks/duckdb_pushdown.py @@ -0,0 +1,103 @@ +"""Benchmark: DuckDB re-scannable stream vs pushdown dataset vs ceiling. + +Times the three ways DuckDB can consume the same 10M-row synthetic +dataset — the re-scannable stream (no pushdown), the default +``register()`` pushdown dataset, and an in-memory ``pyarrow.dataset`` +as the ceiling — and asserts at the end that all three returned the +same answers. Cross-engine comparisons live in +``benchmarks/geospatial/``; this measures the adapter paths within one +engine. + +Usage: python benchmarks/duckdb_pushdown.py (needs duckdb installed) +""" + +import math +import statistics +import time + +import duckdb +import numpy as np +import pandas as pd +import pyarrow.dataset as pads +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.duckdb import XarrayArrowStream + +np.random.seed(0) +N_TIME, N_LAT, N_LON = 1000, 100, 100 # 10M rows +ds = xr.Dataset( + { + "temperature": ( + ["time", "lat", "lon"], + np.random.rand(N_TIME, N_LAT, N_LON), + ), + "humidity": ( + ["time", "lat", "lon"], + np.random.rand(N_TIME, N_LAT, N_LON), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=N_TIME, freq="h"), + "lat": np.linspace(-90, 90, N_LAT), + "lon": np.linspace(-180, 180, N_LON), + }, +).chunk({"time": 50}) # 20 partitions + +con = duckdb.connect() + +QUERIES = { + "full AVG scan": "SELECT AVG(temperature) FROM {t}", + "1pct time filter": ( + "SELECT AVG(temperature) FROM {t} WHERE time < '2020-01-01 10:00:00'" + ), + "bbox filter": ( + "SELECT AVG(temperature) FROM {t} " + "WHERE lat BETWEEN 0 AND 10 AND lon BETWEEN 0 AND 20" + ), + "projection (1 of 2 vars)": "SELECT AVG(humidity) FROM {t}", + "count only": "SELECT COUNT(*) FROM {t}", +} + + +def bench(table, label, n=5): + """Times each query; returns {query: answer} for equivalence checks.""" + print(f"\n== {label} ==") + answers = {} + for qname, q in QUERIES.items(): + sql = q.format(t=table) + times = [] + for _ in range(n): + t0 = time.perf_counter() + r = con.sql(sql).fetchall() + times.append(time.perf_counter() - t0) + answers[qname] = r[0][0] + med = statistics.median(times) + print( + f" {qname:28s} {med:8.3f}s " + f"(min {min(times):.3f} / max {max(times):.3f}) -> {r[0][0]:.6g}" + ) + return answers + + +# re-scannable stream, registered via the stream wrapper explicitly: +# DuckDB scans every row, no filter/projection pushdown +con.register("t_stream", XarrayArrowStream(ds)) +stream = bench("t_stream", "stream (no pushdown)") + +# default register(): the pushdown pyarrow-dataset path +xql.register(con, "t_pushdown", ds) +pushdown = bench("t_pushdown", "register() [pushdown]") + +# ceiling: materialized pa.Table via pyarrow.dataset +table = xql.read_xarray(ds).read_all() +con.register("t_ceiling", pads.dataset(table)) +ceiling = bench("t_ceiling", "ceiling: in-memory pyarrow.dataset") + +# The timings are only meaningful if every path computed the same thing. +for qname in QUERIES: + a, b, c = stream[qname], pushdown[qname], ceiling[qname] + assert math.isclose(a, b, rel_tol=1e-9) and math.isclose( + a, c, rel_tol=1e-9 + ), f"{qname}: paths disagree — stream={a} pushdown={b} ceiling={c}" +print("\nall paths agree") diff --git a/benchmarks/geospatial/01_ndvi.py b/benchmarks/geospatial/01_ndvi.py index e62b400..7b8272f 100644 --- a/benchmarks/geospatial/01_ndvi.py +++ b/benchmarks/geospatial/01_ndvi.py @@ -45,8 +45,7 @@ import xarray as xr -import xarray_sql as xql - +from _engines import EngineContext from _harness import ( CaseSkipped, assert_grid_close, @@ -111,7 +110,8 @@ def main() -> None: f" scene window: {dict(scene.sizes)} ({n:,} pixels, B04=red/B08=NIR)" ) - ctx = xql.XarrayContext() + ctx = EngineContext() + print(f" engine: {ctx.flavor}") ctx.from_dataset("scene", scene, chunks={"y": 256, "x": 256}) sql = """ @@ -122,7 +122,7 @@ def main() -> None: show_sql(sql) for _ in measured("SQL NDVI"): - got = ctx.sql(sql).to_dataset(dims=["y", "x"]).ndvi + got = ctx.sql_to_dataset(sql, dims=["y", "x"]).ndvi # Array reference: the same formula in pure xarray. ``.compute()`` reads the # window and evaluates it here (the scene is lazy), so this measures the same diff --git a/benchmarks/geospatial/02_climatology.py b/benchmarks/geospatial/02_climatology.py index 8099bdc..96135ab 100644 --- a/benchmarks/geospatial/02_climatology.py +++ b/benchmarks/geospatial/02_climatology.py @@ -42,8 +42,7 @@ import xarray as xr -import xarray_sql as xql - +from _engines import EngineContext from _harness import ( CaseSkipped, assert_grid_close, @@ -81,7 +80,8 @@ def main() -> None: except Exception as exc: # noqa: BLE001 — any failure → skip, not crash raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc - ctx = xql.XarrayContext() + ctx = EngineContext() + print(f" engine: {ctx.flavor}") with timed("register full ERA5 (lazy)"): ctx.from_dataset( "era5", @@ -110,8 +110,10 @@ def main() -> None: # A climatology is a gridded product: round-trip the result back to an # xarray Dataset keyed by (latitude, longitude, hour) — how it is used. for _ in measured("SQL diurnal climatology (lazy read)"): - got = ctx.sql(sql, param_values=_PARAMS).to_dataset( - dims=["latitude", "longitude", "hour"] + got = ctx.sql_to_dataset( + sql, + dims=["latitude", "longitude", "hour"], + param_values=_PARAMS, ) # Array reference: the textbook groupby-over-the-cycle reduction, in °C — diff --git a/benchmarks/geospatial/03_zonal_mean.py b/benchmarks/geospatial/03_zonal_mean.py index 4ab9369..581c25f 100644 --- a/benchmarks/geospatial/03_zonal_mean.py +++ b/benchmarks/geospatial/03_zonal_mean.py @@ -36,8 +36,7 @@ import xarray as xr -import xarray_sql as xql - +from _engines import EngineContext from _harness import ( CaseSkipped, assert_grid_close, @@ -75,7 +74,8 @@ def main() -> None: # ERA5 mixes surface (time, lat, lon) and atmospheric (… level …) variables, # so register it as two tables under an ``era5`` schema. - ctx = xql.XarrayContext() + ctx = EngineContext() + print(f" engine: {ctx.flavor}") with timed("register full ERA5"): ctx.from_dataset( "era5", @@ -101,9 +101,11 @@ def main() -> None: # Round-trip the profile back to an xarray Dataset keyed by latitude. for _ in measured("SQL zonal mean (reads one day)"): - got = ctx.sql( - sql, param_values={"start": _START, "end": _END} - ).to_dataset(dims=["latitude"]) + got = ctx.sql_to_dataset( + sql, + dims=["latitude"], + param_values={"start": _START, "end": _END}, + ) # Array reference: reduce the same day over the two un-grouped axes. for _ in measured("xarray reference"): diff --git a/benchmarks/geospatial/04_anomaly.py b/benchmarks/geospatial/04_anomaly.py index 738956a..ff49748 100644 --- a/benchmarks/geospatial/04_anomaly.py +++ b/benchmarks/geospatial/04_anomaly.py @@ -40,8 +40,7 @@ import xarray as xr -import xarray_sql as xql - +from _engines import EngineContext from _harness import ( CaseSkipped, assert_grid_close, @@ -74,7 +73,8 @@ def main() -> None: except Exception as exc: # noqa: BLE001 — any failure → skip, not crash raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc - ctx = xql.XarrayContext() + ctx = EngineContext() + print(f" engine: {ctx.flavor}") with timed("register full ERA5 (lazy)"): ctx.from_dataset( "era5", @@ -113,8 +113,10 @@ def main() -> None: # The anomaly is a gridded field; round-trip it to (time, lat, lon). for _ in measured("SQL anomaly (climatology CTE self-join, lazy read)"): - got = ctx.sql(sql, param_values=_PARAMS).to_dataset( - dims=["time", "latitude", "longitude"] + got = ctx.sql_to_dataset( + sql, + dims=["time", "latitude", "longitude"], + param_values=_PARAMS, ) # Array reference: grouped broadcast-subtract, in pure xarray (lazy window). diff --git a/benchmarks/geospatial/05_forecast_skill.py b/benchmarks/geospatial/05_forecast_skill.py index 2654377..dac6881 100644 --- a/benchmarks/geospatial/05_forecast_skill.py +++ b/benchmarks/geospatial/05_forecast_skill.py @@ -49,8 +49,7 @@ import pandas as pd import xarray as xr -import xarray_sql as xql - +from _engines import EngineContext from _harness import ( CaseSkipped, assert_grid_close, @@ -144,7 +143,8 @@ def main() -> None: f"leads × 2 models)" ) - ctx = xql.XarrayContext() + ctx = EngineContext() + print(f" engine: {ctx.flavor}") # chunks here is the Arrow batch (partition) size each table streams in, not a # filter — no data is dropped. Both windows are small, so one partition each is # fastest (fewer partitions = fewer Python→Arrow round-trips for the same @@ -172,7 +172,7 @@ def main() -> None: show_sql(sql) for _ in measured("SQL RMSE by (model, lead) — lazy JOIN"): - got = ctx.sql(sql).to_dataset(dims=["model", "lead"]).rmse + got = ctx.sql_to_dataset(sql, dims=["model", "lead"]).rmse for _ in measured("xarray reference"): ref = _reference_rmse(forecasts, truth) diff --git a/benchmarks/geospatial/06_zonal_vector.py b/benchmarks/geospatial/06_zonal_vector.py index d1dce28..0ccc7f1 100644 --- a/benchmarks/geospatial/06_zonal_vector.py +++ b/benchmarks/geospatial/06_zonal_vector.py @@ -44,8 +44,7 @@ import numpy as np import xarray as xr -import xarray_sql as xql - +from _engines import EngineContext from _harness import ( CaseSkipped, assert_grid_close, @@ -101,7 +100,8 @@ def main() -> None: f"vector: {len(_REGIONS)} continental boxes" ) - ctx = xql.XarrayContext() + ctx = EngineContext() + print(f" engine: {ctx.flavor}") with timed("register full ERA5 + regions"): ctx.from_dataset( "era5", @@ -131,9 +131,11 @@ def main() -> None: show_sql(sql) for _ in measured("SQL zonal stats (raster × vector range JOIN)"): - got = ctx.sql( - sql, param_values={"start": _START, "end": _END} - ).to_dataset(dims=["region_id"]) + got = ctx.sql_to_dataset( + sql, + dims=["region_id"], + param_values={"start": _START, "end": _END}, + ) # Array reference: one lazy pass — stack the region masks and reduce. No # .load(): the day's field is read inside this timed block (exactly like the diff --git a/benchmarks/geospatial/_engines.py b/benchmarks/geospatial/_engines.py new file mode 100644 index 0000000..66d769b --- /dev/null +++ b/benchmarks/geospatial/_engines.py @@ -0,0 +1,291 @@ +"""Engine-portable SQL layer for the geospatial suite. + +``GEOBENCH_ENGINE`` selects which SQL engine executes each case's query, +so the same case scripts (same SQL, same datasets, same correctness +assertions) can be measured across engines: + +* ``datafusion`` (default) — ``xql.XarrayContext`` over the native + DataFusion table provider, the suite's original path. Requires the + compiled ``xarray_sql._native`` module; raises at startup when it is + missing instead of falling back. +* ``datafusion-arrow`` — a plain ``datafusion.SessionContext`` scanning + ``xql.arrow_dataset`` (pure Python). +* ``duckdb`` — DuckDB over the same pyarrow pushdown datasets. +* ``polars`` — ``polars.SQLContext`` over ``scan_pyarrow_dataset`` frames. + +Every case builds one :class:`EngineContext`, registers datasets exactly +as it always registered them on ``XarrayContext``, and calls +:meth:`EngineContext.sql_to_dataset`. On the ``datafusion`` path this +is byte-for-byte the original behavior (``from_dataset`` + ``sql`` + +``XarrayDataFrame.to_dataset``); the other engines register one pyarrow +dataset per dimension group under flattened table names +(``era5.surface`` → ``era5_surface`` — rewritten in the SQL text) and the +result rows are round-tripped to an ``xr.Dataset`` through pandas. + +The DataFusion-only UDF cases (07 and the UDF half of 09) build +``xql.XarrayContext`` directly rather than through this layer; the suite +runner records them as n/a for every engine except ``datafusion``. +""" + +from __future__ import annotations + +import datetime +import os +import re +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr + + +def engine_name() -> str: + """The engine selected for this process (``GEOBENCH_ENGINE``).""" + engine = os.environ.get("GEOBENCH_ENGINE", "datafusion") + if engine not in _ENGINES: + raise ValueError(f"GEOBENCH_ENGINE={engine!r}; expected {_ENGINES}") + return engine + + +def _group_tables(name, ds, table_names): + """Split ``ds`` into per-dimension-group tables like XarrayContext does. + + Returns ``[(flat_name, dotted_name, sub_dataset)]``; a uniform dataset + keeps its plain name (flat == dotted == name). + """ + groups: dict[tuple, list] = {} + for var, v in ds.data_vars.items(): + groups.setdefault(tuple(v.dims), []).append(var) + if len(groups) == 1: + return [(name, name, ds)] + out = [] + for dims, variables in groups.items(): + sub = (table_names or {}).get(dims) or "_".join(dims) + out.append((f"{name}_{sub}", f"{name}.{sub}", ds[variables])) + return out + + +def _literal(value: Any) -> str: + """Render a parameter value as a SQL literal (for engines without binds).""" + if isinstance(value, (datetime.datetime, pd.Timestamp, np.datetime64)): + return ( + f"TIMESTAMP '{pd.Timestamp(value).strftime('%Y-%m-%d %H:%M:%S')}'" + ) + if isinstance(value, str): + escaped = value.replace("'", "''") + return f"'{escaped}'" + return repr(value) + + +def _to_ns(pdf: pd.DataFrame, dims: list[str]) -> pd.DataFrame: + """Normalize datetime/timedelta dim columns to ns for label alignment.""" + for col in dims: + dtype = pdf[col].dtype + if pd.api.types.is_datetime64_any_dtype(dtype): + pdf[col] = pdf[col].astype("datetime64[ns]") + elif pd.api.types.is_timedelta64_dtype(dtype): + pdf[col] = pdf[col].astype("timedelta64[ns]") + return pdf + + +def _pandas_to_dataset(pdf: pd.DataFrame, dims: list[str]) -> xr.Dataset: + """Round-trip a SQL result table to a gridded ``xr.Dataset`` by ``dims``.""" + pdf = _to_ns(pdf.copy(), dims) + return xr.Dataset.from_dataframe(pdf.set_index(dims).sort_index()) + + +class EngineContext: + """Uniform register-and-query facade over the suite's SQL engines. + + ``EngineContext(engine)`` instantiates the subclass ``_IMPLS`` maps + the engine name to (default: :func:`engine_name`). Subclasses set + ``flavor`` and implement three hooks: ``_connect`` (open the + engine's connection/context), ``_register`` (attach one pyarrow + dataset under a flat table name), and ``_execute`` (run SQL, + returning a ``pandas.DataFrame``). Engines that bypass the shared + pyarrow-dataset path override :meth:`from_dataset` / + :meth:`sql_to_dataset` instead. + """ + + flavor = "" + + def __new__(cls, engine: str | None = None): + if cls is EngineContext: + cls = _IMPLS[engine or engine_name()] + return super().__new__(cls) + + def __init__(self, engine: str | None = None): + self.engine = engine or engine_name() + self._renames: dict[str, str] = {} + self._connect() + + def _connect(self) -> None: + raise NotImplementedError + + def _register(self, flat: str, dataset) -> None: + raise NotImplementedError + + def _execute(self, sql: str, param_values) -> pd.DataFrame: + raise NotImplementedError + + # -- registration ----------------------------------------------------- + + def from_dataset(self, name, ds, *, chunks=None, table_names=None): + """Register ``ds`` as SQL table(s), mirroring XarrayContext naming.""" + import xarray_sql as xql + + for flat, dotted, sub in _group_tables(name, ds, table_names): + if dotted != flat: + self._renames[dotted] = flat + sub_chunks = ( + {d: c for d, c in chunks.items() if d in sub.dims} + if isinstance(chunks, dict) + else chunks + ) or None + self._register(flat, xql.arrow_dataset(sub, sub_chunks)) + + # -- querying ---------------------------------------------------------- + + def _rewrite(self, sql: str, param_values) -> str: + for dotted, flat in self._renames.items(): + sql = re.sub(rf"\b{re.escape(dotted)}\b", flat, sql) + return sql + + def sql_to_dataset( + self, sql: str, *, dims: list[str], param_values=None + ) -> xr.Dataset: + """Run ``sql`` and round-trip the result to an ``xr.Dataset``.""" + pdf = self._execute(self._rewrite(sql, param_values), param_values) + return _pandas_to_dataset(pdf, dims) + + +class _DataFusionNative(EngineContext): + """``xql.XarrayContext`` over the native DataFusion table provider.""" + + flavor = "datafusion (XarrayContext, native)" + + def _connect(self): + try: + import xarray_sql._native # noqa: F401 + except ImportError as exc: + raise RuntimeError( + "GEOBENCH_ENGINE=datafusion requires the compiled " + "xarray_sql._native module (`maturin develop`); use " + "GEOBENCH_ENGINE=datafusion-arrow for the pure-Python " + "pyarrow-dataset path." + ) from exc + import xarray_sql as xql + + self._ctx = xql.XarrayContext() + + def from_dataset(self, name, ds, *, chunks=None, table_names=None): + self._ctx.from_dataset(name, ds, chunks=chunks, table_names=table_names) + + def sql_to_dataset(self, sql, *, dims, param_values=None): + df = ( + self._ctx.sql(sql, param_values=param_values) + if param_values + else self._ctx.sql(sql) + ) + return df.to_dataset(dims=dims) + + +class _DataFusionArrow(EngineContext): + """Plain ``datafusion.SessionContext`` over ``xql.arrow_dataset``.""" + + flavor = "datafusion-arrow (pyarrow dataset, pure Python)" + + def _connect(self): + from datafusion import SessionContext + + self._ctx = SessionContext() + + def _register(self, flat, dataset): + self._ctx.register_dataset(flat, dataset) + + def _execute(self, sql, param_values): + df = ( + self._ctx.sql(sql, param_values=param_values) + if param_values + else self._ctx.sql(sql) + ) + return df.to_pandas() + + +class _DuckDB(EngineContext): + """DuckDB over the same pyarrow pushdown datasets.""" + + flavor = "duckdb" + + def _connect(self): + import duckdb + + self._con = duckdb.connect() + + def _register(self, flat, dataset): + self._con.register(flat, dataset) + + def _execute(self, sql, param_values): + return self._con.execute(sql, param_values or {}).df() + + +class _Polars(EngineContext): + """``polars.SQLContext`` over ``scan_pyarrow_dataset`` frames. + + Keeps the pyarrow datasets and builds the SQLContext per query. + Polars' SQL layer renders TIMESTAMP literals as strptime-plus-cast + expressions it cannot convert to pyarrow filters, so a WHERE over + the full archive would scan everything; the same bounds applied as + native expressions *do* push down. ``_execute`` therefore + pre-filters each frame with the query's window parameters + (identical predicate to the SQL WHERE, which still runs on top). + """ + + flavor = "polars (SQLContext + expression window pushdown)" + + # The window bounds a query passes as parameters, as (column, low + # param, high param); applied per registered frame when the column + # exists — the same inclusive predicate the SQL WHERE states. + _BOUND_PARAMS = ( + ("time", "start", "end"), + ("latitude", "lat_s", "lat_n"), + ("longitude", "lon_w", "lon_e"), + ) + + def _connect(self): + self._tables: dict[str, Any] = {} + + def _register(self, flat, dataset): + self._tables[flat] = dataset + + def _rewrite(self, sql, param_values): + sql = super()._rewrite(sql, param_values) + for key, value in (param_values or {}).items(): + sql = re.sub(rf"\${key}\b", _literal(value), sql) + return sql + + def _execute(self, sql, param_values): + import polars as pl + + ctx = pl.SQLContext() + params = param_values or {} + for flat, dataset in self._tables.items(): + lf = pl.scan_pyarrow_dataset(dataset) + names = set(dataset.schema.names) + for col, lo, hi in self._BOUND_PARAMS: + if col in names and lo in params and hi in params: + lf = lf.filter( + (pl.col(col) >= params[lo]) + & (pl.col(col) <= params[hi]) + ) + ctx.register(flat, lf) + return ctx.execute(sql, eager=True).to_pandas() + + +_IMPLS = { + "datafusion": _DataFusionNative, + "datafusion-arrow": _DataFusionArrow, + "duckdb": _DuckDB, + "polars": _Polars, +} +_ENGINES = tuple(_IMPLS) diff --git a/benchmarks/geospatial/engine_suite.py b/benchmarks/geospatial/engine_suite.py new file mode 100644 index 0000000..b8cb6e4 --- /dev/null +++ b/benchmarks/geospatial/engine_suite.py @@ -0,0 +1,699 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "aiohttp", +# "coiled", +# "datafusion", +# "duckdb", +# "earthengine-api", +# "gcsfs", +# "numpy", +# "pandas", +# "polars", +# "psutil", +# "pyarrow", +# "pyproj", +# "pystac-client", +# "requests", +# "scipy", +# "shapely", +# "xarray", +# "xarray-sql", +# "xee", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""The geospatial suite across engines and VM sizes, via Coiled Functions. + +Runs the nine geospatial cases (``01_ndvi`` … ``09_warp``) under every +SQL engine the suite supports — DataFusion over the native table +provider (``datafusion``, the original path; requires the compiled +native module), DataFusion over the pure-Python pyarrow dataset +(``datafusion-arrow``), DuckDB, and Polars, selected per process +through ``GEOBENCH_ENGINE`` and the ``_engines`` facade — on one reused +Coiled VM per machine size, driven in parallel across sizes. + +``datafusion`` cells need the compiled native module; the first such +cell on a VM provisions it (see :func:`_ensure_native`) and records the +outcome under ``native`` in its result, so a failed build surfaces as +that cell's error rather than a VM startup failure. The driver's own +build is copied in when it imports on that platform; otherwise rustup +(minimal profile) is installed and the shipped crate is built via the +project's maturin build backend, once per source digest. + +The measurement protocol is exactly ``run_perf.sh``'s: every repetition +is a **fresh process** with no warm-up (``GEOBENCH_PROFILE=1 +GEOBENCH_WARMUP=0 GEOBENCH_REPS=1``), so the SQL side and the xarray +reference each pay a cold read on every rep, and each case's own +correctness assertion (SQL answer == array reference) must pass for the +timing to count. The xarray-reference timings are engine-independent; +the tables report the reference column from the DataFusion runs. + +Coverage notes, recorded rather than hidden: cases 07 and 09 build +DataFusion scalar UDFs on ``xql.XarrayContext`` directly, so every +engine except ``datafusion`` is marked n/a; case 08 reads +through Earth Engine and is left on the original context (EE-gated); +cases 07–09 skip cleanly wherever Earth Engine auth is unavailable +(e.g. on the benchmark VMs) with the reason recorded. + +Each (vm, case, engine) cell returns a plain dict; every completed cell +is appended to a local ``--jsonl`` file immediately, and the driver +prints one timestamped line per event. + +Usage:: + + uv run benchmarks/geospatial/engine_suite.py --local --reps 1 \ + --cases 02_climatology --vms local # in-process check + cd /tmp && uv run ~/path/to/xarray-sql/benchmarks/geospatial/engine_suite.py + +(or ``python .../engine_suite.py`` from an environment that already has +the dependencies; the inline metadata above is for ``uv run``.) + +Remote runs (the second form) must be launched from a working directory +outside the repository: Coiled's package sync resolves the repo's +``uv.lock`` when it finds one at the cwd, and that lock does not carry +the driver's dependencies. ``--local`` runs work from anywhere. +""" + +from __future__ import annotations + +import argparse +import csv +import datetime +import io +import json +import os +import statistics +import subprocess +import sys +import tarfile +import tempfile +import threading +import time +from pathlib import Path +from typing import Any + +REGION = "us-central1" + +CASES = [ + "01_ndvi", + "02_climatology", + "03_zonal_mean", + "04_anomaly", + "05_forecast_skill", + "06_zonal_vector", + "07_reproject_udf", + "08_regrid_weights", + "09_warp", +] +ENGINES = ["datafusion", "datafusion-arrow", "duckdb", "polars"] +# Cases whose SQL builds DataFusion scalar UDFs on XarrayContext directly +# (07, and the UDF half of 09): they run only under ``datafusion``. Case +# 08 is portable SQL but Earth-Engine-gated, so it stays on the original +# context. +NOT_PORTABLE = { + "07_reproject_udf": "n/a (DataFusion scalar UDF)", + "09_warp": "n/a (DataFusion scalar UDF)", + "08_regrid_weights": "not ported (Earth-Engine-gated case)", +} +VM_SIZES = ["e2-standard-8", "e2-standard-16", "e2-standard-32"] + + +def cluster_name(vm: str) -> str: + return "xql-geo-" + vm.replace("standard-", "") + + +# -------------------------------------------------------------------------- +# Remote side (runs inside the coiled function, or locally with --local) +# -------------------------------------------------------------------------- + + +def _install_src(src_targz: bytes | None) -> tuple[str, str]: + """Unpack the shipped source tree; returns (sys.path root, geo dir). + + The root is keyed by the tarball's hash so a reused warm VM never + serves a stale tree from an earlier driver run. + """ + import hashlib + + digest = hashlib.md5(src_targz or b"local").hexdigest()[:10] + root = f"/tmp/xql_geo_src_{digest}" + # Written only after extractall returns, so an interrupted + # extraction is retried instead of reused as a corrupt tree. + marker = os.path.join(root, ".extraction-complete") + if src_targz is not None and not os.path.exists(marker): + os.makedirs(root, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(src_targz), mode="r:gz") as tf: + tf.extractall(root) # noqa: S202 — our own tarball + open(marker, "w").close() + return root, os.path.join(root, "benchmarks", "geospatial") + + +def _run_logged(cmd, **kwargs) -> None: + """subprocess.run(check=True) that surfaces stderr on failure.""" + proc = subprocess.run(cmd, capture_output=True, text=True, **kwargs) + if proc.returncode != 0: + raise RuntimeError( + f"{cmd if isinstance(cmd, str) else ' '.join(cmd)} failed:\n" + f"{proc.stderr[-800:]}" + ) + + +def _ensure_native(src_root: str) -> str: + """Make ``xarray_sql._native`` importable from ``src_root``. + + Tries, in order: the module already present in the tree; the one + installed in this interpreter's environment (copied in, when built + for this platform); a from-source build of the shipped crate — + rustup (minimal profile) plus ``pip wheel``, which drives the + project's maturin build backend — installed over the pure-Python + copy. The built module lands in ``src_root``, which is keyed by + source digest, so a warm VM builds at most once per source state. + + Returns a status string for the run log. + """ + import glob + import importlib.util + import shutil + + # cwd well inside the tree, so `-c` resolves xarray_sql only through + # PYTHONPATH=src_root — the same view the case subprocesses get. + geo_dir = os.path.join(src_root, "benchmarks", "geospatial") + env = dict(os.environ, PYTHONPATH=src_root) + + def _importable() -> bool: + return ( + subprocess.run( + [sys.executable, "-c", "import xarray_sql._native"], + env=env, + cwd=geo_dir, + capture_output=True, + ).returncode + == 0 + ) + + if _importable(): + return "importable" + + try: + spec = importlib.util.find_spec("xarray_sql._native") + except ImportError: + spec = None + if spec is not None and spec.origin: + shutil.copy2( + spec.origin, + os.path.join(src_root, "xarray_sql", os.path.basename(spec.origin)), + ) + if _importable(): + return "copied from driver environment" + + t0 = time.monotonic() + build_env = dict(env) + cargo_bin = os.path.expanduser("~/.cargo/bin") + build_env["PATH"] = cargo_bin + os.pathsep + build_env.get("PATH", "") + if shutil.which("cargo", path=build_env["PATH"]) is None: + _run_logged( + "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs " + "| sh -s -- -y --profile minimal --default-toolchain stable", + shell=True, + env=build_env, + ) + wheel_dir = os.path.join(src_root, "wheelhouse") + _run_logged( + [ + sys.executable, + "-m", + "pip", + "wheel", + "--no-deps", + "-w", + wheel_dir, + src_root, + ], + env=build_env, + ) + wheel = sorted(glob.glob(os.path.join(wheel_dir, "xarray_sql-*.whl")))[-1] + _run_logged( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--upgrade", + "--target", + src_root, + wheel, + ], + ) + if not _importable(): + raise RuntimeError(f"built {wheel} but xarray_sql._native still fails") + return f"built from source in {time.monotonic() - t0:.0f}s" + + +def run_case_cell( + case: str, + engine: str, + reps: int, + src_targz: bytes | None = None, + rep_timeout: float = 600.0, +) -> dict: + """One (case, engine) cell: ``reps`` fresh-process cold runs.""" + result: dict[str, Any] = { + "case": case, + "engine": engine, + "status": "ok", + "reps": [], + } + try: + src_root, geo_dir = _install_src(src_targz) + if engine == "datafusion": + try: + result["native"] = _ensure_native(src_root) + except Exception: + result["native"] = "provisioning failed" + raise + env = dict( + os.environ, + GEOBENCH_ENGINE=engine, + GEOBENCH_PROFILE="1", + GEOBENCH_WARMUP="0", + GEOBENCH_REPS="1", + PYTHONUNBUFFERED="1", + PYTHONPATH=src_root, + ) + rows: list[dict] = [] + for rep in range(1, reps + 1): + with tempfile.NamedTemporaryFile(suffix=".csv") as csv_file: + env["GEOBENCH_CSV"] = csv_file.name + t0 = time.perf_counter() + try: + proc = subprocess.run( + [sys.executable, f"{case}.py"], + cwd=geo_dir, + env=env, + capture_output=True, + text=True, + timeout=rep_timeout, + ) + except subprocess.TimeoutExpired: + result["reps"].append({"rep": rep, "status": "timeout"}) + result["status"] = "timeout" + break + wall = round(time.perf_counter() - t0, 3) + out = proc.stdout + # Exit status outranks the skip marker: a case that + # prints SKIPPED and then crashes is an error, not a skip. + if proc.returncode == 0 and "SKIPPED" in out: + reason = next( + ( + line.split("SKIPPED:", 1)[1].strip() + for line in out.splitlines() + if "SKIPPED:" in line + ), + "skipped", + ) + result.update(status="skip", reason=reason[:300]) + break + if proc.returncode != 0: + result.update( + status="error", + error=(proc.stderr.strip() or out.strip())[-600:], + ) + break + flavor = next( + ( + line.split("engine:", 1)[1].strip() + for line in out.splitlines() + if "engine:" in line + ), + engine, + ) + result["flavor"] = flavor + with open(csv_file.name) as fh: + for row in csv.DictReader(fh): + row["rep"] = rep + rows.append(row) + result["reps"].append( + {"rep": rep, "status": "ok", "wall_s": wall} + ) + print(f"[vm] {case} x {engine}: rep {rep} {wall}s", flush=True) + steps: dict[str, dict] = {} + for row in rows: + step = steps.setdefault( + row["step"], {"times_s": [], "peak_mb": 0.0} + ) + step["times_s"].append(float(row["t_median_s"])) + step["peak_mb"] = max(step["peak_mb"], float(row["peak_mb"])) + for step in steps.values(): + times = step["times_s"] + step["median_s"] = round(statistics.median(times), 3) + step["min_s"] = round(min(times), 3) + step["max_s"] = round(max(times), 3) + step["n"] = len(times) + result["steps"] = steps + if result["status"] == "ok" and not steps: + result.update(status="error", error="no perf rows produced") + except Exception as exc: # noqa: BLE001 — cell errors are data + result.update(status="error", error=f"{type(exc).__name__}: {exc}") + return result + + +def probe_environment(src_targz: bytes | None = None) -> dict: + """Machine spec + package versions, gathered where the cells run.""" + import platform + + _install_src(src_targz) + info = { + "platform": platform.platform(), + "python": platform.python_version(), + "cpus": os.cpu_count(), + "node": platform.node(), + } + try: + import psutil + + info["mem_gb"] = round(psutil.virtual_memory().total / 2**30, 1) + except Exception: # noqa: BLE001 + pass + versions = {} + for pkg in ["duckdb", "polars", "datafusion", "pyarrow", "xarray"]: + try: + from importlib import metadata + + versions[pkg] = metadata.version(pkg) + except Exception: # noqa: BLE001 + versions[pkg] = "missing" + return {"machine": info, "versions": versions} + + +# -------------------------------------------------------------------------- +# Driver +# -------------------------------------------------------------------------- + +_PRINT_LOCK = threading.Lock() + + +def log(vm: str, msg: str) -> None: + now = datetime.datetime.now().strftime("%H:%M:%S") + with _PRINT_LOCK: + print(f"[{now}][{vm}] {msg}", flush=True) + + +def _pack_src() -> bytes: + """gzip tar of xarray_sql, benchmarks/geospatial, and the Rust crate. + + Byte-identical for identical file contents (gzip and tar metadata + normalized): _install_src keys its extraction root — and therefore + _ensure_native's build cache on a warm VM — on the digest of these + bytes. + """ + import gzip + + repo = Path(__file__).resolve().parents[2] + + def _normalize(info: tarfile.TarInfo) -> tarfile.TarInfo: + info.mtime = 0 + info.mode = 0o644 + info.uid = info.gid = 0 + info.uname = info.gname = "" + return info + + paths: list[Path] = [] + for rel in ["xarray_sql", "benchmarks/geospatial"]: + paths += [ + p + for p in sorted((repo / rel).rglob("*.py")) + if "__pycache__" not in p.parts + ] + # The crate sources, so `datafusion` cells can build the native + # module where it is not already importable (see _ensure_native). + for rel in [ + "src", + "Cargo.toml", + "Cargo.lock", + "pyproject.toml", + "README.md", + ]: + target = repo / rel + paths += ( + [p for p in sorted(target.rglob("*")) if p.is_file()] + if target.is_dir() + else [target] + ) + buf = io.BytesIO() + with gzip.GzipFile(fileobj=buf, mode="wb", mtime=0) as gz: + # GzipFile is a binary stream at runtime; typeshed wants IO[bytes]. + with tarfile.open(fileobj=gz, mode="w") as tf: # type: ignore[arg-type] + for path in paths: + tf.add( + path, + arcname=str(path.relative_to(repo)), + filter=_normalize, + ) + return buf.getvalue() + + +def _drive_vm(vm, cells, args, src, results, jsonl_lock): + """Run every (case, engine) cell for one VM size, sequentially.""" + if vm == "local": + remote_cell, remote_probe = run_case_cell, probe_environment + submit = None + else: + import coiled + + deco = coiled.function( + name=cluster_name(vm), + vm_type=vm, + region=REGION, + keepalive="10m", + idle_timeout="20 minutes", + spot_policy="on-demand", + package_sync_ignore=["xarray_sql", "xarray-sql"], + environ={"PYTHONUNBUFFERED": "1"}, + ) + remote_cell, remote_probe = deco(run_case_cell), deco(probe_environment) + submit = remote_cell.submit + + log(vm, "probing environment (provisions the VM on first call)...") + meta = None + for attempt in range(1, 4): + try: + meta = remote_probe(src) + break + except Exception as exc: # noqa: BLE001 — transient control plane + log(vm, f"probe attempt {attempt} failed: {exc}"[:200]) + if attempt < 3: + time.sleep(30 * attempt) + if meta is None: + log(vm, "giving up: VM never came up") + return + log(vm, f"machine: {json.dumps(meta['machine'])}") + total = len(cells) + for k, (case, engine) in enumerate(cells, 1): + tag = f"cell {k}/{total} {case} x {engine}" + if case in NOT_PORTABLE and engine != "datafusion": + rec = { + "case": case, + "engine": engine, + "status": "n/a", + "reason": NOT_PORTABLE[case], + } + else: + log(vm, f"{tag}: submitted") + t0 = time.monotonic() + try: + if submit is None: + rec = run_case_cell(case, engine, args.reps, src) + else: + fut = submit(case, engine, args.reps, src) + rec = fut.result(timeout=args.cell_timeout) + except Exception as exc: # noqa: BLE001 + rec = { + "case": case, + "engine": engine, + "status": "error", + "error": f"{type(exc).__name__}: {exc}"[:500], + } + rec["cell_wall_s"] = round(time.monotonic() - t0, 1) + rec["vm"] = vm + if rec.get("native", "importable") != "importable": + log(vm, f"{tag}: native module {rec['native']}") + results.append(rec) + with jsonl_lock, open(args.jsonl, "a") as fh: + fh.write(json.dumps(rec) + "\n") + if rec["status"] == "ok": + sql_step = next( + ( + s + for name, s in rec.get("steps", {}).items() + if name.startswith("SQL") + ), + None, + ) + brief = ( + f"SQL median {sql_step['median_s']}s (n={sql_step['n']})" + if sql_step + else "ok" + ) + log(vm, f"{tag}: ok {brief} [{rec.get('flavor', engine)}]") + else: + detail = rec.get("reason") or rec.get("error", "") + log(vm, f"{tag}: {rec['status']} {detail[:200]}") + meta_rec = {"vm": vm, "case": "_meta", "engine": "", **meta} + results.append(meta_rec) + with jsonl_lock, open(args.jsonl, "a") as fh: + fh.write(json.dumps(meta_rec) + "\n") + if submit is not None: + # Shut the VM down the moment its last cell finishes — don't + # leave the teardown to keepalive expiry. + try: + remote_cell.cluster.shutdown() + log(vm, "cluster shut down") + except Exception as exc: # noqa: BLE001 — teardown best-effort + log(vm, f"cluster shutdown failed: {exc}"[:200]) + + +def _markdown(results: list[dict]) -> str: + """One case x engine table per VM (SQL median s; reference column).""" + out = [] + vms = list(dict.fromkeys(r["vm"] for r in results)) + for vm in vms: + rows = [r for r in results if r["vm"] == vm and r["case"] != "_meta"] + if not rows: + continue + cases = list(dict.fromkeys(r["case"] for r in rows)) + out.append(f"\n### {vm}\n") + out.append("| Case | " + " | ".join(ENGINES) + " | xarray reference |") + out.append("|---|" + "---|" * (len(ENGINES) + 1)) + by = {(r["case"], r["engine"]): r for r in rows} + for case in cases: + cells = [] + for engine in ENGINES: + r = by.get((case, engine)) + if r is None: + cells.append("-") + elif r["status"] != "ok": + detail = r.get("reason") or r.get("error", "") + cells.append(f"{r['status']}: {detail[:40]}") + else: + s = next( + ( + v + for k, v in r["steps"].items() + if k.startswith("SQL") + ), + None, + ) + cells.append( + f"{s['median_s']:.3f}s (n={s['n']}, " + f"{s['peak_mb']:.0f} MB)" + if s + else "?" + ) + df_run = by.get((case, "datafusion"), {}) + ref = (df_run.get("steps") or {}).get("xarray reference") + ref_text = ( + f"{ref['median_s']:.3f}s ({ref['peak_mb']:.0f} MB)" + if ref + else "-" + ) + out.append(f"| {case} | " + " | ".join(cells) + f" | {ref_text} |") + return "\n".join(out) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--local", action="store_true") + ap.add_argument("--reps", type=int, default=5) + ap.add_argument("--cases", default=",".join(CASES)) + ap.add_argument("--engines", default=",".join(ENGINES)) + ap.add_argument("--vms", default=",".join(VM_SIZES)) + ap.add_argument("--cell-timeout", type=float, default=1800.0) + ap.add_argument("--out", default="engine_suite_results.json") + ap.add_argument("--jsonl", default="engine_suite_results.jsonl") + args = ap.parse_args() + + cases = [c for c in args.cases.split(",") if c] + engines = [e for e in args.engines.split(",") if e] + # dict.fromkeys: duplicate VM names would share one cluster and defeat + # the incomplete-run detection, which matches records by VM name. + vms = ( + ["local"] + if args.local + else list(dict.fromkeys(v for v in args.vms.split(",") if v)) + ) + cells = [(c, e) for c in cases for e in engines] + log("plan", f"{len(vms)} VMs x {len(cells)} cells, reps={args.reps}") + for c, e in cells: + note = ( + f" [{NOT_PORTABLE[c]}]" + if c in NOT_PORTABLE and e != "datafusion" + else "" + ) + log("plan", f" {c} x {e}{note}") + src = _pack_src() + log("plan", f"packed source: {len(src) / 1024:.0f} KiB") + + open(args.jsonl, "w").close() + results: list[dict] = [] + jsonl_lock = threading.Lock() + threads = [ + threading.Thread( + target=_drive_vm, + args=(vm, cells, args, src, results, jsonl_lock), + name=vm, + ) + for vm in vms + ] + for i, t in enumerate(threads): + if i: # stagger: concurrent package-sync scans trip the server + time.sleep(20) + t.start() + for t in threads: + t.join() + + payload = { + "meta": { + "region": REGION, + "reps": args.reps, + "protocol": "fresh process per rep, no warmup, cold reads", + }, + "results": results, + } + with open(args.out, "w") as fh: + json.dump(payload, fh, indent=2) + md = _markdown([r for r in results if r.get("case")]) + md_path = os.path.splitext(args.out)[0] + ".md" + with open(md_path, "w") as fh: + fh.write(md + "\n") + print(md) + log("done", f"wrote {args.out}, {md_path}, {args.jsonl}") + # A VM that never produced its _meta record never ran its cells; + # exit nonzero so partial runs cannot pass for complete ones. + incomplete = [ + vm + for vm in vms + if not any( + r.get("vm") == vm and r.get("case") == "_meta" for r in results + ) + ] + if incomplete: + log("done", f"incomplete run: no results from {', '.join(incomplete)}") + sys.exit(1) + failed = [ + f"{r['vm']}/{r['case']}/{r['engine']}" + for r in results + if r.get("status") in ("error", "timeout") + ] + if failed: + log("done", f"failed cells: {', '.join(failed)}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/docs/engines.md b/docs/engines.md new file mode 100644 index 0000000..11dacac --- /dev/null +++ b/docs/engines.md @@ -0,0 +1,259 @@ +# Engines + +xarray-sql translates **data, not queries**. It does not own a SQL +dialect, a query IR, or a transpiler: you pick a query engine and write +that engine's native SQL, using that engine's extension ecosystem +(spatial, H3, …) directly. xarray-sql implements the two seams no engine +builds for itself: + +1. **register** — a lazy `xarray.Dataset` becomes a table on the + engine's own connection, streamed as Arrow record batches only while + a query executes. +2. **round-trip** — the engine's Arrow result plus the source Dataset as + a *template* becomes a labeled `xr.Dataset` again: attrs, non-dim + coordinates, and dtypes recovered. *SQL in, array out.* + +Everything between the seams — geometry functions, dialects, +optimizers — belongs to the engine. + +## DataFusion (default) + +DataFusion is the built-in engine, wrapped in a session: + +```python +import xarray_sql as xql + +ctx = xql.XarrayContext() +ctx.from_dataset("era5", ds, chunks={"time": 24}) +result = ctx.sql("SELECT ... FROM era5").to_dataset() +``` + +This is the deepest integration: the Rust `TableProvider` gives +partition pruning on dimension predicates, projection pushdown to the +storage layer, exact per-partition statistics for the optimizer, and a +lazy chunked round-trip (`to_dataset(chunks=...)`). + +The generic entry point dispatches here too: `xql.register(ctx, "era5", ds)` +works on any `datafusion.SessionContext`. + +### Relation to zarr-datafusion + +[zarr-datafusion](https://crates.io/crates/zarr-datafusion) extends +DataFusion with SQL over Zarr stores natively (early days — a single +0.1.0 release at the time of writing) — for plain-Zarr sources that is +the engine-native path, the same role duckdb-zarr plays for DuckDB. +This library's role is complementary there too: anything xarray can +open (NetCDF, GRIB, Earth Engine via Xee, CF-decoded/virtual datasets, +in-memory arrays), and the round-trip from a query result back to a +labeled Dataset, which no engine extension provides. + + +## DuckDB (adapter) + +```sh +pip install 'xarray-sql[duckdb]' +``` + +```python +import duckdb +import xarray_sql as xql + +con = duckdb.connect() +xql.register(con, "era5", ds) # seam 1 + +con.sql("INSTALL spatial; LOAD spatial;") # DuckDB's own shelf +rel = con.sql(""" + SELECT time, lat, lon, AVG(t2m) AS t2m + FROM era5 + WHERE lat BETWEEN 40 AND 41 + GROUP BY time, lat, lon +""") + +out = xql.to_dataset(rel, template=ds) # seam 2 +``` + +The adapter registers an `XarrayPushdownDataset` — a +`pyarrow.dataset.Dataset` subclass (the same pattern +[Lance](https://github.com/lancedb/lance) uses for `LanceDataset`), so +DuckDB hands each query's column list and pushed predicate to the +source. The scan then loads only the data variables the query mentions, +prunes chunks whose coordinate ranges cannot satisfy the predicate +(via Arrow's own guarantee simplification — sound for every predicate +shape), and prefetches surviving chunks on a thread pool. The table is +lazy, re-queryable, and a bounding-box query over a billions-of-pixels +raster answers in about a second because only the intersecting chunks +are ever read. + +Pushed comparison filters are a correctness contract in DuckDB (it +deletes them from its own plan), so the scanner always applies the +exact expression via pyarrow — pruning is only an optimization on top. +`XarrayArrowStream`, the dependency-light re-scannable C-stream wrapper +without pushdown, remains available as a fallback. + +Details that matter in production: + +- **Finely partitioned axes** (e.g. hourly-chunked reanalysis time with + hundreds of thousands of chunks) prune through a two-level shadow: + a coarse pass over at most 1024 buckets, refined per surviving + bucket — so pruning cost is bounded regardless of chunk count, and + refinement is skipped when a predicate matches most of the axis. +- **Tuning** via `xql.register(con, name, ds, batch_size=..., + prefetch=..., prefetch_bytes=..., coalesce_rows=...)`: `prefetch` + bounds concurrent chunk loads, `prefetch_bytes` caps estimated bytes + in flight, `coalesce_rows` merges runs of consecutive surviving + chunks into single reads, `batch_size` caps rows per Arrow batch. + See the [performance guide](performance.md#the-memory-contract). +- **Source parallelism matters as much as the adapter's**: rioxarray + serializes GDAL tile reads behind a lock by default, which caps any + scan at single-stream speed regardless of `prefetch`. Open rasters + with `rioxarray.open_rasterio(..., lock=False)` — measured 6× on + full scans of a 9-billion-pixel cloud GeoTIFF, making remote reads + as fast as a local copy. + + +### Relation to duckdb-zarr + +[duckdb-zarr](https://github.com/xqlsystems/duckdb-zarr) reads Zarr +stores natively inside DuckDB, with projection pushdown — for +plain-Zarr sources it is the engine-native path and will beat this +adapter. The adapter's role is complementary: anything xarray can open +(NetCDF, GRIB, Earth Engine via Xee, CF-decoded/virtual datasets, +in-memory arrays), and the round-trip from a DuckDB result back to a +labeled Dataset, which no engine extension provides. + +## Polars (via the pyarrow dataset protocol) + +```sh +pip install 'xarray-sql[polars]' +``` + +`xql.arrow_dataset(ds)` returns a real `pyarrow.dataset.Dataset`, so +any engine that consumes that protocol gets the same lazy scan with +projection pushdown and coordinate-range chunk pruning — no adapter +code at all. Polars works today: + +```python +import polars as pl +import xarray_sql as xql + +lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) +out = ( + lf.filter(pl.col("lat") > 0) + .group_by("time") + .agg(pl.col("t2m").mean()) + .collect() +) +xql.to_dataset(out, template=ds) # polars frames speak Arrow PyCapsule +``` + +Polars pushes its predicate and column selection into the dataset scan +(verified: a filtered group-by read 1 of 20 chunks and 3 of 5 columns), +and its results round-trip through `xql.to_dataset` unchanged. The +chunked round-trip is fully supported: windows re-execute on Polars' +streaming engine. + + +## Engine support matrix + +What each integration provides. Known issues and constraints live on +[Known issues & limitations](limitations.md). + +| | DataFusion | DuckDB | Polars | +|---|---|---|---| +| Register | `XarrayContext` / any `SessionContext` | `xql.register(con, name, ds)` | `pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` | +| Projection pushdown | yes | yes | yes | +| Chunk pruning on dim predicates | yes | yes | yes | +| Eager round-trip (`xql.to_dataset`) | yes | yes | yes | +| Chunked round-trip (`chunks=`) | re-execution | `spill=True` [^spill-only] | re-execution (streaming engine) | +| `geometry` column ([geospatial](geospatial.md#geoarrow-point-geometry-columns)) | annotated WKB passes through | native `GEOMETRY` (`"wkb"` encoding) | plain binary/struct | +| Mixed-dimension datasets | one schema, `name.group` tables | `_` tables | filter `data_vars` before `arrow_dataset` | +| Version floor | bundled (core dependency) | `duckdb >= 1.4` (tested on 1.5) | tested on `polars 1.42` | + +[^spill-only]: Why DuckDB relations do not re-execute — and two other + engine-specific issues worth knowing — is explained on + [Known issues & limitations](limitations.md#engine-specific-issues). + +## The lazy round-trip across engines + +`xql.to_dataset(result, chunks=...)` reconstructs a query result as a +*chunked, lazy* `xr.Dataset`: each output chunk re-executes the engine's +query narrowed to that chunk's coordinate window on first access. Over a +table registered through xarray-sql, the window's range predicate flows +back into chunk pruning at the source — accessing one output chunk reads +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 -- "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)"] + HD -- "DuckDB relation" --> NO["NotImplementedError
(upstream deadlock — see
Known issues)"] + HD -- "one-shot Arrow stream" --> NO2["TypeError
(nothing to re-execute)"] + SP -- "True / directory" --> SPL["one-pass spill: stream once
(bounded memory) → temp Parquet →
windows re-execute against the file
(row-group pruning); file deleted
with the Dataset"] +``` + +**Choosing:** re-execution pays per window — right when you'll touch a +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. + +Two knobs matter at scale: + +- `coords="template"` trusts the template's coordinate arrays instead of + running one `DISTINCT` query per dimension — construction then reads + nothing at all. Only valid when the result spans the template's full + extent (an unfiltered scan). On ARCO-ERA5 (1.32M hourly chunks) this + builds a lazy view over a 1.37-trillion-row table in ~0.3 s with zero + source reads; a one-day window then computes in ~2 s reading only the + source chunks under the window. +- Contiguous windows become two-literal range predicates the engine can + push and the source can prune on; stepped or fancy selections fall + back to explicit value lists (exact, just less prunable). + +With `spill=True`, the result is streamed **once** (bounded memory) +into a temporary Parquet file and windows re-execute against that +file — the right shape when most of the result will be touched, the +only chunked option for one-shot Arrow streams, and the required path +for DuckDB relations (see the DuckDB section above). Polars/DataFusion +re-execution remains the default for window-at-a-time access over huge +results. + +## Adding an engine + +An adapter implements one small contract +(`xarray_sql.backends.base.EngineAdapter`): `matches(con)` recognizes +the engine's connection object without importing the engine, and +`register(con, name, ds, chunks=...)` attaches the Dataset as a table. +Arrow C streams are the common wire; pushdown quality is where adapters +differ. The round-trip needs no per-engine work as long as the engine +can hand back Arrow. + +A complete adapter, modeled on the DuckDB one: + +```python +from xarray_sql.backends.base import register_adapter +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + +@register_adapter +class AcmeAdapter: + """Registers Datasets on acme.Connection objects.""" + + @staticmethod + def matches(con) -> bool: + # type inspection only, so `acme` stays an optional dependency + return type(con).__module__.split(".")[0] == "acme" + + @staticmethod + def register(con, name, ds, *, chunks=None, **kwargs): + dataset = XarrayPushdownDataset(ds, chunks, **kwargs) + con.register_arrow(name, dataset) # the engine's own API + return con +``` + +`xql.register(con, "t", ds)` then dispatches here whenever `matches` +recognizes the connection. Engines that consume the pyarrow dataset +protocol (DuckDB, Polars) get projection pushdown and chunk pruning for +free; an engine that only accepts Arrow streams can register +`XarrayArrowStream(ds)` instead, trading pushdown away. diff --git a/docs/examples.md b/docs/examples.md index f80c5b5..483fb2a 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -142,6 +142,17 @@ ctx.sql('SELECT * FROM goes.scalar').to_pandas().shape # -> (1, 89) Override the default name like any other group with `table_names={(): 'metadata'}`. A runnable version of the ERA5 example lives at -[`perf_tests/era5_temp_profile.py`](../perf_tests/era5_temp_profile.py). +[`perf_tests/era5_temp_profile.py`](https://github.com/xqlsystems/xarray-sql/blob/main/perf_tests/era5_temp_profile.py). [arco-era5]: https://github.com/google-research/arco-era5 + + +## The same tables on DuckDB and Polars + +Every example above registers through an `XarrayContext`, 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 +splitting rules for mixed-dimension Datasets, same round-trip through +`xql.to_dataset(result, template=ds)`. See [Engines](engines.md) for the +support matrix and per-engine details. diff --git a/docs/geospatial.md b/docs/geospatial.md index b2bd119..2ae8625 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -10,15 +10,16 @@ scalar UDF. The array paradigm (NumPy, Xarray, Dask) is a wonderful *interface* for these operations. But it is not the only one, and for a large and growing audience — the people fluent in SQL rather than in `apply_ufunc` and rechunking — it is not -the most accessible one. [`xarray-sql`](../README.md) lets you pose these -questions in SQL and answers them with a real query engine (DataFusion). The +the most accessible one. [`xarray-sql`](index.md) lets you pose these +questions in SQL and answers them with a real query engine (DataFusion here; +[the same tables serve DuckDB and Polars](engines.md)). The datasets are opened *lazily*, so a query against the whole archive reads only the variable and the slice it actually needs. And because a gridded result is still gridded data, every query here round-trips its answer straight back to an `xarray.Dataset` — SQL in, an array out, ready to plot or save. This page makes the argument case by case. Every claim below is backed by a -runnable script in [`benchmarks/geospatial/`](../benchmarks/geospatial/) that +runnable script in [`benchmarks/geospatial/`](https://github.com/xqlsystems/xarray-sql/tree/main/benchmarks/geospatial/) that poses the operation in SQL and **asserts the answer matches an xarray/array reference** to floating-point tolerance. The point is not that "SQL is faster"; the point is that the SQL reads like the *definition* of the operation and @@ -53,15 +54,15 @@ through SQL one by one, turns out to be — almost entirely — queries. | Operation | The "array" framing | The relational reality | Script | |-----------|---------------------|------------------------|--------| -| Spectral index (NDVI) | `apply_ufunc` over a raster | column arithmetic | [`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) | -| Climatology | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | [`02_climatology.py`](../benchmarks/geospatial/02_climatology.py) | -| Zonal mean | reduce over lon/time axes | `GROUP BY lat` | [`03_zonal_mean.py`](../benchmarks/geospatial/03_zonal_mean.py) | -| Anomaly | grouped broadcast-subtract | climatology CTE self-`JOIN` | [`04_anomaly.py`](../benchmarks/geospatial/04_anomaly.py) | -| Forecast skill (RMSE) | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` | [`05_forecast_skill.py`](../benchmarks/geospatial/05_forecast_skill.py) | -| Zonal stats over regions | rasterize polygons + mask | raster × vector range `JOIN` | [`06_zonal_vector.py`](../benchmarks/geospatial/06_zonal_vector.py) | -| Reprojection | per-pixel CRS transform | scalar **UDF** (`ST_Transform`-style) | [`07_reproject_udf.py`](../benchmarks/geospatial/07_reproject_udf.py) | -| Regridding | interpolation to a new grid | sparse-weight table `JOIN` | [`08_regrid_weights.py`](../benchmarks/geospatial/08_regrid_weights.py) | -| Warp (reproject + resample) | CRS transform *and* interpolation | reproject **UDF** → weight-table `JOIN` | [`09_warp.py`](../benchmarks/geospatial/09_warp.py) | +| Spectral index (NDVI) | `apply_ufunc` over a raster | column arithmetic | [`01_ndvi.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/01_ndvi.py) | +| Climatology | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | [`02_climatology.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/02_climatology.py) | +| Zonal mean | reduce over lon/time axes | `GROUP BY lat` | [`03_zonal_mean.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/03_zonal_mean.py) | +| Anomaly | grouped broadcast-subtract | climatology CTE self-`JOIN` | [`04_anomaly.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/04_anomaly.py) | +| Forecast skill (RMSE) | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` | [`05_forecast_skill.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/05_forecast_skill.py) | +| Zonal stats over regions | rasterize polygons + mask | raster × vector range `JOIN` | [`06_zonal_vector.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/06_zonal_vector.py) | +| Reprojection | per-pixel CRS transform | scalar **UDF** (`ST_Transform`-style) | [`07_reproject_udf.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/07_reproject_udf.py) | +| Regridding | interpolation to a new grid | sparse-weight table `JOIN` | [`08_regrid_weights.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/08_regrid_weights.py) | +| Warp (reproject + resample) | CRS transform *and* interpolation | reproject **UDF** → weight-table `JOIN` | [`09_warp.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/09_warp.py) | ## 1. A pixel-wise formula is a column expression @@ -79,7 +80,7 @@ Invalid pixels need no special handling: xarray decodes the band's `_FillValue` to `NaN` on open, and `NaN` propagates through the arithmetic on both sides, so the masking is free. -[`01_ndvi.py`](../benchmarks/geospatial/01_ndvi.py) runs this against a **real +[`01_ndvi.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/01_ndvi.py) runs this against a **real Sentinel-2 L2A scene in Zarr** — discovered with `pystac-client` and opened the canonical way with `xr.open_datatree` (ESA's EOPF sample service) — and matches xarray's `apply_ufunc`-style result over a million pixels. @@ -98,12 +99,12 @@ FROM era5 GROUP BY latitude, longitude, date_part('hour', time) ``` The grouping keys are the dimensions you keep; everything else is reduced. No -layout to reason about. [`02_climatology.py`](../benchmarks/geospatial/02_climatology.py) +layout to reason about. [`02_climatology.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/02_climatology.py) computes the **diurnal cycle** of ERA5 2m-temperature over a region — averaging each cell by hour of day — and matches `da.groupby("time.hour").mean()` across ~500k cells. -A **zonal mean** ([`03_zonal_mean.py`](../benchmarks/geospatial/03_zonal_mean.py)) +A **zonal mean** ([`03_zonal_mean.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/03_zonal_mean.py)) is the same idea with fewer keys: the axes you "reduce over" are simply the columns you don't `GROUP BY`. @@ -128,7 +129,7 @@ FROM era5 a JOIN clim c AND date_part('hour', a.time) = c.hour ``` -[`04_anomaly.py`](../benchmarks/geospatial/04_anomaly.py) computes the +[`04_anomaly.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/04_anomaly.py) computes the climatology once (the CTE) and joins it back to every observation. ## 4. Forecast evaluation is a `JOIN` on valid time + aggregate @@ -158,10 +159,10 @@ Both models are stacked along a `model` dimension into one forecast table, so a single query scores them together, grouped by the `model` column. The entire evaluation — temporal alignment across three time axes, spatial matching, and the score — is one JOIN and one aggregate. -[`05_forecast_skill.py`](../benchmarks/geospatial/05_forecast_skill.py) runs it +[`05_forecast_skill.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/05_forecast_skill.py) runs it for both models, matches an xarray reference, and reproduces the published result that GraphCast edges out Pangu at every lead — the classic "error grows with -horizon" curve (≈0.3 K at 6 h rising to ≈2.5 K at 9 days): +horizon" curve (≈0.3 K at 6 h rising to ≈2.5 K at 9 days). The result round-trips to a `pandas` table directly (`got.to_pandas()`), RMSE in kelvin by lead time: @@ -200,7 +201,7 @@ GROUP BY r.region This is the README's promise — *joining tabular data with raster data* — made literal: the raster is the full ERA5 archive (the `WHERE` prunes it to a day), the regions are a second SQL table, and the spatial relationship is an ordinary -`BETWEEN`. See [`06_zonal_vector.py`](../benchmarks/geospatial/06_zonal_vector.py) +`BETWEEN`. See [`06_zonal_vector.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/06_zonal_vector.py) — it reports e.g. Sahara 33 °C vs Greenland −8 °C for a June day. (Rectangular regions keep this simple; arbitrary polygons would follow the same shape, with a point-in-polygon test in the join.) @@ -225,7 +226,7 @@ SELECT x, y, FROM grid ``` -[`07_reproject_udf.py`](../benchmarks/geospatial/07_reproject_udf.py) validates +[`07_reproject_udf.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/07_reproject_udf.py) validates this against **Earth Engine itself**: it opens a UTM grid through [Xee](https://github.com/google/Xee) carrying `ee.Image.pixelLonLat()`, so EE's own geodesy engine reports the true lon/lat of every pixel — an *independent* @@ -251,7 +252,7 @@ FROM weights w JOIN src s ON s.cell_id = w.src_id GROUP BY w.dst_id ``` -[`08_regrid_weights.py`](../benchmarks/geospatial/08_regrid_weights.py) regrids +[`08_regrid_weights.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/08_regrid_weights.py) regrids real **SRTM elevation** (Sierra Nevada terrain, opened from the Earth Engine catalog through [Xee](https://github.com/google/Xee)) coarse → fine and matches xarray's bilinear `.interp()` exactly. So regridding does not weaken the thesis — @@ -259,7 +260,7 @@ it is the most relational operation of all. **A warp is just the two composed.** The full operation a GIS calls *warp* (GDAL and rasterio's `reproject`) does both at once: change the CRS *and* resample onto -the new grid. [`09_warp.py`](../benchmarks/geospatial/09_warp.py) writes it as the +the new grid. [`09_warp.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/09_warp.py) writes it as the two cases above run back to back — the 07 reproject UDF carries the target lon/lat grid back into the source UTM space, arrays turn those reprojected points into bilinear weights, and the 08 `JOIN` applies them: @@ -284,6 +285,37 @@ warp lands exactly where the split predicts: the row-independent half is a UDF, the many-to-many half is a `JOIN`, and the only genuinely geometric step — turning the reprojected points into weights — is the array work the next section is about. +## GeoArrow point-geometry columns + +`register(..., geometry=("x", "y"))` derives a `geometry` point column +from two coordinate dims. With the default `"wkb"` encoding DuckDB +(spatial loaded) ingests it as a native `GEOMETRY` with the CRS +attached, so geometry predicates need no `ST_Point(x, y)` construction: + +```sql +SELECT avg(risk) FROM eri +WHERE y BETWEEN -29 AND -28 AND x BETWEEN -58 AND -57 -- prunes chunks + AND ST_Within(geometry, ST_GeomFromText('POLYGON (...)')) -- refines +``` + +**Always pair geometry predicates with bbox conjuncts on the coordinate +columns.** Engines do not push functions like `ST_Within` into the +scan, so a geometry-only predicate scans (and encodes) every chunk — +measured ~29x slower than the paired form on a 10M-row grid, where the +bbox prunes first and the exact polygon test is nearly free. +`xql.bbox_conjuncts(geom, x=..., y=...)` renders the conjuncts from any +geometry's envelope (shapely objects or plain +`(xmin, ymin, xmax, ymax)` tuples), with `pad=` for +`ST_DWithin`-style margins — so the idiom is one f-string. The full +reasoning lives in [Known issues & limitations](limitations.md#geometry-predicates-alone-cannot-prune). + +`geometry_encoding="point"` emits GeoArrow-native separated coordinates +instead (the struct children *are* the coordinate arrays): zero-parse +for GeoPandas 1.x (`GeoDataFrame.from_arrow`), lonboard, geoarrow-rs +and SedonaDB. DuckDB does not consume this encoding — pick per +destination. The CRS tag defaults to `OGC:CRS84`; pass +`geometry_crs=...` for anything else. + ## Where the array paradigm still earns its keep The boundary is **weight generation**. Applying a regridding is a join; @@ -314,13 +346,13 @@ uv run benchmarks/geospatial/02_climatology.py # standalone (PEP 723 deps) ``` Each script prints its SQL, runs the array reference, and asserts the two agree. -See [`benchmarks/geospatial/README.md`](../benchmarks/geospatial/README.md) for +See [`benchmarks/geospatial/README.md`](https://github.com/xqlsystems/xarray-sql/tree/main/benchmarks/geospatial) for the full list and dataset notes. ## Results Correctness is the headline, but every case is also profiled. The numbers below -come from [`run_perf.sh`](../benchmarks/geospatial/run_perf.sh) on a single Google +come from [`run_perf.sh`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/run_perf.sh) on a single Google Compute Engine `e2-standard-8` (8 vCPU, 32 GB) in `us-central1` — in-region with the ARCO-ERA5 and WeatherBench 2 buckets, so the cloud read is fast — with Earth Engine reached from the same VM, so all nine cases share one machine and one release build. @@ -383,6 +415,85 @@ machine — across three `e2-standard-8` runs it has measured ≈10.7 s, ≈12 s ≈23 s, while the read-bound *reference* stays near 0.25 s. So read the 05 ratio as "the relational form costs real CPU here," not as a fixed multiplier. +### The suite across engines and machine sizes + +The table above measures the DataFusion-native path. The same cases also run +through the suite's engine layer +([`_engines.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/_engines.py), selected per process +with `GEOBENCH_ENGINE`) as four engines: `datafusion` (the native table +provider — the harness builds the compiled module from the shipped crate on +each VM), `datafusion-arrow` (the pure-Python pyarrow-dataset path, through +`SessionContext.register_dataset(xql.arrow_dataset(ds))`), DuckDB, and +Polars. We ran the portable cases across all four on an `e2-standard-8` in +`us-central1`, in-region with the data, under the same protocol: **fresh +process per repetition, no warmup, five cold reps**, and every engine's +answer asserted against the xarray reference before its timing counts +([`engine_suite.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/engine_suite.py) drives it). + +Scope, stated plainly. The `datafusion` column is the same code path as the +headline table, re-measured on this VM and day, so the two DataFusion columns +compare the native table provider and the pyarrow-dataset path on identical +hardware; DuckDB and Polars consume the same pyarrow pushdown datasets. +Polars executes the identical SQL via `polars.SQLContext`, with the query's +window bounds also applied as native scan-level expressions (its SQL +`TIMESTAMP` literals compile to `strptime` casts that never reach the pyarrow +scanner as filters); its SQL dialect cannot express case 06's range `JOIN` (a +`BETWEEN` join constraint), recorded as unsupported rather than worked +around. Cases 07 and 09 build DataFusion scalar UDFs (n/a on the other +engines), case 08 is Earth-Engine-gated, and 07–09 skip on these VMs (no +Earth Engine auth) — the skip reasons ride along in the results. + +The same grid was also measured on `e2-standard-16` and `e2-standard-32` in +the same run, with no practical difference: every case is dominated by a +single-stream cold cloud read plus a mostly single-threaded row pipeline, so +extra vCPUs buy nothing, and the spread across sizes is shared-core `e2` and +network variance, not engine behavior (the 16-vCPU VM measured *slower* than +the 8-vCPU one on most cells in this run). The full per-size tables live in +the `xarray-sql-notes` repository (`engine-matrix-results.md`). + +Software: CPython 3.12.8, duckdb 1.5.5, polars 1.42.1, datafusion 54.0.0, +pyarrow 25.0.0, xarray 2026.7.0, Linux (glibc 2.41). Medians of 5 cold reps; +peak is the Python-allocator peak per process. + +| Case | DataFusion (native) | DataFusion (pyarrow) | DuckDB | Polars | xarray reference | +|---|--:|--:|--:|--:|--:| +| 01 · NDVI | 4.308 s (114 MB) | 4.380 s (97 MB) | 5.325 s (106 MB) | 5.725 s (106 MB) | 0.444 s (42 MB) | +| 02 · Climatology | 7.235 s (1116 MB) | 8.746 s (1142 MB) | 4.622 s (627 MB) | 4.936 s (637 MB) | 2.367 s (44 MB) | +| 03 · Zonal mean | 4.763 s (406 MB) | 3.614 s (403 MB) | 2.958 s (403 MB) | 3.705 s (413 MB) | 0.829 s (250 MB) | +| 04 · Anomaly | 10.416 s (1117 MB) | 16.262 s (3003 MB) | 9.817 s (627 MB) | 13.837 s (936 MB) | 4.410 s (76 MB) | +| 05 · Forecast skill | 1.791 s (170 MB) | 1.814 s (175 MB) | 1.797 s (170 MB) | 2.405 s (184 MB) | 0.247 s (2 MB) | +| 06 · Zonal stats | 2.390 s (515 MB) | 6.131 s (513 MB) | 9.112 s (503 MB) | unsupported (range `JOIN`) | 1.813 s (1262 MB) | + +Cases 07–09 could not run on these VMs (no Earth Engine auth), so their rows +come from the headline run instead: the original `e2-standard-8` GCE VM with +Earth Engine access, DataFusion **native** path (07 and 09 are DataFusion +scalar UDFs, n/a on the other engines): + +| Case | DataFusion (native path) | xarray reference | +|---|--:|--:| +| 07 · Reprojection (PROJ scalar UDF) | 0.029 s (0.3 MB) | — | +| 08 · Regridding (weight-table `JOIN`) | 0.875 s (11.9 MB) | 0.850 s (13.3 MB) | +| 09 · Warp (reproject UDF → regrid `JOIN`) | 0.281 s (0.8 MB) | 0.817 s (11.2 MB) | + +The engine story, in three observations. **Native vs pyarrow on the same +engine:** the native table provider wins where rows are consumed in bulk — +case 06's range `JOIN` (2.4 vs 6.1 s) and the anomaly self-`JOIN` 04 (10.4 vs +16.3 s, at a third of the peak memory), ~1.2× on climatology 02 — while on +the read-bound cases (01, 05) the two are at parity, and on the plain zonal +mean 03 the pyarrow path is the faster one (3.6 vs 4.8 s). **Across +engines**, DuckDB is the fastest consumer of the shared pushdown scan on the +plain group-bys (02, 03) and narrowly beats native DataFusion on 04 (9.8 vs +10.4 s), while native DataFusion leads the join-heavy 06 outright; Polars +stays close to DuckDB on 02 and falls back on 04. The spread between engines is much smaller +on cases whose cost is the read itself (01, 05), which is the same lesson as +the headline table: the paradigm and the I/O set the floor, the engine sets +the constant. And a benchmark side-effect worth keeping: streaming case 05's +full window through the pyarrow protocol surfaced a real library bug — +`pa.array` returns a `ChunkedArray` for a large string dimension coordinate, +which the pivot's fast path passed straight into `RecordBatch.from_arrays`; +fixed, and pinned by a regression test in `tests/test_df.py` +(`test_iter_record_batches_large_string_dim_coord`). + ## Analysis: how a relational operation spends its time Why is SQL slower, and where does the time actually go? Profiling case 05 — the diff --git a/docs/limitations.md b/docs/limitations.md new file mode 100644 index 0000000..8aab000 --- /dev/null +++ b/docs/limitations.md @@ -0,0 +1,157 @@ +# Known issues & limitations + +What does not work, why, and what to do instead. How the machinery +*works* — scan pipeline, tuning, cost model — lives in +[Engines](engines.md) and the [performance guide](performance.md); this +page is only the sharp edges. Everything here is pinned by tests. + +## Engine-specific issues + +Pick your engine: + +=== "DataFusion" + + No engine-specific known issues. DataFusion is the deepest + integration (native table provider, chunked round-trip via + re-execution); the constraints below apply as everywhere. + +=== "DuckDB" + + **Re-executing relations from worker threads deadlocks.** + + - *Symptom:* a chunked round-trip (`chunks=`) of a DuckDB relation + would hang intermittently (~50% of runs) when dask workers + re-execute the relation, whenever the query scans a + Python-backed table. + - *Scope:* duckdb-python 1.4–1.5 with CPython 3.12 (observed on + macOS); unaffected by `SET threads=1`, connection-level + serialization, or thread-pool pre-warming. The identical + topology through Polars never hangs. The deadlock is in the + interpreter/engine thread-state interaction, not in xarray-sql. + - *What the library does:* `chunks=` on a DuckDB relation raises + `NotImplementedError` immediately rather than hanging. + `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. + + **Derived relations break under concurrent materialization.** + + - *Symptom:* materializing two relations derived from the same + base concurrently raises `InvalidInputException` (they share + pending-query state upstream). + - *What the library does:* the round-trip handle serializes every + engine call on one dedicated thread. Nothing to do on your + side — documented so the serialization is not mistaken for a + missing optimization. + + **GeoArrow-native points are not consumed.** + + - *Symptom:* a `geometry` column registered with + `geometry_encoding="point"` binds as a plain struct; `ST_*` + functions reject it. + - *What to do:* use the default `"wkb"` encoding for DuckDB — it + binds as a native `GEOMETRY` with the CRS attached. See + [Geospatial in SQL](geospatial.md#geoarrow-point-geometry-columns). + +=== "Polars" + + **Float `is_in` literals lose precision.** + + - *Symptom:* `is_in` with **float** literals can silently match + nothing (reproducible without xarray-sql; integer and timestamp + value sets are unaffected). + - *What to do:* prefer `is_between` for float coordinates in your + own queries. + - *What the library does:* the lazy round-trip's window queries + render float value lists as degenerate ranges internally, so + reconstruction is immune. + + **No geometry types.** + + - *Symptom:* a registered `geometry` column arrives as plain + binary (WKB) or a plain struct; there are no `ST_*` functions. + - *What to do:* filter on the coordinate columns instead, and do + geometry work in DuckDB, DataFusion, or GeoPandas. + + **Single-threaded source pull.** + + Polars pulls the scan sequentially; source-side parallelism comes + from the adapter's prefetch pool (`prefetch`, `prefetch_bytes`), + not from the consumer. Not a bug — worth knowing when sizing + scans. + +## Constraints in any engine + +These follow from the data model — no engine or configuration avoids +them. + +### Geometry predicates alone cannot prune + +Engines never push function calls (`ST_Within`, casts, arithmetic) into +a scan — only plain column-vs-constant comparisons, `IN`, `IS NULL`, +and boolean combinations. A geometry-only `WHERE` therefore scans and +encodes every chunk (measured ~29x slower than the paired form on a +10M-row grid). Pair every geometry predicate with range conjuncts on +the coordinate columns; [`bbox_conjuncts`][xarray_sql.bbox_conjuncts] +renders them from the geometry's envelope. See +[Geospatial in SQL](geospatial.md#geoarrow-point-geometry-columns). + +### Filters on data variables always scan + +Chunk pruning and arithmetic counting rest on per-chunk *coordinate* +ranges. A predicate on a data variable (`t2m > 300`) carries no such +guarantee: every surviving chunk is scanned, and the filter is applied +row-exactly. No configuration changes this; it is what the data model +can prove. + +### NaN coordinates disable pruning for their chunks + +A NaN/NaT anywhere in a chunk's coordinate span poisons its min/max +guarantee, so that chunk is kept for every predicate. This is the +correct trade: a range that pretended to cover NaN would let engines +whose NaN ordering differs (DuckDB sorts NaN greatest) silently lose +rows. Chunks without NaN prune normally. + +### String, object, and cftime dimensions never prune + +Chunk guarantees are built for numeric and datetime coordinates only; +predicates on other dimension types conservatively scan every chunk +(row-exactly, as always). + +### Scan-path pruning is per-dimension + +On the scan path, each dimension's surviving chunks are computed +independently and combined as a product, so a predicate pairing +*specific* ranges across dims — `(t < a AND lat < b) OR (t > c AND +lat > d)` — also reads the cross combinations (sound, conservative; +per-dim indexes are what keep million-chunk axes cheap). `count_rows` +refines the crosses away with cross-dimension bucket analysis; ordinary +scans accept the extra reads. + +### Sparse results can explode the dense grid + +The eager 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. + +### One-shot Arrow streams cannot re-execute + +A materialized table or bare C-stream has no query behind it, so the +re-execution form of `chunks=` cannot serve it — `spill=True` (one +pass to a temporary Parquet file) is the chunked path for these. + +### Mixed-dimension datasets split into one DuckDB table per dim group + +DuckDB registration has no schema namespace, so variables with +different dims land in suffixed tables (`_`), sharing one +set of coordinate reads. DataFusion registers the same layout as +`name.group` tables inside one schema. + +### Pointwise indexers on lazy round-trip arrays are slower + +Vectorized (pointwise) selection goes through xarray's +outer-then-gather fallback — correct, but slower than slice windows. diff --git a/docs/performance.md b/docs/performance.md new file mode 100644 index 0000000..349589d --- /dev/null +++ b/docs/performance.md @@ -0,0 +1,245 @@ +# Performance guide + +How to get engine-limited speed out of registered xarray tables. Every +number below was measured on real cloud rasters (billions of pixels); +your mileage scales with network and core count, but the *ratios* are +structural. + +## How a scan decides what to read + +Every engine query over a registered table flows through one pipeline; +each tuning knob on this page acts on one of its stages: + +```mermaid +flowchart TB + Q["engine calls scanner(columns, filter)"] --> P["prune chunks
per-dim coordinate ranges +
Arrow guarantee simplification"] + Q --> J["project
only referenced variables are read"] + P --> C["coalesce (opt-in)
merge consecutive surviving chunks
into single reads"] + C --> F["prefetch pool
bounded by prefetch (threads)
and prefetch_bytes (memory)"] + J --> F + F --> X["exact filter
the pushed expression is applied
row-exactly — pruning is only
ever an optimization"] + X --> B["Arrow batches → engine"] +``` + +Two invariants hold everywhere: pruning never decides correctness (the +exact expression is always applied — engines delete pushed conjuncts +from their own plans), and only what reaches `scanner()` can prune +(engines push plain comparisons, never function calls). + +Everything on this page up to [Per-engine notes](#per-engine-notes) +applies whichever engine you query with. + +## Make the source read in parallel + +The single biggest lever is usually the reader, not the engine. + +**GeoTIFF / rioxarray**: `rioxarray.open_rasterio` serializes GDAL tile +reads behind a lock by default, capping every scan at single-stream +speed no matter how many threads the adapter runs. On GDAL ≥ 3.11 use +the natively thread-safe LIBERTIFF driver; on older GDAL pass +`lock=False`: + +```python +da = rioxarray.open_rasterio( + url, chunks={"x": 2048, "y": 2048}, + driver="LIBERTIFF", # GDAL >= 3.11; else keep lock=False only + lock=False, +) +``` + +Measured on a 9-billion-pixel public cloud GeoTIFF, full-table +aggregation: default open **277 s** → `lock=False` **43 s** → +LIBERTIFF + `GDAL_NUM_THREADS=ALL_CPUS` **24 s**. With parallel reads, +remote (`/vsicurl/`) matched a local copy of the same file — the +network was never the bottleneck, the lock was. + +Remote-read environment preset worth exporting for `/vsicurl/` sources: + +```python +os.environ.update( + GDAL_NUM_THREADS="ALL_CPUS", + GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", + VSI_CACHE="TRUE", +) +``` + +**Zarr**: zarr-python 3's async store defaults to only 10 concurrent +requests; raise it before opening remote stores: + +```python +zarr.config.set({"async.concurrency": 128}) +``` + +On a moderately sized windowed query (~40 chunks of 4096² uint8 per +variable, GCS) this was a modest gain (4.2 s → 3.7 s); it matters more +as chunk counts grow and chunks shrink. The obstore-backed +`zarr.storage.ObjectStore` is worth benchmarking for high-concurrency +workloads, but was not faster at this scale in our tests — measure +before switching. + +## Choose chunk sizes for the scan, not just the store + +Every chunk costs one prefetch task, one pivot call, and one shadow +fragment. Aim for **1–8 M rows per chunk** (e.g. 2048²–4096² pixels for +2-D grids). The same 10 M-row scan ran 1.7× faster in 4 chunks than in +20. Axes with hundreds of thousands of chunks still prune in +milliseconds (the shadow index is bucketed), but scanning them pays +per-chunk overhead. + +## Tune the adapter knobs + +```python +xql.register(con, "t", ds, prefetch=12, batch_size=262_144) +``` + +- `prefetch`: chunk loads kept in flight ahead of the engine. The + default (4) saturates local CPU work; raise to 8–12 for remote + sources where latency dominates. Memory scales with + `prefetch × pivoted chunk size`. +- `batch_size`: rows per Arrow batch. The default (64 Ki) is fine; + values between 64 Ki and 1 Mi measured within a few percent of each + other. + +## The memory contract + +Peak scan memory is bounded by `prefetch × pivoted-block-size` plus the +engine's own aggregation state — it does not grow with the amount of +data scanned. Measured on ARCO-ERA5 over anonymous GCS: a one-month +full-globe aggregation (772M rows) peaks at the same resident set size +as the one-week scan (174M rows), ~0.75 GB with the defaults. + +`prefetch_bytes` caps *estimated bytes* in flight instead of block +count — set it when `coalesce_rows` makes blocks large or ragged. +The block size is the source chunk size unless `coalesce_rows` is set, +in which case in-flight units are merged blocks: raising +`coalesce_rows` buys fewer round-trips at proportionally higher peak +memory (`prefetch=16, coalesce_rows=8_000_000` peaked at ~1.2 GB on the +same scan while cutting wall time ~1.5-2x). Size the two together. + +`count(*)`-shaped queries never pay scan memory at all: unfiltered +counts are pure chunk arithmetic, and filtered counts scan only the +boundary chunks the filter cannot prove — at any filter breadth; see +[What counting costs](#what-counting-costs). + +## Let pushdown do its job + +Selective queries are fast *because of their predicates*: bounding-box +`WHERE` clauses on dimension columns prune to intersecting chunks, and +only the variables a query references are read. Corollaries: + +- Prefer explicit column lists over `SELECT *` on wide datasets. +- Spatial functions (`ST_Within`, ...) are not pushed down — pair them + with a bounding-box predicate that is: the box prunes, the geometry + refines. +- A query with no `WHERE` on dimension columns is a full scan on any + engine; that's physics, not a missing optimization. + +## What counting costs + +`count(*)` never pays scan memory, and usually no I/O either: + +```mermaid +flowchart TB + C["count_rows(filter)"] --> U{"filter?"} + U -- none --> A["pure arithmetic
0 reads"] + U -- "coordinate ranges" --> H["hierarchical strictness:
bucket-products proven or pruned
whole; only mixed cells recurse"] + H --> E["boundary chunks scanned exactly
(usually 0-2 per range edge,
at any axis size)"] + U -- "data variables" --> S["every surviving chunk scanned
(values carry no coordinate
guarantee — see Known issues)"] +``` + +Coordinate-range counts stay arithmetic at any breadth (a +near-universal filter over a million single-row chunks counts with +zero reads), and the strictness pass applies cross-dimension +information, so paired-range predicates count without reading the +cross combinations. + +## Stop re-scanning: cache derived tables + +Registered tables are virtual — every query re-streams the source. +Statistics you ask repeatedly should pay the scan once: create a native +table from your query, sorted by the coordinate columns so the engine's +storage compresses the repetitive coordinates (DuckDB picks ALP/RLE on +sorted runs) and zone maps prune range predicates. + +```sql +CREATE OR REPLACE TABLE grid_cube AS +SELECT FLOOR(y) AS lat, FLOOR(x) AS lon, klass, COUNT(*) AS n +FROM grid GROUP BY 1, 2, 3 +ORDER BY lat, lon; + +SELECT * FROM grid_cube WHERE lat = -32; -- native speed +``` + +One engine quirk to know: on DataFusion, DDL is a lazy plan — collect +it or nothing happens: + +```python +ctx.sql("CREATE OR REPLACE TABLE grid_cube AS ...").collect() +``` + +## Round-trip faster with ORDER BY + +Results that arrive **grid-ordered** — sorted by the dimension columns, +outermost first — reconstruct with a single reshape; unordered results +(DuckDB's parallel scans return chunk order, not grid order) pay a +per-row positional scatter instead, measured ~2x slower on large +windows. When you will round-trip a large result, add +`ORDER BY ` to the query. + +And if what you want is a raw sub-array of a registered Dataset rather +than a relational answer, plain `ds.sel(...)` is the direct path — SQL +adds value when the question is relational. + +## Per-engine notes + +=== "DataFusion" + + **Two registration paths.** `XarrayContext.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 + (`ctx.register_dataset(xql.arrow_dataset(ds))`), not to the native + provider. + + **DDL/DML is lazy.** `CREATE TABLE`/`INSERT` statements are plans — + `.collect()` them or nothing executes (the caching recipe above + shows this). + +=== "DuckDB" + + **Connections and threads.** Registered Python objects are + connection-local: `con.cursor()` does not inherit them, and one + connection's result slot is not thread-safe. For multithreaded + querying, give each thread its own cursor and register the *same* + dataset object on it: + + ```python + dataset = xql.arrow_dataset(ds) + def worker(): + cur = con.cursor() + cur.register("t", dataset) # cheap; shares the pruning index + ... + ``` + + The dataset object itself is safe to share across threads + (verified under concurrent query load). + + **Row order.** DuckDB's parallel scans return results in chunk + order, not grid order — add `ORDER BY ` before round-tripping + large results (see [above](#round-trip-faster-with-order-by)). + + **Geometry.** Register with the default `"wkb"` encoding; pair + `ST_*` predicates with bbox conjuncts so pruning still applies. + +=== "Polars" + + **Parallelism.** Polars pulls the scan single-threaded; source-side + parallelism comes entirely from the adapter's `prefetch` / + `prefetch_bytes`, so tune those rather than Polars settings. + + **Batch sizing.** `scan_pyarrow_dataset` passes its `batch_size` + through to the scanner (honored), so Polars morsel sizing works as + documented on their side. + + **Large results.** Collect with `engine="streaming"` to keep + memory bounded; the lazy round-trip's windows already do this. diff --git a/pyproject.toml b/pyproject.toml index 7df9a51..8a35be4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "maturin" [project] name = "xarray_sql" dynamic = ["version"] -description = "Querry Xarray with SQL." +description = "Query Xarray with SQL." readme = "README.md" requires-python = ">=3.10" # Python 3.9 EOL is October 31, 2025. license = {text = "Apache-2.0"} @@ -36,12 +36,21 @@ dependencies = [ ] [project.optional-dependencies] +duckdb = [ + "duckdb>=1.4.0", +] +polars = [ + # collect(engine="streaming") landed in polars 1.25; + # LazyFrame.collect_batches (streamed max_result_bytes enforcement) + # in 1.33. PolarsHandle degrades gracefully below both floors. + "polars>=1.33", +] geo = [ "pyproj", ] test = [ "cftime", - "xarray-sql[geo]", + "xarray-sql[duckdb,polars,geo]", "pytest", "xarray[io]", "gcsfs", @@ -115,3 +124,6 @@ cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "** [tool.pytest.ini_options] testpaths = ["tests"] +markers = [ + "integration: needs network (anonymous GCS); excluded from the CI unit run", +] diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py new file mode 100644 index 0000000..34aebc7 --- /dev/null +++ b/tests/test_arrow_dataset.py @@ -0,0 +1,542 @@ +"""Contract tests for the engine-neutral pyarrow dataset view. + +``xql.arrow_dataset`` returns a real ``pyarrow.dataset.Dataset`` for +consumers of ``schema``, ``scanner``, ``get_fragments`` and the scan +conveniences — pyarrow itself and Polars are exercised here; DuckDB has +its own suite in ``test_duckdb_backend.py``. DataFusion's native Rust +``TableProvider`` (``XarrayContext``) fills the same role through +DataFusion's own extension trait and sits outside this contract. + +The contract, one section of this file per clause: + +1. Exactness — the pushed filter is applied row-exactly after pruning; + pruning never decides correctness. +2. Projection — exactly the requested columns come back; what gets read + is the projected columns plus the filter's, nothing else. +3. Pruning/counting — provably impossible regions are never read, + provable counts are pure arithmetic, and anything uncertain + (boundary chunks, NaN/NaT coordinates, opaque expressions) is + scanned conservatively. +4. Laziness — construction reads dimension coordinates only; scans + repeat, survive mid-scan abandonment, and run concurrently. +5. Tuning (``batch_size``, ``prefetch``, ``prefetch_bytes``, + ``coalesce_rows``) changes the shape of the work, never the result; + non-positive ``batch_size`` is rejected at construction; fragments + stay one per source chunk. +6. Schema stays on offset types (a view type disables DuckDB's + pushdown); the prefetch pool is fully started at construction and + shuts down when the dataset is collected. +""" + +import threading + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.compute as pc +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + + +@pytest.fixture +def ds() -> xr.Dataset: + np.random.seed(3) + return xr.Dataset( + { + "temperature": ( + ["time", "lat"], + 20 + 5 * np.random.randn(20, 6), + ), + "humidity": (["time", "lat"], np.random.rand(20, 6)), + }, + coords={ + "time": pd.date_range("2022-01-01", periods=20, freq="D"), + "lat": np.linspace(-25.0, 25.0, 6), + }, + ).chunk({"time": 5}) + + +class _ChunkCounter: + def __init__(self): + self.blocks = [] + self.column_sets = [] + + def __call__(self, block, names): + self.blocks.append(block) + self.column_sets.append(tuple(names)) + + +def _hourly_grid() -> xr.Dataset: + """100 hourly steps x 4 latitudes with sequential values.""" + return xr.Dataset( + {"t2m": (["time", "lat"], np.arange(100.0 * 4).reshape(100, 4))}, + coords={ + "time": pd.date_range("2020-01-01", periods=100, freq="h"), + "lat": np.linspace(-30.0, 30.0, 4), + }, + ) + + +@pytest.fixture +def counted(): + """A pushdown dataset over hourly data with a chunk-read counter.""" + source = _hourly_grid() + counter = _ChunkCounter() + dataset = XarrayPushdownDataset( + source, {"time": 10}, _iteration_callback=counter + ) + return dataset, counter + + +# -- Dataset protocol surface ------------------------------------------------ + + +def test_to_table_projects_and_filters(ds): + table = xql.arrow_dataset(ds).to_table( + columns=["time", "temperature"], + filter=pc.field("lat") > 0, + ) + assert table.column_names == ["time", "temperature"] + assert table.num_rows == 20 * 3 # lat > 0 keeps 3 of 6 latitudes + + +def test_count_rows_and_head(ds): + dataset = xql.arrow_dataset(ds) + assert dataset.count_rows() == 20 * 6 + assert dataset.head(7).num_rows == 7 + + +def test_get_fragments_prunes_and_scans(ds): + dataset = xql.arrow_dataset(ds) + assert len(dataset.get_fragments()) == 4 # time chunked by 5 + + # A time predicate covering the first chunk keeps one fragment. + early = pc.field("time") < pa.scalar( + pd.Timestamp("2022-01-06"), type=pa.timestamp("ns") + ) + kept = dataset.get_fragments(filter=early) + assert len(kept) == 1 + assert kept[0].to_table().num_rows == 5 * 6 + + # An unsatisfiable predicate prunes everything. + assert dataset.get_fragments(filter=pc.field("lat") > 100) == [] + + +def test_scanner_honors_batch_size(ds): + dataset = xql.arrow_dataset(ds) + batches = list(dataset.scanner(batch_size=7).to_batches()) + assert sum(b.num_rows for b in batches) == 20 * 6 + assert max(b.num_rows for b in batches) <= 7 + + # The kwarg travels through the inherited to_batches path, which is + # how Polars sizes its morsels. + sizes = [b.num_rows for b in dataset.to_batches(batch_size=11)] + assert sum(sizes) == 20 * 6 + assert max(sizes) <= 11 + + +def test_schema_never_uses_view_types(ds): + # A single view-typed column disables DuckDB's filter pushdown for + # the whole table; pin the schema to offset layouts so a pyarrow + # upgrade cannot regress this silently. + for field in xql.arrow_dataset(ds).schema: + assert field.type not in (pa.string_view(), pa.binary_view()) + + +# -- Consumer integrations --------------------------------------------------- + + +def test_datafusion_register_dataset_round_trips(ds): + from datafusion import SessionContext + + ctx = SessionContext() + ctx.register_dataset("t", xql.arrow_dataset(ds)) + out = ctx.sql( + "SELECT time, AVG(temperature) AS temperature FROM t " + "WHERE lat > 0 GROUP BY time ORDER BY time" + ).to_pandas() + expected = ds.temperature.sel(lat=ds.lat[ds.lat > 0]).mean("lat").compute() + np.testing.assert_allclose(out["temperature"].values, expected.values) + + +def test_dask_from_map_over_fragments(ds): + dd = pytest.importorskip("dask.dataframe") + + frags = xql.arrow_dataset(ds).get_fragments() + ddf = dd.from_map(lambda f: f.to_table().to_pandas(), frags) + assert len(ddf.compute()) == 20 * 6 + + +def test_polars_scan_pushdown_round_trip(ds): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) + out = ( + lf.filter(pl.col("lat") > 0) + .group_by("time") + .agg(pl.col("temperature").mean()) + .sort("time") + .collect() + ) + expected = ( + ds.temperature.sel(lat=ds.lat[ds.lat > 0]).mean("lat").compute().values + ) + np.testing.assert_allclose(out["temperature"].to_numpy(), expected) + + +def test_polars_result_round_trips_to_xarray(ds): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) + frame = ( + lf.group_by("time") + .agg(pl.col("temperature").mean().alias("temperature")) + .sort("time") + .collect() + ) + # Polars DataFrames export Arrow via the PyCapsule protocol, so the + # engine-agnostic round-trip works unchanged. + out = xql.to_dataset(frame, template=ds) + assert list(out.dims) == ["time"] + assert out.sizes["time"] == 20 + + +# -- Projection: what is returned vs what is read ---------------------------- + + +def test_filter_only_columns_are_read_but_not_returned(): + src = xr.Dataset( + {"a": (["i"], np.arange(100.0)), "b": (["i"], np.arange(100.0) * 2)}, + coords={"i": np.arange(100.0)}, + ) + counter = _ChunkCounter() + dataset = XarrayPushdownDataset(src, {"i": 10}, _iteration_callback=counter) + table = dataset.to_table(columns=["a"], filter=pc.field("b") >= 100.0) + assert table.column_names == ["a"] + assert table.num_rows == 50 + # The filter column is read alongside the projected one, nothing else. + assert set(counter.column_sets) == {("a", "b")} + + +def test_empty_projection_is_a_real_projection(counted): + dataset, counter = counted + table = dataset.scanner(columns=[]).to_table() + assert table.num_columns == 0 + assert table.num_rows == 100 * 4 + + # With a filter, only the filter's column is read, still zero returned. + counter.column_sets.clear() + table = dataset.scanner( + columns=[], filter=pc.field("t2m") < 40.0 + ).to_table() + assert table.num_columns == 0 + assert table.num_rows == 40 # values 0..39: ten hours x four latitudes + assert set(counter.column_sets) == {("t2m",)} + + +# -- Counting and pruning ---------------------------------------------------- + +_T0 = pd.Timestamp("2020-01-01 03:00") +_T1 = pd.Timestamp("2020-01-02 03:00") + + +def _ts(value): + return pa.scalar(value, type=pa.timestamp("ns")) + + +@pytest.mark.parametrize( + "predicate, rows, reads", + [ + # No filter: pure arithmetic, no chunk is read. + (None, 100 * 4, 0), + # [03:00, 27:00): chunks 0 and 2 are boundary, chunk 1 is provably + # inside the range and must be counted arithmetically. + ( + (pc.field("time") >= _ts(_T0)) & (pc.field("time") < _ts(_T1)), + 24 * 4, + 2, + ), + # A data-variable filter carries no coordinate guarantee: every + # chunk is a boundary chunk, and the count must still be row-exact. + (pc.field("t2m") >= 200.0, 200, 10), + # Unsatisfiable: everything pruned, nothing read. + (pc.field("lat") > 100.0, 0, 0), + ], + ids=["unfiltered", "strict-chunks", "data-variable", "unsatisfiable"], +) +def test_count_rows_contract(counted, predicate, rows, reads): + dataset, counter = counted + assert dataset.count_rows(filter=predicate) == rows + assert len(counter.blocks) == reads + + +def test_count_rows_broad_filter_stays_arithmetic(): + # 100k single-element chunks with a filter keeping nearly all of + # them: the hierarchical strictness analysis must prove whole + # buckets at once instead of scanning every survivor. + reads: list = [] + dataset = XarrayPushdownDataset( + xr.Dataset( + {"v": (["step"], np.arange(100_000.0))}, + coords={"step": np.arange(100_000.0)}, + ), + {"step": 1}, + _iteration_callback=lambda b, n: reads.append(b), + ) + assert dataset.count_rows(filter=pc.field("step") >= 100.0) == 99_900 + assert len(reads) <= 2 # at most the bucket-edge chunk + + +def test_count_rows_cross_dimension_refinement(): + # Paired ranges across two chunked dims: per-dim pruning keeps the + # union of each dim's survivors (so the cross combinations too); + # the strictness pass must prune the crosses and count exactly. + t = np.arange(200.0) + lat = np.linspace(-45.0, 45.0, 20) + reads: list = [] + dataset = XarrayPushdownDataset( + xr.Dataset( + {"v": (["t", "lat"], np.arange(200.0 * 20).reshape(200, 20))}, + coords={"t": t, "lat": lat}, + ), + {"t": 10, "lat": 10}, + _iteration_callback=lambda b, n: reads.append(b), + ) + predicate = ((pc.field("t") < 5.0) & (pc.field("lat") < -40.0)) | ( + (pc.field("t") >= 190.0) & (pc.field("lat") > 40.0) + ) + n = dataset.count_rows(filter=predicate) + expected = int( + ( + ((t[:, None] < 5) & (lat[None, :] < -40)) + | ((t[:, None] >= 190) & (lat[None, :] > 40)) + ).sum() + ) + assert n == expected + # Per-dim pruning alone keeps 4 chunk combos (2 t-chunks x 2 + # lat-chunks); cross-dim refinement drops the 2 crosses. + assert len(reads) <= 2 + + +def test_poisoned_coordinates_prune_conservatively(): + # A NaN (or NaT) inside a coordinate chunk voids its range guarantee: + # that chunk must be scanned, never pruned or counted arithmetically, + # and the result must match the oracle (NaN compares False). + x = np.array([0.0, 1.0, np.nan, 3.0, 4.0, 5.0]) + reads: list = [] + dataset = XarrayPushdownDataset( + xr.Dataset({"v": (["x"], np.arange(6.0))}, coords={"x": x}), + {"x": 2}, + _iteration_callback=lambda b, n: reads.append(b), + ) + assert dataset.count_rows(filter=pc.field("x") > 0.5) == 4 + assert any(b["x"] == slice(2, 4) for b in reads) # the NaN chunk + + t = pd.to_datetime(["2020-01-01", "2020-01-02", "NaT", "2020-01-04"]) + nat = XarrayPushdownDataset( + xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": t}), + {"time": 2}, + ) + lo = _ts(pd.Timestamp("2020-01-02")) + assert nat.count_rows(filter=pc.field("time") >= lo) == 2 + + +# -- Scan scheduling knobs: shape of the work, never the result --------------- + + +@pytest.mark.parametrize("coalesce_rows", [None, 10 * 4, 30 * 4, 10_000]) +def test_coalesce_results_identical(coalesce_rows): + source = _hourly_grid() + dataset = XarrayPushdownDataset( + source, {"time": 10}, coalesce_rows=coalesce_rows + ) + predicate = (pc.field("time") >= _ts(_T0)) & ( + pc.field("time") < _ts(pd.Timestamp("2020-01-03 07:00")) + ) + table = dataset.to_table(filter=predicate) + assert table.num_rows == 52 * 4 + expected = source.t2m.isel(time=slice(3, 55)).values.ravel() + np.testing.assert_array_equal( + np.sort(table["t2m"].to_numpy()), np.sort(expected) + ) + + +def test_coalesce_merges_consecutive_chunk_runs(): + source = _hourly_grid() + reads: list[dict] = [] + dataset = XarrayPushdownDataset( + source, + {"time": 10}, + coalesce_rows=30 * 4, # up to 3 source chunks per read + _iteration_callback=lambda b, n: reads.append(b), + ) + # An unfiltered scan of 10 chunks arrives as ceil(10/3) = 4 reads. + assert dataset.to_table().num_rows == 400 + assert len(reads) == 4 + spans = sorted((b["time"].start, b["time"].stop) for b in reads) + assert spans == [(0, 30), (30, 60), (60, 90), (90, 100)] + + # Pruning still applies before merging: a filter keeping chunks + # 0-2 and 7-9 yields one merged read per consecutive run. + reads.clear() + keep = (pc.field("time") < _ts(pd.Timestamp("2020-01-02 06:00"))) | ( + pc.field("time") >= _ts(pd.Timestamp("2020-01-03 22:00")) + ) + table = dataset.to_table(filter=keep) + assert table.num_rows == (30 + 30) * 4 + spans = sorted((b["time"].start, b["time"].stop) for b in reads) + assert spans == [(0, 30), (70, 100)] + + +def test_coalesce_only_affects_scanner_not_fragments(): + dataset = XarrayPushdownDataset( + _hourly_grid(), {"time": 10}, coalesce_rows=10_000 + ) + # Fragment consumers (DataFusion, dask) keep one fragment per source + # chunk for their own parallelism. + assert len(dataset.get_fragments()) == 10 + + +def test_prefetch_bytes_scan_reads_every_block_once(): + source = xr.Dataset( + {"v": (["step"], np.arange(10_000.0))}, + coords={"step": np.arange(10_000.0)}, + ) + reads: list = [] + # 100 chunks of 100 rows x 16 bytes/row = 1600 bytes per block; a + # 4000-byte budget throttles admission well below the 8 threads + # prefetch allows. The scan must still visit every block exactly + # once and return the full table. + dataset = XarrayPushdownDataset( + source, + {"step": 100}, + prefetch=8, + prefetch_bytes=4_000, + _iteration_callback=lambda b, n: reads.append(b), + ) + table = dataset.to_table() + assert table.num_rows == 10_000 + assert len(reads) == 100 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"prefetch": 0}, # at most one worker: the pool-less path + {"prefetch": -1}, + {"coalesce_rows": 0}, + {"prefetch_bytes": 0}, + ], + ids=lambda kw: ( + next(iter(kw.items()))[0] + "=" + str(next(iter(kw.values()))) + ), +) +def test_degenerate_tuning_values_still_scan_exactly(kwargs): + # Degenerate knob values may degrade the schedule (prefetch <= 1 + # takes the pool-less path) but never the answer, and never hang. + src = xr.Dataset( + {"v": (["i"], np.arange(100.0))}, coords={"i": np.arange(100.0)} + ) + dataset = XarrayPushdownDataset(src, {"i": 10}, **kwargs) + assert dataset.to_table().num_rows == 100 + assert dataset.count_rows(filter=pc.field("i") >= 50.0) == 50 + + +@pytest.mark.parametrize("batch_size", [0, -5]) +def test_non_positive_batch_size_fails_at_construction(batch_size): + # batch_size cannot degrade gracefully: a zero size never advances + # the zero-column scan's row loop, so it is rejected eagerly. + src = xr.Dataset( + {"v": (["i"], np.arange(100.0))}, coords={"i": np.arange(100.0)} + ) + with pytest.raises(ValueError, match="batch_size"): + XarrayPushdownDataset(src, {"i": 10}, batch_size=batch_size) + + +# -- Laziness, re-scannability, concurrency ----------------------------------- + + +def test_abandoned_scanner_does_not_wedge_later_scans(counted): + dataset, counter = counted + batches = dataset.scanner().to_batches() + next(batches) + del batches # LIMIT-style early stop: consumer walks away mid-scan + counter.blocks.clear() + assert dataset.count_rows() == 400 + table = dataset.to_table(columns=["t2m"]) + assert table.num_rows == 400 + + +def test_concurrent_scans_are_isolated_and_exact(): + # Engines scan from their own worker threads; simultaneous filtered + # scans over one dataset must not cross-talk. + src = xr.Dataset( + {"v": (["i"], np.arange(20_000.0))}, coords={"i": np.arange(20_000.0)} + ) + dataset = XarrayPushdownDataset(src, {"i": 500}, prefetch=4) + results: list = [None] * 8 + errors: list = [] + + def worker(k): + try: + lo = k * 1000.0 + predicate = (pc.field("i") >= lo) & (pc.field("i") < lo + 3000.0) + results[k] = dataset.to_table(filter=predicate).num_rows + except Exception as exc: # noqa: BLE001 — reported by the assert + errors.append(f"{k}: {exc}") + + threads = [threading.Thread(target=worker, args=(k,)) for k in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(60) + assert not any(t.is_alive() for t in threads), "a scan wedged" + assert not errors + assert results == [3000] * 8 + + # Two batch iterators over the same dataset, consumed alternately, + # stay exact and independent. + a = dataset.scanner().to_batches() + b = dataset.scanner().to_batches() + rows_a = rows_b = 0 + exhausted_a = exhausted_b = False + while not (exhausted_a and exhausted_b): + batch = next(a, None) + if batch is None: + exhausted_a = True + else: + rows_a += batch.num_rows + batch = next(b, None) + if batch is None: + exhausted_b = True + else: + rows_b += batch.num_rows + assert rows_a == rows_b == 20_000 + + +# -- Pool lifecycle ------------------------------------------------------------ + + +def test_prefetch_pool_threads_all_started_at_construction(ds): + dataset = XarrayPushdownDataset(ds, {"time": 5}, prefetch=6) + # Every pool thread must exist before the first scan: a thread + # spawned later, from inside an engine's scan callback, is exactly + # the deadlock the pre-spawn exists to prevent. Thread accounting + # is only visible on the executor's private state. + assert len(dataset._pool._threads) == 6 + + +def test_prefetch_pool_shut_down_when_dataset_dies(ds): + import gc + + dataset = XarrayPushdownDataset(ds, {"time": 5}, prefetch=4) + pool = dataset._pool + del dataset + gc.collect() + # A shut-down executor refuses new work — the observable contract + # that its threads have been told to exit. + with pytest.raises(RuntimeError): + pool.submit(lambda: None) diff --git a/tests/test_arrow_dataset_integration.py b/tests/test_arrow_dataset_integration.py new file mode 100644 index 0000000..f7880e2 --- /dev/null +++ b/tests/test_arrow_dataset_integration.py @@ -0,0 +1,405 @@ +"""Integration tests: the arrow-dataset contract against real cloud stores. + +The same contract ``test_arrow_dataset.py`` pins on synthetic data, +exercised at scale: real Zarr stores read over the network, consumed +through the real engines (DuckDB, Polars, DataFusion). Assertions are +plan-shape — exactly which source chunks each query reads, exact row +counts, values matched against a direct xarray read of the same window — +so a pruning or fast-path regression fails long before it shows up in +wall-clock noise. + +Two axes, both extensible: + +* ``CASES`` — one :class:`StoreCase` per dataset. Expectations are + computed from the case's declared cadence and windows plus the store's + own coordinates. A new dataset is a config entry, provided its + temporal dimension is named ``time``, its cadence is regular, and its + grid is dense; anything else needs test changes, not just a case. +* ``ENGINES`` — engine name to runner function; every runner executes + the same windowed aggregation through its engine's idiomatic path + (DuckDB SQL, Polars lazy expressions, DataFusion SQL). + +Reads anonymously from public buckets. Excluded from the CI unit run +(``pytest -m "not integration"``); run deliberately with +``pytest -m integration tests/test_arrow_dataset_integration.py``. +""" + +import threading +from dataclasses import dataclass, field + +import pandas as pd +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + +pytestmark = pytest.mark.integration + + +@dataclass(frozen=True) +class StoreCase: + id: str + url: str + variable: str # the variable every scan queries + other_variable: str # registered alongside; must never be read + time_chunk: int # registration granularity, steps per chunk + steps_per_day: int # from the store's cadence + window_start: str # a day-window anchor, chunk-aligned + month: str # a chunk-aligned month for the arithmetic count + bbox: dict = field(default_factory=dict) # dim -> (lo, hi), inclusive + storage_options: dict = field(default_factory=dict) + + +CASES = [ + StoreCase( + id="arco-era5", + url="gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", + variable="2m_temperature", + other_variable="10m_u_component_of_wind", + time_chunk=1, + steps_per_day=24, + window_start="2020-01-01", + month="2020-01", + bbox={"latitude": (36, 44), "longitude": (350, 360)}, + storage_options={"token": "anon"}, + ), +] + + +# -- Engine runners ------------------------------------------------------------ +# One function per engine: run count(*) + avg(variable) over a half-open +# time window (plus an optional bbox) through the engine's idiomatic +# path, returning (rows, mean). Adding an engine is one function and one +# registry entry; tests parametrize over the registry. + + +def _sql(case, t0, t1, bbox) -> str: + conds = [f"time >= TIMESTAMP '{t0}'", f"time < TIMESTAMP '{t1}'"] + for dim, (lo, hi) in bbox.items(): + conds.append(f'"{dim}" BETWEEN {lo} AND {hi}') + return ( + f'SELECT count(*), avg("{case.variable}") FROM t ' + f"WHERE {' AND '.join(conds)}" + ) + + +def _duckdb_scan(case, dataset, t0, t1, bbox): + duckdb = pytest.importorskip("duckdb") + con = duckdb.connect() + con.register("t", dataset) + n, mean = con.execute(_sql(case, t0, t1, bbox)).fetchone() + return int(n), float(mean) + + +def _datafusion_scan(case, dataset, t0, t1, bbox): + from datafusion import SessionContext + + ctx = SessionContext() + ctx.register_dataset("t", dataset) + row = ctx.sql(_sql(case, t0, t1, bbox)).to_pandas().iloc[0] + return int(row.iloc[0]), float(row.iloc[1]) + + +def _polars_scan(case, dataset, t0, t1, bbox): + pl = pytest.importorskip("polars") + lf = pl.scan_pyarrow_dataset(dataset).filter( + (pl.col("time") >= t0.to_pydatetime()) + & (pl.col("time") < t1.to_pydatetime()) + ) + for dim, (lo, hi) in bbox.items(): + lf = lf.filter((pl.col(dim) >= lo) & (pl.col(dim) <= hi)) + out = lf.select( + pl.len().alias("n"), pl.col(case.variable).mean().alias("mean") + ).collect() + return int(out["n"][0]), float(out["mean"][0]) + + +# DataFusion consumes the dataset through get_fragments() — one +# fragment per source chunk, which scanner-level coalescing +# deliberately leaves untouched (see the contract in +# test_arrow_dataset.py). The other engines scan the whole dataset. +_datafusion_scan.consumes_fragments = True # type: ignore[attr-defined] + +ENGINES = { + "duckdb": _duckdb_scan, + "polars": _polars_scan, + "datafusion": _datafusion_scan, +} + + +@pytest.fixture(params=sorted(ENGINES), ids=str) +def scan(request): + """The engine runner under test.""" + return ENGINES[request.param] + + +# -- Dataset cases ------------------------------------------------------------- + + +@pytest.fixture(scope="module", params=CASES, ids=lambda c: c.id) +def case(request) -> StoreCase: + c: StoreCase = request.param + # The expectations below assume windows land on chunk boundaries; + # reject a miswritten case loudly instead of failing tests obscurely. + assert c.steps_per_day % c.time_chunk == 0, ( + f"{c.id}: steps_per_day must be a multiple of time_chunk" + ) + steps_from_midnight = ( + pd.Timestamp(c.window_start) - pd.Timestamp(c.window_start).normalize() + ) / (pd.Timedelta(days=1) / c.steps_per_day) + assert steps_from_midnight % c.time_chunk == 0, ( + f"{c.id}: window_start is not chunk-aligned" + ) + return c + + +@pytest.fixture(scope="module") +def source(case) -> xr.Dataset: + ds = xr.open_zarr( + case.url, + chunks=None, + storage_options=case.storage_options, + consolidated=True, + )[[case.variable, case.other_variable]] + # Chunk boundaries are laid out from the store's own time origin; + # both declared anchors must land on one or the fast-path/read + # expectations below are silently wrong for this case. + origin = pd.Timestamp(ds.time.values[0]) + chunk_span = pd.Timedelta(days=1) / case.steps_per_day * case.time_chunk + for name in ("window_start", "month"): + anchor = pd.Timestamp(getattr(case, name)) + assert (anchor - origin) % chunk_span == pd.Timedelta(0), ( + f"{case.id}: {name} is not aligned to the store's chunk grid" + ) + return ds + + +def _tracked(case, source, variables=None, **kwargs): + """A pushdown dataset over ``variables`` recording every block read.""" + reads: list = [] + column_sets: list = [] + dataset = XarrayPushdownDataset( + source[variables or [case.variable]], + {"time": case.time_chunk}, + prefetch=16, + _iteration_callback=lambda b, names: ( + reads.append(b), + column_sets.append(tuple(names)), + ), + **kwargs, + ) + return dataset, reads, column_sets + + +def _grid_cells(case, source, use_bbox=True) -> int: + """Cells per time step, inside the case's bbox unless disabled.""" + cells = 1 + for dim in source[case.variable].dims: + if dim == "time": + continue + vals = source[dim].values + if use_bbox and dim in case.bbox: + lo, hi = case.bbox[dim] + cells *= int(((vals >= lo) & (vals <= hi)).sum()) + else: + cells *= len(vals) + return cells + + +def _day_window(case, day=0, days=1) -> tuple[pd.Timestamp, pd.Timestamp]: + start = pd.Timestamp(case.window_start) + pd.Timedelta(days=day) + return start, start + pd.Timedelta(days=days) + + +def _day_chunks(case: StoreCase) -> int: + return case.steps_per_day // case.time_chunk + + +# -- The contract, engine by engine --------------------------------------------- + + +def test_windowed_scan_prunes_and_is_exact(scan, case, source): + # One day + bbox: every engine must push the window down so only the + # day's chunks are read, and row count and mean must match a direct + # xarray read of the same window. + dataset, reads, _ = _tracked(case, source) + t0, t1 = _day_window(case) + n, mean = scan(case, dataset, t0, t1, case.bbox) + + assert len(reads) == _day_chunks(case), "the engine did not prune" + assert n == case.steps_per_day * _grid_cells(case, source) + + window = source[case.variable].sel(time=slice(t0, t1 - pd.Timedelta("1ns"))) + for dim, (lo, hi) in case.bbox.items(): + keep = window[dim][(window[dim] >= lo) & (window[dim] <= hi)] + window = window.sel({dim: keep}) + # rel=1e-5: the store is float32, so the two sides accumulate the + # mean in different orders and dtypes. + assert mean == pytest.approx(float(window.mean()), rel=1e-5) + + +def test_projection_reads_only_the_referenced_variable(scan, case, source): + # Two variables registered; a query touching one must never read the + # other, whichever engine decides the column set. + dataset, _, column_sets = _tracked( + case, source, variables=[case.variable, case.other_variable] + ) + t0, t1 = _day_window(case) + scan(case, dataset, t0, t1, {}) + read = {name for cols in column_sets for name in cols} + assert case.variable in read + assert case.other_variable not in read + + +def test_dataset_is_rescannable_across_queries(scan, case, source): + # One wrapper, two queries of different shapes: the second scan must + # see fresh state, not a consumed stream or stale pruning (a bbox + # left over) from the first. + dataset, reads, _ = _tracked(case, source) + t0, t1 = _day_window(case) + n, _ = scan(case, dataset, t0, t1, case.bbox) + assert n == case.steps_per_day * _grid_cells(case, source) + assert len(reads) == _day_chunks(case) + + reads.clear() + t0, t1 = _day_window(case, day=1) + n, _ = scan(case, dataset, t0, t1, {}) + assert n == case.steps_per_day * _grid_cells(case, source, use_bbox=False) + assert len(reads) == _day_chunks(case) + + +def test_coalescing_merges_consecutive_reads(scan, case, source): + # The same day window in a handful of merged reads instead of one + # per chunk; the answer must not change. + cells = _grid_cells(case, source, use_bbox=False) + merge_chunks = 6 # chunks per merged read + dataset, reads, _ = _tracked( + case, + source, + coalesce_rows=merge_chunks * case.time_chunk * cells, + ) + t0, t1 = _day_window(case) + n, _ = scan(case, dataset, t0, t1, {}) + assert n == case.steps_per_day * cells + if getattr(scan, "consumes_fragments", False): + # Fragment consumers read one source chunk per fragment. + assert len(reads) == _day_chunks(case) + else: + assert len(reads) == -(-_day_chunks(case) // merge_chunks) # ceil + + +def test_concurrent_queries_stay_exact(scan, case, source): + # Engines scan from worker threads; two simultaneous queries over + # disjoint days of one wrapper must both come back exact. + dataset, _, _ = _tracked(case, source) + cells = _grid_cells(case, source, use_bbox=False) + results: list = [None, None] + errors: list = [] + + def worker(day): + try: + t0, t1 = _day_window(case, day) + results[day] = scan(case, dataset, t0, t1, {})[0] + except Exception as exc: # noqa: BLE001 — reported by the assert + errors.append(str(exc)) + + # daemon: a genuinely wedged scan must fail the assert, not keep + # the interpreter alive after pytest reports it. + threads = [ + threading.Thread(target=worker, args=(d,), daemon=True) for d in (0, 1) + ] + for t in threads: + t.start() + for t in threads: + t.join(120) + assert not any(t.is_alive() for t in threads), "a scan wedged" + assert not errors + assert results == [case.steps_per_day * cells] * 2 + + +# -- Dataset-level fast paths (no engine in the loop) --------------------------- + + +def test_multiday_global_window_prunes_to_its_chunks(case, source): + # Windows wider than one day prune exactly (7 days of hourly chunks: + # 168 reads on arco-era5). A property of the dataset's own scanner, + # so it runs once here instead of once per engine; batches are + # streamed and dropped to keep the 100M-row scan out of memory. + import pyarrow as pa + import pyarrow.compute as pc + + dataset, reads, _ = _tracked(case, source) + days = 7 + t0, t1 = _day_window(case, days=days) + predicate = (pc.field("time") >= pa.scalar(t0, type=pa.timestamp("ns"))) & ( + pc.field("time") < pa.scalar(t1, type=pa.timestamp("ns")) + ) + scanner = dataset.scanner(columns=[case.variable], filter=predicate) + rows = sum(batch.num_rows for batch in scanner.to_batches()) + assert len(reads) == days * _day_chunks(case), "wide window mispruned" + assert rows == days * case.steps_per_day * _grid_cells( + case, source, use_bbox=False + ) + + +def test_count_rows_fast_path_reads_nothing(case, source): + # A chunk-aligned month: every surviving chunk is provably inside + # the range, so the count is pure arithmetic. + import pyarrow as pa + import pyarrow.compute as pc + + dataset, reads, _ = _tracked(case, source) + lo = pd.Timestamp(case.month) + hi = lo + pd.offsets.MonthBegin(1) + predicate = (pc.field("time") >= pa.scalar(lo, type=pa.timestamp("ns"))) & ( + pc.field("time") < pa.scalar(hi, type=pa.timestamp("ns")) + ) + steps = (hi - lo) / pd.Timedelta(days=1) * case.steps_per_day + count = dataset.count_rows(filter=predicate) + assert reads == [], "count_rows fast path must not read data" + assert count == int(steps) * _grid_cells(case, source, use_bbox=False) + + +def test_polars_lazy_roundtrip_window_reads_only_its_blocks(case, source): + # Lazy round-trip: construction reads nothing with template coords; + # a one-day window reads only its own coalesced blocks. Polars is + # the one engine whose results re-execute over this dataset — + # DuckDB relations cannot (see limitations.md) and DataFusion's + # chunked round-trip lives on its native path. + pl = pytest.importorskip("polars") + if tuple(int(p) for p in pl.__version__.split(".")[:2]) >= (1, 43): + # polars 1.43 regressed streaming re-execution over pyarrow + # datasets: this window takes >10 minutes against 7s on 1.42. + pytest.skip("polars >= 1.43 streaming re-execution regression") + cells = _grid_cells(case, source, use_bbox=False) + merge_chunks = 6 + dataset, reads, _ = _tracked( + case, + source, + variables=[case.variable], + coalesce_rows=merge_chunks * case.time_chunk * cells, + ) + lf = pl.scan_pyarrow_dataset(dataset) + reads.clear() + lazy = xql.to_dataset( + lf, + template=source[[case.variable]], + chunks={"time": case.steps_per_day}, + coords="template", + ) + assert reads == [], "lazy construction must not read the source" + t0, t1 = _day_window(case) + value = float( + lazy[case.variable] + .sel(time=slice(t0, t1 - pd.Timedelta("1ns"))) + .mean() + .compute() + ) + assert len(reads) == -(-_day_chunks(case) // merge_chunks) + oracle = float( + source[case.variable] + .sel(time=slice(t0, t1 - pd.Timedelta("1ns"))) + .mean() + ) + assert value == pytest.approx(oracle, rel=1e-6) diff --git a/tests/test_df.py b/tests/test_df.py index aec174a..5185b57 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -16,6 +16,7 @@ explode, from_map, from_map_batched, + group_vars_by_dims, iter_record_batches, partition_metadata, pivot, @@ -429,20 +430,22 @@ def test_read_xarray_loads_one_chunk_at_a_time(large_ds): peaks.append(cur_peak) for size in sizes: - # Observed range: 1.59–1.83× on macOS, up to ~2.7× on Linux - # (glibc + Arrow allocate more intermediate buffers). - # iter_record_batches holds data-variable arrays (≈1× chunk) while - # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk). + # iter_record_batches' whole-partition fast path holds the + # data-variable arrays (≈1× chunk) plus repeat/tile-expanded + # coordinate columns (n_dims × 8 bytes × rows, ≈1.5× chunk + # for this 3-dim float64 dataset) for the partition being + # streamed; batches themselves are zero-copy slices. assert chunk_size * 1.3 < size, f"size {size} unexpectedly low" - assert chunk_size * 3.5 > size, f"size {size} unexpectedly high" + assert chunk_size * 4.0 > size, f"size {size} unexpectedly high" for peak in peaks: - # Observed range: 1.84–3.28× on macOS, up to ~4.15× on Linux - # (glibc + Arrow hold more intermediate buffers at peak). - # Peak includes data arrays + Arrow batch + temporary coordinate index - # arrays; the first batch of each chunk is highest (Dask compute overhead). + # Peak adds transient buffers on top of the steady state: + # np.repeat/np.tile intermediates for the coordinate columns + # and Arrow's from_pandas null scan; the first batch of each + # chunk is highest (Dask compute overhead). Observed ~5.04× + # on macOS. assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low" - assert chunk_size * 5.0 > peak, f"peak {peak} unexpectedly high" + assert chunk_size * 6.5 > peak, f"peak {peak} unexpectedly high" assert max(peaks) < large_ds.nbytes finally: @@ -560,6 +563,27 @@ def test_compute_chunks_tuples_sum_to_dim_size(): assert sum(tup) == ds.sizes[dim] +def test_iter_record_batches_large_string_dim_coord(): + """A string dim coord big enough that pa.array returns a ChunkedArray. + + Pivoting tiles a string dimension coordinate across every row of the + partition; for a few million rows pyarrow's numpy-unicode conversion + returns a ChunkedArray, which RecordBatch.from_arrays rejects. + Regression: found by the forecast-skill benchmark (a 2-model x 3.3M-row + window) streaming through the pyarrow dataset protocol into DuckDB. + """ + n_x = 1_700_000 # 2 * n_x rows: comfortably past the chunking threshold + ds = xr.Dataset( + {"value": (("model", "x"), np.zeros((2, n_x), dtype="float32"))}, + coords={"model": ["pangu", "graphcast"], "x": np.arange(n_x)}, + ) + schema = _parse_schema(ds) + got_rows = 0 + for batch in iter_record_batches(ds, schema, batch_size=DEFAULT_BATCH_SIZE): + got_rows += batch.num_rows + assert got_rows == 2 * n_x + + # -- Object-dtype and out-of-ns-range coordinate support -------------------- @@ -693,3 +717,53 @@ def test_partition_metadata_in_range_datetime_still_pruned(): for m in meta: _, _, tag = m["time"] assert tag == "timestamp_ns" + + +class TestGroupVarsByDims: + def test_single_dim_group(self): + ds = xr.Dataset( + { + "a": (["x", "y"], np.zeros((2, 3))), + "b": (["x", "y"], np.ones((2, 3))), + } + ) + groups = group_vars_by_dims(ds) + assert groups == {("x", "y"): ["a", "b"]} + + def test_multiple_dim_groups(self): + ds = xr.Dataset( + { + "surface": (["time", "lat", "lon"], np.zeros((2, 3, 4))), + "upper": ( + ["time", "lat", "lon", "level"], + np.zeros((2, 3, 4, 5)), + ), + } + ) + groups = group_vars_by_dims(ds) + assert set(groups.keys()) == { + ("time", "lat", "lon"), + ("time", "lat", "lon", "level"), + } + assert groups[("time", "lat", "lon")] == ["surface"] + assert groups[("time", "lat", "lon", "level")] == ["upper"] + + def test_empty_dataset(self): + assert group_vars_by_dims(xr.Dataset()) == {} + + def test_includes_scalar_group(self): + """Scalar (0-dim) variables group under the empty dims tuple.""" + ds = xr.Dataset( + {"band": (["y", "x"], np.zeros((2, 3))), "projection": ((), 0)} + ) + groups = group_vars_by_dims(ds) + assert groups == {("y", "x"): ["band"], (): ["projection"]} + + def test_ignores_coords(self): + """Coordinate variables shouldn't be returned as groups.""" + ds = xr.Dataset( + {"v": (["x"], np.arange(3))}, + coords={"x": np.arange(3), "label": ("x", ["a", "b", "c"])}, + ) + groups = group_vars_by_dims(ds) + assert groups == {("x",): ["v"]} diff --git a/tests/test_duckdb_backend.py b/tests/test_duckdb_backend.py new file mode 100644 index 0000000..0a00450 --- /dev/null +++ b/tests/test_duckdb_backend.py @@ -0,0 +1,387 @@ +"""Tests for the DuckDB engine adapter and the engine-agnostic round-trip. + +Covers the two seams of the multi-engine design: ``xql.register`` puts a +lazy Dataset on a DuckDB connection, DuckDB executes its own SQL dialect +(including extensions), and ``xql.to_dataset`` rebuilds a labeled +Dataset from the Arrow result. +""" + +import duckdb +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.duckdb import ( + XarrayArrowStream, + XarrayPushdownDataset, +) + + +@pytest.fixture +def ds() -> xr.Dataset: + np.random.seed(7) + time = pd.date_range("2021-01-01", periods=8, freq="h") + lat = np.linspace(-10.0, 10.0, 5) + lon = np.linspace(0.0, 40.0, 6) + temperature = 15 + 8 * np.random.randn(8, 5, 6) + precipitation = 10 * np.random.rand(8, 5, 6) + return xr.Dataset( + data_vars=dict( + temperature=(["time", "lat", "lon"], temperature), + precipitation=(["time", "lat", "lon"], precipitation), + ), + coords=dict(time=time, lat=lat, lon=lon), + attrs=dict(description="Synthetic weather."), + ).chunk({"time": 4}) + + +@pytest.fixture +def con(ds) -> duckdb.DuckDBPyConnection: + connection = duckdb.connect() + xql.register(connection, "weather", ds) + return connection + + +def test_full_scan_round_trips(con, ds): + rel = con.sql( + "SELECT time, lat, lon, temperature, precipitation FROM weather " + "ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + xr.testing.assert_allclose(out, ds.compute()) + assert out.attrs == ds.attrs + + +def test_aggregation_round_trips_on_surviving_dims(con, ds): + rel = con.sql( + "SELECT time, AVG(temperature) AS temperature FROM weather " + "GROUP BY time ORDER BY time" + ) + out = xql.to_dataset(rel, template=ds) + + expected = ds["temperature"].mean(["lat", "lon"]).compute() + assert list(out.dims) == ["time"] + np.testing.assert_allclose(out["temperature"].values, expected.values) + + +def test_registered_table_is_requeryable(con): + first = con.sql("SELECT COUNT(*) AS n FROM weather").fetchone()[0] + second = con.sql("SELECT COUNT(*) AS n FROM weather").fetchone()[0] + assert first == second == 8 * 5 * 6 + + +def test_registration_is_lazy(ds): + reads: list = [] + stream = XarrayArrowStream( + ds, _iteration_callback=lambda b, p: reads.append(b) + ) + + con = duckdb.connect() + con.register("weather", stream) + assert reads == [] # registration reads no data + + con.sql("SELECT AVG(temperature) FROM weather").fetchall() + assert len(reads) > 0 # data was read during query execution + + +def test_where_filter_yields_sparse_result(con, ds): + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather " + "WHERE lat > 0 ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + expected = ds[["temperature"]].sel(lat=ds.lat[ds.lat > 0]).compute() + xr.testing.assert_allclose(out, expected) + + +def test_template_sparsity_reindexes_to_full_extent(con, ds): + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather WHERE lat > 0" + ) + out = xql.to_dataset(rel, template=ds, sparsity="template") + + assert out.sizes == {"time": 8, "lat": 5, "lon": 6} + assert out["temperature"].isnull().sum() == 8 * 3 * 6 # lat <= 0 cells + + +def test_duckdb_dialect_and_join(con, ds): + # Engine-native SQL: DuckDB's date_part plus a join against a local + # relation — nothing xarray-sql has to understand. + con.sql("CREATE TABLE labels AS SELECT 0 AS h, 'midnight' AS label") + rel = con.sql( + "SELECT w.time, AVG(w.temperature) AS temperature, ANY_VALUE(l.label) AS label " + "FROM weather w JOIN labels l ON date_part('hour', w.time) = l.h " + "GROUP BY w.time" + ) + out = xql.to_dataset(rel, dims=["time"]) + assert out.sizes == {"time": 1} + + +def test_to_dataset_accepts_plain_arrow_table(ds): + table = pa.table( + { + "time": pd.date_range("2021-01-01", periods=3, freq="h"), + "temperature": [1.0, 2.0, 3.0], + } + ) + out = xql.to_dataset(table, dims=["time"]) + np.testing.assert_allclose(out["temperature"].values, [1.0, 2.0, 3.0]) + + +def test_to_dataset_requires_dims_or_template(): + table = pa.table({"a": [1, 2], "b": [3.0, 4.0]}) + with pytest.raises(ValueError, match="dims cannot be inferred"): + xql.to_dataset(table) + + +def test_to_dataset_rejects_missing_dim_column(): + table = pa.table({"a": [1, 2], "b": [3.0, 4.0]}) + with pytest.raises(ValueError, match="not columns of the result"): + xql.to_dataset(table, dims=["z"]) + + +def test_register_splits_mixed_dimension_variables(ds): + mixed = ds.assign(surface=ds["temperature"].isel(time=0, drop=True)) + con = duckdb.connect() + xql.register(con, "weather", mixed) + + n_full = con.sql("SELECT COUNT(*) FROM weather_time_lat_lon").fetchone()[0] + n_surface = con.sql("SELECT COUNT(*) FROM weather_lat_lon").fetchone()[0] + assert n_full == 8 * 5 * 6 + assert n_surface == 5 * 6 + + +def test_pushdown_dataset_rejects_mixed_dimension_variables(ds): + mixed = ds.assign(surface=ds["temperature"].isel(time=0, drop=True)) + with pytest.raises(ValueError, match="dimensions must be equal"): + XarrayPushdownDataset(mixed) + + +def _tracked_connection(ds): + """Register ds with an iteration callback; returns (con, reads).""" + reads: list = [] + dataset = XarrayPushdownDataset( + ds, _iteration_callback=lambda block, cols: reads.append((block, cols)) + ) + con = duckdb.connect() + con.register("weather", dataset) + return con, reads + + +def test_projection_pushdown_skips_unrequested_variables(ds): + con, reads = _tracked_connection(ds) + con.sql("SELECT AVG(temperature) FROM weather").fetchall() + assert reads # data was read + for _, cols in reads: + assert "precipitation" not in cols + + +def test_filter_pushdown_prunes_chunks(ds): + # ds is chunked {"time": 4} -> 2 chunks; this predicate covers only + # the first chunk, so the second is never loaded. + con, reads = _tracked_connection(ds) + n = con.sql( + "SELECT COUNT(*) FROM weather WHERE time < '2021-01-01 04:00:00'" + ).fetchone()[0] + assert n == 4 * 5 * 6 + assert len(reads) == 1 + + +def test_pushed_filter_is_applied_exactly(ds): + # DuckDB trusts pushed comparison filters and does not re-apply + # them, so the scan itself must enforce the predicate row-exactly — + # including inside chunks that pruning keeps. + con, _ = _tracked_connection(ds) + out = con.sql( + "SELECT COUNT(*) FROM weather " + "WHERE time = '2021-01-01 02:00:00' AND lat > 0" + ).fetchone()[0] + expected = int( + (ds.time == np.datetime64("2021-01-01T02:00:00")).sum() + * (ds.lat > 0).sum() + * ds.sizes["lon"] + ) + assert out == expected + + +def test_filter_on_variable_outside_projection(ds): + # The filter references `temperature`, the projection only `lat`; + # the scan must widen its columns to evaluate the predicate. + con, _ = _tracked_connection(ds) + got = con.sql( + "SELECT COUNT(DISTINCT lat) FROM weather WHERE temperature > 20" + ).fetchone()[0] + expected = len( + np.unique( + ds.lat.values[np.where((ds.temperature > 20).any(["time", "lon"]))] + ) + ) + assert got == expected + + +def test_or_and_in_filters_round_trip(ds): + con, _ = _tracked_connection(ds) + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather " + "WHERE lat < -5 OR lat > 5 ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + mask = (ds.lat < -5) | (ds.lat > 5) + expected = ds[["temperature"]].sel(lat=ds.lat[mask]).compute() + 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_fully_pruned_scan_returns_empty(con, ds): + n = con.sql( + "SELECT COUNT(*) FROM weather WHERE time >= '2022-01-01'" + ).fetchone()[0] + assert n == 0 + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather " + "WHERE time >= '2022-01-01'" + ) + out = xql.to_dataset(rel, template=ds) + assert out.sizes.get("time", 0) == 0 + + +def test_limit_terminates_early(con): + rows = con.sql("SELECT time, temperature FROM weather LIMIT 5").fetchall() + assert len(rows) == 5 + + +def test_descending_coordinate_pruning_is_correct(ds): + # Latitude stored north→south, like most rasters and ERA5. + flipped = ds.isel(lat=slice(None, None, -1)).chunk({"lat": 2}) + con = duckdb.connect() + xql.register(con, "weather", flipped) + got = con.sql("SELECT COUNT(*) FROM weather WHERE lat > 4").fetchone()[0] + expected = int((ds.lat > 4).sum()) * ds.sizes["time"] * ds.sizes["lon"] + assert got == expected + + +def test_integer_and_string_variables_round_trip(): + ds = xr.Dataset( + { + "klass": (["y", "x"], np.arange(12, dtype=np.uint8).reshape(3, 4)), + "label": ( + ["y", "x"], + np.array([["a"] * 4, ["b"] * 4, ["c"] * 4]), + ), + }, + coords={"y": np.arange(3), "x": np.arange(4)}, + ).chunk({"y": 2}) + con = duckdb.connect() + xql.register(con, "grid", ds) + rows = con.sql( + "SELECT label, SUM(klass) AS total FROM grid " + "WHERE klass >= 4 GROUP BY label ORDER BY label" + ).fetchall() + assert rows == [("b", 22), ("c", 38)] + + +def test_finely_chunked_dimension_uses_bucketed_pruning(): + # 5000 single-step time chunks exceeds the shadow fanout (1024), so + # pruning goes through the coarse-then-refine path; an equality in + # the middle of the axis must load exactly one chunk. + n = 5000 + ds = xr.Dataset( + {"v": (["time", "x"], np.random.rand(n, 2))}, + coords={ + "time": pd.date_range("2000-01-01", periods=n, freq="h"), + "x": np.arange(2), + }, + ).chunk({"time": 1}) + reads: list = [] + dataset = XarrayPushdownDataset( + ds, _iteration_callback=lambda block, cols: reads.append(block) + ) + con = duckdb.connect() + con.register("t", dataset) + + got = con.sql( + "SELECT COUNT(*) FROM t WHERE time = '2000-03-15 07:00:00'" + ).fetchone()[0] + assert got == 2 + assert len(reads) == 1 + + # A range spanning most of the axis stays correct (refinement is + # skipped when it cannot pay for itself). + reads.clear() + got = con.sql( + "SELECT COUNT(*) FROM t WHERE time >= '2000-01-01 12:00:00'" + ).fetchone()[0] + assert got == (n - 12) * 2 + + +def test_register_kwargs_are_forwarded(ds): + con = duckdb.connect() + xql.register(con, "weather", ds, prefetch=1, batch_size=7) + n = con.sql("SELECT COUNT(*) FROM weather").fetchone()[0] + assert n == 8 * 5 * 6 + + +def test_register_dispatches_to_datafusion(): + # The same entry point serves the default engine. + ctx = xql.XarrayContext() + small = xr.Dataset( + {"v": (["x"], np.arange(4.0))}, coords={"x": np.arange(4)} + ).chunk({"x": 2}) + xql.register(ctx, "t", small) + out = ctx.sql("SELECT x, v FROM t ORDER BY x").to_dataset() + np.testing.assert_allclose(out["v"].values, np.arange(4.0)) + + +def test_register_rejects_unknown_connection(ds): + with pytest.raises(TypeError, match="No xarray-sql engine adapter"): + xql.register(object(), "weather", ds) + + +def test_nan_coordinate_chunk_is_not_pruned(): + # NaN in a chunk's coordinate must disable pruning for that span, + # never poison the range guarantee (which would silently drop rows). + ds = xr.Dataset( + {"v": (["lat"], np.arange(6.0))}, + coords={"lat": [np.nan, 5.0, 10.0, 20.0, 30.0, 40.0]}, + ).chunk({"lat": 2}) + con = duckdb.connect() + xql.register(con, "t", ds) + assert con.sql("SELECT v FROM t WHERE lat = 5.0").fetchall() == [(1.0,)] + + +def test_cftime_dataset_aggregates_under_projection(): + cftime = pytest.importorskip("cftime") + + times = xr.date_range( + "2000-01-01", periods=6, calendar="360_day", use_cftime=True + ) + ds = xr.Dataset( + {"v": (["time"], np.arange(6.0))}, coords={"time": times} + ).chunk({"time": 3}) + con = duckdb.connect() + xql.register(con, "t", ds) + # The scan projects only `v`; the cftime dim column is absent from + # the scan schema but still shapes the iteration. + assert con.sql("SELECT SUM(v) FROM t").fetchone()[0] == 15.0 + + +def test_null_dimension_value_round_trips_positionally(): + # A NULL in a result's dim column must reject the affine fast path + # and fall back to positional scatter. + table = pa.table( + { + "lat": pa.array([0.0, None, 1.0, 3.0], type=pa.float64()), + "v": [10.0, 99.0, 11.0, 13.0], + } + ) + out = xql.to_dataset(table, dims=["lat"]) + np.testing.assert_allclose(out["v"].values, [10.0, 99.0, 11.0, 13.0]) diff --git a/tests/test_geometry.py b/tests/test_geometry.py new file mode 100644 index 0000000..5e3238d --- /dev/null +++ b/tests/test_geometry.py @@ -0,0 +1,170 @@ +"""GeoArrow point-geometry columns derived at registration.""" + +import json + +import numpy as np +import pyarrow as pa +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + + +@pytest.fixture +def grid() -> xr.Dataset: + return xr.Dataset( + {"risk": (["y", "x"], np.arange(8.0 * 6).reshape(8, 6))}, + coords={ + "y": np.linspace(-28.0, -29.4, 8), # descending, like rasters + "x": np.linspace(-58.0, -57.0, 6), + }, + ) + + +def test_geometry_field_annotation(grid): + dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y")) + field = dataset.schema.field("geometry") + assert field.type == pa.binary() + assert field.metadata[b"ARROW:extension:name"] == b"geoarrow.wkb" + meta = json.loads(field.metadata[b"ARROW:extension:metadata"]) + assert meta == {"crs": "OGC:CRS84"} + + +def test_wkb_points_decode_exactly(grid): + dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y")) + table = dataset.to_table(columns=["geometry", "x", "y"]) + blob = table["geometry"][0].as_py() + assert len(blob) == 21 and blob[0] == 1 + x = np.frombuffer(blob, "= -28.7) & (grid.y <= -28.0) + expected = grid.risk.values[inside.values, :] + assert got == (expected.size, round(float(expected.mean()), 3)) + + +def test_geopandas_consumes_native_points(grid): + gpd = pytest.importorskip("geopandas") + + dataset = xql.arrow_dataset( + grid, {"y": 4}, geometry=("x", "y"), geometry_encoding="point" + ) + gdf = gpd.GeoDataFrame.from_arrow(dataset.to_table()) + assert gdf.geometry.iloc[0].x == float(grid.x[0]) + assert str(gdf.crs).endswith("CRS84") + + +def test_geometry_name_collision_raises(): + clash = xr.Dataset( + {"geometry": (["x"], np.arange(3.0))}, + coords={"x": np.arange(3.0)}, + ) + with pytest.raises(ValueError, match="shadow"): + xql.arrow_dataset(clash, {"x": 3}, geometry=("x", "x")) + + +def test_bbox_conjuncts_prunes_and_pairs_with_st_within(grid, spatial_con): + con, reads = spatial_con + + bounds = (-58.1, -28.75, -56.9, -27.9) # xmin, ymin, xmax, ymax + conjuncts = xql.bbox_conjuncts(bounds, x="x", y="y") + assert '"x" BETWEEN' in conjuncts and '"y" BETWEEN' in conjuncts + reads.clear() + got = con.execute( + f"SELECT count(*) FROM t WHERE {conjuncts} " + "AND ST_Within(geometry, ST_GeomFromText(" + "'POLYGON ((-58.1 -28.75, -56.9 -28.75, -56.9 -27.9, " + "-58.1 -27.9, -58.1 -28.75))'))" + ).fetchone() + assert len(reads) == 1 # the y-range pruned to one chunk + inside = (grid.y >= -28.75) & (grid.y <= -27.9) + assert got[0] == int(inside.sum()) * grid.sizes["x"] + + +def test_bbox_conjuncts_accepts_bounds_objects(): + class Boxy: + bounds = (1.0, 2.0, 3.0, 4.0) + + sql = xql.bbox_conjuncts(Boxy(), x="lon", y="lat", pad=0.5) + assert sql == '"lon" BETWEEN 0.5 AND 3.5 AND "lat" BETWEEN 1.5 AND 4.5' + + +def test_wkb_points_guards_int32_offset_overflow(): + from xarray_sql.geometry import _wkb_points + + # Stride-0 broadcast views: len() reports ~103M points without + # allocating them, and the guard must fire before any buffer is + # built (pa.binary() offsets are int32; n * 21 would overflow). + n = 103_000_000 + x = np.broadcast_to(np.float64(0.0), (n,)) + with pytest.raises(ValueError, match="int32"): + _wkb_points(x, x) diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py new file mode 100644 index 0000000..185191e --- /dev/null +++ b/tests/test_lazy_roundtrip.py @@ -0,0 +1,377 @@ +"""Lazy chunked round-trip through engines beyond DataFusion. + +``xql.to_dataset(result, chunks=...)`` re-executes the engine's query +per accessed window. These tests verify the reconstruction is correct +on Polars frames (DuckDB chunked reconstruction fails fast — see +DuckDBHandle.supports_chunked — while its eager path works), that +laziness is real, and that one-shot +streams are rejected with a clear error. +""" + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + + +@pytest.fixture +def source() -> xr.Dataset: + np.random.seed(7) + return xr.Dataset( + { + "t2m": ( + ["time", "lat"], + np.random.rand(100, 6).astype(np.float64), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=100, freq="h"), + "lat": np.linspace(-25.0, 25.0, 6), + }, + attrs={"title": "synthetic"}, + ) + + +@pytest.fixture +def registered(source): + """A DuckDB connection with the source registered + a read counter.""" + duckdb = pytest.importorskip("duckdb") + + reads: list[dict] = [] + dataset = XarrayPushdownDataset( + source, {"time": 10}, _iteration_callback=lambda b, n: reads.append(b) + ) + con = duckdb.connect() + con.register("t", dataset) + return con, reads + + +def test_duckdb_chunked_fails_fast_with_guidance(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + # Re-executing a DuckDB relation from dask worker threads + # intermittently deadlocks inside duckdb-python when the query + # scans a Python-backed table; the library refuses instead of + # hanging (see DuckDBHandle.supports_chunked). + with pytest.raises(NotImplementedError, match="Polars"): + xql.to_dataset(rel, template=source, chunks={"time": 10}) + + +def test_duckdb_eager_round_trip_through_handle(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + out = xql.to_dataset(rel, template=source) + assert not out.chunks + assert out.attrs == source.attrs + xr.testing.assert_allclose(out, source) + + +def test_duckdb_eager_filtered_and_aggregated(source, registered): + con, _ = registered + rel = con.sql( + "SELECT time, avg(t2m) AS t2m FROM t " + "WHERE lat > 0 GROUP BY time ORDER BY time" + ) + out = xql.to_dataset(rel, template=source) + expected = source.t2m.sel(lat=source.lat[source.lat > 0]).mean("lat") + np.testing.assert_allclose(out.t2m.values, expected.values) + + +def test_polars_lazyframe_chunked_round_trip(source): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + out = xql.to_dataset(lf, template=source, chunks={"time": 20}) + assert out.chunks + xr.testing.assert_allclose(out.compute(), source) + + # Eager path through the same handle (LazyFrame has no stream + # protocol; the handle executes it once). + eager = xql.to_dataset(lf, template=source) + xr.testing.assert_allclose(eager, source) + + +def test_polars_eager_frame_is_reexecutable(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, template=source, chunks={"time": 50}) + xr.testing.assert_allclose(out.compute(), source) + + +def test_one_shot_stream_with_chunks_raises(source, registered): + con, _ = registered + table = con.sql("SELECT * FROM t").to_arrow_table() + with pytest.raises(TypeError, match="re-executable"): + xql.to_dataset(table, template=source, chunks={"time": 10}) + + +def test_inherit_without_chunked_source_falls_back_to_eager(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + out = xql.to_dataset(rel, template=source, chunks="inherit") + # The in-memory template has no multi-chunk dim: eager, dense. + assert not out.chunks + xr.testing.assert_allclose(out, source) + + +def test_stepped_indexer_uses_value_lists(source): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + out = xql.to_dataset(lf, template=source, chunks={"time": 10}) + # A step-2 selection is not a contiguous coordinate range; the + # values path must return exactly the requested rows. + stepped = out.t2m.isel(time=slice(10, 30, 2)).compute() + np.testing.assert_allclose( + stepped.values, source.t2m.isel(time=slice(10, 30, 2)).values + ) + + +def test_descending_coordinate_windows(): + pl = pytest.importorskip("polars") + + desc = xr.Dataset( + {"v": (["lat"], np.arange(8.0))}, + coords={"lat": np.linspace(70.0, 0.0, 8)}, # descending, like ERA5 + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(desc, {"lat": 4})) + out = xql.to_dataset( + lf, + template=desc, + chunks={"lat": 4}, + coords="template", + ) + xr.testing.assert_allclose(out.compute(), desc) + window = out.v.isel(lat=slice(2, 6)).compute() + np.testing.assert_allclose(window.values, desc.v.isel(lat=slice(2, 6))) + + +def test_unsorted_template_coords_window_exactly(): + pl = pytest.importorskip("polars") + + # Template coords are used verbatim, so the backend can see a + # non-monotonic coordinate array. A contiguous positional window + # like 1:3 then has monotonic values [7, 55], but the value range + # [7, 55] also admits 23 at position 3 — the scatter would write + # that unrequested row over a requested cell. Windows over a + # non-monotonic coordinate must use explicit value lists. + src = xr.Dataset( + {"v": (["x"], np.arange(4.0))}, + coords={"x": np.array([102.0, 7.0, 55.0, 23.0])}, + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(src, {"x": 4})) + out = xql.to_dataset(lf, template=src, chunks={"x": 4}, coords="template") + window = out.v.isel(x=slice(1, 3)).compute() + np.testing.assert_array_equal(window.values, [1.0, 2.0]) + xr.testing.assert_allclose(out.compute(), src) + + +def test_polars_float_value_windows_are_exact(): + pl = pytest.importorskip("polars") + + # Non-representable float coordinates: upstream Polars is_in drops + # them (silently matching nothing); the handle's degenerate-range + # translation must return exactly the requested rows. + src = xr.Dataset( + {"v": (["lat", "t"], np.arange(38.0).reshape(19, 2))}, + coords={"lat": np.linspace(-45.0, 45.0, 19), "t": [0.0, 1.0]}, + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(src, {"lat": 5})) + out = xql.to_dataset(lf, template=src, chunks={"lat": 5}) + # A stepped (non-contiguous) selection forces the value-list path. + picked = out.v.isel(lat=slice(1, 12, 2)).compute() + np.testing.assert_allclose( + picked.values, src.v.isel(lat=slice(1, 12, 2)).values + ) + + +def test_max_result_bytes_guards_stream_collection(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + with pytest.raises(ValueError, match="max_result_bytes"): + xql.to_dataset(rel, template=source, max_result_bytes=1_000) + # A generous budget passes untouched. + out = xql.to_dataset(rel, template=source, max_result_bytes=10**9) + xr.testing.assert_allclose(out, source) + + +def test_polars_large_float_value_lists_stay_flat(): + pl = pytest.importorskip("polars") + + # 5000 stepped float values in a single window: a left-deep OR + # chain plans quadratically at this size (seconds per window); the + # flat any_horizontal translation must stay exact and quick. + n = 10_000 + src = xr.Dataset( + {"v": (["x"], np.arange(float(n)))}, + coords={"x": np.linspace(-45.0, 45.0, n)}, + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(src, {"x": n})) + out = xql.to_dataset(lf, template=src, chunks={"x": n}) + picked = out.v.isel(x=slice(1, None, 2)).compute() + np.testing.assert_allclose( + picked.values, src.v.isel(x=slice(1, None, 2)).values + ) + + +def test_collect_streaming_falls_back_on_older_polars(): + from xarray_sql.lazyscan import _collect_streaming + + class OldLazyFrame: + # Pre-1.25 collect(): no ``engine`` keyword. + def collect(self): + return "collected" + + assert _collect_streaming(OldLazyFrame()) == "collected" + + +def test_max_result_bytes_guards_polars_lazyframe(source): + pl = pytest.importorskip("polars") + + # The LazyFrame eager fallback collects inside the engine before + # any batch surfaces; with a budget set it must stream through + # collect_batches so the guard fires before full materialization. + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + with pytest.raises(ValueError, match="max_result_bytes"): + xql.to_dataset(lf, template=source, max_result_bytes=1_000) + out = xql.to_dataset(lf, template=source, max_result_bytes=10**9) + xr.testing.assert_allclose(out, source) + + +def test_max_result_bytes_guards_table_only_results(source, registered): + con, _ = registered + table = con.sql("SELECT * FROM t").to_arrow_table() + + class TableOnly: + # The narrowest result surface: to_arrow_table() materializes + # in full before the budget can see a batch, so the guard runs + # on the materialized size. + def __init__(self, t): + self._t = t + + def to_arrow_table(self): + return self._t + + with pytest.raises(ValueError, match="max_result_bytes"): + xql.to_dataset( + TableOnly(table), template=source, max_result_bytes=1_000 + ) + out = xql.to_dataset( + TableOnly(table), template=source, max_result_bytes=10**9 + ) + xr.testing.assert_allclose(out, source) + + +def test_max_result_bytes_guards_dense_blowup(registered): + con, _ = registered + # A sparse diagonal: tiny Arrow payload, huge dense grid (the + # coordinate product), so the dense-size check must fire even + # though the stream fits the budget. + diag = pa.table( + { + "a": np.arange(3000.0), + "b": np.arange(3000.0), + "v": np.ones(3000), + } + ) + with pytest.raises(ValueError, match="dense reconstruction"): + xql.to_dataset(diag, dims=["a", "b"], max_result_bytes=10_000_000) + + +def test_duckdb_spill_chunked_round_trip(source, registered, tmp_path): + con, reads = registered + rel = con.sql("SELECT * FROM t") + reads.clear() + out = xql.to_dataset( + rel, template=source, chunks={"time": 10}, spill=tmp_path + ) + spilled = list(tmp_path.glob("*.parquet")) + assert len(spilled) == 1 + # The source was streamed exactly once (10 chunks), during the spill. + assert len(reads) == 10 + assert out.chunks + reads.clear() + xr.testing.assert_allclose(out.compute(), source) + # Windows re-execute against the Parquet file, not the source. + assert reads == [] + + +def test_duckdb_spill_filtered_aggregation(source, registered, tmp_path): + con, _ = registered + rel = con.sql( + "SELECT time, avg(t2m) AS t2m FROM t " + "WHERE lat > 0 GROUP BY time ORDER BY time" + ) + out = xql.to_dataset( + rel, template=source, chunks={"time": 25}, spill=tmp_path + ) + expected = source.t2m.sel(lat=source.lat[source.lat > 0]).mean("lat") + np.testing.assert_allclose(out.t2m.compute().values, expected.values) + + +def test_one_shot_table_spill_chunked(source, registered, tmp_path): + con, _ = registered + table = con.sql("SELECT * FROM t").to_arrow_table() + out = xql.to_dataset( + table, template=source, chunks={"time": 10}, spill=tmp_path + ) + assert out.chunks + xr.testing.assert_allclose(out.compute(), source) + + +def test_spill_file_removed_when_dataset_dies(source, registered, tmp_path): + import gc + + con, _ = registered + rel = con.sql("SELECT * FROM t") + out = xql.to_dataset( + rel, template=source, chunks={"time": 10}, spill=tmp_path + ) + assert list(tmp_path.glob("*.parquet")) + del out + gc.collect() + assert list(tmp_path.glob("*.parquet")) == [] + + +def test_spill_requires_chunks(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + with pytest.raises(ValueError, match="spill= only applies"): + xql.to_dataset(rel, template=source, spill=True) + + +def test_duckdb_handle_runner_stopped_when_handle_dies(source, registered): + import gc + + from xarray_sql.lazyscan import DuckDBHandle + + con, _ = registered + handle = DuckDBHandle(con.sql("SELECT * FROM t")) + runner = handle._runner + del handle + gc.collect() + # A shut-down executor refuses new work — the observable contract + # that the handle's dedicated engine thread has been told to exit. + with pytest.raises(RuntimeError): + runner.submit(lambda: None) + + +def test_polars_spill_uses_streaming_sink(source, tmp_path): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + out = xql.to_dataset( + lf, template=source, chunks={"time": 20}, spill=tmp_path + ) + xr.testing.assert_allclose(out.compute(), source) diff --git a/tests/test_sql.py b/tests/test_sql.py index 746dfb3..66c6dd2 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -6,7 +6,6 @@ import xarray as xr from xarray_sql import XarrayContext -from xarray_sql.sql import _group_vars_by_dims def test_sanity(air_dataset_small): @@ -325,58 +324,6 @@ def test_gregorian_like_no_cftime_udf(self): ).collect() -class TestGroupVarsByDims: - """Unit tests for the private _group_vars_by_dims helper.""" - - def test_single_dim_group(self): - ds = xr.Dataset( - { - "a": (["x", "y"], np.zeros((2, 3))), - "b": (["x", "y"], np.ones((2, 3))), - } - ) - groups = _group_vars_by_dims(ds) - assert groups == {("x", "y"): ["a", "b"]} - - def test_multiple_dim_groups(self): - ds = xr.Dataset( - { - "surface": (["time", "lat", "lon"], np.zeros((2, 3, 4))), - "upper": ( - ["time", "lat", "lon", "level"], - np.zeros((2, 3, 4, 5)), - ), - } - ) - groups = _group_vars_by_dims(ds) - assert set(groups.keys()) == { - ("time", "lat", "lon"), - ("time", "lat", "lon", "level"), - } - assert groups[("time", "lat", "lon")] == ["surface"] - assert groups[("time", "lat", "lon", "level")] == ["upper"] - - def test_empty_dataset(self): - assert _group_vars_by_dims(xr.Dataset()) == {} - - def test_includes_scalar_group(self): - """Scalar (0-dim) variables group under the empty dims tuple.""" - ds = xr.Dataset( - {"band": (["y", "x"], np.zeros((2, 3))), "projection": ((), 0)} - ) - groups = _group_vars_by_dims(ds) - assert groups == {("y", "x"): ["band"], (): ["projection"]} - - def test_ignores_coords(self): - """Coordinate variables shouldn't be returned as groups.""" - ds = xr.Dataset( - {"v": (["x"], np.arange(3))}, - coords={"x": np.arange(3), "label": ("x", ["a", "b", "c"])}, - ) - groups = _group_vars_by_dims(ds) - assert groups == {("x",): ["v"]} - - class TestFromDatasetMultiDims: """from_dataset should split datasets with mixed dims into multiple tables.""" diff --git a/tests/test_sql_recipes.py b/tests/test_sql_recipes.py new file mode 100644 index 0000000..5e654be --- /dev/null +++ b/tests/test_sql_recipes.py @@ -0,0 +1,61 @@ +"""The performance guide's SQL recipes, pinned on both engines. + +The guide documents caching as plain engine SQL rather than wrapping +it in a helper; this test runs the documented statement on DuckDB and +DataFusion so the recipe cannot rot. +""" + +import duckdb +import numpy as np +import pytest +import xarray as xr + +import xarray_sql as xql + + +def _grid() -> xr.Dataset: + np.random.seed(11) + return xr.Dataset( + { + "klass": ( + ["y", "x"], + np.random.randint(1, 6, (64, 64), dtype=np.uint8), + ) + }, + coords={ + "y": np.linspace(-34.0, -30.0, 64), + "x": np.linspace(-66.0, -62.0, 64), + }, + ).chunk({"y": 32}) + + +@pytest.fixture(params=["duckdb", "datafusion"]) +def con(request): + connection = ( + duckdb.connect() if request.param == "duckdb" else xql.XarrayContext() + ) + xql.register(connection, "grid", _grid()) + return connection + + +def _rows(con, sql) -> list[tuple]: + result = con.sql(sql) + if hasattr(result, "fetchall"): + return list(result.fetchall()) + frame = result.to_pandas() + return [tuple(r) for r in frame.itertuples(index=False)] + + +def test_documented_caching_recipe(con): + # The performance guide documents caching as plain engine SQL; this + # pins the recipe on both engines — including that DataFusion DDL + # is a lazy plan that must be collected to execute. + ctas = ( + "CREATE OR REPLACE TABLE cube AS " + "SELECT FLOOR(y) AS lat, klass, COUNT(*) AS n FROM grid " + "GROUP BY 1, 2 ORDER BY lat, klass" + ) + result = con.sql(ctas) + if hasattr(result, "collect"): + result.collect() + assert _rows(con, "SELECT SUM(n) FROM cube")[0][0] == 64 * 64 diff --git a/uv.lock b/uv.lock index 97391a1..cf72eb8 100644 --- a/uv.lock +++ b/uv.lock @@ -484,6 +484,48 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/d5/c5db1ea3394c6e1732fb3286b3bd878b59507a8f77d32a2cebda7d7b7cd4/donfig-0.8.1.post1-py3-none-any.whl", hash = "sha256:2a3175ce74a06109ff9307d90a230f81215cbac9a751f4d1c6194644b8204f9d", size = 21592, upload-time = "2024-05-23T14:13:55.283Z" }, ] +[[package]] +name = "duckdb" +version = "1.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/19/e57151753576373c6696a12022648546cca6038e8833fda2908ee2342d9b/duckdb-1.5.5.tar.gz", hash = "sha256:72f33ee57ca7595b23957671a2cc7f7fe2be0ecc2d68f63abedcfcaa3a5c1238", size = 18066741, upload-time = "2026-07-22T10:55:17.819Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/d4/298acf9331a80b3ce6ac64dd940e7e13f4058fb69d18914445f02e3c7bfe/duckdb-1.5.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:3b805507f88171b428b21c966c30e9a3d54e30b24528918a44ed0032542bc26f", size = 32702934, upload-time = "2026-07-22T10:53:19.069Z" }, + { url = "https://files.pythonhosted.org/packages/d5/90/c489fb63d64b2e7ee109ce8460bdede003a0f256e5b41a03a2a1c4764058/duckdb-1.5.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b08e19cc856220d8a26fa62abc2264b349aff67255e9373c6a3f607addd56dc6", size = 17343604, upload-time = "2026-07-22T10:53:22.767Z" }, + { url = "https://files.pythonhosted.org/packages/0b/27/effa80a15b1f0c61c235622f797868485359e8c9ad6a8e358e7a0c479151/duckdb-1.5.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a17e6a922e42a5c06ed2353fe78c5dff2610f6632d603836f9606ad0bf754079", size = 15488179, upload-time = "2026-07-22T10:53:25.945Z" }, + { url = "https://files.pythonhosted.org/packages/5d/07/21212345c8d24ba62dceaa20be3b21f5c46f1510b1b42ce93bb058afe0c4/duckdb-1.5.5-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1bdc38922c365c37720149f90d90b1e9823eb82dad6830855b5f87537fa6fc0c", size = 19367323, upload-time = "2026-07-22T10:53:30.23Z" }, + { url = "https://files.pythonhosted.org/packages/3f/d0/10371ae875fb4b5ef61bb892743b4b2e90c512b371fdf29317deb744857d/duckdb-1.5.5-cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e238060db5ca59879882a6e9b015e2c65d5c64ddf281ba1d7a9a2033764152cf", size = 21476568, upload-time = "2026-07-22T10:53:33.486Z" }, + { url = "https://files.pythonhosted.org/packages/b2/35/09568ce617dd7bc0757b3d7b6a981660b9e4f0b7594de8ed776755eae740/duckdb-1.5.5-cp310-cp310-win_amd64.whl", hash = "sha256:4acc72798ba1885a9c17d1242903d2cd502f13b1271c7677f7cab25d8578eceb", size = 13156129, upload-time = "2026-07-22T10:53:37.55Z" }, + { url = "https://files.pythonhosted.org/packages/9c/c2/b62ec24d57bb8df4e24b0b58f7f8facb32f5fdb9f1895aed9e9fcdded168/duckdb-1.5.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:1b543841b0ae18a9c982345cfa3987e9c065d3a4b0f067daa473d92d1e65f528", size = 32708371, upload-time = "2026-07-22T10:53:41.642Z" }, + { url = "https://files.pythonhosted.org/packages/8a/ce/769171ba45f0b73632dc3bc3108d891e81dd6c6bbfba630a34a75b4dcc0f/duckdb-1.5.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1a925d06c2a4c3b64553d6cc1aced5028d376d4479bed689a7d47e9b1dccd80a", size = 17343979, upload-time = "2026-07-22T10:53:44.951Z" }, + { url = "https://files.pythonhosted.org/packages/46/59/a8e3384ee916e00d5dcf985194c1511d61978540778a1e96fa47f9fb3e0d/duckdb-1.5.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0c42757cb34722144bd4dfb94b6f336339e7b2468f6813fa7fa9a319ba07bab4", size = 15493704, upload-time = "2026-07-22T10:53:47.912Z" }, + { url = "https://files.pythonhosted.org/packages/6f/1d/9840179c2607b90523a2884a129c4d4e6dbdc1178ba62a976c1043beba88/duckdb-1.5.5-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e72f9e1a4f90a5c8483ad4d540e495bf0834ba61c360b52499a573d7ed62a3f", size = 19366574, upload-time = "2026-07-22T10:53:51.876Z" }, + { url = "https://files.pythonhosted.org/packages/b5/55/f9641a4eebcc2f4df631287d6c3b9ed2eea3b92644f93acbad825e3972b6/duckdb-1.5.5-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9b6f86ed85d4ef5e0211eaebf75d057bd8bb520bba438a95dd0f4e42234bbfe", size = 21477952, upload-time = "2026-07-22T10:53:55.575Z" }, + { url = "https://files.pythonhosted.org/packages/3c/3a/07c3556e37a5c97b95917b029c8fdde4a25fbd76a660bacdac195cf20dcb/duckdb-1.5.5-cp311-cp311-win_amd64.whl", hash = "sha256:9f4287f97ccf0c1f3d471e7115be2b067cbf99627e2d34bffd462dd64703cddc", size = 13156986, upload-time = "2026-07-22T10:53:58.823Z" }, + { url = "https://files.pythonhosted.org/packages/4f/ff/07b48eef2078ca033847e9caa46cc7633b714c5f91ad1ce091c8ca89d792/duckdb-1.5.5-cp311-cp311-win_arm64.whl", hash = "sha256:179633a3fc6296c75d57c69c1e239fa9e5cdcb670fd1dbff88a02663f932905c", size = 14001317, upload-time = "2026-07-22T10:54:01.724Z" }, + { url = "https://files.pythonhosted.org/packages/d6/40/2e05d324400fdaa5656c9f48d6298da421cb034d85e509fa0e6e325cf04b/duckdb-1.5.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d4dd65f8941a604b947e0b9b4b4f7165988e29a23ec0b69b4038520956d9933e", size = 32753858, upload-time = "2026-07-22T10:54:05.514Z" }, + { url = "https://files.pythonhosted.org/packages/79/15/5ceb58ffb5bb8a62b3fd7abb39c41467cdf94850ece02e6d88664dfc75ce/duckdb-1.5.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33db46679b071f108d57139493dee2d37e1f5efcf5c5c039c2969eed11a6c8a7", size = 17368293, upload-time = "2026-07-22T10:54:09.139Z" }, + { url = "https://files.pythonhosted.org/packages/bf/5c/bf02da0b354fe83cca4f95a4fbf762181af466f7d551ab2a093f7698882a/duckdb-1.5.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f0b88535a5d86fdd63dba6ea02ab68c003dfb9e4892b11256ef24c4da208baae", size = 15509131, upload-time = "2026-07-22T10:54:12.228Z" }, + { url = "https://files.pythonhosted.org/packages/ea/a9/5f1f09da421d8e930e0b063d11c1b3f90363f40ede74438cd188afdd13a2/duckdb-1.5.5-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f316eae2323d9a851883fdf2dee91c1f9efe251ab33e14a2272f82a913422ed6", size = 19391959, upload-time = "2026-07-22T10:54:15.551Z" }, + { url = "https://files.pythonhosted.org/packages/4f/98/6549769f158126fa64fd6c1ac2eb59a18282146c939867a3eb31b7c1db07/duckdb-1.5.5-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7a6d2d11859d82a936ebdcb30ce3d8a1cbb3e990bff05c12abb9b54c44fa7bd1", size = 21510909, upload-time = "2026-07-22T10:54:19.681Z" }, + { url = "https://files.pythonhosted.org/packages/af/b7/5753b41d3124838f868f9f523362812d9fc45409e9e4dd70dcbb0a25826e/duckdb-1.5.5-cp312-cp312-win_amd64.whl", hash = "sha256:ddfbdb096c11d51ee22492397d342c90a82e62c5d09961477895934d0a25372f", size = 13168544, upload-time = "2026-07-22T10:54:22.789Z" }, + { url = "https://files.pythonhosted.org/packages/5c/28/44b679c7d46245f8398feae7edac959d1b83d4eb143e25b3fce0630b78bd/duckdb-1.5.5-cp312-cp312-win_arm64.whl", hash = "sha256:2725d2b9ace3a4e75d72fc5a239f6a44b502c580edadb8fb2676db772c5f9282", size = 13988684, upload-time = "2026-07-22T10:54:26.003Z" }, + { url = "https://files.pythonhosted.org/packages/47/37/4a38116e7700720fd152c666292214fd3abdf916496991296d8d1f66efbf/duckdb-1.5.5-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:cd98829b67788609017e65c761bd42a5dd0f9129441bed8bda4d6881ccf819f0", size = 32754294, upload-time = "2026-07-22T10:54:29.822Z" }, + { url = "https://files.pythonhosted.org/packages/66/42/7d392f1ba1eee0eaf4ab4c8c7a604bfe3536cd63f979cf5c98798664f807/duckdb-1.5.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:feead93c56679b79592d437c62975d39cb67adedffa7592c763baf8160ac7366", size = 17368211, upload-time = "2026-07-22T10:54:33.359Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a5/0a6f4fa60562faa615e55e15bd1953a2f2b17a8edd8105e5cda215e43457/duckdb-1.5.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:49c963d9469373d7aba8d750d9ea565ab823e94166efed953f184dd9b169b98c", size = 15509136, upload-time = "2026-07-22T10:54:36.369Z" }, + { url = "https://files.pythonhosted.org/packages/e4/cb/023c89f51978545b9fab318581bba0c457a58e7530d2d933e54ae7d8647c/duckdb-1.5.5-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a736217825461732b5442d05a220f3da2e23a0dae114efbf08c9bf171b53098a", size = 19392147, upload-time = "2026-07-22T10:54:39.551Z" }, + { url = "https://files.pythonhosted.org/packages/3e/c5/41bef391fb8b23dbc133c9f2ba016e7a7a8124513d2cc1b430f1897d87e4/duckdb-1.5.5-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:078e6a60dd8eedde5832f45422ca5c4a6b8c837aeabd8a56ca0b7d933f588053", size = 21511060, upload-time = "2026-07-22T10:54:42.788Z" }, + { url = "https://files.pythonhosted.org/packages/07/9f/c44dfc1f924ac29b3252dc1b91393c01d009dbfe9f8ed33f10b986151bd1/duckdb-1.5.5-cp313-cp313-win_amd64.whl", hash = "sha256:6826504277dba513c0c5d71d828456c94d729c9d2482f94b2e289f90a9167e28", size = 13168028, upload-time = "2026-07-22T10:54:46.127Z" }, + { url = "https://files.pythonhosted.org/packages/ca/88/591384b2cd59abddd6f5dc175e60374f9abae6064429f0c4402854c10f44/duckdb-1.5.5-cp313-cp313-win_arm64.whl", hash = "sha256:baa9c5702002fabb559ded2a39008f9f421fcbc7237d388b8213eff1e08858de", size = 13989955, upload-time = "2026-07-22T10:54:49.262Z" }, + { url = "https://files.pythonhosted.org/packages/3e/56/12c65bfa2d2605b81981b264788891bcf11ec72227889554cead5d8d13b9/duckdb-1.5.5-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8e6413dd40facb7b8ab21bd844450cd8f549b29e138635be9cf090ef4d2049e2", size = 32761946, upload-time = "2026-07-22T10:54:53.412Z" }, + { url = "https://files.pythonhosted.org/packages/b9/46/682ce155f17e0d2822d4f13ee3db9ca4b5b7c2da61b841b2629035e1f4bc/duckdb-1.5.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:64078acfd16541132ac6e191eb81b2845554444a0305cc1aa581ba107e514aa8", size = 17375069, upload-time = "2026-07-22T10:54:57.269Z" }, + { url = "https://files.pythonhosted.org/packages/39/ce/a24bcbd3289c8f305a430759c5fc12242740b4af3e17f7593f3a34e333d2/duckdb-1.5.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8c11775cc99a447618d5f1840126db17f2652f3eae05529df4f81f40e2df7151", size = 15519791, upload-time = "2026-07-22T10:55:00.681Z" }, + { url = "https://files.pythonhosted.org/packages/d9/76/3a01afbc615c1d418c0de58a6b68ac5ce2a8563232c0464bfbc2ce552398/duckdb-1.5.5-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77bbc1e6ba12e1e06f9020117bdf848627ecfdf36f907550e62e008e6109dece", size = 19398251, upload-time = "2026-07-22T10:55:04.168Z" }, + { url = "https://files.pythonhosted.org/packages/a1/43/3a5e81d1728f4d234c79bfe385808ee7c04834f7c37a4b5c257459c25614/duckdb-1.5.5-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fbf0f2d48b43c6c304d00463b463c27ead6c4b01c3c1816b750f728decf71afe", size = 21513851, upload-time = "2026-07-22T10:55:07.864Z" }, + { url = "https://files.pythonhosted.org/packages/91/41/fc7c829172c60ca22485251eab285f4f1a0d87b486a024c726f21471d86e/duckdb-1.5.5-cp314-cp314-win_amd64.whl", hash = "sha256:9dc826c4b50e64f6c4e4d07a3a9cb075ef70ba3899dc43ec5493dc3d7b04b353", size = 13691858, upload-time = "2026-07-22T10:55:11.181Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2c/95d9216b79e9273689d7ebce125a54503ed0c9bd7da931f0265888e99779/duckdb-1.5.5-cp314-cp314-win_arm64.whl", hash = "sha256:63e48d4b74b15aeacd688976432a7225163df8c226eddeb8536bba2d4d4ff433", size = 14470180, upload-time = "2026-07-22T10:55:14.445Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.0" @@ -1651,6 +1693,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "polars" +version = "1.43.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "polars-runtime-32" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/13/3873f213304bcbaaf39e63c8b905ceb460a0524448d57f86a829f6d4d0fd/polars-1.43.2.tar.gz", hash = "sha256:c699671b99eb71ff53334d237917aaa3db5ad4dda480abcb6c80e0eaee7b677b", size = 750312, upload-time = "2026-08-01T06:28:30.872Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/fe/0888040a24e4504098b85d8ad486b14cb01cf6b030bbe479dfc2dcffc2ac/polars-1.43.2-py3-none-any.whl", hash = "sha256:22aa0cb92a1ee2d60d6a15a638b2e8e0dd99aea21ac0cd8fb29da8e382e075a9", size = 847150, upload-time = "2026-08-01T06:27:15.543Z" }, +] + +[[package]] +name = "polars-runtime-32" +version = "1.43.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/06/11b578eeef05f867e3ee31b2a2fdd8e7684c2aa47822c49935d1be789c38/polars_runtime_32-1.43.2.tar.gz", hash = "sha256:d7b7c486bccee75a6af0158b87077da3d054657e3c60036b28644f4e1c7fdbf7", size = 3095669, upload-time = "2026-08-01T06:28:32.315Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/fc/12e6d4ca34d820297651134cfa35f86c33e898539fc6629cbb35d0089697/polars_runtime_32-1.43.2-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:91abf205d4ec93f92ba95386b7f8776559ae3dfce425ed2e527efa75d117d04a", size = 53088908, upload-time = "2026-08-01T06:27:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/32/81/833b0853551deb810854f96b43dea342b6e6c9b0ea1afcccf774157d519d/polars_runtime_32-1.43.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:2cc3ff96fd44789b02eb5c15b98dfcb000101636b177d3034fae2feec19b118f", size = 47540529, upload-time = "2026-08-01T06:27:21.391Z" }, + { url = "https://files.pythonhosted.org/packages/83/55/7b2a75af14c9294d97f3bec132dd3018ddcd988bef32b5d28322150b8c11/polars_runtime_32-1.43.2-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:10ed36e615ab362feb7406e6d084e124b445ad284caa73bd93ae7e65745ed894", size = 51366340, upload-time = "2026-08-01T06:27:24.776Z" }, + { url = "https://files.pythonhosted.org/packages/62/60/64deacb3abc70c52e2d88a808a052d1621c86a48fe9194f2c065579ab1cd/polars_runtime_32-1.43.2-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6d5a7ae004a2723ebf4427f6d6a639f30f86af4cf077075f6b35d04711154fc3", size = 57304599, upload-time = "2026-08-01T06:27:27.875Z" }, + { url = "https://files.pythonhosted.org/packages/52/95/d6e3a236d7630e17c40d0ddee839bf2be9acf548fdc0e5ad65ed9ff0cac6/polars_runtime_32-1.43.2-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:09339eacc6d392206e78aabbaaa37d7276eb969b798f46cb1f367fd718798c60", size = 51520580, upload-time = "2026-08-01T06:27:30.913Z" }, + { url = "https://files.pythonhosted.org/packages/b6/5a/2deb8eac70e9a2ac26d88a66ae7cf52612865026f4f4a5e7ab11ad9d52bf/polars_runtime_32-1.43.2-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:452b400e59e7f56e4c6437f435e796903272a9388feee12de2bea049ae87025e", size = 55204471, upload-time = "2026-08-01T06:27:33.886Z" }, + { url = "https://files.pythonhosted.org/packages/29/9e/647401ae8a607bc0cc40ed7b8592d5b1be90ded0dc9b9d6d3aeb03f9524b/polars_runtime_32-1.43.2-cp310-abi3-win_amd64.whl", hash = "sha256:00e33c28e321410c8d66e814a90043101e3bdd9ed2c6dabda07565aa8adbbdf1", size = 52572176, upload-time = "2026-08-01T06:27:37.048Z" }, + { url = "https://files.pythonhosted.org/packages/96/8d/60a50c3f36c85218a7ffcb48c6fe2ce1f7bec799152d68b8658ebed2179c/polars_runtime_32-1.43.2-cp310-abi3-win_arm64.whl", hash = "sha256:350a4868cae85bf8b3f81b33ba47927c15256bd9264dfc8c0753f1b927eac9d3", size = 46582513, upload-time = "2026-08-01T06:27:40.025Z" }, +] + [[package]] name = "pooch" version = "1.8.2" @@ -1929,6 +1999,123 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6f/2c/5b079febdc65e1c3fb2729bf958d18b45be7113828528e8a0b5850dd819a/pymdown_extensions-10.21-py3-none-any.whl", hash = "sha256:91b879f9f864d49794c2d9534372b10150e6141096c3908a455e45ca72ad9d3f", size = 268877, upload-time = "2026-02-15T20:44:05.464Z" }, ] +[[package]] +name = "pyproj" +version = "3.7.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "certifi", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/10/a8480ea27ea4bbe896c168808854d00f2a9b49f95c0319ddcbba693c8a90/pyproj-3.7.1.tar.gz", hash = "sha256:60d72facd7b6b79853f19744779abcd3f804c4e0d4fa8815469db20c9f640a47", size = 226339, upload-time = "2025-02-16T04:28:46.621Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/a3/c4cd4bba5b336075f145fe784fcaf4ef56ffbc979833303303e7a659dda2/pyproj-3.7.1-cp310-cp310-macosx_13_0_x86_64.whl", hash = "sha256:bf09dbeb333c34e9c546364e7df1ff40474f9fddf9e70657ecb0e4f670ff0b0e", size = 6262524, upload-time = "2025-02-16T04:27:19.725Z" }, + { url = "https://files.pythonhosted.org/packages/40/45/4fdf18f4cc1995f1992771d2a51cf186a9d7a8ec973c9693f8453850c707/pyproj-3.7.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:6575b2e53cc9e3e461ad6f0692a5564b96e7782c28631c7771c668770915e169", size = 4665102, upload-time = "2025-02-16T04:27:24.428Z" }, + { url = "https://files.pythonhosted.org/packages/0c/d2/360eb127380106cee83569954ae696b88a891c804d7a93abe3fbc15f5976/pyproj-3.7.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8cb516ee35ed57789b46b96080edf4e503fdb62dbb2e3c6581e0d6c83fca014b", size = 9432667, upload-time = "2025-02-16T04:27:27.04Z" }, + { url = "https://files.pythonhosted.org/packages/76/a5/c6e11b9a99ce146741fb4d184d5c468446c6d6015b183cae82ac822a6cfa/pyproj-3.7.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1e47c4e93b88d99dd118875ee3ca0171932444cdc0b52d493371b5d98d0f30ee", size = 9259185, upload-time = "2025-02-16T04:27:30.35Z" }, + { url = "https://files.pythonhosted.org/packages/41/56/a3c15c42145797a99363fa0fdb4e9805dccb8b4a76a6d7b2cdf36ebcc2a1/pyproj-3.7.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:3e8d276caeae34fcbe4813855d0d97b9b825bab8d7a8b86d859c24a6213a5a0d", size = 10469103, upload-time = "2025-02-16T04:27:33.542Z" }, + { url = "https://files.pythonhosted.org/packages/ef/73/c9194c2802fefe2a4fd4230bdd5ab083e7604e93c64d0356fa49c363bad6/pyproj-3.7.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f173f851ee75e54acdaa053382b6825b400cb2085663a9bb073728a59c60aebb", size = 10401391, upload-time = "2025-02-16T04:27:36.051Z" }, + { url = "https://files.pythonhosted.org/packages/c5/1d/ce8bb5b9251b04d7c22d63619bb3db3d2397f79000a9ae05b3fd86a5837e/pyproj-3.7.1-cp310-cp310-win32.whl", hash = "sha256:f550281ed6e5ea88fcf04a7c6154e246d5714be495c50c9e8e6b12d3fb63e158", size = 5869997, upload-time = "2025-02-16T04:27:38.302Z" }, + { url = "https://files.pythonhosted.org/packages/09/6a/ca145467fd2e5b21e3d5b8c2b9645dcfb3b68f08b62417699a1f5689008e/pyproj-3.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:3537668992a709a2e7f068069192138618c00d0ba113572fdd5ee5ffde8222f3", size = 6278581, upload-time = "2025-02-16T04:27:41.051Z" }, + { url = "https://files.pythonhosted.org/packages/ab/0d/63670fc527e664068b70b7cab599aa38b7420dd009bdc29ea257e7f3dfb3/pyproj-3.7.1-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:a94e26c1a4950cea40116775588a2ca7cf56f1f434ff54ee35a84718f3841a3d", size = 6264315, upload-time = "2025-02-16T04:27:44.539Z" }, + { url = "https://files.pythonhosted.org/packages/25/9d/cbaf82cfb290d1f1fa42feb9ba9464013bb3891e40c4199f8072112e4589/pyproj-3.7.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:263b54ba5004b6b957d55757d846fc5081bc02980caa0279c4fc95fa0fff6067", size = 4666267, upload-time = "2025-02-16T04:27:47.019Z" }, + { url = "https://files.pythonhosted.org/packages/79/53/24f9f9b8918c0550f3ff49ad5de4cf3f0688c9f91ff191476db8979146fe/pyproj-3.7.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f6d6a2ccd5607cd15ef990c51e6f2dd27ec0a741e72069c387088bba3aab60fa", size = 9680510, upload-time = "2025-02-16T04:27:49.239Z" }, + { url = "https://files.pythonhosted.org/packages/3c/ac/12fab74a908d40b63174dc704587febd0729414804bbfd873cabe504ff2d/pyproj-3.7.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8c5dcf24ede53d8abab7d8a77f69ff1936c6a8843ef4fcc574646e4be66e5739", size = 9493619, upload-time = "2025-02-16T04:27:52.65Z" }, + { url = "https://files.pythonhosted.org/packages/c4/45/26311d6437135da2153a178125db5dfb6abce831ce04d10ec207eabac70a/pyproj-3.7.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3c2e7449840a44ce860d8bea2c6c1c4bc63fa07cba801dcce581d14dcb031a02", size = 10709755, upload-time = "2025-02-16T04:27:55.239Z" }, + { url = "https://files.pythonhosted.org/packages/99/52/4ecd0986f27d0e6c8ee3a7bc5c63da15acd30ac23034f871325b297e61fd/pyproj-3.7.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0829865c1d3a3543f918b3919dc601eea572d6091c0dd175e1a054db9c109274", size = 10642970, upload-time = "2025-02-16T04:27:58.343Z" }, + { url = "https://files.pythonhosted.org/packages/3f/a5/d3bfc018fc92195a000d1d28acc1f3f1df15ff9f09ece68f45a2636c0134/pyproj-3.7.1-cp311-cp311-win32.whl", hash = "sha256:6181960b4b812e82e588407fe5c9c68ada267c3b084db078f248db5d7f45d18a", size = 5868295, upload-time = "2025-02-16T04:28:01.712Z" }, + { url = "https://files.pythonhosted.org/packages/92/39/ef6f06a5b223dbea308cfcbb7a0f72e7b506aef1850e061b2c73b0818715/pyproj-3.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:5ad0ff443a785d84e2b380869fdd82e6bfc11eba6057d25b4409a9bbfa867970", size = 6279871, upload-time = "2025-02-16T04:28:04.988Z" }, + { url = "https://files.pythonhosted.org/packages/e6/c9/876d4345b8d17f37ac59ebd39f8fa52fc6a6a9891a420f72d050edb6b899/pyproj-3.7.1-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:2781029d90df7f8d431e29562a3f2d8eafdf233c4010d6fc0381858dc7373217", size = 6264087, upload-time = "2025-02-16T04:28:09.036Z" }, + { url = "https://files.pythonhosted.org/packages/ff/e6/5f8691f8c90e7f402cc80a6276eb19d2ec1faa150d5ae2dd9c7b0a254da8/pyproj-3.7.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:d61bf8ab04c73c1da08eedaf21a103b72fa5b0a9b854762905f65ff8b375d394", size = 4669628, upload-time = "2025-02-16T04:28:10.944Z" }, + { url = "https://files.pythonhosted.org/packages/42/ec/16475bbb79c1c68845c0a0d9c60c4fb31e61b8a2a20bc18b1a81e81c7f68/pyproj-3.7.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:04abc517a8555d1b05fcee768db3280143fe42ec39fdd926a2feef31631a1f2f", size = 9721415, upload-time = "2025-02-16T04:28:13.342Z" }, + { url = "https://files.pythonhosted.org/packages/b3/a3/448f05b15e318bd6bea9a32cfaf11e886c4ae61fa3eee6e09ed5c3b74bb2/pyproj-3.7.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:084c0a475688f934d386c2ab3b6ce03398a473cd48adfda70d9ab8f87f2394a0", size = 9556447, upload-time = "2025-02-16T04:28:15.818Z" }, + { url = "https://files.pythonhosted.org/packages/6a/ae/bd15fe8d8bd914ead6d60bca7f895a4e6f8ef7e3928295134ff9a7dad14c/pyproj-3.7.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a20727a23b1e49c7dc7fe3c3df8e56a8a7acdade80ac2f5cca29d7ca5564c145", size = 10758317, upload-time = "2025-02-16T04:28:18.338Z" }, + { url = "https://files.pythonhosted.org/packages/9d/d9/5ccefb8bca925f44256b188a91c31238cae29ab6ee7f53661ecc04616146/pyproj-3.7.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bf84d766646f1ebd706d883755df4370aaf02b48187cedaa7e4239f16bc8213d", size = 10771259, upload-time = "2025-02-16T04:28:20.822Z" }, + { url = "https://files.pythonhosted.org/packages/2a/7d/31dedff9c35fa703162f922eeb0baa6c44a3288469a5fd88d209e2892f9e/pyproj-3.7.1-cp312-cp312-win32.whl", hash = "sha256:5f0da2711364d7cb9f115b52289d4a9b61e8bca0da57f44a3a9d6fc9bdeb7274", size = 5859914, upload-time = "2025-02-16T04:28:23.303Z" }, + { url = "https://files.pythonhosted.org/packages/3e/47/c6ab03d6564a7c937590cff81a2742b5990f096cce7c1a622d325be340ee/pyproj-3.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:aee664a9d806612af30a19dba49e55a7a78ebfec3e9d198f6a6176e1d140ec98", size = 6273196, upload-time = "2025-02-16T04:28:25.227Z" }, + { url = "https://files.pythonhosted.org/packages/ef/01/984828464c9960036c602753fc0f21f24f0aa9043c18fa3f2f2b66a86340/pyproj-3.7.1-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:5f8d02ef4431dee414d1753d13fa82a21a2f61494737b5f642ea668d76164d6d", size = 6253062, upload-time = "2025-02-16T04:28:27.861Z" }, + { url = "https://files.pythonhosted.org/packages/68/65/6ecdcdc829811a2c160cdfe2f068a009fc572fd4349664f758ccb0853a7c/pyproj-3.7.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:0b853ae99bda66cbe24b4ccfe26d70601d84375940a47f553413d9df570065e0", size = 4660548, upload-time = "2025-02-16T04:28:29.526Z" }, + { url = "https://files.pythonhosted.org/packages/67/da/dda94c4490803679230ba4c17a12f151b307a0d58e8110820405ca2d98db/pyproj-3.7.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:83db380c52087f9e9bdd8a527943b2e7324f275881125e39475c4f9277bdeec4", size = 9662464, upload-time = "2025-02-16T04:28:31.437Z" }, + { url = "https://files.pythonhosted.org/packages/6f/57/f61b7d22c91ae1d12ee00ac4c0038714e774ebcd851b9133e5f4f930dd40/pyproj-3.7.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b35ed213892e211a3ce2bea002aa1183e1a2a9b79e51bb3c6b15549a831ae528", size = 9497461, upload-time = "2025-02-16T04:28:33.848Z" }, + { url = "https://files.pythonhosted.org/packages/b7/f6/932128236f79d2ac7d39fe1a19667fdf7155d9a81d31fb9472a7a497790f/pyproj-3.7.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a8b15b0463d1303bab113d1a6af2860a0d79013c3a66fcc5475ce26ef717fd4f", size = 10708869, upload-time = "2025-02-16T04:28:37.34Z" }, + { url = "https://files.pythonhosted.org/packages/1d/0d/07ac7712994454a254c383c0d08aff9916a2851e6512d59da8dc369b1b02/pyproj-3.7.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:87229e42b75e89f4dad6459200f92988c5998dfb093c7c631fb48524c86cd5dc", size = 10729260, upload-time = "2025-02-16T04:28:40.639Z" }, + { url = "https://files.pythonhosted.org/packages/b0/d0/9c604bc72c37ba69b867b6df724d6a5af6789e8c375022c952f65b2af558/pyproj-3.7.1-cp313-cp313-win32.whl", hash = "sha256:d666c3a3faaf3b1d7fc4a544059c4eab9d06f84a604b070b7aa2f318e227798e", size = 5855462, upload-time = "2025-02-16T04:28:42.827Z" }, + { url = "https://files.pythonhosted.org/packages/98/df/68a2b7f5fb6400c64aad82d72bcc4bc531775e62eedff993a77c780defd0/pyproj-3.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:d3caac7473be22b6d6e102dde6c46de73b96bc98334e577dfaee9886f102ea2e", size = 6266573, upload-time = "2025-02-16T04:28:44.727Z" }, +] + +[[package]] +name = "pyproj" +version = "3.7.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +dependencies = [ + { name = "certifi", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/04/90/67bd7260b4ea9b8b20b4f58afef6c223ecb3abf368eb4ec5bc2cdef81b49/pyproj-3.7.2.tar.gz", hash = "sha256:39a0cf1ecc7e282d1d30f36594ebd55c9fae1fda8a2622cee5d100430628f88c", size = 226279, upload-time = "2025-08-14T12:05:42.18Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/bd/f205552cd1713b08f93b09e39a3ec99edef0b3ebbbca67b486fdf1abe2de/pyproj-3.7.2-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:2514d61f24c4e0bb9913e2c51487ecdaeca5f8748d8313c933693416ca41d4d5", size = 6227022, upload-time = "2025-08-14T12:03:51.474Z" }, + { url = "https://files.pythonhosted.org/packages/75/4c/9a937e659b8b418ab573c6d340d27e68716928953273e0837e7922fcac34/pyproj-3.7.2-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:8693ca3892d82e70de077701ee76dd13d7bca4ae1c9d1e739d72004df015923a", size = 4625810, upload-time = "2025-08-14T12:03:53.808Z" }, + { url = "https://files.pythonhosted.org/packages/c0/7d/a9f41e814dc4d1dc54e95b2ccaf0b3ebe3eb18b1740df05fe334724c3d89/pyproj-3.7.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:5e26484d80fea56273ed1555abaea161e9661d81a6c07815d54b8e883d4ceb25", size = 9638694, upload-time = "2025-08-14T12:03:55.669Z" }, + { url = "https://files.pythonhosted.org/packages/ad/ab/9bdb4a6216b712a1f9aab1c0fcbee5d3726f34a366f29c3e8c08a78d6b70/pyproj-3.7.2-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:281cb92847814e8018010c48b4069ff858a30236638631c1a91dd7bfa68f8a8a", size = 9493977, upload-time = "2025-08-14T12:03:57.937Z" }, + { url = "https://files.pythonhosted.org/packages/c9/db/2db75b1b6190f1137b1c4e8ef6a22e1c338e46320f6329bfac819143e063/pyproj-3.7.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9c8577f0b7bb09118ec2e57e3babdc977127dd66326d6c5d755c76b063e6d9dc", size = 10841151, upload-time = "2025-08-14T12:04:00.271Z" }, + { url = "https://files.pythonhosted.org/packages/89/f7/989643394ba23a286e9b7b3f09981496172f9e0d4512457ffea7dc47ffc7/pyproj-3.7.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a23f59904fac3a5e7364b3aa44d288234af267ca041adb2c2b14a903cd5d3ac5", size = 10751585, upload-time = "2025-08-14T12:04:02.228Z" }, + { url = "https://files.pythonhosted.org/packages/53/6d/ad928fe975a6c14a093c92e6a319ca18f479f3336bb353a740bdba335681/pyproj-3.7.2-cp311-cp311-win32.whl", hash = "sha256:f2af4ed34b2cf3e031a2d85b067a3ecbd38df073c567e04b52fa7a0202afde8a", size = 5908533, upload-time = "2025-08-14T12:04:04.821Z" }, + { url = "https://files.pythonhosted.org/packages/79/e0/b95584605cec9ed50b7ebaf7975d1c4ddeec5a86b7a20554ed8b60042bd7/pyproj-3.7.2-cp311-cp311-win_amd64.whl", hash = "sha256:0b7cb633565129677b2a183c4d807c727d1c736fcb0568a12299383056e67433", size = 6320742, upload-time = "2025-08-14T12:04:06.357Z" }, + { url = "https://files.pythonhosted.org/packages/b7/4d/536e8f93bca808175c2d0a5ac9fdf69b960d8ab6b14f25030dccb07464d7/pyproj-3.7.2-cp311-cp311-win_arm64.whl", hash = "sha256:38b08d85e3a38e455625b80e9eb9f78027c8e2649a21dec4df1f9c3525460c71", size = 6245772, upload-time = "2025-08-14T12:04:08.365Z" }, + { url = "https://files.pythonhosted.org/packages/8d/ab/9893ea9fb066be70ed9074ae543914a618c131ed8dff2da1e08b3a4df4db/pyproj-3.7.2-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:0a9bb26a6356fb5b033433a6d1b4542158fb71e3c51de49b4c318a1dff3aeaab", size = 6219832, upload-time = "2025-08-14T12:04:10.264Z" }, + { url = "https://files.pythonhosted.org/packages/53/78/4c64199146eed7184eb0e85bedec60a4aa8853b6ffe1ab1f3a8b962e70a0/pyproj-3.7.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:567caa03021178861fad27fabde87500ec6d2ee173dd32f3e2d9871e40eebd68", size = 4620650, upload-time = "2025-08-14T12:04:11.978Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ac/14a78d17943898a93ef4f8c6a9d4169911c994e3161e54a7cedeba9d8dde/pyproj-3.7.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:c203101d1dc3c038a56cff0447acc515dd29d6e14811406ac539c21eed422b2a", size = 9667087, upload-time = "2025-08-14T12:04:13.964Z" }, + { url = "https://files.pythonhosted.org/packages/b8/be/212882c450bba74fc8d7d35cbd57e4af84792f0a56194819d98106b075af/pyproj-3.7.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:1edc34266c0c23ced85f95a1ee8b47c9035eae6aca5b6b340327250e8e281630", size = 9552797, upload-time = "2025-08-14T12:04:16.624Z" }, + { url = "https://files.pythonhosted.org/packages/ba/c0/c0f25c87b5d2a8686341c53c1792a222a480d6c9caf60311fec12c99ec26/pyproj-3.7.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:aa9f26c21bc0e2dc3d224cb1eb4020cf23e76af179a7c66fea49b828611e4260", size = 10837036, upload-time = "2025-08-14T12:04:18.733Z" }, + { url = "https://files.pythonhosted.org/packages/5d/37/5cbd6772addde2090c91113332623a86e8c7d583eccb2ad02ea634c4a89f/pyproj-3.7.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f9428b318530625cb389b9ddc9c51251e172808a4af79b82809376daaeabe5e9", size = 10775952, upload-time = "2025-08-14T12:04:20.709Z" }, + { url = "https://files.pythonhosted.org/packages/69/a1/dc250e3cf83eb4b3b9a2cf86fdb5e25288bd40037ae449695550f9e96b2f/pyproj-3.7.2-cp312-cp312-win32.whl", hash = "sha256:b3d99ed57d319da042f175f4554fc7038aa4bcecc4ac89e217e350346b742c9d", size = 5898872, upload-time = "2025-08-14T12:04:22.485Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a6/6fe724b72b70f2b00152d77282e14964d60ab092ec225e67c196c9b463e5/pyproj-3.7.2-cp312-cp312-win_amd64.whl", hash = "sha256:11614a054cd86a2ed968a657d00987a86eeb91fdcbd9ad3310478685dc14a128", size = 6312176, upload-time = "2025-08-14T12:04:24.736Z" }, + { url = "https://files.pythonhosted.org/packages/5d/68/915cc32c02a91e76d02c8f55d5a138d6ef9e47a0d96d259df98f4842e558/pyproj-3.7.2-cp312-cp312-win_arm64.whl", hash = "sha256:509a146d1398bafe4f53273398c3bb0b4732535065fa995270e52a9d3676bca3", size = 6233452, upload-time = "2025-08-14T12:04:27.287Z" }, + { url = "https://files.pythonhosted.org/packages/be/14/faf1b90d267cea68d7e70662e7f88cefdb1bc890bd596c74b959e0517a72/pyproj-3.7.2-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:19466e529b1b15eeefdf8ff26b06fa745856c044f2f77bf0edbae94078c1dfa1", size = 6214580, upload-time = "2025-08-14T12:04:28.804Z" }, + { url = "https://files.pythonhosted.org/packages/35/48/da9a45b184d375f62667f62eba0ca68569b0bd980a0bb7ffcc1d50440520/pyproj-3.7.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c79b9b84c4a626c5dc324c0d666be0bfcebd99f7538d66e8898c2444221b3da7", size = 4615388, upload-time = "2025-08-14T12:04:30.553Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e7/d2b459a4a64bca328b712c1b544e109df88e5c800f7c143cfbc404d39bfb/pyproj-3.7.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:ceecf374cacca317bc09e165db38ac548ee3cad07c3609442bd70311c59c21aa", size = 9628455, upload-time = "2025-08-14T12:04:32.435Z" }, + { url = "https://files.pythonhosted.org/packages/f8/85/c2b1706e51942de19076eff082f8495e57d5151364e78b5bef4af4a1d94a/pyproj-3.7.2-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:5141a538ffdbe4bfd157421828bb2e07123a90a7a2d6f30fa1462abcfb5ce681", size = 9514269, upload-time = "2025-08-14T12:04:34.599Z" }, + { url = "https://files.pythonhosted.org/packages/34/38/07a9b89ae7467872f9a476883a5bad9e4f4d1219d31060f0f2b282276cbe/pyproj-3.7.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f000841e98ea99acbb7b8ca168d67773b0191de95187228a16110245c5d954d5", size = 10808437, upload-time = "2025-08-14T12:04:36.485Z" }, + { url = "https://files.pythonhosted.org/packages/12/56/fda1daeabbd39dec5b07f67233d09f31facb762587b498e6fc4572be9837/pyproj-3.7.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8115faf2597f281a42ab608ceac346b4eb1383d3b45ab474fd37341c4bf82a67", size = 10745540, upload-time = "2025-08-14T12:04:38.568Z" }, + { url = "https://files.pythonhosted.org/packages/0d/90/c793182cbba65a39a11db2ac6b479fe76c59e6509ae75e5744c344a0da9d/pyproj-3.7.2-cp313-cp313-win32.whl", hash = "sha256:f18c0579dd6be00b970cb1a6719197fceecc407515bab37da0066f0184aafdf3", size = 5896506, upload-time = "2025-08-14T12:04:41.059Z" }, + { url = "https://files.pythonhosted.org/packages/be/0f/747974129cf0d800906f81cd25efd098c96509026e454d4b66868779ab04/pyproj-3.7.2-cp313-cp313-win_amd64.whl", hash = "sha256:bb41c29d5f60854b1075853fe80c58950b398d4ebb404eb532536ac8d2834ed7", size = 6310195, upload-time = "2025-08-14T12:04:42.974Z" }, + { url = "https://files.pythonhosted.org/packages/82/64/fc7598a53172c4931ec6edf5228280663063150625d3f6423b4c20f9daff/pyproj-3.7.2-cp313-cp313-win_arm64.whl", hash = "sha256:2b617d573be4118c11cd96b8891a0b7f65778fa7733ed8ecdb297a447d439100", size = 6230748, upload-time = "2025-08-14T12:04:44.491Z" }, + { url = "https://files.pythonhosted.org/packages/aa/f0/611dd5cddb0d277f94b7af12981f56e1441bf8d22695065d4f0df5218498/pyproj-3.7.2-cp313-cp313t-macosx_13_0_x86_64.whl", hash = "sha256:d27b48f0e81beeaa2b4d60c516c3a1cfbb0c7ff6ef71256d8e9c07792f735279", size = 6241729, upload-time = "2025-08-14T12:04:46.274Z" }, + { url = "https://files.pythonhosted.org/packages/15/93/40bd4a6c523ff9965e480870611aed7eda5aa2c6128c6537345a2b77b542/pyproj-3.7.2-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:55a3610d75023c7b1c6e583e48ef8f62918e85a2ae81300569d9f104d6684bb6", size = 4652497, upload-time = "2025-08-14T12:04:48.203Z" }, + { url = "https://files.pythonhosted.org/packages/1b/ae/7150ead53c117880b35e0d37960d3138fe640a235feb9605cb9386f50bb0/pyproj-3.7.2-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:8d7349182fa622696787cc9e195508d2a41a64765da9b8a6bee846702b9e6220", size = 9942610, upload-time = "2025-08-14T12:04:49.652Z" }, + { url = "https://files.pythonhosted.org/packages/d8/17/7a4a7eafecf2b46ab64e5c08176c20ceb5844b503eaa551bf12ccac77322/pyproj-3.7.2-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:d230b186eb876ed4f29a7c5ee310144c3a0e44e89e55f65fb3607e13f6db337c", size = 9692390, upload-time = "2025-08-14T12:04:51.731Z" }, + { url = "https://files.pythonhosted.org/packages/c3/55/ae18f040f6410f0ea547a21ada7ef3e26e6c82befa125b303b02759c0e9d/pyproj-3.7.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:237499c7862c578d0369e2b8ac56eec550e391a025ff70e2af8417139dabb41c", size = 11047596, upload-time = "2025-08-14T12:04:53.748Z" }, + { url = "https://files.pythonhosted.org/packages/e6/2e/d3fff4d2909473f26ae799f9dda04caa322c417a51ff3b25763f7d03b233/pyproj-3.7.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8c225f5978abd506fd9a78eaaf794435e823c9156091cabaab5374efb29d7f69", size = 10896975, upload-time = "2025-08-14T12:04:55.875Z" }, + { url = "https://files.pythonhosted.org/packages/f2/bc/8fc7d3963d87057b7b51ebe68c1e7c51c23129eee5072ba6b86558544a46/pyproj-3.7.2-cp313-cp313t-win32.whl", hash = "sha256:2da731876d27639ff9d2d81c151f6ab90a1546455fabd93368e753047be344a2", size = 5953057, upload-time = "2025-08-14T12:04:58.466Z" }, + { url = "https://files.pythonhosted.org/packages/cc/27/ea9809966cc47d2d51e6d5ae631ea895f7c7c7b9b3c29718f900a8f7d197/pyproj-3.7.2-cp313-cp313t-win_amd64.whl", hash = "sha256:f54d91ae18dd23b6c0ab48126d446820e725419da10617d86a1b69ada6d881d3", size = 6375414, upload-time = "2025-08-14T12:04:59.861Z" }, + { url = "https://files.pythonhosted.org/packages/5b/f8/1ef0129fba9a555c658e22af68989f35e7ba7b9136f25758809efec0cd6e/pyproj-3.7.2-cp313-cp313t-win_arm64.whl", hash = "sha256:fc52ba896cfc3214dc9f9ca3c0677a623e8fdd096b257c14a31e719d21ff3fdd", size = 6262501, upload-time = "2025-08-14T12:05:01.39Z" }, + { url = "https://files.pythonhosted.org/packages/42/17/c2b050d3f5b71b6edd0d96ae16c990fdc42a5f1366464a5c2772146de33a/pyproj-3.7.2-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:2aaa328605ace41db050d06bac1adc11f01b71fe95c18661497763116c3a0f02", size = 6214541, upload-time = "2025-08-14T12:05:03.166Z" }, + { url = "https://files.pythonhosted.org/packages/03/68/68ada9c8aea96ded09a66cfd9bf87aa6db8c2edebe93f5bf9b66b0143fbc/pyproj-3.7.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:35dccbce8201313c596a970fde90e33605248b66272595c061b511c8100ccc08", size = 4617456, upload-time = "2025-08-14T12:05:04.563Z" }, + { url = "https://files.pythonhosted.org/packages/81/e4/4c50ceca7d0e937977866b02cb64e6ccf4df979a5871e521f9e255df6073/pyproj-3.7.2-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:25b0b7cb0042444c29a164b993c45c1b8013d6c48baa61dc1160d834a277e83b", size = 9615590, upload-time = "2025-08-14T12:05:06.094Z" }, + { url = "https://files.pythonhosted.org/packages/05/1e/ada6fb15a1d75b5bd9b554355a69a798c55a7dcc93b8d41596265c1772e3/pyproj-3.7.2-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:85def3a6388e9ba51f964619aa002a9d2098e77c6454ff47773bb68871024281", size = 9474960, upload-time = "2025-08-14T12:05:07.973Z" }, + { url = "https://files.pythonhosted.org/packages/51/07/9d48ad0a8db36e16f842f2c8a694c1d9d7dcf9137264846bef77585a71f3/pyproj-3.7.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b1bccefec3875ab81eabf49059e2b2ea77362c178b66fd3528c3e4df242f1516", size = 10799478, upload-time = "2025-08-14T12:05:14.102Z" }, + { url = "https://files.pythonhosted.org/packages/85/cf/2f812b529079f72f51ff2d6456b7fef06c01735e5cfd62d54ffb2b548028/pyproj-3.7.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d5371ca114d6990b675247355a801925814eca53e6c4b2f1b5c0a956336ee36e", size = 10710030, upload-time = "2025-08-14T12:05:16.317Z" }, + { url = "https://files.pythonhosted.org/packages/99/9b/4626a19e1f03eba4c0e77b91a6cf0f73aa9cb5d51a22ee385c22812bcc2c/pyproj-3.7.2-cp314-cp314-win32.whl", hash = "sha256:77f066626030f41be543274f5ac79f2a511fe89860ecd0914f22131b40a0ec25", size = 5991181, upload-time = "2025-08-14T12:05:19.492Z" }, + { url = "https://files.pythonhosted.org/packages/04/b2/5a6610554306a83a563080c2cf2c57565563eadd280e15388efa00fb5b33/pyproj-3.7.2-cp314-cp314-win_amd64.whl", hash = "sha256:5a964da1696b8522806f4276ab04ccfff8f9eb95133a92a25900697609d40112", size = 6434721, upload-time = "2025-08-14T12:05:21.022Z" }, + { url = "https://files.pythonhosted.org/packages/ae/ce/6c910ea2e1c74ef673c5d48c482564b8a7824a44c4e35cca2e765b68cfcc/pyproj-3.7.2-cp314-cp314-win_arm64.whl", hash = "sha256:e258ab4dbd3cf627809067c0ba8f9884ea76c8e5999d039fb37a1619c6c3e1f6", size = 6363821, upload-time = "2025-08-14T12:05:22.627Z" }, + { url = "https://files.pythonhosted.org/packages/e4/e4/5532f6f7491812ba782a2177fe9de73fd8e2912b59f46a1d056b84b9b8f2/pyproj-3.7.2-cp314-cp314t-macosx_13_0_x86_64.whl", hash = "sha256:bbbac2f930c6d266f70ec75df35ef851d96fdb3701c674f42fd23a9314573b37", size = 6241773, upload-time = "2025-08-14T12:05:24.577Z" }, + { url = "https://files.pythonhosted.org/packages/20/1f/0938c3f2bbbef1789132d1726d9b0e662f10cfc22522743937f421ad664e/pyproj-3.7.2-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:b7544e0a3d6339dc9151e9c8f3ea62a936ab7cc446a806ec448bbe86aebb979b", size = 4652537, upload-time = "2025-08-14T12:05:26.391Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a8/488b1ed47d25972f33874f91f09ca8f2227902f05f63a2b80dc73e7b1c97/pyproj-3.7.2-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:f7f5133dca4c703e8acadf6f30bc567d39a42c6af321e7f81975c2518f3ed357", size = 9940864, upload-time = "2025-08-14T12:05:27.985Z" }, + { url = "https://files.pythonhosted.org/packages/c7/cc/7f4c895d0cb98e47b6a85a6d79eaca03eb266129eed2f845125c09cf31ff/pyproj-3.7.2-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:5aff3343038d7426aa5076f07feb88065f50e0502d1b0d7c22ddfdd2c75a3f81", size = 9688868, upload-time = "2025-08-14T12:05:30.425Z" }, + { url = "https://files.pythonhosted.org/packages/b2/b7/c7e306b8bb0f071d9825b753ee4920f066c40fbfcce9372c4f3cfb2fc4ed/pyproj-3.7.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:b0552178c61f2ac1c820d087e8ba6e62b29442debddbb09d51c4bf8acc84d888", size = 11045910, upload-time = "2025-08-14T12:05:32.507Z" }, + { url = "https://files.pythonhosted.org/packages/42/fb/538a4d2df695980e2dde5c04d965fbdd1fe8c20a3194dc4aaa3952a4d1be/pyproj-3.7.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:47d87db2d2c436c5fd0409b34d70bb6cdb875cca2ebe7a9d1c442367b0ab8d59", size = 10895724, upload-time = "2025-08-14T12:05:35.465Z" }, + { url = "https://files.pythonhosted.org/packages/e8/8b/a3f0618b03957de9db5489a04558a8826f43906628bb0b766033aa3b5548/pyproj-3.7.2-cp314-cp314t-win32.whl", hash = "sha256:c9b6f1d8ad3e80a0ee0903a778b6ece7dca1d1d40f6d114ae01bc8ddbad971aa", size = 6056848, upload-time = "2025-08-14T12:05:37.553Z" }, + { url = "https://files.pythonhosted.org/packages/bc/56/413240dd5149dd3291eda55aa55a659da4431244a2fd1319d0ae89407cfb/pyproj-3.7.2-cp314-cp314t-win_amd64.whl", hash = "sha256:1914e29e27933ba6f9822663ee0600f169014a2859f851c054c88cf5ea8a333c", size = 6517676, upload-time = "2025-08-14T12:05:39.126Z" }, + { url = "https://files.pythonhosted.org/packages/15/73/a7141a1a0559bf1a7aa42a11c879ceb19f02f5c6c371c6d57fd86cefd4d1/pyproj-3.7.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d9d25bae416a24397e0d85739f84d323b55f6511e45a522dd7d7eae70d10c7e4", size = 6391844, upload-time = "2025-08-14T12:05:40.745Z" }, +] + [[package]] name = "pytest" version = "8.4.1" @@ -2581,9 +2768,23 @@ docs = [ { name = "mkdocstrings", extra = ["python"] }, { name = "zensical" }, ] +duckdb = [ + { name = "duckdb" }, +] +geo = [ + { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +polars = [ + { name = "polars" }, +] test = [ { name = "cftime" }, + { name = "duckdb" }, { name = "gcsfs" }, + { name = "polars" }, + { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pytest" }, { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, extra = ["io"], marker = "python_full_version < '3.11'" }, { name = "xarray", version = "2025.7.0", source = { registry = "https://pypi.org/simple" }, extra = ["io"], marker = "python_full_version >= '3.11'" }, @@ -2602,18 +2803,22 @@ requires-dist = [ { name = "cftime", marker = "extra == 'test'" }, { name = "dask", specifier = ">=2024.8.0" }, { name = "datafusion", 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 = "pyproj", marker = "extra == 'geo'" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "pytest", marker = "extra == 'test'" }, { name = "watchfiles", marker = "extra == 'dev'" }, { name = "xarray", specifier = ">=2024.7.0" }, { name = "xarray", extras = ["io"], marker = "extra == 'test'" }, { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" }, + { name = "xarray-sql", extras = ["duckdb", "polars", "geo"], marker = "extra == 'test'" }, { name = "zensical", marker = "extra == 'docs'" }, ] -provides-extras = ["dev", "docs", "test"] +provides-extras = ["dev", "docs", "duckdb", "geo", "polars", "test"] [package.metadata.requires-dev] dev = [ diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py index d1e5984..0f98e1e 100644 --- a/xarray_sql/__init__.py +++ b/xarray_sql/__init__.py @@ -1,6 +1,9 @@ from . import cftime +from .backends import arrow_dataset, register +from .geometry import bbox_conjuncts from .df import from_map from .reader import read_xarray, read_xarray_table +from .roundtrip import to_dataset from .sql import XarrayContext __all__ = [ @@ -8,5 +11,9 @@ "XarrayContext", "read_xarray_table", "read_xarray", + "arrow_dataset", + "bbox_conjuncts", + "register", + "to_dataset", "from_map", # deprecated ] diff --git a/xarray_sql/backends/__init__.py b/xarray_sql/backends/__init__.py new file mode 100644 index 0000000..82aad8e --- /dev/null +++ b/xarray_sql/backends/__init__.py @@ -0,0 +1,34 @@ +"""Engine adapters — the *register* seam of xarray-sql. + +xarray-sql translates data, not queries, across two seams — the two +boundaries between xarray and a query engine that neither side builds +for itself: *register* (a lazy ``xarray.Dataset`` becomes a table on +the engine's own connection; this package) and *round-trip* (an Arrow +result becomes a labeled Dataset again; [xarray_sql.to_dataset][]). +SQL dialects, geometry, H3, and optimizers belong to each engine and +its extension ecosystem. + +Adapters register themselves on import via +[register_adapter][xarray_sql.backends.base.register_adapter]; +[register][xarray_sql.backends.base.register] dispatches on the connection +type. +""" + +from .base import EngineAdapter, get_adapter, register, register_adapter +from . import datafusion as _datafusion # noqa: F401 (self-registers) +from . import duckdb as _duckdb # noqa: F401 (self-registers) +from .pyarrow import ( + XarrayArrowStream, + XarrayPushdownDataset, + arrow_dataset, +) + +__all__ = [ + "EngineAdapter", + "XarrayArrowStream", + "XarrayPushdownDataset", + "arrow_dataset", + "get_adapter", + "register", + "register_adapter", +] diff --git a/xarray_sql/backends/base.py b/xarray_sql/backends/base.py new file mode 100644 index 0000000..59eb488 --- /dev/null +++ b/xarray_sql/backends/base.py @@ -0,0 +1,118 @@ +"""Engine-adapter dispatch for [xarray_sql.register][]. + +An *engine adapter* implements the register seam: given an engine's native +connection object and a lazy ``xarray.Dataset``, register the Dataset as +a queryable table on that connection. The Arrow C-stream protocol is the +common wire between xarray and every engine; adapters differ only in how +a stream is attached to the connection and in what pushdown the engine +can do against it. + +Adapters self-describe which connections they accept via ``matches``, +which must not require the engine's package to be importable (detection +is by type inspection), so optional engines stay optional. +""" + +from __future__ import annotations + +from typing import Any, Protocol, TypeGuard, TypeVar, cast + +import xarray as xr + +from ..df import Chunks + +ConT = TypeVar("ConT") +"""An engine's native connection type (e.g. ``duckdb.DuckDBPyConnection``).""" + + +class EngineAdapter(Protocol[ConT]): + """One engine's implementation of the register seam.""" + + @staticmethod + def matches(con: object) -> TypeGuard[ConT]: + """Whether *con* is a connection this adapter can register into.""" + ... + + @staticmethod + def register( + con: ConT, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> ConT: + """Register *ds* as table *name* on *con*; returns *con*.""" + ... + + +_ADAPTERS: list[type[EngineAdapter[Any]]] = [] + +_A = TypeVar("_A", bound=type[EngineAdapter[Any]]) + + +def register_adapter(cls: _A) -> _A: + """Class decorator adding an adapter to the dispatch list.""" + _ADAPTERS.append(cls) + return cls + + +def get_adapter(con: object) -> type[EngineAdapter[Any]]: + """Return the first adapter whose ``matches(con)`` is true.""" + for adapter in _ADAPTERS: + if adapter.matches(con): + return adapter + raise TypeError( + f"No xarray-sql engine adapter for connection of type " + f"{type(con).__module__}.{type(con).__qualname__}. " + f"Supported: DataFusion SessionContext and DuckDB connections." + ) + + +def register( + con: ConT, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, +) -> ConT: + """Register a lazy xarray Dataset as a table on an engine connection. + + The engine is inferred from the connection type. Data is not read at + registration time; the engine pulls Arrow record batches lazily during + query execution. Write your SQL in the engine's own dialect and use + the engine's extension ecosystem directly — xarray-sql translates the + data, not the queries. + + Example (DuckDB):: + + import duckdb + import xarray_sql as xql + + con = duckdb.connect() + xql.register(con, "era5", ds) + rel = con.sql("SELECT time, AVG(t2m) AS t2m FROM era5 GROUP BY time") + result = xql.to_dataset(rel, template=ds) + + Args: + con: An engine connection: a ``datafusion.SessionContext`` (or + [xarray_sql.XarrayContext][]) or a + ``duckdb.DuckDBPyConnection``. + name: The table name to register the Dataset under. Datasets + whose variables have differing dimensions are split into one + table per dimension group (a SQL schema ``name.group`` on + DataFusion; ``name_group`` tables on DuckDB). + ds: An xarray Dataset. + chunks: Xarray-like chunks specification controlling partition + granularity. Defaults to the Dataset's existing chunks. + **kwargs: Adapter-specific options, forwarded as-is — e.g. + ``table_names`` on DataFusion, ``batch_size`` / ``prefetch`` + on DuckDB. + + Returns: + The connection, to allow chaining. + """ + # The connection type is erased by the runtime dispatch; every adapter + # returns the connection it was given. + adapter: Any = get_adapter(con) + return cast(ConT, adapter.register(con, name, ds, chunks=chunks, **kwargs)) diff --git a/xarray_sql/backends/datafusion.py b/xarray_sql/backends/datafusion.py new file mode 100644 index 0000000..3d7c29e --- /dev/null +++ b/xarray_sql/backends/datafusion.py @@ -0,0 +1,46 @@ +"""DataFusion engine adapter. + +DataFusion is xarray-sql's default engine and the richest adapter: the +Rust ``LazyArrowStreamTable`` table provider gives partition pruning on +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. +""" + +from __future__ import annotations + +from typing import Any, TypeGuard + +import xarray as xr +from datafusion import SessionContext + +from ..df import Chunks +from ..reader import read_xarray_table +from ..sql import XarrayContext +from .base import register_adapter + + +@register_adapter +class DataFusionAdapter: + """Registers Datasets on ``datafusion.SessionContext`` connections.""" + + @staticmethod + def matches(con: object) -> TypeGuard[SessionContext]: + return isinstance(con, SessionContext) + + @staticmethod + def register( + con: SessionContext, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> SessionContext: + # XarrayContext.from_dataset adds dim-group splitting, cftime UDF + # registration, and round-trip metadata tracking on top of the + # plain table registration; use it when available. + if isinstance(con, XarrayContext): + return con.from_dataset(name, ds, chunks=chunks, **kwargs) + con.register_table(name, read_xarray_table(ds, chunks, **kwargs)) + return con diff --git a/xarray_sql/backends/duckdb.py b/xarray_sql/backends/duckdb.py new file mode 100644 index 0000000..1651ac0 --- /dev/null +++ b/xarray_sql/backends/duckdb.py @@ -0,0 +1,89 @@ +"""DuckDB engine adapter. + +Registers a lazy ``xarray.Dataset`` on a ``duckdb.DuckDBPyConnection`` +as an [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]: +DuckDB classifies it with a real ``isinstance`` check against +``pyarrow.dataset.Dataset`` and calls ``scanner(columns=[...], +filter=)`` once per query, giving +projection pushdown, coordinate-range chunk pruning, and prefetched +parallel production (see [xarray_sql.backends.pyarrow][]). + +This adapter never imports the ``duckdb`` package at runtime — detection +is by connection type, and registration is a method call on the +connection — so DuckDB stays a purely optional dependency +(``pip install xarray-sql[duckdb]``). + +Zarr-native scanning inside DuckDB is what the [duckdb-zarr](https://github.com/xqlsystems/duckdb-zarr) extension provides; this +adapter instead covers everything xarray can open (NetCDF, GRIB, Xee, CF +decoding, in-memory) and pairs with [xarray_sql.to_dataset][] for +the labeled round-trip. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, TypeGuard + +import xarray as xr + +from ..df import Chunks, group_vars_by_dims +from .base import register_adapter +from .pyarrow import XarrayArrowStream, XarrayPushdownDataset + +if TYPE_CHECKING: + import duckdb + +__all__ = ["DuckDBAdapter", "XarrayArrowStream", "XarrayPushdownDataset"] + + +@register_adapter +class DuckDBAdapter: + """Registers Datasets on ``duckdb.DuckDBPyConnection`` connections.""" + + @staticmethod + def matches(con: object) -> TypeGuard[duckdb.DuckDBPyConnection]: + # The connection class lives in ``duckdb`` or, in newer releases, + # the ``_duckdb`` C-extension module. + root = type(con).__module__.split(".")[0] + return root in ("duckdb", "_duckdb") + + @staticmethod + def register( + con: duckdb.DuckDBPyConnection, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> duckdb.DuckDBPyConnection: + """Register ``ds`` on a DuckDB connection. + + Datasets whose variables all share the same dimensions become a + single table named ``name``. Mixed-dimension datasets are split + into one table per dimension group, named + ``___...`` (DuckDB registration has no schema + namespace to mirror the DataFusion adapter's ``name.group`` + layout). Extra keyword arguments (``batch_size``, ``prefetch``) + are forwarded to [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. + """ + groups = group_vars_by_dims(ds) + if len(groups) <= 1: + con.register(name, XarrayPushdownDataset(ds, chunks, **kwargs)) + return con + # Materialise dim coordinates once and share across sub-tables. + coord_arrays = { + str(dim): ds.coords[dim].values + for dim in ds.dims + if dim in ds.coords + } + for dims, var_names in groups.items(): + suffix = "_".join(dims) or "scalar" + con.register( + f"{name}_{suffix}", + XarrayPushdownDataset( + ds[var_names], + chunks, + coord_arrays=coord_arrays, + **kwargs, + ), + ) + return con diff --git a/xarray_sql/backends/pyarrow.py b/xarray_sql/backends/pyarrow.py new file mode 100644 index 0000000..dd0f254 --- /dev/null +++ b/xarray_sql/backends/pyarrow.py @@ -0,0 +1,1143 @@ +"""Engine-neutral pyarrow views of lazy xarray Datasets. + +Two ways to hand a lazy ``xarray.Dataset`` to an Arrow-speaking query +engine: + +* [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset] — a real ``pyarrow.dataset.Dataset`` + subclass (the pattern Lance uses for ``LanceDataset``). Consumers of + the pyarrow dataset protocol — DuckDB via ``con.register``, Polars via + ``pl.scan_pyarrow_dataset``, or pyarrow itself — call + [scanner][xarray_sql.backends.pyarrow.XarrayPushdownDataset.scanner] with the columns a query needs + and the predicate it pushed down, so the scan loads only the needed + data variables from only the chunks whose coordinate ranges can + satisfy the predicate. Construct one with [arrow_dataset][xarray_sql.backends.pyarrow.arrow_dataset]. +* [XarrayArrowStream][xarray_sql.backends.pyarrow.XarrayArrowStream] — a re-scannable Arrow C-stream + (PyCapsule) view. No source-level pushdown, but works with any + PyCapsule consumer; the dependency-light fallback. + +Correctness contract shared by all consumers of the pushdown dataset: +engines may delete the filter conjuncts they push down (DuckDB does), +so the returned scanner applies the expression exactly via +``pyarrow.dataset.Scanner``; chunk pruning is only ever an optimization +on top. +""" + +from __future__ import annotations + +import itertools +import math +import re +import threading +import weakref +from collections import deque +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.dataset as pads +import pyarrow.fs as pafs +import xarray as xr + +from ..df import ( + Block, + Chunks, + DEFAULT_BATCH_SIZE, + _ensure_default_indexes, + _parse_schema, + iter_record_batches, + resolve_chunks, +) +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.""" + +_SHADOW_FANOUT = 1024 +"""Maximum fragments per shadow level. + +A dimension with more chunks than this gets a two-level shadow: a coarse +level of at most this many buckets, refined per surviving bucket. This +bounds shadow construction cost for finely partitioned datasets (e.g. +hundreds of thousands of single-step time chunks) at registration and +query time alike. +""" + +_REFINE_MAX_FRACTION = 0.25 +"""Skip fine-level pruning when the coarse pass kept more buckets. + +Refinement builds one sub-shadow per surviving bucket; when a predicate +matches most of the axis that cost cannot pay for itself, so the scan +falls back to the (sound) coarse answer. +""" + +_STRICT_LEVEL_BUDGET = 4096 +"""Maximum guarantee fragments per level of the strictness analysis. + +Strictness classifies bucket-products of surviving chunks recursively: +whole buckets prove (or prune) at once, and only mixed cells refine. +Each level builds at most this many fragments, so grids of millions of +chunks resolve in a handful of vectorized passes. +""" + +_STRICT_MAX_DEPTH = 6 +"""Recursion bound for the strictness analysis (a backstop: realistic +grids terminate in 2-3 levels).""" + + +def _guarantee_shadow( + guarantees: list[tuple[str, pc.Expression]], schema: pa.Schema +) -> pads.FileSystemDataset: + """A path-only dataset whose fragments carry the given guarantees. + + Each ``(path, guarantee)`` pair becomes one fragment: the path is a + label (never opened) and the guarantee is its + ``partition_expression``, so ``get_fragments(filter=...)`` delegates + satisfiability to Arrow's guarantee simplification. + """ + fmt = pads.IpcFileFormat() + fs = pafs.LocalFileSystem() + fragments = [ + fmt.make_fragment(path, fs, partition_expression=guarantee) + for path, guarantee in guarantees + ] + return pads.FileSystemDataset(fragments, schema, fmt, fs) + + +class _DimShadow: + """Chunk-pruning index for one dimension of the source grid. + + Fragment ``i`` of a shadow ``FileSystemDataset`` carries the + guarantee ``dim ∈ [min, max]`` of chunk-span ``i`` as its + ``partition_expression``; ``get_fragments(filter=...)`` then lets + Arrow's guarantee simplification decide which spans can satisfy a + predicate — sound for every predicate shape, conservative on columns + the guarantee does not mention, and the fragments' paths are never + opened. + + Axes with more than ``_SHADOW_FANOUT`` chunks use two levels: a + coarse shadow over buckets of consecutive chunks, plus per-bucket + fine shadows built lazily for the buckets a query keeps. + """ + + def __init__( + self, + name: str, + schema: pa.Schema, + coord: np.ndarray, + bounds: np.ndarray, + ): + self._name = name + # The full table schema, not just this dimension's field: the + # pushed predicate may reference any column, and get_fragments + # must be able to bind all of them (guarantees stay per-dim; + # unmentioned columns are conservatively unconstrained). + self._schema = schema + self._field_type = schema.field(name).type + self._coord = coord + self._bounds = bounds + self._n = len(bounds) - 1 + self._step = max(1, math.ceil(self._n / _SHADOW_FANOUT)) + self._n_buckets = math.ceil(self._n / self._step) + self._coarse = self._build( + [ + (b * self._step, min((b + 1) * self._step, self._n)) + for b in range(self._n_buckets) + ] + ) + self._fine: dict[int, pads.FileSystemDataset] = {} + + def _build(self, spans: list[tuple[int, int]]) -> pads.FileSystemDataset: + """A shadow whose fragment ``i`` guarantees chunk-span ``spans[i]``.""" + guarantees: list[tuple[str, pc.Expression]] = [] + for i, (lo_chunk, hi_chunk) in enumerate(spans): + vals = self._coord[self._bounds[lo_chunk] : self._bounds[hi_chunk]] + if (vals.dtype.kind == "f" and np.isnan(vals).any()) or ( + vals.dtype.kind == "M" and np.isnat(vals).any() + ): + # NaN/NaT poisons min/max into a (dim >= NaN) guarantee + # that Arrow simplifies every predicate against as false, + # silently pruning rows. An always-true guarantee keeps + # the span unprunable instead. + guarantee = pc.scalar(True) + else: + # min/max (not first/last) so descending axes like + # latitude 90→-90 carry correct ranges. + lo = pa.scalar(vals.min(), type=self._field_type) + hi = pa.scalar(vals.max(), type=self._field_type) + guarantee = (pc.field(self._name) >= lo) & ( + pc.field(self._name) <= hi + ) + guarantees.append((str(i), guarantee)) + return _guarantee_shadow(guarantees, self._schema) + + @staticmethod + def _kept_indices( + shadow: pads.FileSystemDataset, filter: pc.Expression + ) -> list[int]: + return sorted( + int(frag.path) for frag in shadow.get_fragments(filter=filter) + ) + + def kept(self, filter: pc.Expression) -> list[int] | None: + """Chunk indices that can satisfy ``filter``; ``None`` means all.""" + try: + buckets = self._kept_indices(self._coarse, filter) + if self._step == 1: + return buckets if len(buckets) < self._n else None + if len(buckets) > _REFINE_MAX_FRACTION * self._n_buckets: + # Refining most of the axis costs more than it saves; + # answer with the coarse buckets, which is still sound. + return ( + None + if len(buckets) == self._n_buckets + else [ + i + for b in buckets + for i in range( + b * self._step, + min((b + 1) * self._step, self._n), + ) + ] + ) + kept: list[int] = [] + for b in buckets: + fine = self._fine.get(b) + if fine is None: + start = b * self._step + stop = min((b + 1) * self._step, self._n) + fine = self._build([(i, i + 1) for i in range(start, stop)]) + self._fine[b] = fine + start = b * self._step + kept.extend(start + i for i in self._kept_indices(fine, filter)) + return kept + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError): + return None # conservative: scan every chunk of this dim + + +class XarrayArrowStream: + """A re-scannable Arrow C-stream view over a lazy xarray Dataset. + + Arrow PyCapsule consumers (DuckDB among them) call + ``__arrow_c_stream__`` once per scan. Each call constructs a fresh + [XarrayRecordBatchReader][xarray_sql.reader.XarrayRecordBatchReader] over the same + lazy Dataset, so — unlike registering a ``pyarrow.RecordBatchReader`` + directly, which is exhausted after one query — the same registered + table supports any number of queries, and data is only read while a + query is executing. + + The PyCapsule scan path gets no source-level pushdown (the producer + never sees the query's columns or filters), so + [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset] is the default registration object; + this class remains as the dependency-light fallback. + """ + + def __init__( + self, + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, + ): + # Validate eagerly (same checks XarrayRecordBatchReader runs) so + # registration fails fast instead of erroring mid-query. + probe = XarrayRecordBatchReader(ds, chunks, batch_size=batch_size) + self._ds = ds + self._chunks = chunks + self._batch_size = batch_size + self._schema = probe.schema + self._iteration_callback = _iteration_callback + + def __arrow_c_stream__( + self, requested_schema: object | None = None + ) -> object: + reader = XarrayRecordBatchReader( + self._ds, + self._chunks, + batch_size=self._batch_size, + _iteration_callback=self._iteration_callback, + ) + return reader.__arrow_c_stream__(requested_schema) + + def __arrow_c_schema__(self) -> object: + return self._schema.__arrow_c_schema__() + + +class XarrayPushdownDataset(pads.Dataset): + """A pushdown-capable ``pyarrow.dataset.Dataset`` view of a Dataset. + + Consumers that speak the pyarrow dataset protocol (DuckDB, Polars, + ...) call [scanner][xarray_sql.backends.pyarrow.XarrayPushdownDataset.scanner] with the columns a query needs and the + predicate it pushed down; the scan then loads only the needed data + variables from only the chunks whose coordinate ranges can satisfy + the predicate. + + The base class is never initialized (there is no C++ dataset behind + this object — the same construction Lance uses for ``LanceDataset``); + every entry point consumers touch is overridden in Python, and the + few inherited members that would read uninitialized native state are + stubbed out. + + References: + Lance's ``LanceDataset``, a ``pyarrow.dataset.Dataset`` subclass + built the same way: https://github.com/lancedb/lance + (``python/python/lance/dataset.py``). + """ + + def __init__( + self, + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + prefetch: int = DEFAULT_PREFETCH, + prefetch_bytes: int | None = None, + coalesce_rows: int | None = None, + geometry: tuple[str, str] | None = None, + geometry_encoding: str = "wkb", + geometry_crs: str | None = "OGC:CRS84", + coord_arrays: dict[str, np.ndarray] | None = None, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, + ): + # Deliberately no super().__init__() — see class docstring. + ds = _ensure_default_indexes(ds) + if ds.data_vars: + fst = next(iter(ds.values())).dims + if not all(da.dims == fst for da in ds.values()): + raise ValueError( + "All dimensions must be equal. " + "Please filter data_vars in the Dataset." + ) + self._ds = ds + self._schema = _parse_schema(ds) + self._geometry = tuple(geometry) if geometry else None + self._geometry_encoding = geometry_encoding + if self._geometry: + if GEOMETRY_COLUMN in self._schema.names: + raise ValueError( + f"geometry= would shadow an existing column named " + f"{GEOMETRY_COLUMN!r}." + ) + missing = [d for d in self._geometry if d not in self._schema.names] + if missing: + raise ValueError( + f"geometry dims {missing} are not columns of the " + f"table; available: {self._schema.names}" + ) + self._schema = self._schema.append( + 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() + } + # Reuse pre-materialised coordinate arrays where the caller has + # them (e.g. shared across the tables of a dim-group split); each + # missing dim costs one read, a network round-trip for Zarr. + self._coord_arrays = dict(coord_arrays or {}) + for d in ds.dims: + if str(d) not in self._coord_arrays: + self._coord_arrays[str(d)] = ds.coords[d].values + if batch_size <= 0: + # A zero size would never advance the zero-column scan's + # row loop; fail here rather than mid-scan. + raise ValueError(f"batch_size must be positive, got {batch_size}") + self._batch_size = batch_size + self._prefetch = prefetch + self._prefetch_bytes = prefetch_bytes + self._coalesce_rows = coalesce_rows + self._iteration_callback = _iteration_callback + self._shadows: dict[str, _DimShadow] | None = None + self._span_cache: dict[ + str, tuple[np.ndarray, np.ndarray, np.ndarray] + ] = {} + # One long-lived pool shared by every scan, its threads spawned + # NOW — never from inside an engine's scan callback. Creating a + # pool (and its OS threads) per scan deadlocks when the scan is + # driven from an engine executing under another thread pool + # (dask computing chunks of a lazy round-trip): thread startup + # and concurrent.futures' global shutdown lock interleave with + # the engine's callback needing the GIL. Scans that stop early + # cancel their queued loads instead of tearing the pool down. + # Like any thread state, the pool does not survive fork(); use + # from forked workers (e.g. PyTorch DataLoader) is tracked in + # https://github.com/xqlsystems/xarray-sql/issues/145. + self._pool: ThreadPoolExecutor | None = None + if self._prefetch > 1: + self._pool = ThreadPoolExecutor(max_workers=self._prefetch) + # Each pre-spawn task parks on the barrier, so no thread can + # take a second task and the executor is forced to start all + # ``prefetch`` OS threads before __init__ returns (submitting + # plain no-ops lets one idle thread absorb several of them, + # leaving the rest to spawn later inside an engine's scan + # callback — the deadlock this pre-spawn exists to prevent). + barrier = threading.Barrier(self._prefetch + 1) + spawn = [ + self._pool.submit(barrier.wait) for _ in range(self._prefetch) + ] + barrier.wait() + for f in spawn: + f.result() + # Stop the pool's threads when the dataset dies; live scans + # keep the dataset alive through their generator closures, + # so nothing in flight is cut short. The callback is bound + # to the executor, not the dataset, so the finalizer holds + # no reference cycle back to self. + weakref.finalize( + self, self._pool.shutdown, wait=False, cancel_futures=True + ) + + # ------------------------------------------------------------------ + # The consumer-facing surface + # ------------------------------------------------------------------ + + @property + def schema(self) -> pa.Schema: + return self._schema + + def scanner( + self, + columns: list[str] | None = None, + filter: pc.Expression | None = None, + batch_size: int | None = None, + **kwargs: Any, + ) -> pads.Scanner: + """Build a scanner for the requested columns and predicate. + + ``filter`` is applied exactly by the returned scanner (DuckDB + deletes the conjuncts it pushes down and trusts the source to + enforce them); chunk pruning and column selection only reduce + how much data is read to get there. ``batch_size`` caps rows per + emitted batch (Polars passes it through ``to_batches``). Extra + keyword arguments from other pyarrow-dataset consumers are + accepted and ignored. + """ + kept = None if filter is None else self._prune(filter) + blocks = ( + self._coalesced_blocks(kept) + if self._coalesce_rows + else self._blocks(kept) + ) + return self._scanner_for_blocks(blocks, columns, filter, batch_size) + + def _scanner_for_blocks( + self, + blocks: Iterator[Block] | list[Block], + columns: list[str] | None, + filter: pc.Expression | None, + batch_size: int | None = None, + ) -> pads.Scanner: + """A scanner over the given blocks; shared by dataset and fragments.""" + # ``None`` means every column (the pyarrow convention); an + # explicitly empty list is a real projection ("no payload"), not + # a request for the full schema. + proj = list(self._schema.names) if columns is None else list(columns) + scan_names = self._scan_columns(proj, filter) + scan_schema = pa.schema([self._schema.field(n) for n in scan_names]) + size = batch_size or self._batch_size + if self._geometry and GEOMETRY_COLUMN in scan_names: + batches = self._batches_with_geometry(scan_schema, blocks, size) + else: + batches = self._batch_generator(scan_schema, blocks, size) + return pads.Scanner.from_batches( + batches, schema=scan_schema, columns=proj, filter=filter + ) + + def _batches_with_geometry( + self, + scan_schema: pa.Schema, + blocks: Iterator[Block] | list[Block], + batch_size: int, + ) -> Iterator[pa.RecordBatch]: + """Emit ``scan_schema`` batches, synthesizing the geometry column. + + The pivot never materializes geometry: batches are produced with + the coordinate dims the geometry derives from, and the geometry + column is built per batch from those columns (for the native + encoding the point struct's children *are* the coordinate + arrays — a schema annotation, not a copy). + """ + assert self._geometry is not None + x_dim, y_dim = self._geometry + base_names = [n for n in scan_schema.names if n != GEOMETRY_COLUMN] + for d in (x_dim, y_dim): + if d not in base_names: + base_names.append(d) + base_schema = pa.schema([self._schema.field(n) for n in base_names]) + for batch in self._batch_generator(base_schema, blocks, batch_size): + geom = build_geometry( + self._geometry_encoding, + batch.column(base_names.index(x_dim)), + batch.column(base_names.index(y_dim)), + ) + arrays = [ + geom + if name == GEOMETRY_COLUMN + else batch.column(base_names.index(name)) + for name in scan_schema.names + ] + yield pa.RecordBatch.from_arrays(arrays, schema=scan_schema) + + def _batch_generator( + self, + scan_schema: pa.Schema, + blocks: Iterator[Block] | list[Block], + batch_size: int, + ) -> Iterator[pa.RecordBatch]: + names = list(scan_schema.names) + data_vars = [n for n in names if n in self._ds.data_vars] + # Select only the needed variables before slicing so unrequested + # variables are never loaded (dimension coords come via coords). + base = ( + self._ds[data_vars] + if data_vars + else self._ds.drop_vars(list(self._ds.data_vars)) + ) + + def load(block: Block) -> list[pa.RecordBatch]: + if self._iteration_callback is not None: + self._iteration_callback(block, names) + if not names: + # Zero-column projection: row counts are chunk + # arithmetic; no coordinate or variable data is read. + out = [] + rows = self._block_rows(block) + while rows > 0: + n = min(rows, batch_size) + out.append( + pa.table({"_": np.empty(n, np.int8)}) + .select([]) + .to_batches()[0] + ) + rows -= n + return out + return list( + iter_record_batches(base.isel(block), scan_schema, batch_size) + ) + + # Estimated pivoted bytes per row: gates admission when a + # byte budget is set, so peak memory tracks bytes in flight + # rather than block count (blocks vary in size under + # coalesce_rows). + row_width = 0 + for field in scan_schema: + try: + row_width += np.dtype(field.type.to_pandas_dtype()).itemsize + except (TypeError, NotImplementedError): + row_width += 8 + + def generate() -> Iterator[pa.RecordBatch]: + block_iter = iter(blocks) + first = next(block_iter, None) + if first is None: + return + second = next(block_iter, None) + if self._pool is None or second is None: + # Single-block scans (a lazy round-trip window that maps + # onto one source chunk) skip the pool entirely. + yield from load(first) + if second is not None: + yield from load(second) + for block in block_iter: + yield from load(block) + return + pool = self._pool + budget = self._prefetch_bytes + pending: deque = deque() + inflight = 0 + + def submit(block: Block) -> None: + nonlocal inflight + estimate = self._block_rows(block) * row_width + pending.append((pool.submit(load, block), estimate)) + inflight += estimate + + def drain_one() -> Iterator[pa.RecordBatch]: + nonlocal inflight + future, estimate = pending.popleft() + inflight -= estimate + yield from future.result() + + try: + submit(first) + submit(second) + for block in block_iter: + submit(block) + while len(pending) > 1 and ( + len(pending) >= self._prefetch + or (budget is not None and inflight > budget) + ): + yield from drain_one() + while pending: + yield from drain_one() + finally: + # Consumer may stop early (e.g. LIMIT): drop queued work + # without waiting for in-flight loads. The pool itself is + # shared across scans and stays up. + for future, _ in pending: + future.cancel() + + return generate() + + def get_fragments( + self, filter: pc.Expression | None = None + ) -> list["_XarrayFragment"]: + """One fragment per chunk of the source grid, pruned by ``filter``. + + This is how DataFusion consumes the dataset + (``SessionContext.register_dataset`` plans one partition per + fragment and scans them in parallel), and enables the Dask + pattern ``from_map(lambda f: f.to_table().to_pandas(), + ds.get_fragments())``. + """ + kept = None if filter is None else self._prune(filter) + return [_XarrayFragment(self, block) for block in self._blocks(kept)] + + def count_rows( + self, filter: pc.Expression | None = None, **kwargs: Any + ) -> int: + """Count rows, reading as little data as possible. + + Without a filter the count is pure chunk arithmetic — no I/O at + all. With a filter, chunks are split three ways: pruned chunks + contribute nothing, chunks whose coordinate ranges *prove* the + filter true contribute their exact size arithmetically, and only + the undecided boundary chunks are scanned (reading just the + columns the filter references). + """ + if not self._ds.sizes: + return int(self.scanner(columns=[], filter=filter).count_rows()) + if filter is None: + return int(np.prod([self._ds.sizes[d] for d in self._ds.dims])) + kept = self._prune(filter) + proven, boundary = self._strict_partition(kept, filter) + return proven + int( + self._scanner_for_blocks(boundary, [], filter).count_rows() + ) + + # Inherited convenience methods (to_table, head, to_batches, take) + # route through scanner() and keep working; the members below would + # touch the uninitialized native dataset. + + @property + def partition_expression(self) -> pc.Expression: + # The dataset-level guarantee: trivially true. The base class + # getter reads native state this object does not have. + return pc.scalar(True) + + def filter(self, expression: pc.Expression): + # A lazily-composed filter view is implementable; tracked in + # https://github.com/xqlsystems/xarray-sql/issues/239. + raise NotImplementedError( + "Use scanner(filter=...) or the engine's WHERE clause." + ) + + def replace_schema(self, schema: pa.Schema): + raise NotImplementedError + + # Guarded delegation for sort_by/join/join_asof (engine-level SQL + # joins already work through scanner()) is tracked in + # https://github.com/xqlsystems/xarray-sql/issues/240. + + def sort_by(self, sorting, **kwargs): + raise NotImplementedError + + def join(self, *args, **kwargs): + raise NotImplementedError + + def join_asof(self, *args, **kwargs): + raise NotImplementedError + + def __reduce__(self): + raise TypeError("XarrayPushdownDataset is not picklable.") + + # ------------------------------------------------------------------ + # Projection: which columns must be read + # ------------------------------------------------------------------ + + def _scan_columns( + self, proj: list[str], filter: pc.Expression | None + ) -> list[str]: + """Columns to read: the projection plus any the filter references. + + The consumer's column list need not include filter-only columns + (DuckDB drops pushed conjuncts from its plan and has no upstream + use for them). Rather than parsing the expression, probe it + against an empty table and grow the column set from the "no match + for field" errors until it evaluates; on anything unexpected fall + back to scanning every column, which is always correct. + """ + if filter is None: + return proj + wanted = set(proj) + for _ in range(len(self._schema.names) + 1): + probe = pa.table( + { + n: pa.array([], type=self._schema.field(n).type) + for n in self._schema.names + if n in wanted + } + ) + try: + probe.filter(filter) + except pa.lib.ArrowInvalid as exc: + match = re.search(r"FieldRef\.Name\((.*?)\)", str(exc)) + name = match.group(1) if match else None + if ( + name is not None + and name in set(self._schema.names) - wanted + ): + wanted.add(name) + continue + return list(self._schema.names) + else: + return [n for n in self._schema.names if n in wanted] + return list(self._schema.names) + + # ------------------------------------------------------------------ + # Pruning: which chunks can satisfy the predicate + # ------------------------------------------------------------------ + + def _dim_shadows(self) -> dict[str, _DimShadow]: + """One pruning index per prunable dimension, built lazily. + + Keeping one shadow per dimension (Σ n_d fragments) instead of one + per chunk (Π n_d) is what keeps this cheap for finely partitioned + datasets, and is sound: a chunk is dropped only when the full + predicate is provably false given that single dimension's range. + """ + if self._shadows is not None: + return self._shadows + shadows: dict[str, _DimShadow] = {} + for dim in self._resolved: + name = str(dim) + if name not in self._schema.names: + continue + coord = self._coord_arrays[name] + if coord.dtype.kind not in ("i", "u", "f", "M"): + continue # strings/objects/cftime: never prune this dim + try: + shadows[name] = _DimShadow( + name, self._schema, coord, self._chunk_bounds[dim] + ) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError): + continue # conservative: no pruning on this dim + self._shadows = shadows + return shadows + + def _prune(self, filter: pc.Expression) -> dict[str, list[int]]: + """Per-dimension chunk indices that can satisfy ``filter``. + + Satisfiability is delegated to Arrow's guarantee simplification + (see ``_DimShadow``) — no expression decoding here, and + predicates on columns a shadow knows nothing about are + conservatively kept. Dimensions without a shadow, or where every + chunk survives, are absent from the result. + """ + kept: dict[str, list[int]] = {} + for name, shadow in self._dim_shadows().items(): + indices = shadow.kept(filter) + if indices is not None: + kept[name] = indices + return kept + + # ------------------------------------------------------------------ + # Strictness: which surviving chunks satisfy the filter entirely + # ------------------------------------------------------------------ + + def _chunk_spans( + self, name: str + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Cached vectorized per-chunk ``(lo, hi, poisoned)`` for a dim.""" + cached = self._span_cache.get(name) + if cached is not None: + return cached + coord = self._coord_arrays[name] + starts = self._chunk_bounds[name][:-1] + lo = np.minimum.reduceat(coord, starts) + hi = np.maximum.reduceat(coord, starts) + if coord.dtype.kind == "f": + bad_values = np.isnan(coord) + elif coord.dtype.kind == "M": + bad_values = np.isnat(coord) + else: + bad_values = np.zeros(len(coord), dtype=bool) + bad = np.bitwise_or.reduceat(bad_values, starts) + self._span_cache[name] = (lo, hi, bad) + return self._span_cache[name] + + def _strict_partition( + self, + kept: dict[str, list[int]] | None, + filter: pc.Expression, + ) -> tuple[int, list[Block] | Iterator[Block]]: + """``(rows proven to satisfy filter, boundary blocks to scan)``. + + A cell of the surviving chunk grid with conjunctive coordinate + guarantee ``G`` satisfies ``filter`` everywhere iff + ``G ∧ ¬filter`` is unsatisfiable, and can be *dropped* entirely + iff ``G ∧ filter`` is — both decided by Arrow's guarantee + simplification. Cells are classified hierarchically: each level + buckets the surviving indices into at most + ``_STRICT_LEVEL_BUDGET`` products, proves/prunes whole buckets + at once (the prune side refines the per-dimension pruning with + cross-dimension information), and recurses only into mixed + cells — so million-chunk axes resolve in two or three levels. + Anything undecidable (NaN spans, non-numeric dims, expression + shapes the simplifier rejects) conservatively lands in the + boundary set, which the caller scans exactly. + """ + dims = list(self._resolved.keys()) + lists = {d: list(self._surviving(kept, d)) for d in dims} + if not dims or any(not v for v in lists.values()): + return 0, [] + lens = {d: np.diff(self._chunk_bounds[d]) for d in dims} + outer = self._outer_rows() + usable = { + d: str(d) in self._schema.names + and self._coord_arrays[str(d)].dtype.kind in ("i", "u", "f", "M") + for d in dims + } + proven = 0 + boundary: list[Block] = [] + + def bucket_guarantee( + d: Any, indices: np.ndarray + ) -> pc.Expression | None: + """Conjunctive [min, max] guarantee over one dimension's index + bucket, or None when unprovable (NaN spans, non-numeric dims).""" + if not usable[d]: + return None + lo, hi, bad = self._chunk_spans(str(d)) + if bad[indices].any(): + return None + field_type = self._schema.field(str(d)).type + return ( + pc.field(str(d)) >= pa.scalar(lo[indices].min(), field_type) + ) & (pc.field(str(d)) <= pa.scalar(hi[indices].max(), field_type)) + + def rows_of(cell: dict) -> int: + """Rows spanned by ``cell``: the product of its chunks' lengths + per dimension, times the rows of dimensions outside the grid.""" + rows = outer + for d in dims: + rows *= int(lens[d][np.asarray(cell[d])].sum()) + return rows + + def classify(cell: dict, depth: int) -> None: + """Prove, prune, or split one cell of the surviving chunk grid. + + A cell is a hyper-rectangle of chunk indices per dimension. + Its indices are bucketed into at most + ``_STRICT_LEVEL_BUDGET`` products; Arrow's guarantee + simplification then decides each bucket-product wholesale: + proven (every row satisfies ``filter``, counted into + ``proven`` without reading), pruned (provably empty, + dropped), or mixed (recurse). Undecidable single-chunk + cells land in ``boundary`` for exact scanning. + """ + nonlocal proven + ks = {d: len(cell[d]) for d in dims} + while int(np.prod(list(ks.values()))) > _STRICT_LEVEL_BUDGET: + widest = max(ks, key=lambda d: ks[d]) + if ks[widest] == 1: + break + ks[widest] = max(1, ks[widest] // 2) + buckets = { + d: np.array_split(np.asarray(cell[d]), ks[d]) for d in dims + } + combos = list(itertools.product(*(range(ks[d]) for d in dims))) + guarantees: list[pc.Expression | None] = [] + for combo in combos: + g: pc.Expression | None = None + complete = True + for d, b in zip(dims, combo): + gd = bucket_guarantee(d, buckets[d][b]) + if gd is None: + if usable[d]: + complete = False # NaN span: never provable + continue + g = gd if g is None else g & gd + guarantees.append(g if (g is not None and complete) else None) + decidable = [i for i, g in enumerate(guarantees) if g is not None] + satisfiable = set(decidable) + unstrict = set(decidable) + if decidable: + shadow = _guarantee_shadow( + [(str(i), guarantees[i]) for i in decidable], self._schema + ) + satisfiable = { + int(f.path) for f in shadow.get_fragments(filter=filter) + } + unstrict = { + int(f.path) for f in shadow.get_fragments(filter=~filter) + } + for i, combo in enumerate(combos): + subcell = {d: list(buckets[d][b]) for d, b in zip(dims, combo)} + if guarantees[i] is not None: + if i not in satisfiable: + continue # provably empty: cross-dim refinement + if i not in unstrict: + proven += rows_of(subcell) + continue + if all(len(v) == 1 for v in subcell.values()): + boundary.append( + self._block_for_combo( + tuple(v[0] for v in subcell.values()) + ) + ) + elif depth < _STRICT_MAX_DEPTH: + classify(subcell, depth + 1) + else: + # Depth backstop; unreachable for realistic grids. + for c in itertools.product(*subcell.values()): + boundary.append(self._block_for_combo(c)) + + try: + classify(lists, 0) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError): + return 0, self._blocks(kept) # scan every survivor, exactly + return proven, boundary + + def _block_rows(self, block: Block) -> int: + rows = 1 + for d, sl in block.items(): + size = self._ds.sizes[d] + start, stop, _ = sl.indices(size) + rows *= stop - start + return rows + + # ------------------------------------------------------------------ + # Scan: load surviving chunks, prefetching ahead of the consumer + # ------------------------------------------------------------------ + + def _surviving( + self, kept: dict[str, list[int]] | None, dim: Any + ) -> list[int] | range: + """One dim's surviving chunk indices; every chunk when unpruned.""" + return (kept or {}).get(str(dim), range(len(self._resolved[dim]))) + + def _outer_rows(self) -> int: + """Rows contributed per grid cell by dims outside the chunk grid. + + Unresolved dims span their full extent in every block, so they + multiply every block's row count uniformly. + """ + rows = 1 + for d in self._ds.dims: + if d not in self._resolved: + rows *= self._ds.sizes[d] + return rows + + def _combos( + self, kept: dict[str, list[int]] | None + ) -> Iterator[tuple[int, ...]]: + """Surviving chunk-index combinations, in grid order.""" + dims = list(self._resolved.keys()) + if not dims: + return + yield from itertools.product(*(self._surviving(kept, d) for d in dims)) + + def _block_for_combo(self, combo: tuple[int, ...]) -> Block: + block: Block = {d: slice(None) for d in self._ds.dims} + for d, i in zip(self._resolved.keys(), combo): + bounds = self._chunk_bounds[d] + block[d] = slice(int(bounds[i]), int(bounds[i + 1])) + return block + + def _blocks(self, kept: dict[str, list[int]] | None) -> Iterator[Block]: + """Yield isel-able block slices for the surviving chunk grid.""" + if not self._resolved: + yield {} + return + for combo in self._combos(kept): + yield self._block_for_combo(combo) + + def _coalesced_blocks( + self, kept: dict[str, list[int]] | None + ) -> Iterator[Block]: + """Blocks with runs of consecutive chunks merged along one dim. + + Runs are merged along the most finely chunked dimension while + the merged block stays under ``coalesce_rows`` rows. One merged + block is one ``isel`` — on Zarr sources its member chunks are + fetched by the store's own concurrent batch read instead of one + request per chunk through the prefetch pool. + """ + if not self._resolved: + yield {} + return + dims = list(self._resolved.keys()) + merge_dim = max(dims, key=lambda d: len(self._resolved[d])) + others = [d for d in dims if d != merge_dim] + ranges = {d: list(self._surviving(kept, d)) for d in dims} + merge_bounds = self._chunk_bounds[merge_dim] + outer_rows = self._outer_rows() + + def flush(prefix: tuple[int, ...], run: list[int]) -> Block: + block: Block = {d: slice(None) for d in self._ds.dims} + for d, i in zip(others, prefix): + bounds = self._chunk_bounds[d] + block[d] = slice(int(bounds[i]), int(bounds[i + 1])) + block[merge_dim] = slice( + int(merge_bounds[run[0]]), int(merge_bounds[run[-1] + 1]) + ) + return block + + coalesce_rows = self._coalesce_rows + assert coalesce_rows is not None # only reached with coalescing on + for prefix in itertools.product(*(ranges[d] for d in others)): + per_row = outer_rows + for d, i in zip(others, prefix): + bounds = self._chunk_bounds[d] + per_row *= int(bounds[i + 1] - bounds[i]) + run: list[int] = [] + run_rows = 0 + for i in ranges[merge_dim]: + rows = int(merge_bounds[i + 1] - merge_bounds[i]) * per_row + if run and ( + i != run[-1] + 1 or run_rows + rows > coalesce_rows + ): + yield flush(prefix, run) + run, run_rows = [], 0 + run.append(i) + run_rows += rows + if run: + yield flush(prefix, run) + + +class _XarrayFragment: + """One chunk of the source grid, presented as a dataset fragment. + + Fragment consumers (DataFusion's ``DatasetExec`` plans one partition + per fragment; Dask maps over them) call [scanner][xarray_sql.backends.pyarrow.XarrayPushdownDataset.scanner] with the + columns and predicate for this piece; the pushed filter is applied + row-exactly, same as the parent dataset's scanner. + """ + + def __init__(self, dataset: XarrayPushdownDataset, block: Block): + self._dataset = dataset + self._block = block + + @property + def physical_schema(self) -> pa.Schema: + return self._dataset.schema + + def scanner( + self, + schema: pa.Schema | None = None, + columns: list[str] | None = None, + filter: pc.Expression | None = None, + batch_size: int | None = None, + **kwargs: Any, + ) -> pads.Scanner: + return self._dataset._scanner_for_blocks( + [self._block], columns, filter, batch_size + ) + + def to_batches(self, **kwargs: Any) -> Iterator[pa.RecordBatch]: + return iter(self.scanner(**kwargs).to_batches()) + + def to_table(self, **kwargs: Any) -> pa.Table: + return self.scanner(**kwargs).to_table() + + def count_rows(self, **kwargs: Any) -> int: + return int(self.scanner(**kwargs).count_rows()) + + def __dask_tokenize__(self) -> tuple: + # Dask hashes from_map inputs; the parent dataset is unpicklable, + # so provide a deterministic token from the fragment's identity. + return ( + "xarray_sql._XarrayFragment", + repr(self._block), + self._dataset.schema.to_string(), + ) + + def __repr__(self) -> str: + return f"_XarrayFragment({self._block!r})" + + +def arrow_dataset( + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + prefetch: int = DEFAULT_PREFETCH, + prefetch_bytes: int | None = None, + coalesce_rows: int | None = None, + geometry: tuple[str, str] | None = None, + geometry_encoding: str = "wkb", + geometry_crs: str | None = "OGC:CRS84", +) -> XarrayPushdownDataset: + """A pushdown-capable ``pyarrow.dataset.Dataset`` view of ``ds``. + + The returned object works anywhere a pyarrow dataset does, keeping + projection pushdown and coordinate-range chunk pruning:: + + import polars as pl + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) + + import duckdb + duckdb.connect().register("t", xql.arrow_dataset(ds)) + + xql.arrow_dataset(ds).to_table(columns=["t2m"], filter=...) + + Args: + 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. + 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). + prefetch_bytes: Optional cap on *estimated pivoted bytes* in + flight; admission then tracks bytes rather than block count, + which keeps peak memory steady when ``coalesce_rows`` makes + blocks large or uneven. ``prefetch`` still bounds + concurrency (thread count). + coalesce_rows: When set, merge runs of consecutive surviving + chunks along the most finely chunked dimension into single + reads of at most this many rows. Fewer, larger source + requests — the win on remote stores, where each merged read + fetches its member chunks through the store's own concurrent + batching. Memory scales with ``prefetch`` x the *merged* + block size, so size accordingly (e.g. ``8_000_000``). + geometry: ``(x_dim, y_dim)`` coordinate dims to derive a + ``geometry`` point column from (see + [xarray_sql.geometry][]). With the default ``"wkb"`` + encoding, DuckDB (spatial loaded) sees a native ``GEOMETRY`` + column, so ``ST_Within(geometry, ...)`` works directly. + geometry_encoding: ``"wkb"`` (default; DuckDB-consumable) or + ``"point"`` (GeoArrow native separated coordinates — the + struct children are the coordinate arrays; for GeoPandas, + lonboard, geoarrow-rs consumers). + geometry_crs: CRS tag carried in the extension metadata. + Defaults to ``OGC:CRS84`` (plain longitude/latitude); pass + ``None`` to omit, or an authority code / PROJJSON string. + + Returns: + An [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. + """ + return XarrayPushdownDataset( + ds, + chunks, + batch_size=batch_size, + prefetch=prefetch, + prefetch_bytes=prefetch_bytes, + coalesce_rows=coalesce_rows, + geometry=geometry, + geometry_encoding=geometry_encoding, + geometry_crs=geometry_crs, + ) diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 9acaa94..1f7c49c 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -29,7 +29,6 @@ # Calendar classification # --------------------------------------------------------------------------- -#: Calendars close enough to proleptic Gregorian for ``pa.timestamp('us')``. GREGORIAN_LIKE_CALENDARS: frozenset[str] = frozenset( { "standard", @@ -41,10 +40,13 @@ "366_day", } ) +"""Calendars close enough to proleptic Gregorian for ``pa.timestamp('us')``.""" -#: Default CF-convention units when no encoding is available on the coordinate. -#: Microseconds give sub-second precision and fit int64 for ±292 k years. DEFAULT_UNITS: str = "microseconds since 1970-01-01T00:00:00" +"""Default CF-convention units when no encoding is available on the coordinate. + +Microseconds give sub-second precision and fit int64 for ±292 k years. +""" def is_gregorian_like(calendar: str) -> bool: @@ -111,7 +113,7 @@ def encoding(ds: xr.Dataset, coord_name: str) -> tuple[str, str]: """Return ``(units, calendar)`` for a cftime coordinate. Reads xarray ``.encoding`` metadata (from the originating NetCDF file) - first, falling back to :data:`DEFAULT_UNITS`. + first, falling back to [DEFAULT_UNITS][xarray_sql.cftime.DEFAULT_UNITS]. """ cal = calendar(ds, coord_name) or "standard" enc = ds.coords[coord_name].encoding diff --git a/xarray_sql/df.py b/xarray_sql/df.py index cd5e432..83879f7 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,4 +1,5 @@ import itertools +from collections import defaultdict from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping from typing import Any @@ -131,6 +132,23 @@ def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]: yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks)) +def group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: + """Group a Dataset's data variables by their exact dimension tuple. + + Variables that share dimensions can share a table; each distinct + dimension tuple becomes its own table when a mixed-dimension Dataset + is registered:: + + ("time", "lat", "lon"): ["temperature_2m", "wind_speed"], + ("time", "lat", "lon", "level"): ["pressure", "humidity"] + """ + groups = defaultdict(list) + for var_name, var in ds.data_vars.items(): + dims = var.dims + groups[dims].append(var_name) + return groups + + def _block_len(block: Block) -> int: return int(np.prod([v.stop - v.start for v in block.values()])) @@ -293,9 +311,39 @@ def dataset_to_record_batch( return pa.RecordBatch.from_arrays(arrays, schema=schema) -#: Default number of rows per emitted Arrow RecordBatch. -#: 64 K rows balances DataFusion pipeline depth against per-batch overhead. DEFAULT_BATCH_SIZE: int = 65_536 +"""Default number of rows per emitted Arrow RecordBatch. + +64 K rows balances DataFusion pipeline depth against per-batch overhead. +""" + +_FULL_PIVOT_MAX_ROWS: int = 8_388_608 +"""Row cap for the whole-partition coordinate fast path in +iter_record_batches. + +Below this, coordinate columns are materialised for the full partition +with repeat/tile (sequential writes, ~3x faster than per-batch index +arithmetic) and batches are zero-copy slices; the cost is holding every +coordinate column of the partition in memory at once (rows x 8 bytes x +n_dims). Above it — e.g. single-time-step reanalysis partitions with +tens of millions of rows — the per-batch path keeps peak memory at +O(batch_size) per coordinate instead. +""" + + +def _as_single_array(values, type: pa.DataType, *, from_pandas: bool = False): + """``pa.array`` that always returns a contiguous ``pa.Array``. + + ``pa.array`` may return a ``ChunkedArray`` instead of an ``Array`` for + large inputs (observed for numpy fixed-width unicode columns of a few + million rows — e.g. a string dimension coordinate tiled across a full + partition). ``RecordBatch.from_arrays`` rejects chunked input, so + flatten it back to one contiguous array. + """ + arr = pa.array(values, type=type, from_pandas=from_pandas) + if isinstance(arr, pa.ChunkedArray): + arr = arr.combine_chunks() + return arr def iter_record_batches( @@ -338,6 +386,8 @@ def iter_record_batches( # Preload small 1-D coordinate arrays (negligible memory). # Convert cftime objects to numeric values matching the schema type. + # Projected scans may omit dimension columns from the schema; those + # dims still shape the iteration but never emit a column. coord_values = {} schema_names = set(schema.names) for name in dim_names: @@ -368,6 +418,36 @@ def iter_record_batches( else: data_arrays[field.name] = raw.ravel() + if 0 < total_rows <= _FULL_PIVOT_MAX_ROWS: + # Fast path: build each coordinate column once for the whole + # partition. In C order, dim k's flat column is its coord values + # each repeated prod(shape[k+1:]) times, with that pattern tiled + # prod(shape[:k]) times — two sequential-write kernels, much + # faster than per-batch division/modulo plus gather. Batches are + # then zero-copy slices of the full-partition Arrow arrays. + full_arrays = [] + for field in schema: + name = field.name + if name in ds.coords and name in ds.dims: + k = dim_names.index(name) + outer = int(np.prod(shape[:k])) + col = np.repeat(coord_values[name], strides[k]) + if outer > 1: + col = np.tile(col, outer) + full_arrays.append(_as_single_array(col, field.type)) + else: + full_arrays.append( + _as_single_array( + data_arrays[name], field.type, from_pandas=True + ) + ) + for row_start in range(0, total_rows, batch_size): + yield pa.RecordBatch.from_arrays( + [a.slice(row_start, batch_size) for a in full_arrays], + schema=schema, + ) + return + for row_start in range(0, total_rows, batch_size): row_end = min(row_start + batch_size, total_rows) row_idx = np.arange(row_start, row_end) @@ -379,13 +459,13 @@ def iter_record_batches( k = dim_names.index(name) coord_idx = (row_idx // strides[k]) % shape[k] arrays.append( - pa.array(coord_values[name][coord_idx], type=field.type) + _as_single_array(coord_values[name][coord_idx], field.type) ) else: arrays.append( - pa.array( + _as_single_array( data_arrays[name][row_start:row_end], - type=field.type, + field.type, from_pandas=True, ) ) @@ -412,7 +492,7 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: Only *dimension coordinates* become dimension columns, so a dimension without a coordinate would be dropped. Callers must run the Dataset through - :func:`_ensure_default_indexes` first (the readers do) so every dimension + ``_ensure_default_indexes`` first (the readers do) so every dimension has a coordinate and appears as a column. Uses the xarray index type to detect cftime coordinates without diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index cd1db39..99ef7b3 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -1,13 +1,13 @@ """Reconstruct xarray Datasets from SQL query results. The inverse of the forward Dataset-to-table pivot done by -:func:`xarray_sql.df.pivot`. Internally defines an :class:`XarrayDataFrame` +[xarray_sql.df.pivot][]. Internally defines an [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapper around the DataFusion ``DataFrame`` returned by -:meth:`XarrayContext.sql`, with a :meth:`XarrayDataFrame.to_dataset` +[XarrayContext.sql][xarray_sql.sql.XarrayContext.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 -:meth:`XarrayDataFrame.to_dataset` -- the xarray idiom for tuning how a +[XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] -- the xarray idiom for tuning how a result is partitioned -- rather than by reflecting on the query plan: * **Eager** (``chunks=None``, or the default ``"inherit"`` when the @@ -17,7 +17,7 @@ (aggregations), whose results are small, and it never re-executes. * **Lazy / chunked** (``chunks`` is a mapping, ``"auto"``, or ``"inherit"`` over a multi-chunk source dimension): data variables are - backed by :class:`SQLBackendArray` wrapped in + backed by [SQLBackendArray][xarray_sql.ds.SQLBackendArray] wrapped in ``xarray.core.indexing.LazilyIndexedArray`` and chunked via xarray's configured chunk manager (dask, cubed, ...). Each chunk maps onto the source partitions and reads its coordinate range on access by @@ -38,7 +38,8 @@ import pandas as pd import pyarrow as pa import xarray as xr -from datafusion import col, literal + +from .lazyscan import DataFusionHandle, DimSpec, LazyResultHandle Sparsity = Literal["result", "template"] """Output coordinate extent for a filtered round-trip. @@ -147,6 +148,36 @@ def _apply_template(ds: xr.Dataset, template: xr.Dataset) -> xr.Dataset: return out +def _axis_numeric(values: np.ndarray) -> np.ndarray: + """View an axis as float64 for affine position arithmetic.""" + if values.dtype.kind == "M": + return values.astype("datetime64[ns]").view("int64").astype("float64") + return values.astype("float64", copy=False) + + +def _affine_axis(requested: np.ndarray) -> tuple[float, float] | None: + """``(origin, step)`` when *requested* is uniformly spaced, else None. + + Uniform spacing must hold exactly enough that ``rint((v - origin) / + step)`` recovers every index: the deviation of each element from its + affine prediction is checked against a quarter step. Non-numeric + axes (strings, cftime objects) never qualify. + """ + if requested.dtype.kind not in ("i", "u", "f", "M") or len(requested) < 2: + return None + numeric = _axis_numeric(requested) + step = (numeric[-1] - numeric[0]) / (len(numeric) - 1) + if step == 0 or not np.isfinite(step): + return None + predicted = numeric[0] + step * np.arange(len(numeric)) + # Written as a <= comparison so a NaN anywhere in the axis (e.g. a + # NULL dim value in the result) fails the check and falls back to + # the searchsorted path, which handles it positionally. + if not (np.abs(numeric - predicted) <= 0.25 * abs(step)).all(): + return None + return float(numeric[0]), float(step) + + def _scatter_batches_to_ndarray( batches: list[pa.RecordBatch], dimension_columns: list[str], @@ -182,8 +213,17 @@ def _scatter_batches_to_ndarray( # positions, and template coords like air_temperature.lat are descending). # ``np.searchsorted`` requires ascending input, so we sort each requested # array once, search there, and remap back to the original positions. - sorted_idx = {d: np.argsort(requested[d]) for d in dimension_columns} - sorted_req = {d: requested[d][sorted_idx[d]] for d in dimension_columns} + # Uniformly spaced axes (the norm for rasters and regular time steps, + # ascending or descending) skip the search entirely: the position is + # ``rint((value - origin) / step)``, a fused vector op several times + # faster than a per-row binary search. + affine = {d: _affine_axis(requested[d]) for d in dimension_columns} + sorted_idx = { + d: np.argsort(requested[d]) + for d in dimension_columns + if affine[d] is None + } + sorted_req = {d: requested[d][sorted_idx[d]] for d in sorted_idx} for batch in batches: if batch.num_rows == 0: @@ -195,8 +235,16 @@ def _scatter_batches_to_ndarray( for d in dimension_columns: col_arr = batch.column(schema_names.index(d)) vals = col_arr.to_numpy(zero_copy_only=False) - pos_in_sorted = np.searchsorted(sorted_req[d], vals) - positions.append(sorted_idx[d][pos_in_sorted]) + pair = affine[d] + if pair is not None: + origin, step = pair + pos = np.rint((_axis_numeric(vals) - origin) / step).astype( + np.intp + ) + positions.append(pos) + else: + pos_in_sorted = np.searchsorted(sorted_req[d], vals) + positions.append(sorted_idx[d][pos_in_sorted]) value_arr = batch.column(schema_names.index(var_name)).to_numpy( zero_copy_only=False ) @@ -208,28 +256,31 @@ def _scatter_batches_to_ndarray( class SQLBackendArray(xr.backends.BackendArray): - """Read-only lazy N-D array view over a DataFusion DataFrame. + """Read-only lazy N-D array view over a re-executable SQL result. Bridges xarray's lazy-indexing interface - (:class:`xarray.backends.BackendArray`) to a DataFusion query result, + (``xarray.backends.BackendArray``) to an engine query result, so an xarray ``Dataset`` can present a SQL query as if it were a materialized N-D array without actually loading any data until the caller asks for it. This is the workhorse that lets - :meth:`XarrayDataFrame.to_dataset` return a Dataset cheaply. + [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] (and the engine-agnostic + ``xql.to_dataset(chunks=...)``) return a Dataset cheaply. On each ``__getitem__`` call, the requested xarray indexer is - translated into a DataFusion filter expression (``df.filter(expr)``) - and a column projection (``df.select(*cols)``). The filtered - DataFrame is consumed via ``execute_stream`` as a sequence of Arrow - ``RecordBatch`` es and scattered into a preallocated numpy buffer, - so only the requested data is materialized. + translated into per-dimension coordinate windows and a column + projection, executed through a + [LazyResultHandle][xarray_sql.lazyscan.LazyResultHandle] (DataFusion, DuckDB, + or Polars — each renders the windows with its own typed expression + API). The resulting Arrow ``RecordBatch`` es are scattered into a + preallocated numpy buffer, so only the requested data is + materialized. Constraints and caveats: - Read-only: there is no write path; the backend exists to surface query results, not to round-trip writes into a SQL store. - - The underlying DataFusion ``DataFrame`` holds a reference to its - originating ``SessionContext``, which is not picklable. The class + - The underlying engine object may hold non-picklable references + (DataFusion's ``SessionContext``, a DuckDB connection). The class therefore overrides ``__copy__`` and ``__deepcopy__`` to return ``self`` -- this is safe because the backend is read-only. - ``IndexingSupport.OUTER``: ``BasicIndexer`` and ``OuterIndexer`` @@ -238,33 +289,39 @@ class SQLBackendArray(xr.backends.BackendArray): works, just less efficiently. Raises: - ValueError, datafusion exceptions: propagated from the - underlying ``df.filter().select().execute_stream()`` chain - if a predicate refers to a missing column, the dtype of a - literal is incompatible, or the execution itself fails. + ValueError, engine exceptions: propagated from the underlying + filter/project/execute chain if a predicate refers to a + missing column, the dtype of a literal is incompatible, or + the execution itself fails. AssertionError: from ``np.searchsorted`` mis-alignment, which indicates the result contains coordinate values not present in the wrapper's pre-computed coord arrays -- usually a symptom of a filtered query whose coord discovery missed a value. - Constructed by :func:`_build_lazy_scan`; users should not instantiate + Constructed by ``_build_lazy_scan``; users should not instantiate this class directly. """ def __init__( self, - inner_df: Any, + handle: LazyResultHandle, var_name: str, dimension_columns: list[str], coord_arrays: dict[str, np.ndarray], shape: tuple[int, ...], dtype: np.dtype, ) -> None: - self._inner_df = inner_df + self._handle = handle self._var_name = var_name self._dimension_columns = list(dimension_columns) self._coord_arrays = coord_arrays + # Computed once per dim: whether the whole coordinate array is + # strictly monotonic, the precondition for translating contiguous + # positional windows into value ranges (see _dim_spec). + self._monotonic = { + d: _strictly_monotonic(coord_arrays[d]) for d in dimension_columns + } self.shape = tuple(shape) self.dtype = np.dtype(dtype) @@ -291,28 +348,30 @@ def __deepcopy__(self, memo: dict) -> "SQLBackendArray": # ------------------------------------------------------------------ def _raw_getitem(self, key: tuple) -> np.ndarray: - """Materialize the indexed region described by *key* via DataFusion + Arrow. + """Materialize the indexed region described by *key* via the engine. ``key`` is a tuple of ``int``/``slice``/1-D integer-array, one per - dim, in :attr:`_dimension_columns` order. + dim, in ``_dimension_columns`` order. """ requested: dict[str, np.ndarray] = {} - # Dims whose indexer covers the full extent (slice(None) or - # equivalent). For these we omit the filter predicate entirely - # so DataFusion doesn't have to evaluate a tautology. - full_dims: set[str] = set() + # Per-dim windows for the engine. Dims whose indexer covers the + # full extent are omitted entirely so the engine doesn't have to + # evaluate a tautology. + specs: dict[str, DimSpec] = {} drop_axes: list[int] = [] for axis, (dim, k) in enumerate( zip(self._dimension_columns, key, strict=True) ): coord = self._coord_arrays[dim] + contiguous = False if isinstance(k, slice): start = 0 if k.start is None else k.start stop = len(coord) if k.stop is None else k.stop step = 1 if k.step is None else k.step requested[dim] = np.asarray(coord[start:stop:step]) + contiguous = step == 1 if start == 0 and stop >= len(coord) and step == 1: - full_dims.add(dim) + continue elif isinstance(k, (int, np.integer)): requested[dim] = np.asarray([coord[int(k)]]) drop_axes.append(axis) @@ -323,7 +382,11 @@ def _raw_getitem(self, key: tuple) -> np.ndarray: len(arr) == len(coord) and (arr == np.arange(len(coord))).all() ): - full_dims.add(dim) + continue + contiguous = len(arr) > 1 and bool((np.diff(arr) == 1).all()) + specs[dim] = _dim_spec( + requested[dim], contiguous, self._monotonic[dim] + ) out_shape = tuple(len(requested[d]) for d in self._dimension_columns) if any(n == 0 for n in out_shape): @@ -333,38 +396,9 @@ def _raw_getitem(self, key: tuple) -> np.ndarray: ) return cast(np.ndarray, squeezed) - # Build a single DataFusion filter expression as the AND of per-dim - # predicates. For a single requested value: equality. For multiple: - # OR-chain of equalities (DataFusion 52.0.0 does not expose a clean - # ``Expr.in_list`` from Python; OR-chained equalities constant-fold - # equivalently and stay typed). - predicates = [] - for dim in self._dimension_columns: - if dim in full_dims: - continue - vals = requested[dim] - if len(vals) == 1: - predicates.append(col(f'"{dim}"') == literal(vals[0])) - else: - eq = col(f'"{dim}"') == literal(vals[0]) - for v in vals[1:]: - eq = eq | (col(f'"{dim}"') == literal(v)) - predicates.append(eq) - - filtered = self._inner_df - if predicates: - combined = predicates[0] - for p in predicates[1:]: - combined = combined & p - filtered = filtered.filter(combined) - projected = filtered.select( - *(col(f'"{c}"') for c in self._dimension_columns + [self._var_name]) + batches = self._handle.fetch( + specs, self._dimension_columns + [self._var_name] ) - - # Consume the projected DataFrame as Arrow RecordBatches. The - # DataFusion wrapper exposes ``.to_pyarrow()`` to convert each - # batch into a true ``pyarrow.RecordBatch``. - batches = [b.to_pyarrow() for b in projected.execute_stream()] return _scatter_batches_to_ndarray( batches=batches, dimension_columns=self._dimension_columns, @@ -376,26 +410,97 @@ def _raw_getitem(self, key: tuple) -> np.ndarray: ) -def _materialize( - inner_df: Any, +def _strictly_monotonic(coord: np.ndarray) -> bool: + """Whether ``coord`` is strictly increasing or strictly decreasing. + + Strict monotonicity of the whole coordinate array is the + precondition for translating a contiguous positional window into a + value range: with duplicated or unsorted values, ``[min, max]`` of a + window admits coordinate values at positions outside the window. + NaN/NaT (whose comparisons are all false) and non-comparable object + arrays report ``False``, which safely falls back to value lists. + """ + if len(coord) < 2: + return True + head, tail = coord[:-1], coord[1:] + try: + return bool((tail > head).all() or (tail < head).all()) + except TypeError: + return False + + +def _dim_spec( + vals: np.ndarray, contiguous: bool, coord_monotonic: bool +) -> DimSpec: + """The engine window for one dim's requested coordinate values. + + A contiguous run of positions over a strictly monotonic coordinate + array is exactly the value range ``[min, max]`` — a two-literal + predicate engines can push into range pruning. Monotonicity must + hold for the *entire* coordinate array (``coord_monotonic``), not + just the requested window: template coords are used verbatim, and + over a non-monotonic array a window's ``[min, max]`` admits values + at positions outside the window, which the scatter would then write + to wrong cells. Anything else (stepped slices, fancy indexers, + non-monotonic or duplicated coords) must be an explicit value list: + a range would admit rows the scatter did not request. + """ + if contiguous and coord_monotonic and len(vals) > 1: + return ("range", vals.min(), vals.max()) + return ("values", vals, None) + + +def _c_order_grid( + dim_cols: dict[str, np.ndarray], + coord_arrays: dict[str, np.ndarray], + dimension_columns: list[str], + total_rows: int, +) -> bool: + """Whether the result rows form the complete grid in C order. + + True iff the row count is exactly the coordinate product and every + dimension column is its coordinates repeated/tiled in C order — the + shape any unfiltered or bbox-windowed scan produces. When it holds, + data variables are dense row-major arrays already and can be + reshaped instead of scatter-written (one memcpy versus a + ``searchsorted`` per dimension per row). + """ + shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + if total_rows != int(np.prod(shape)) or total_rows == 0: + return False + for k, d in enumerate(dimension_columns): + inner = int(np.prod(shape[k + 1 :])) + outer = int(np.prod(shape[:k])) + view = dim_cols[d].reshape(outer, shape[k], inner) + if not (view == coord_arrays[d][None, :, None]).all(): + return False + return True + + +def _dataset_from_batches( + batches: list[pa.RecordBatch], 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. + """Build a dense in-memory Dataset from Arrow ``RecordBatch`` es. + + The engine-agnostic core of the eager round-trip: derives the + coordinates and every data variable from a single already-executed + result, whichever engine produced it. ``field_types`` values only + need a ``to_pandas_dtype()`` method (both ``pyarrow.DataType`` and + DataFusion's Arrow type wrappers qualify). + + Complete grid-ordered results (unfiltered scans, bbox windows) are + reshaped directly; anything else — sparse results from filtered + queries, engine-reordered rows — falls back to the positional + scatter, which handles arbitrary row order. """ - batches = [b.to_pyarrow() for b in inner_df.execute_stream()] - + dim_cols: dict[str, np.ndarray] = {} coord_arrays: dict[str, np.ndarray] = {} for d in dimension_columns: if not batches: + dim_cols[d] = np.asarray([]) coord_arrays[d] = np.asarray([]) continue vals = np.concatenate( @@ -404,6 +509,7 @@ def _materialize( for b in batches ] ) + dim_cols[d] = vals # Preserve the order coordinate values first appear in the result so an # ORDER BY direction (e.g. ``ORDER BY level DESC``) carries through to # the Dataset dimension instead of being force-sorted ascending. @@ -411,27 +517,64 @@ def _materialize( # internally, so arbitrarily-ordered coordinates are placed correctly. coord_arrays[d] = np.asarray(pd.unique(vals)) shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + total_rows = sum(b.num_rows for b in batches) + + grid_ordered = _c_order_grid( + dim_cols, coord_arrays, dimension_columns, total_rows + ) data_vars: dict[str, xr.Variable] = {} for name in field_names: if name in dimension_columns: continue np_dtype = np.dtype(field_types[name].to_pandas_dtype()) - dense = _scatter_batches_to_ndarray( - batches=batches, - dimension_columns=dimension_columns, - requested=coord_arrays, - var_name=name, - out_shape=shape, - dtype=np_dtype, - drop_axes=[], - ) + if grid_ordered: + flat = np.concatenate( + [ + b.column(b.schema.names.index(name)).to_numpy( + zero_copy_only=False + ) + for b in batches + ] + ) + dense = flat.astype(np_dtype, copy=False).reshape(shape) + else: + dense = _scatter_batches_to_ndarray( + batches=batches, + dimension_columns=dimension_columns, + requested=coord_arrays, + var_name=name, + out_shape=shape, + dtype=np_dtype, + drop_axes=[], + ) data_vars[name] = xr.Variable(dimension_columns, dense) coords_arg = {d: coord_arrays[d] for d in dimension_columns} return xr.Dataset(data_vars=data_vars, coords=coords_arg) +def _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"} @@ -486,7 +629,7 @@ def _maybe_template_coords( table and the registered Dataset carries all requested dims. Returns ``None`` otherwise so the caller falls back to per-dim discovery. Skipping discovery avoids one full plan execution per dim and - preserves the source's coordinate order (xarray-sql#171). + preserves the source's coordinate order. Coord values come from the **scanned** registered Dataset, not from any user-supplied ``template=`` (which is for metadata recovery @@ -506,41 +649,29 @@ def _maybe_template_coords( def _build_lazy_scan( - inner_df: Any, + handle: LazyResultHandle, dimension_columns: list[str], field_names: list[str], field_types: dict[str, Any], - templates: dict[str, xr.Dataset] | None = None, + coord_arrays: dict[str, np.ndarray] | None = None, ) -> xr.Dataset: - """Build a lazy Dataset whose data vars are :class:`SQLBackendArray`. + """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 - scanned table's registered Dataset (fast path, for unfiltered scans -- see - :func:`_maybe_template_coords`) or from per-dim - ``inner_df.select(col(d)).distinct().sort(...).execute_stream()``; the table - provider projects to that single coordinate column and skips data variables, - so discovery reads coordinate values only (no data-variable I/O). + once wrapped by ``Dataset.chunk``, every chunk reads its coordinate range + via a pushdown filter on first access. 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 + registered pushdown table the engine projects to that single coordinate + column, so discovery reads coordinate values only (no data-variable I/O). """ - coord_arrays = _maybe_template_coords( - templates, dimension_columns, inner_df - ) if coord_arrays is None: coord_arrays = {} for d in dimension_columns: - dim_only = ( - inner_df.select(col(f'"{d}"')) - .distinct() - .sort(col(f'"{d}"').sort()) - ) - chunks = [b.to_pyarrow() for b in dim_only.execute_stream()] - if not chunks: - coord_arrays[d] = np.asarray([]) - continue - coord_arrays[d] = np.concatenate( - [c.column(0).to_numpy(zero_copy_only=False) for c in chunks] - ) + # ``distinct`` returns engine order; sort ascending so + # 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) data_vars: dict[str, xr.Variable] = {} @@ -549,7 +680,7 @@ def _build_lazy_scan( continue np_dtype = field_types[name].to_pandas_dtype() backend = SQLBackendArray( - inner_df=inner_df, + handle=handle, var_name=name, dimension_columns=dimension_columns, coord_arrays=coord_arrays, @@ -636,14 +767,14 @@ def _result_to_xarray( ) -> xr.Dataset: """Reconstruct an ``xr.Dataset`` from a SQL result. - ``chunks`` (already resolved by :meth:`XarrayDataFrame._resolve_chunks`) + ``chunks`` (already resolved by ``XarrayDataFrame._resolve_chunks``) selects the execution strategy: * ``None`` -> eager: execute once and materialize a dense Dataset - (:func:`_materialize`). Correct for any query and the right default for + (``_materialize``). Correct for any query and the right default for reductions, whose results are small. - * a mapping (or ``"auto"``) -> lazy/chunked: build :class:`SQLBackendArray` - data variables (:func:`_build_lazy_scan`) and wrap them with + * 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 pushdown. The chunk grid maps onto the source partitions. Chunking goes through xarray's configured chunk manager (dask, cubed, ...), so no @@ -666,13 +797,35 @@ def _result_to_xarray( ds = _materialize(inner_df, dimension_columns, field_names, field_types) else: ds = _build_lazy_scan( - inner_df, + DataFusionHandle(inner_df), dimension_columns, field_names, field_types, - templates=templates, + coord_arrays=_maybe_template_coords( + templates, dimension_columns, inner_df + ), ) + return _finish_dataset( + ds, + dimension_columns, + template, + sparsity, + fill_value, + chunks, + field_types, + ) + +def _finish_dataset( + ds: xr.Dataset, + dimension_columns: list[str], + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str | None, + field_types: dict[str, Any], +) -> xr.Dataset: + """Shared reconstruction tail: sparsity, template metadata, chunking.""" if sparsity == "template": assert template is not None indexers = { @@ -708,17 +861,17 @@ def _result_to_xarray( class XarrayDataFrame: """Wrapper around a DataFusion ``DataFrame`` with xarray-aware helpers. - Returned by :meth:`xarray_sql.XarrayContext.sql`. Forwards every + Returned by [xarray_sql.XarrayContext.sql][]. Forwards every attribute it does not define itself to the wrapped DataFrame, so ``.collect()``, ``.schema()``, ``.show()``, ``.count()`` all work unchanged. Carries a private snapshot of the context's registered Datasets so - :meth:`to_dataset` can default ``dims`` and recover metadata + ``to_dataset`` can default ``dims`` and recover metadata dropped by the forward pivot. Users should not construct this class directly; let - :meth:`XarrayContext.sql` produce it. + [XarrayContext.sql][xarray_sql.sql.XarrayContext.sql] produce it. """ def __init__( @@ -730,13 +883,13 @@ def __init__( Args: inner: The underlying ``datafusion.DataFrame`` returned by - :meth:`XarrayContext.sql`. + [XarrayContext.sql][xarray_sql.sql.XarrayContext.sql]. templates: Snapshot of the registered Datasets on the producing context, keyed by the SQL identifier each was registered - under. Used by :meth:`to_dataset` to recover metadata that + under. Used by ``to_dataset`` to recover metadata that the forward pivot strips. ``None`` means no metadata recovery is possible from registrations alone; callers may - still pass ``template=`` to :meth:`to_dataset` explicitly. + still pass ``template=`` to ``to_dataset`` explicitly. """ object.__setattr__(self, "_inner", inner) object.__setattr__(self, "_templates", dict(templates or {})) @@ -840,7 +993,7 @@ def _resolve_chunks( (a single full chunk is not "chunked"), so reductions that drop the chunked dimension resolve to ``None`` (eager) automatically. Mappings pass through unchanged; ``"auto"`` passes through here and is snapped to - source partition boundaries later (see :func:`_auto_chunks`). + source partition boundaries later (see ``_auto_chunks``). """ if chunks is None: return None @@ -885,7 +1038,7 @@ def _infer_dimension_columns( become the dimensions, so aggregations that drop dims (e.g. ``GROUP BY time`` over a ``(time, lat, lon)`` grid) round-trip on the surviving dim(s). Uses the data variable's dim order (via - :func:`_ds_var_dims`) so the original axis order is preserved. + ``_ds_var_dims``) so the original axis order is preserved. """ result_cols = set(self._result_columns()) diff --git a/xarray_sql/geometry.py b/xarray_sql/geometry.py new file mode 100644 index 0000000..6e3e640 --- /dev/null +++ b/xarray_sql/geometry.py @@ -0,0 +1,131 @@ +"""GeoArrow point-geometry columns derived from coordinate dimensions. + +A regular grid's pivot already materializes per-row x/y coordinate +columns; a point-geometry column is those same values under a GeoArrow +extension annotation. Two encodings: + +* ``"wkb"`` (default) — 21-byte WKB points under the ``geoarrow.wkb`` + extension name. DuckDB (>= 1.2, spatial loaded) ingests the column as + a native ``GEOMETRY`` with the CRS attached, so ``ST_Within(geometry, + ...)`` works with no ``ST_Point(x, y)`` construction in user SQL. +* ``"point"`` — GeoArrow native points with *separated* coordinates + (``struct`` under ``geoarrow.point``): the + child arrays are the coordinate columns themselves, no per-row + parsing for consumers that execute on native layouts (GeoPandas 1.x, + geoarrow-rs, lonboard, SedonaDB). DuckDB does not consume this + encoding. + +The CRS rides in the extension metadata (GeoArrow 0.2 allows +authority:code strings alongside PROJJSON). ``OGC:CRS84`` is the +correct tag for plain longitude/latitude grids. +""" + +from __future__ import annotations + +import json +from typing import Any + +import numpy as np +import pyarrow as pa + +GEOMETRY_COLUMN = "geometry" + +_ENCODINGS = ("wkb", "point") + + +def geometry_field(encoding: str, crs: str | None) -> pa.Field: + """The schema field for the derived geometry column.""" + if encoding not in _ENCODINGS: + raise ValueError( + f"geometry_encoding must be one of {_ENCODINGS}, got {encoding!r}" + ) + metadata = { + b"ARROW:extension:name": f"geoarrow.{encoding}".encode(), + } + if crs is not None: + metadata[b"ARROW:extension:metadata"] = json.dumps( + {"crs": crs} + ).encode() + storage = ( + pa.binary() + if encoding == "wkb" + else pa.struct([("x", pa.float64()), ("y", pa.float64())]) + ) + return pa.field(GEOMETRY_COLUMN, storage, metadata=metadata) + + +def build_geometry(encoding: str, x: pa.Array, y: pa.Array) -> pa.Array: + """Point geometries for one batch's x/y coordinate columns.""" + if encoding == "point": + return pa.StructArray.from_arrays( + [x.cast(pa.float64()), y.cast(pa.float64())], ["x", "y"] + ) + return _wkb_points( + np.ascontiguousarray(x.to_numpy(zero_copy_only=False), " pa.Array: + """Vectorized 21-byte little-endian WKB point encoding.""" + n = len(x) + # ``pa.binary()`` carries int32 offsets, which the final offset + # (n * 21) overflows past ~102M points; the buffers would build + # silently corrupt. Unreachable through the pivot (batch_size caps + # rows per batch well below this), so guard rather than widen the + # storage to large_binary, which DuckDB's ingestion expects not to + # see. + if n * 21 > np.iinfo(np.int32).max: + raise ValueError( + f"cannot WKB-encode {n:,} points in a single batch: " + "pa.binary() offsets are int32 and n * 21 bytes would " + "overflow them. Use a smaller batch_size." + ) + buf = np.empty((n, 21), dtype=np.uint8) + buf[:, 0] = 1 # little-endian byte order mark + buf[:, 1:5] = np.array([1, 0, 0, 0], dtype=np.uint8) # WKB type 1: Point + buf[:, 5:13] = x.view(np.uint8).reshape(n, 8) + buf[:, 13:21] = y.view(np.uint8).reshape(n, 8) + offsets = pa.py_buffer( + np.arange(0, (n + 1) * 21, 21, dtype=np.int32).tobytes() + ) + return pa.Array.from_buffers( + pa.binary(), n, [None, offsets, pa.py_buffer(buf.tobytes())] + ) + + +def bbox_conjuncts( + bounds: Any, x: str = "x", y: str = "y", pad: float = 0.0 +) -> str: + """SQL bbox conjuncts for a geometry's envelope — the pruning half. + + Engines do not push ``ST_*`` functions into the scan, so a + geometry-only predicate reads every chunk; pairing it with range + conjuncts on the coordinate columns restores pruning. This helper + renders those conjuncts from a geometry's envelope:: + + poly = shapely.from_wkt("POLYGON (...)") + sql = ( + f"SELECT avg(risk) FROM eri " + f"WHERE {xql.bbox_conjuncts(poly, x='x', y='y')} " + f"AND ST_Within(geometry, ST_GeomFromText('{poly.wkt}'))" + ) + + Args: + bounds: ``(xmin, ymin, xmax, ymax)``, or any object with a + ``.bounds`` attribute in that convention (shapely + geometries qualify). + x: The x/longitude column name. + y: The y/latitude column name. + pad: Optional margin added on every side (e.g. to be safe + around ``ST_DWithin``-style predicates). + + Returns: + A SQL snippet ``"x" BETWEEN a AND b AND "y" BETWEEN c AND d``. + """ + values = getattr(bounds, "bounds", bounds) + xmin, ymin, xmax, ymax = (float(v) for v in values) + return ( + f'"{x}" BETWEEN {xmin - pad!r} AND {xmax + pad!r} ' + f'AND "{y}" BETWEEN {ymin - pad!r} AND {ymax + pad!r}' + ) diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py new file mode 100644 index 0000000..cbf0551 --- /dev/null +++ b/xarray_sql/lazyscan.py @@ -0,0 +1,369 @@ +"""Re-executable engine handles behind the lazy chunked round-trip. + +The lazy path of ``to_dataset(chunks=...)`` re-executes the engine's +query per accessed chunk, narrowed to that chunk's coordinate window and +columns. That requires the engine result to be *re-executable* — a +handle onto the query, not a one-shot stream of its rows. Each handle +here adapts one engine's native lazy surface to the three operations the +reconstruction needs: + +* [schema][xarray_sql.lazyscan.LazyResultHandle.schema] — result column names/types, without + executing the query; +* [distinct][xarray_sql.lazyscan.LazyResultHandle.distinct] — one column's distinct values + (coordinate discovery; the caller sorts); +* [fetch][xarray_sql.lazyscan.LazyResultHandle.fetch] — the result narrowed by per-dimension + windows and projected to the requested columns, as Arrow batches. + +Windows are passed as [DimSpec][xarray_sql.lazyscan.DimSpec] values instead of rendered SQL so +each engine can express them with its own *typed* expression API — +strings would re-open every literal-formatting pitfall (timestamps, +floats, quoting) per dialect. + +Handles compose with the registration seam: when the wrapped query +scans a Dataset registered through xarray-sql's pushdown machinery, the +per-chunk range filter flows back through the engine into +[XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset], so each +output chunk's access reads only the source chunks it maps onto. +""" + +from __future__ import annotations + +import weakref +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Literal, Protocol, cast + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from datafusion import col, literal + +DimSpec = tuple[Literal["range", "values"], Any, Any] +"""One dimension's window: ``("range", lo, hi)`` (inclusive bounds; the +requested coordinate positions are contiguous) or ``("values", array, +None)`` (explicit value list, for stepped/fancy indexers). + +The payload stays ``Any``: the values are coordinate scalars handed to +the engine's typed expression API, which does the comparing — Python +never orders them, and a concrete union over coordinate dtypes +(timestamps, cftime, numerics, strings) would stay incomplete.""" + + +def _collect_streaming(lf: Any) -> Any: + """Collect a Polars LazyFrame on the streaming engine, if available. + + ``collect(engine="streaming")`` needs polars >= 1.25 (the test + extra pins higher); an older installed polars raises TypeError on + the unknown keyword, and the plain in-memory collect is a correct, + if less memory-frugal, fallback. + """ + try: + return lf.collect(engine="streaming") + except TypeError: + return lf.collect() + + +def _plain(value: Any) -> Any: + """A plain-Python literal (numpy scalars don't travel to engines).""" + if isinstance(value, np.datetime64): + return pd.Timestamp(value) + if isinstance(value, np.timedelta64): + return pd.Timedelta(value) + if isinstance(value, np.generic): + return value.item() + return value + + +class LazyResultHandle(Protocol): + """A re-executable query result (see module docstring).""" + + supports_chunked: bool = True + """Whether fetch() may be driven from consumer worker threads (the + chunked reconstruction). Handles for engines that cannot safely + re-execute under foreign threads set this False; the eager path + remains available.""" + + def schema(self) -> pa.Schema: ... + + def distinct(self, column: str) -> np.ndarray: ... + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: ... + + def spill_parquet(self, path: str) -> None: ... + + # Handles may additionally offer ``stream(columns)``, yielding the + # unfiltered result's Arrow batches incrementally instead of + # materializing it first. The eager round-trip uses it to enforce + # ``max_result_bytes`` while collecting; handles without it refuse + # the budget rather than blow past it after collecting. + + +class DataFusionHandle: + """Handle over a ``datafusion.DataFrame``.""" + + supports_chunked = True + + def __init__(self, df: Any) -> None: + self._df = df + + def schema(self) -> pa.Schema: + return self._df.schema() + + def distinct(self, column: str) -> np.ndarray: + dim_only = self._df.select(col(f'"{column}"')).distinct() + batches = [b.to_pyarrow() for b in dim_only.execute_stream()] + if not batches: + return np.asarray([]) + return np.concatenate( + [b.column(0).to_numpy(zero_copy_only=False) for b in batches] + ) + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + predicate = None + for dim, (kind, a, b) in specs.items(): + c = col(f'"{dim}"') + if kind == "range": + p = (c >= literal(a)) & (c <= literal(b)) + else: + # DataFusion 52.0.0 exposes no clean ``Expr.in_list`` + # from Python; OR-chained equalities constant-fold + # equivalently and stay typed. + p = c == literal(a[0]) + for v in a[1:]: + p = p | (c == literal(v)) + predicate = p if predicate is None else predicate & p + out = self._df if predicate is None else self._df.filter(predicate) + out = out.select(*(col(f'"{n}"') for n in columns)) + return [b.to_pyarrow() for b in out.execute_stream()] + + def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: + """Execute once, yielding Arrow batches as the plan produces them.""" + out = self._df.select(*(col(f'"{n}"') for n in columns)) + return (b.to_pyarrow() for b in out.execute_stream()) + + def spill_parquet(self, path: str) -> None: + with pq.ParquetWriter(path, self.schema()) as writer: + for batch in self._df.execute_stream(): + writer.write_batch(batch.to_pyarrow()) + + +class DuckDBHandle: + """Handle over a ``duckdb.DuckDBPyRelation``. + + Relations are lazy relational algebra: ``filter``/``project`` derive + new relations and every materialization runs the query again, so a + single relation can serve any number of per-chunk fetches, each + narrowed to its own window — the *re-executable* property the + module docstring requires of every handle. Predicates are built + with DuckDB's typed expression API, never rendered SQL text. + + Every engine call runs on one dedicated thread owned by the handle. + A relation is bound to one connection, and a query over a table + registered through xarray-sql re-enters Python from DuckDB's + execution threads (the Arrow scan callback); driving such queries + directly from several consumer threads at once (dask computing + output chunks of a lazy round-trip) deadlocks between the + connection's serialization, the callback's need for the GIL, and + the consumer pool's own thread management. Funnelling execution + through a single pre-started thread reproduces the topology that is + known safe — one thread inside the engine, every other thread + parked on a GIL-releasing wait. + """ + + supports_chunked = False + """Chunked (lazy) reconstruction is disabled for DuckDB relations. + + Windows of a chunked round-trip re-execute the relation from the + consumer's worker threads (dask). A DuckDB query whose source is a + Python-callback Arrow scan (any table registered through xarray-sql) + intermittently deadlocks inside duckdb-python/CPython when other + Python threads start or stop during execution — reproduced on + duckdb 1.4-1.5 / CPython 3.12 / macOS at ~50% of runs, regardless + of ``SET threads=1``, connection serialization, or pool pre-warming. + Until that upstream race is fixed, chunked DuckDB round-trips fail + fast instead of hanging; the eager path (and every other handle + operation) runs on one dedicated thread and is unaffected. + """ + + def __init__(self, rel: Any) -> None: + self._rel = rel + 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 + # would otherwise linger for the life of the process, one per + # discarded handle. The callback is bound to the executor, not + # the handle, so the finalizer holds no reference back to self. + weakref.finalize( + self, self._runner.shutdown, wait=False, cancel_futures=True + ) + + def _run(self, fn: Any) -> Any: + return self._runner.submit(fn).result() + + @staticmethod + def _to_arrow_table(rel: Any) -> pa.Table: + if hasattr(rel, "to_arrow_table"): + return rel.to_arrow_table() + return rel.fetch_arrow_table() # duckdb < 1.5 + + @staticmethod + def _to_arrow_reader(rel: Any) -> pa.RecordBatchReader: + if hasattr(rel, "to_arrow_reader"): + return rel.to_arrow_reader() + return rel.fetch_record_batch() # duckdb < 1.5 + + def schema(self) -> pa.Schema: + return self._run( + lambda: self._to_arrow_table(self._rel.limit(0)).schema + ) + + def distinct(self, column: str) -> np.ndarray: + import duckdb + + table = self._run( + lambda: self._to_arrow_table( + self._rel.project(duckdb.ColumnExpression(column)).distinct() + ) + ) + return np.asarray(table.column(0).to_numpy(zero_copy_only=False)) + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + import duckdb + + predicate = None + for dim, (kind, a, b) in specs.items(): + c = duckdb.ColumnExpression(dim) + if kind == "range": + p = (c >= duckdb.ConstantExpression(_plain(a))) & ( + c <= duckdb.ConstantExpression(_plain(b)) + ) + else: + p = c.isin(*(duckdb.ConstantExpression(_plain(v)) for v in a)) + predicate = p if predicate is None else predicate & p + rel = self._rel if predicate is None else self._rel.filter(predicate) + rel = rel.project(*(duckdb.ColumnExpression(n) for n in columns)) + + return cast( + list[pa.RecordBatch], + self._run(lambda: list(self._to_arrow_reader(rel))), + ) + + def spill_parquet(self, path: str) -> None: + def run() -> None: + reader = self._to_arrow_reader(self._rel) + with pq.ParquetWriter(path, reader.schema) as writer: + for batch in reader: + writer.write_batch(batch) + + self._run(run) + + +class PolarsHandle: + """Handle over a ``polars.LazyFrame``. + + Per-window fetches run on the streaming engine, so a window read + never materializes more than the window even when the frame scans + an out-of-core source. + """ + + supports_chunked = True + + def __init__(self, lf: Any) -> None: + self._lf = lf + + def schema(self) -> pa.Schema: + import polars as pl + + return pl.DataFrame(schema=self._lf.collect_schema()).to_arrow().schema + + def distinct(self, column: str) -> np.ndarray: + import polars as pl + + out = _collect_streaming(self._lf.select(pl.col(column).unique())) + return out.to_series().to_numpy() + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + import polars as pl + + exprs = [] + for dim, (kind, a, b) in specs.items(): + if kind == "range": + exprs.append(pl.col(dim).is_between(_plain(a), _plain(b))) + elif getattr(a, "dtype", None) is not None and a.dtype.kind == "f": + # Upstream Polars translates float ``is_in`` literals + # imprecisely (silently matching nothing); degenerate + # ranges compare exactly. Reproduced on polars 1.42. + # ``any_horizontal`` keeps the disjunction flat — a + # left-deep OR chain plans quadratically in the number + # of values. + exprs.append( + pl.any_horizontal( + [pl.col(dim).is_between(*(_plain(v),) * 2) for v in a] + ) + ) + else: + exprs.append(pl.col(dim).is_in([_plain(v) for v in a])) + lf = self._lf.filter(*exprs) if exprs else self._lf + out = _collect_streaming(lf.select([pl.col(n) for n in columns])) + return cast(list[pa.RecordBatch], out.to_arrow().to_batches()) + + def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: + """Execute once, yielding Arrow batches incrementally. + + Unlike [fetch][xarray_sql.lazyscan.PolarsHandle.fetch], whose ``collect`` materializes the whole + result inside the engine before any batch surfaces, this yields + batches as the streaming engine produces them, so a byte budget + can fire before the result is fully in memory. Requires + ``LazyFrame.collect_batches`` (polars >= 1.33); older Polars + raises here, before anything is collected. + """ + import polars as pl + + lf = self._lf.select([pl.col(n) for n in columns]) + if not hasattr(lf, "collect_batches"): + raise ValueError( + "streaming collection of a Polars LazyFrame requires " + "polars >= 1.33 (LazyFrame.collect_batches)." + ) + + def generate() -> Iterator[pa.RecordBatch]: + for frame in lf.collect_batches(engine="streaming"): + yield from frame.to_arrow().to_batches() + + return generate() + + def spill_parquet(self, path: str) -> None: + self._lf.sink_parquet(path) + + +def resolve_lazy_handle(result: Any) -> LazyResultHandle | None: + """Adapt an engine result to a handle, or ``None`` if it is one-shot. + + Recognizes DuckDB relations, Polars lazy *and* eager frames (an + eager frame re-executes trivially over its in-memory data), and + DataFusion DataFrames. ``pyarrow`` tables/readers and bare + ``__arrow_c_stream__`` objects are one-shot streams: there is no + query to re-execute, so the lazy path cannot serve them. + """ + root = type(result).__module__.split(".")[0] + if root in ("duckdb", "_duckdb") and hasattr(result, "filter"): + return DuckDBHandle(result) + if root == "polars": + import polars as pl + + if isinstance(result, pl.LazyFrame): + return PolarsHandle(result) + if isinstance(result, pl.DataFrame): + return PolarsHandle(result.lazy()) + if hasattr(result, "execute_stream") and hasattr(result, "logical_plan"): + return DataFusionHandle(result) + return None diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py index 2d7f74f..2470d1b 100644 --- a/xarray_sql/proj.py +++ b/xarray_sql/proj.py @@ -42,8 +42,8 @@ would return ``inf``); NULL CRS arguments yield NaN as well. Requires ``pyproj`` (``pip install xarray-sql[geo]``). When pyproj is -installed, :class:`xarray_sql.XarrayContext` registers ``reproject()`` -automatically; :func:`register` is the explicit hook for plain +installed, [xarray_sql.XarrayContext][] registers ``reproject()`` +automatically; [register][xarray_sql.proj.register] is the explicit hook for plain DataFusion ``SessionContext`` objects or custom UDF names. """ @@ -61,10 +61,10 @@ __all__ = ["register"] -#: Arrow type returned by ``reproject()``: destination coordinates in -#: ``always_xy`` order — ``x`` is easting/longitude, ``y`` is -#: northing/latitude. RETURN_TYPE = pa.struct([("x", pa.float64()), ("y", pa.float64())]) +"""Arrow type returned by ``reproject()``: destination coordinates in +``always_xy`` order — ``x`` is easting/longitude, ``y`` is +northing/latitude.""" # --------------------------------------------------------------------------- diff --git a/xarray_sql/roundtrip.py b/xarray_sql/roundtrip.py new file mode 100644 index 0000000..e342828 --- /dev/null +++ b/xarray_sql/roundtrip.py @@ -0,0 +1,495 @@ +"""Engine-agnostic round-trip: Arrow query results → labeled ``xr.Dataset``. + +The second seam of xarray-sql. Any engine's result — a DuckDB relation, +a ``pyarrow.Table``, a ``pyarrow.RecordBatchReader``, or any object +implementing the Arrow PyCapsule stream protocol — plus the registered +Dataset as a *template* is enough to rebuild a labeled, metadata-carrying +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 +[DuckDBHandle][xarray_sql.lazyscan.DuckDBHandle]); pair them with +``spill=True`` instead. +""" + +from __future__ import annotations + +import os +import tempfile +import weakref +from collections.abc import Mapping +from typing import Any, Literal + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq +import xarray as xr + +from .ds import ( + Sparsity, + XarrayDataFrame, + _build_lazy_scan, + _dataset_from_batches, + _ds_var_dims, + _finish_dataset, +) +from .lazyscan import LazyResultHandle, PolarsHandle, resolve_lazy_handle + + +def _guarded(batches: Any, max_bytes: int | None) -> list[pa.RecordBatch]: + """Collect a batch iterable, erroring cleanly past ``max_bytes``. + + A result that would blow past the budget raises with the running + size instead of exhausting memory, before the (larger) dense + reconstruction is even attempted. + """ + if max_bytes is None: + return list(batches) + out: list[pa.RecordBatch] = [] + total = 0 + for batch in batches: + total += batch.nbytes + if total > max_bytes: + raise ValueError( + f"result exceeded max_result_bytes={max_bytes:,} while " + f"materializing (>= {total:,} bytes after " + f"{sum(b.num_rows for b in out) + batch.num_rows:,} rows). " + "Aggregate further, or reconstruct lazily with chunks=." + ) + out.append(batch) + return out + + +def _open_stream(result: Any) -> tuple[pa.Schema, Any] | None: + """The result's Arrow batches as ``(schema, iterable)``, or ``None``. + + Probes, in order: ``pyarrow.Table`` / ``pyarrow.RecordBatch``, + ``pyarrow.RecordBatchReader``, ``__arrow_c_stream__`` (the Arrow + PyCapsule protocol — DuckDB relations qualify on duckdb >= 1.1), and + a ``fetch_record_batch()`` method (DuckDB relations on older + versions). + """ + if isinstance(result, pa.RecordBatch): + return result.schema, [result] + if isinstance(result, pa.Table): + return result.schema, result.to_batches() + if isinstance(result, pa.RecordBatchReader): + return result.schema, result + if hasattr(result, "__arrow_c_stream__"): + reader = pa.RecordBatchReader.from_stream(result) + return reader.schema, reader + if hasattr(result, "fetch_record_batch"): + reader = result.fetch_record_batch() + return reader.schema, reader + return None + + +def _result_to_batches( + result: Any, max_bytes: int | None = None +) -> tuple[pa.Schema, list[pa.RecordBatch]]: + """Normalize an engine result into ``(schema, record batches)``. + + Accepts everything ``_open_stream`` recognizes, then objects + with a ``to_arrow_table()`` method (DataFusion DataFrames and the + [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapper), then re-executable + results without a stream protocol (a Polars LazyFrame), executed + once through their lazy handle. + """ + opened = _open_stream(result) + if opened is not None: + schema, batches = opened + if isinstance(result, (pa.RecordBatch, pa.Table)): + # Already in memory: nothing left for the budget to bound + # (the dense-size check still applies downstream). + return schema, list(batches) + return schema, _guarded(batches, max_bytes) + handle = resolve_lazy_handle(result) + if hasattr(result, "to_arrow_table") and ( + max_bytes is None or handle is None + ): + # This branch materializes the whole result in one call before + # the budget can observe a single batch, so with a budget set a + # re-executable result streams through its handle below instead; + # the post-materialization nbytes check is the fallback guard + # for one-shot results whose only surface is to_arrow_table(). + table = result.to_arrow_table() + if max_bytes is not None and table.nbytes > max_bytes: + raise ValueError( + f"result materialized to {table.nbytes:,} bytes, over " + f"max_result_bytes={max_bytes:,}. Aggregate further, or " + "reconstruct lazily with chunks=." + ) + return table.schema, table.to_batches() + if handle is not None: + schema = handle.schema() + names = list(schema.names) + if max_bytes is None: + return schema, handle.fetch({}, names) + # fetch() may materialize the whole result inside the engine + # before any batch surfaces (Polars collect()), which would + # defeat the budget; enforce it on a true batch stream, or + # refuse up front instead of erroring after the memory is spent. + stream = getattr(handle, "stream", None) + if stream is None: + raise ValueError( + "max_result_bytes cannot be enforced for " + f"{type(result).__qualname__}: the result materializes " + "fully before batches surface. Drop max_result_bytes=, " + "or reconstruct lazily with chunks=." + ) + return schema, _guarded(stream(names), max_bytes) + raise TypeError( + f"Cannot read an Arrow stream from {type(result).__qualname__}; " + "expected a pyarrow Table/RecordBatch/RecordBatchReader, an object " + "implementing __arrow_c_stream__, or an engine result exposing " + "fetch_record_batch()/to_arrow_table()." + ) + + +def to_dataset( + result: Any, + dims: list[str] | None = None, + template: xr.Dataset | None = None, + sparsity: Sparsity = "result", + fill_value: Any = np.nan, + chunks: Mapping[str, int] | str | None = None, + coords: Literal["discover", "template"] = "discover", + max_result_bytes: int | None = None, + spill: bool | str | os.PathLike = False, +) -> xr.Dataset: + """Convert an engine's Arrow result into a labeled ``xr.Dataset``. + + The engine-agnostic counterpart of + [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset]: SQL in, array out, for engines + xarray-sql does not wrap in a session of its own. + + Example (DuckDB):: + + con = duckdb.connect() + xql.register(con, "era5", ds) + rel = con.sql( + "SELECT time, lat, lon, AVG(t2m) AS t2m FROM era5 " + "GROUP BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + Args: + result: The engine's query result: a ``pyarrow.Table``, + ``RecordBatch`` or ``RecordBatchReader``, any object + implementing ``__arrow_c_stream__`` (DuckDB relations), or an + object with ``fetch_record_batch()`` / ``to_arrow_table()``. + The result is consumed once. + dims: Result columns to use as Dataset dimensions. When ``None``, + defaults to the ``template``'s dimensions that survive into + the result columns (so aggregations that drop dims round-trip + on the remaining ones). Either ``dims`` or ``template`` must + be given. + template: The source Dataset registered with the engine. Recovers + metadata the tabular pivot strips (attrs, encoding, non-dim + coordinates, dim-coord dtype) and provides the ``dims`` + default. + sparsity: ``"result"`` (default) keeps only dim values present in + 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 + (e.g. ``{"time": 100}``), ``"auto"``, or ``"inherit"`` + selects the lazy/chunked path: data variables are + reconstructed window by window on access, each window + re-executing the engine's query narrowed to its coordinate + range (over a table registered through xarray-sql, that + filter flows back into chunk pruning at the source). + Requires a re-executable ``result`` — a Polars + LazyFrame/DataFrame or a DataFusion DataFrame; DuckDB + relations refuse the chunked path (add ``spill=True``). + coords: How the lazy path learns each dimension's coordinate + values. ``"discover"`` (default) runs one ``DISTINCT`` query + per dim — correct for any query. ``"template"`` trusts the + 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 + 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 — + instead of exhausting memory. ``None`` (default) means + unlimited. Results whose only surface is + ``to_arrow_table()`` necessarily materialize in full before + the budget can be checked (the check then runs on the + materialized size); re-executable results stream instead, + so the budget fires before full materialization. + spill: Chunked reconstruction from a one-pass on-disk spill + instead of per-window re-execution: the result is streamed + *once* (bounded memory) into a temporary Parquet file, and + windows re-execute against that file. This serves the two + results the re-execution path cannot — DuckDB relations and + one-shot Arrow streams — and trades per-window narrowness + for a single full pass plus temporary disk. ``True`` spills + to the system temp dir; a path spills into that directory. + The file is removed when the returned Dataset is garbage + collected. Requires Polars; only valid with ``chunks=``. + + Returns: + An ``xr.Dataset`` with ``dims`` as dimensions and the remaining + result columns as data variables — dense and in-memory by + default, lazily chunked when ``chunks`` is given. + + Raises: + ValueError: When neither ``dims`` nor ``template`` resolves the + dimension columns, a requested dim is missing from the result, + or ``sparsity="template"`` is used without a template. + TypeError: When ``result`` exposes no readable Arrow stream, or + ``chunks`` is requested for a one-shot stream that cannot be + re-executed. + """ + if sparsity not in ("result", "template"): + raise ValueError( + f"sparsity must be 'result' or 'template', got {sparsity!r}" + ) + if sparsity == "template" and template is None: + raise ValueError("sparsity='template' requires template= to be given") + if coords not in ("discover", "template"): + raise ValueError( + f"coords must be 'discover' or 'template', got {coords!r}" + ) + if coords == "template" and template is None: + raise ValueError("coords='template' requires template= to be given") + if spill and chunks is None: + raise ValueError( + "spill= only applies to chunked reconstruction; pass chunks=." + ) + + if chunks is not None: + if spill: + return _to_dataset_spilled( + result, + dims, + template, + sparsity, + fill_value, + chunks, + coords, + spill, + ) + return _to_dataset_lazy( + result, dims, template, sparsity, fill_value, chunks, coords + ) + + schema, batches = _result_to_batches(result, max_result_bytes) + field_names = [f.name for f in schema] + field_types = {f.name: f.type for f in schema} + + dims = _resolve_dims(dims, template, field_names) + + if max_result_bytes is not None: + _check_dense_size( + batches, dims, field_names, field_types, max_result_bytes + ) + ds = _dataset_from_batches(batches, dims, field_names, field_types) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, None, field_types + ) + + +def _check_dense_size( + batches: list[pa.RecordBatch], + dims: list[str], + field_names: list[str], + field_types: dict[str, Any], + max_bytes: int, +) -> None: + """Error before allocating dense arrays larger than the budget. + + The dense grid is the coordinate product, which for sparse results + can dwarf the Arrow input; check it against the same budget before + a single output array is allocated. + """ + sizes = [] + for d in dims: + # Vectorized distinct count: a per-row Python set (to_pylist) + # costs orders of magnitude more on wide results. + arrays = [b.column(b.schema.names.index(d)) for b in batches] + sizes.append(len(pc.unique(pa.chunked_array(arrays))) if arrays else 0) + cells = int(np.prod(sizes)) if sizes else 0 + total = sum( + cells * np.dtype(field_types[n].to_pandas_dtype()).itemsize + for n in field_names + if n not in dims + ) + if total > max_bytes: + raise ValueError( + f"dense reconstruction needs {total:,} bytes " + f"({cells:,} grid cells), over max_result_bytes=" + f"{max_bytes:,}. Aggregate further, or reconstruct lazily " + "with chunks=." + ) + + +def _resolve_dims( + dims: list[str] | None, + template: xr.Dataset | None, + field_names: list[str], +) -> list[str]: + """Dimension columns, inferred from the template when not given.""" + if dims is None: + if template is None: + raise ValueError( + "dims cannot be inferred without a template; pass " + "dims=[...] or template=." + ) + dims = [d for d in _ds_var_dims(template) if d in field_names] + if not dims: + raise ValueError( + "dims cannot be inferred: no template dimension survives " + "in the result columns. Pass dims=[...] explicitly." + ) + missing = [d for d in dims if d not in field_names] + if missing: + raise ValueError( + f"dims {missing} are not columns of the result {field_names}." + ) + return dims + + +def _to_dataset_lazy( + result: Any, + dims: list[str] | None, + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str, + coords: Literal["discover", "template"], + _handle: LazyResultHandle | None = None, +) -> xr.Dataset: + """The chunked reconstruction behind ``to_dataset(chunks=...)``.""" + 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 " + 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." + ) + schema = handle.schema() + field_names = [f.name for f in schema] + field_types = {f.name: f.type for f in schema} + dims = _resolve_dims(dims, template, field_names) + + coord_arrays = None + if coords == "template": + assert template is not None + missing = [d for d in dims if d not in template.coords] + if missing: + raise ValueError( + f"coords='template' requires the template to carry coords " + f"for every dim; missing {missing}." + ) + coord_arrays = {d: np.asarray(template.coords[d].values) for d in dims} + + resolved = XarrayDataFrame._resolve_chunks(chunks, template, dims) + if resolved is None: + # "inherit" with no chunked source dimension to inherit from: + # eager is the right execution, exactly as on the wrapper path. + batches = handle.fetch({}, field_names) + ds = _dataset_from_batches(batches, dims, field_names, field_types) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, None, field_types + ) + if not getattr(handle, "supports_chunked", True): + raise NotImplementedError( + "Chunked reconstruction is not supported for " + f"{type(result).__qualname__}: re-executing a DuckDB " + "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 " + "Polars (pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))) " + "or a DataFusion context." + ) + ds = _build_lazy_scan( + handle, dims, field_names, field_types, coord_arrays=coord_arrays + ) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, resolved, field_types + ) + + +def _to_dataset_spilled( + result: Any, + dims: list[str] | None, + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str, + coords: Literal["discover", "template"], + spill: bool | str | os.PathLike, +) -> xr.Dataset: + """Chunked reconstruction from a one-pass temporary Parquet spill. + + The result is streamed exactly once with bounded memory — through + the engine handle where one exists (DuckDB spills on its dedicated + engine thread; Polars uses its streaming sink), or straight from + the Arrow stream for one-shot results — and the ordinary lazy + reconstruction then runs against a Polars scan of the file, whose + per-window predicates enjoy Parquet row-group pruning. The file is + removed when the reconstruction handle is garbage collected. + """ + import polars as pl + + directory = os.fspath(spill) if not isinstance(spill, bool) else None + fd, path = tempfile.mkstemp(suffix=".parquet", dir=directory) + os.close(fd) + try: + handle = resolve_lazy_handle(result) + if handle is not None: + handle.spill_parquet(path) + else: + _stream_to_parquet(result, path) + except BaseException: + os.unlink(path) + raise + spilled = PolarsHandle(pl.scan_parquet(path)) + weakref.finalize(spilled, _unlink_quietly, path) + return _to_dataset_lazy( + result, + dims, + template, + sparsity, + fill_value, + chunks, + coords, + _handle=spilled, + ) + + +def _unlink_quietly(path: str) -> None: + try: + os.unlink(path) + except OSError: + pass + + +def _stream_to_parquet(result: Any, path: str) -> None: + """Write a one-shot Arrow result to Parquet, batch by batch.""" + opened = _open_stream(result) + if opened is None: + raise TypeError( + f"cannot spill {type(result).__qualname__}: no readable " + "Arrow stream." + ) + schema, batches = opened + with pq.ParquetWriter(path, schema) as writer: + for batch in batches: + writer.write_batch(batch) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index eb3655e..a5df4fb 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,11 +1,10 @@ import xarray as xr from datafusion import SessionContext from datafusion.catalog import Schema -from collections import defaultdict from types import ModuleType from . import cftime as cft -from .df import Chunks +from .df import Chunks, group_vars_by_dims from .ds import XarrayDataFrame from .reader import read_xarray_table @@ -100,7 +99,7 @@ def from_dataset( Returns: self, to allow chaining. """ - groups = _group_vars_by_dims(input_table) + groups = group_vars_by_dims(input_table) # Materialise dim coordinates once and share across every sub-table. # For Zarr-backed parents (e.g. ARCO-ERA5 on GCS) this saves one @@ -170,7 +169,7 @@ def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None: break # One UDF per context is enough. def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: - """Run a SQL query, returning an :class:`XarrayDataFrame` wrapper. + """Run a SQL query, returning an [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapper. Identical to ``datafusion.SessionContext.sql`` except the returned object wraps the DataFusion DataFrame. The wrapper exposes @@ -185,20 +184,7 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: **kwargs: Forwarded to ``SessionContext.sql``. Returns: - An :class:`XarrayDataFrame` wrapping the DataFusion DataFrame. + An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame. """ inner = super().sql(query, *args, **kwargs) return XarrayDataFrame(inner, templates=self._registered_datasets) - - -def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: - """Group variables in the dataset based on shared dims. - - ("time", "lat", "lon"): ["temperature_2m", "wind_speed"], - ("time", "lat", "lon", "level"): ["pressure", "humidity"] - """ - groups = defaultdict(list) - for var_name, var in ds.data_vars.items(): - dims = var.dims - groups[dims].append(var_name) - return groups diff --git a/zensical.toml b/zensical.toml index f182a70..b80208b 100644 --- a/zensical.toml +++ b/zensical.toml @@ -2,16 +2,25 @@ site_name = "xarray-sql" site_description = "Query Xarray with SQL" site_author = "Alexander Merose" -site_url = "https://alxmrs.github.io/xarray-sql" -repo_url = "https://github.com/alxmrs/xarray-sql" -repo_name = "alxmrs/xarray-sql" +site_url = "https://xqlsystems.github.io/xarray-sql" +repo_url = "https://github.com/xqlsystems/xarray-sql" +repo_name = "xqlsystems/xarray-sql" edit_uri = "edit/main/docs/" nav = [ {"Home" = "index.md"}, - {"Examples" = "examples.md"}, - {"Geospatial in SQL" = "geospatial.md"}, - {"Contributing" = "contributing.md"}, - {"Reference" = "reference/xarray_sql.md"} + {"Getting started" = "examples.md"}, + {"Concepts" = [ + {"Engines" = "engines.md"}, + ]}, + {"Guides" = [ + {"Geospatial in SQL" = "geospatial.md"}, + {"Performance" = "performance.md"}, + ]}, + {"Reference" = [ + {"API" = "reference/xarray_sql.md"}, + {"Known issues" = "limitations.md"}, + {"Contributing" = "contributing.md"}, + ]} ] # Theme configuration @@ -71,6 +80,10 @@ line_spans = "__span" pygments_lang_class = true [project.markdown_extensions.pymdownx.superfences] +[[project.markdown_extensions.pymdownx.superfences.custom_fences]] +name = "mermaid" +class = "mermaid" +format = "pymdownx.superfences.fence_code_format" [project.markdown_extensions.pymdownx.tasklist] custom_checkbox = true @@ -79,6 +92,8 @@ custom_checkbox = true [project.markdown_extensions.admonition] +[project.markdown_extensions.footnotes] + [project.markdown_extensions.pymdownx.snippets] url_download = true base_path = ["."] @@ -116,7 +131,7 @@ filters = ["!^_"] [[project.extra.social]] icon = "fontawesome/brands/github" -link = "https://github.com/alxmrs/xarray-sql" +link = "https://github.com/xqlsystems/xarray-sql" [[project.extra.social]] icon = "fontawesome/brands/python"