diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cc4799d..cd47c21 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/README.md b/README.md index a254d29..6a444b0 100644 --- a/README.md +++ b/README.md @@ -2,10 +2,10 @@ _Query [Xarray](https://xarray.dev/) with SQL_ -[![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) ```shell pip install xarray-sql @@ -15,7 +15,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 @@ -58,6 +62,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, @@ -130,6 +149,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 @@ -137,8 +158,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) @@ -155,7 +174,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? @@ -188,6 +207,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. @@ -217,8 +239,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? @@ -226,15 +248,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? @@ -247,7 +271,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: @@ -257,18 +281,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..22932e8 --- /dev/null +++ b/benchmarks/duckdb_pushdown.py @@ -0,0 +1,89 @@ +"""Benchmark: DuckDB adapter v1 (stream) vs v2 (pushdown) vs ceilings. + +Usage: .venv-duckdb/bin/python bench_v2.py +""" + +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=3): + print(f"\n== {label} ==") + 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) + print(f" {qname:28s} {min(times):8.3f}s -> {r[0][0]:.6g}") + + +# v1: re-scannable stream (registered via the stream wrapper explicitly) +con.register("t_v1", XarrayArrowStream(ds)) +bench("t_v1", "v1 stream (no pushdown)") + +# v2: default register() — pushdown path (once implemented) +xql.register(con, "t_v2", ds) +bench("t_v2", "v2 register() [pushdown]") + +# ceiling: materialized pa.Table via pyarrow.dataset +table = xql.read_xarray(ds).read_all() +con.register("t_ceiling", pads.dataset(table)) +bench("t_ceiling", "ceiling: in-memory pyarrow.dataset") + +# reference: DataFusion engine on the same dataset +ctx = xql.XarrayContext() +xql.register(ctx, "t_df", ds) +print("\n== DataFusion reference ==") +for qname, q in QUERIES.items(): + sql = q.format(t="t_df") + times = [] + for _ in range(3): + t0 = time.perf_counter() + r = ctx.sql(sql).to_pandas() + times.append(time.perf_counter() - t0) + print(f" {qname:28s} {min(times):8.3f}s -> {float(r.iloc[0, 0]):.6g}") diff --git a/benchmarks/era5_out_of_core.py b/benchmarks/era5_out_of_core.py new file mode 100644 index 0000000..c38c92b --- /dev/null +++ b/benchmarks/era5_out_of_core.py @@ -0,0 +1,147 @@ +"""Out-of-core benchmark + plan-shape assertions on public ARCO-ERA5. + +Registers gs://gcp-public-data-arco-era5 (1,323,648 hourly time chunks; +nominally 1.37 trillion rows for one surface variable) without dask and +drives DuckDB, Polars, and the lazy round-trip against it, asserting +the *shape of the work* — exactly which source chunks each query reads +(via the scanner's iteration callback) and exact row counts — not just +the answers. A pruning or fast-path regression that silently falls back +to scanning everything fails these assertions long before it shows up +in wall-clock noise. + +Needs network (anonymous GCS). Run: python benchmarks/era5_out_of_core.py +""" + +import time + +import duckdb +import pandas as pd +import polars as pl +import pyarrow as pa +import pyarrow.compute as pc +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + +URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" + +GRID = 721 * 1440 + +Q_DAY_BBOX = ( + 'SELECT round(avg("2m_temperature") - 273.15, 2), count(*) FROM era5 ' + "WHERE time >= TIMESTAMP '2020-01-01' AND time < TIMESTAMP '2020-01-02' " + "AND latitude BETWEEN 36 AND 44 AND longitude BETWEEN 350 AND 360" +) +Q_WEEK_GLOBE = ( + 'SELECT round(avg("2m_temperature") - 273.15, 2), count(*) FROM era5 ' + "WHERE time >= TIMESTAMP '2020-01-03' AND time < TIMESTAMP '2020-01-10'" +) + + +def timed(label, fn, reads, expect_reads, expect_rows=None): + reads.clear() + t0 = time.time() + out = fn() + wall = time.time() - t0 + assert len(reads) == expect_reads, ( + f"{label}: read {len(reads)} blocks, expected {expect_reads} — " + "pruning/coalescing regressed" + ) + if expect_rows is not None: + assert out[-1] == expect_rows, (label, out, expect_rows) + print(f"{label:34s} {wall:6.2f}s reads={len(reads):3d} {out}") + return out + + +def main() -> None: + ds = xr.open_zarr( + URL, chunks=None, storage_options={"token": "anon"}, consolidated=True + )[["2m_temperature"]] + + reads: list = [] + dataset = XarrayPushdownDataset( + ds, + {"time": 1}, + prefetch=16, + _iteration_callback=lambda b, n: reads.append(b), + ) + con = duckdb.connect() + con.register("era5", dataset) + + # 24 hourly chunks of 1,323,648; bbox trims rows exactly. + timed( + "duckdb day+bbox", + lambda: con.execute(Q_DAY_BBOX).fetchone(), + reads, + expect_reads=24, + expect_rows=24 * 33 * 40, + ) + timed( + "duckdb week globe", + lambda: con.execute(Q_WEEK_GLOBE).fetchone(), + reads, + expect_reads=168, + expect_rows=168 * GRID, + ) + # count(*) over a chunk-aligned window: every surviving chunk is + # provably inside the range, so the count is pure arithmetic. + jan = ( + pc.field("time") + >= pa.scalar(pd.Timestamp("2020-01-01"), type=pa.timestamp("ns")) + ) & ( + pc.field("time") + < pa.scalar(pd.Timestamp("2020-02-01"), type=pa.timestamp("ns")) + ) + result = timed( + "count_rows fast path (January)", + lambda: (dataset.count_rows(filter=jan),), + reads, + expect_reads=0, + ) + assert result[0] == 744 * GRID, result + + # Coalescing: same day+bbox in 4 merged reads instead of 24. + coalesced = XarrayPushdownDataset( + ds, + {"time": 1}, + prefetch=16, + coalesce_rows=8_000_000, + _iteration_callback=lambda b, n: reads.append(b), + ) + con2 = duckdb.connect() + con2.register("era5", coalesced) + timed( + "duckdb day+bbox coalesced", + lambda: con2.execute(Q_DAY_BBOX).fetchone(), + reads, + expect_reads=4, + expect_rows=24 * 33 * 40, + ) + + # Lazy round-trip through Polars: construction reads nothing with + # template coords; a 1-day window reads only its 4 coalesced blocks. + lf = pl.scan_pyarrow_dataset(coalesced) + reads.clear() + lazy = xql.to_dataset( + lf, template=ds, chunks={"time": 24}, coords="template" + ) + assert reads == [], "lazy construction must not read the source" + timed( + "polars lazy 1-day window", + lambda: ( + float( + lazy["2m_temperature"] + .sel(time=slice("2020-01-01", "2020-01-01 23:00")) + .mean() + .compute() + ), + ), + reads, + expect_reads=4, + ) + print("all plan-shape assertions passed") + + +if __name__ == "__main__": + main() 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..f9d9c4d --- /dev/null +++ b/benchmarks/geospatial/_engines.py @@ -0,0 +1,224 @@ +"""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``, the suite's original + path, when the native module is importable; otherwise a plain + ``datafusion.SessionContext`` over ``xql.arrow_dataset`` (pure Python). + Which one ran is recorded in :attr:`EngineContext.flavor`. +* ``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-native 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) do not use +this layer; the suite runner records them as n/a for other engines. +""" + +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 + +_ENGINES = ("datafusion", "duckdb", "polars") + + +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 _native_available() -> bool: + if os.environ.get("GEOBENCH_NO_NATIVE"): # test hook: force fallback + return False + try: + import xarray_sql._native # noqa: F401 + + return True + except Exception: # noqa: BLE001 — pure-Python source tree + return False + + +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.""" + + def __init__(self, engine: str | None = None): + self.engine = engine or engine_name() + self.flavor = self.engine + self._native = False + self._renames: dict[str, str] = {} + if self.engine == "datafusion": + if _native_available(): + import xarray_sql as xql + + self.flavor = "datafusion (XarrayContext, native)" + self._native = True + self._ctx = xql.XarrayContext() + else: + from datafusion import SessionContext + + self.flavor = "datafusion (pyarrow dataset, no native)" + self._ctx = SessionContext() + elif self.engine == "duckdb": + import duckdb + + self._con = duckdb.connect() + else: + # Polars: keep the pyarrow datasets and build 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. sql_to_dataset therefore pre-filters each + # frame with the query's window parameters (identical + # predicate to the SQL WHERE, which still runs on top). + self.flavor = "polars (SQLContext + expression window pushdown)" + self._polars_tables: dict[str, Any] = {} + + # -- registration ----------------------------------------------------- + + def from_dataset(self, name, ds, *, chunks=None, table_names=None): + """Register ``ds`` as SQL table(s), mirroring XarrayContext naming.""" + if self._native: + self._ctx.from_dataset( + name, ds, chunks=chunks, table_names=table_names + ) + return + 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 + dataset = xql.arrow_dataset(sub, sub_chunks) + if self.engine == "duckdb": + self._con.register(flat, dataset) + elif self.engine == "polars": + self._polars_tables[flat] = dataset + else: + self._ctx.register_dataset(flat, dataset) + + # -- 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) + if self.engine == "polars" and param_values: + for key, value in param_values.items(): + sql = re.sub(rf"\${key}\b", _literal(value), 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``.""" + sql = self._rewrite(sql, param_values) + if self.engine == "datafusion": + df = ( + self._ctx.sql(sql, param_values=param_values) + if param_values + else self._ctx.sql(sql) + ) + if self._native: + return df.to_dataset(dims=dims) + return _pandas_to_dataset(df.to_pandas(), dims) + if self.engine == "duckdb": + pdf = self._con.execute(sql, param_values or {}).df() + return _pandas_to_dataset(pdf, dims) + out = self._polars_execute(sql, param_values) + return _pandas_to_dataset(out.to_pandas(), dims) + + # 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 _polars_execute(self, sql: str, param_values): + import polars as pl + + ctx = pl.SQLContext() + params = param_values or {} + for flat, dataset in self._polars_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) diff --git a/benchmarks/geospatial/engine_suite.py b/benchmarks/geospatial/engine_suite.py new file mode 100644 index 0000000..4034ae2 --- /dev/null +++ b/benchmarks/geospatial/engine_suite.py @@ -0,0 +1,460 @@ +"""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 (the original path), 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. + +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, so DuckDB/Polars are 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:: + + python benchmarks/geospatial/engine_suite.py --local --reps 1 \ + --cases 02_climatology --vms local # in-process check + python benchmarks/geospatial/engine_suite.py # 3 VMs, full suite +""" + +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 + +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", "duckdb", "polars"] +# Cases whose SQL builds DataFusion scalar UDFs (07, and the UDF half of +# 09): not expressible on the other engines. 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}" + marker = os.path.join(root, "benchmarks", "geospatial", "_engines.py") + 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 + return root, os.path.join(root, "benchmarks", "geospatial") + + +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 = {"case": case, "engine": engine, "status": "ok", "reps": []} + try: + src_root, geo_dir = _install_src(src_targz) + 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 + if "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 (pure Python).""" + repo = Path(__file__).resolve().parents[2] + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for rel in ["xarray_sql", "benchmarks/geospatial"]: + for path in sorted((repo / rel).rglob("*.py")): + if "__pycache__" in path.parts: + continue + tf.add(path, arcname=str(path.relative_to(repo))) + 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]) + 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 + 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]}") + results.append({"vm": vm, "case": "_meta", "engine": "", **meta}) + with jsonl_lock, open(args.jsonl, "a") as fh: + fh.write(json.dumps(results[-1]) + "\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] + vms = ["local"] if args.local else [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}") + + +if __name__ == "__main__": + main() diff --git a/docs/engines.md b/docs/engines.md new file mode 100644 index 0000000..224531b --- /dev/null +++ b/docs/engines.md @@ -0,0 +1,220 @@ +# 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`. + + +## 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 uses), 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. + 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 085a1b0..37eb03c 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.) @@ -221,7 +222,7 @@ SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat 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* @@ -242,7 +243,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 — @@ -250,7 +251,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: @@ -274,6 +275,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; @@ -304,13 +336,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. @@ -373,6 +405,81 @@ 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 +under DuckDB and Polars 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`), so we ran the portable cases across all three engines +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. These runs exercise the **pyarrow-dataset backend** for +every engine — including DataFusion, which here reads through +`SessionContext.register_dataset(xql.arrow_dataset(ds))` rather than the +native `XarrayContext` the headline table used (the VMs run the pure-Python +tree; the flavor that executed is recorded per cell in the raw results). So +the DataFusion column below is *not* the same code path as the table above — +on the ERA5 group-by cases the pyarrow path runs ~1.5–2× behind its native +sibling, and DuckDB is the fastest consumer of the shared scan. 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`, +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 full per-size tables live in the +`xarray-sql-notes` repository (`engine-matrix-results.md`). + +Software: CPython 3.12.8, duckdb 1.5.4, 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 (pyarrow path) | DuckDB | Polars | xarray reference | +|---|--:|--:|--:|--:| +| 01 · NDVI | 3.840 s (109 MB) | 4.248 s (106 MB) | 4.672 s (100 MB) | 0.263 s (42 MB) | +| 02 · Climatology | 7.505 s (1135 MB) | 4.043 s (627 MB) | 4.360 s (637 MB) | 2.448 s (44 MB) | +| 03 · Zonal mean | 3.229 s (403 MB) | 2.758 s (403 MB) | 3.351 s (413 MB) | 0.890 s (250 MB) | +| 04 · Anomaly | 12.857 s (3003 MB) | 6.094 s (627 MB) | 11.472 s (862 MB) | 5.103 s (80 MB) | +| 05 · Forecast skill | 1.854 s (172 MB) | 1.722 s (170 MB) | 2.295 s (182 MB) | 0.213 s (2 MB) | +| 06 · Zonal stats | 5.311 s (513 MB) | 4.706 s (503 MB) | unsupported (range `JOIN`) | 1.624 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 is consistent. **DuckDB is the +fastest SQL consumer of the shared pushdown scan on every ARCO-ERA5 case** — +on the group-by and join cases (02, 04, 06) it runs ~1.5–2× ahead of the +DataFusion pyarrow path, and it is the only engine that keeps the anomaly +self-`JOIN` (04) within ~1.2× of the array reference. Polars matches DuckDB +on the plain group-bys (02, 03) but falls back toward DataFusion on the +join-heavy 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` — now fixed with a +regression test. + ## Analysis: how a relational operation spends its time Why is SQL slower, and where does the time actually go? Profiling case 05 — the @@ -472,3 +579,4 @@ scale. The point of this suite is not to crown a winner but to show that the lin between the two is exactly where the operation is dense versus where it is relational, and that for a surprising share of geoscience, the operation is relational. + 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..5b08b81 --- /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 RSS 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 4273d31..965a28f 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,8 +36,22 @@ 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", +] test = [ "cftime", + "duckdb>=1.4.0", + # 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", "pytest", "xarray[io]", "gcsfs", diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py new file mode 100644 index 0000000..6ff6d49 --- /dev/null +++ b/tests/test_arrow_dataset.py @@ -0,0 +1,369 @@ +"""Tests for the engine-neutral pyarrow dataset view. + +``xql.arrow_dataset`` returns a real ``pyarrow.dataset.Dataset``, so it +serves any consumer of that protocol — pyarrow itself and Polars are +exercised here; DuckDB has its own suite in ``test_duckdb_backend.py``. +""" + +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}) + + +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_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 + + +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()) + + +class _ChunkCounter: + def __init__(self): + self.blocks = [] + + def __call__(self, block, names): + self.blocks.append(block) + + +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 + + +def test_count_rows_unfiltered_is_pure_arithmetic(counted): + dataset, counter = counted + assert dataset.count_rows() == 100 * 4 + assert counter.blocks == [] # no chunk was read + + +def test_count_rows_strict_chunks_counted_without_reading(counted): + dataset, counter = counted + # [03:00, 27:00): chunks 0 and 2 are boundary, chunk 1 is provably + # inside the range and must be counted arithmetically. + lo = pa.scalar(pd.Timestamp("2020-01-01 03:00"), type=pa.timestamp("ns")) + hi = pa.scalar(pd.Timestamp("2020-01-02 03:00"), type=pa.timestamp("ns")) + predicate = (pc.field("time") >= lo) & (pc.field("time") < hi) + assert dataset.count_rows(filter=predicate) == 24 * 4 + assert len(counter.blocks) == 2 # only the two boundary chunks + + +def test_count_rows_data_variable_filter_is_exact(counted): + dataset, counter = counted + # A filter on a data variable carries no coordinate guarantee: every + # chunk is a boundary chunk, and the count must still be row-exact. + assert dataset.count_rows(filter=pc.field("t2m") >= 200.0) == 200 + assert len(counter.blocks) == 10 + + +def test_count_rows_unsatisfiable_filter(counted): + dataset, counter = counted + assert dataset.count_rows(filter=pc.field("lat") > 100.0) == 0 + assert counter.blocks == [] + + +def test_empty_projection_is_a_real_projection(counted): + dataset, _ = counted + table = dataset.scanner(columns=[]).to_table() + assert table.num_columns == 0 + assert table.num_rows == 100 * 4 + + +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 + + +@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 + ) + lo = pa.scalar(pd.Timestamp("2020-01-01 03:00"), type=pa.timestamp("ns")) + hi = pa.scalar(pd.Timestamp("2020-01-03 07:00"), type=pa.timestamp("ns")) + predicate = (pc.field("time") >= lo) & (pc.field("time") < hi) + 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") + < pa.scalar(pd.Timestamp("2020-01-02 06:00"), type=pa.timestamp("ns")) + ) | ( + pc.field("time") + >= pa.scalar(pd.Timestamp("2020-01-03 22:00"), type=pa.timestamp("ns")) + ) + 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_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_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 + + +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_df.py b/tests/test_df.py index 2c81144..283cae3 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -403,20 +403,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: @@ -532,3 +534,24 @@ def test_compute_chunks_tuples_sum_to_dim_size(): result = compute_chunks(ds, {"a": 3, "b": 4, "c": 5}) for dim, tup in result.items(): 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 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_recipes.py b/tests/test_sql_recipes.py new file mode 100644 index 0000000..5216ac2 --- /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 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/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..cc2c177 --- /dev/null +++ b/xarray_sql/backends/__init__.py @@ -0,0 +1,33 @@ +"""Engine adapters — the *register* seam of xarray-sql. + +xarray-sql translates data, not queries: it registers lazy +``xarray.Dataset`` objects as tables on a query engine's own connection +(seam 1, this package) and turns Arrow results back into labeled +Datasets (seam 2, :func:`xarray_sql.to_dataset`). SQL dialects, +geometry, H3, and optimizers belong to each engine and its extension +ecosystem. + +Adapters register themselves on import via +:func:`~xarray_sql.backends.base.register_adapter`; +:func:`~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..16b80db --- /dev/null +++ b/xarray_sql/backends/base.py @@ -0,0 +1,111 @@ +"""Engine-adapter dispatch for :func:`xarray_sql.register`. + +An *engine adapter* implements one 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, runtime_checkable + +import xarray as xr + +from ..df import Chunks + + +@runtime_checkable +class EngineAdapter(Protocol): + """One engine's implementation of the register seam.""" + + @staticmethod + def matches(con: Any) -> bool: + """Whether *con* is a connection this adapter can register into.""" + ... + + @staticmethod + def register( + con: Any, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> Any: + """Register *ds* as table *name* on *con*; returns *con*.""" + ... + + +_ADAPTERS: list[type[EngineAdapter]] = [] + + +def register_adapter(cls: type) -> type: + """Class decorator adding an adapter to the dispatch list.""" + _ADAPTERS.append(cls) + return cls + + +def get_adapter(con: Any) -> type[EngineAdapter]: + """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: Any, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, +) -> Any: + """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 + :class:`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. + """ + return get_adapter(con).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..ec61ed9 --- /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 +:func:`xarray_sql.register` seam onto that existing machinery. +""" + +from __future__ import annotations + +from typing import Any + +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: Any) -> bool: + 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..ca36d4c --- /dev/null +++ b/xarray_sql/backends/duckdb.py @@ -0,0 +1,88 @@ +"""DuckDB engine adapter. + +Registers a lazy ``xarray.Dataset`` on a ``duckdb.DuckDBPyConnection`` +as an :class:`~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 :mod:`xarray_sql.backends.pyarrow`). + +This adapter never imports the ``duckdb`` package — 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 +`_ extension provides; this +adapter instead covers everything xarray can open (NetCDF, GRIB, Xee, CF +decoding, in-memory) and pairs with :func:`xarray_sql.to_dataset` for +the labeled round-trip. +""" + +from __future__ import annotations + +from typing import Any + +import xarray as xr + +from ..df import Chunks +from ..sql import _group_vars_by_dims +from .base import register_adapter +from .pyarrow import XarrayArrowStream, XarrayPushdownDataset + +__all__ = ["DuckDBAdapter", "XarrayArrowStream", "XarrayPushdownDataset"] + + +@register_adapter +class DuckDBAdapter: + """Registers Datasets on ``duckdb.DuckDBPyConnection`` connections.""" + + @staticmethod + def matches(con: Any) -> bool: + # 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: Any, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> Any: + """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 :class:`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..b8d01fa --- /dev/null +++ b/xarray_sql/backends/pyarrow.py @@ -0,0 +1,1104 @@ +"""Engine-neutral pyarrow views of lazy xarray Datasets. + +Two ways to hand a lazy ``xarray.Dataset`` to an Arrow-speaking query +engine: + +* :class:`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 + :meth:`~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 :func:`arrow_dataset`. +* :class:`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 + :class:`~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 + :class:`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 :meth:`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. + """ + + 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 + 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. + 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._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 _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 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 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 + 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): + raise NotImplementedError( + "Use scanner(filter=...) or the engine's WHERE clause." + ) + + def replace_schema(self, schema: pa.Schema): + raise NotImplementedError + + 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 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 :class:`_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: + 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 = outer + for d in dims: + rows *= int(lens[d][np.asarray(cell[d])].sum()) + return rows + + def classify(cell: dict, depth: int) -> None: + 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 + + 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 > self._coalesce_rows + ): + yield flush(prefix, run) + run, run_rows = [], 0 + run.append(i) + run_rows += rows + if run: + yield flush(prefix, run) + + 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 + 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((self._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() + + +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 :meth:`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 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 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 + :mod:`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 :class:`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/df.py b/xarray_sql/df.py index ab80056..5a82934 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -297,6 +297,32 @@ def dataset_to_record_batch( #: 64 K rows balances DataFusion pipeline depth against per-batch overhead. DEFAULT_BATCH_SIZE: int = 65_536 +#: 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. +_FULL_PIVOT_MAX_ROWS: int = 8_388_608 + + +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( ds: xr.Dataset, @@ -338,8 +364,12 @@ 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 = {} for name in dim_names: + if name not in schema.names: + continue vals = ds.coords[name].values if cft.is_cftime(vals): coord_values[name] = cft.convert_for_field(vals, schema.field(name)) @@ -361,6 +391,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) @@ -372,13 +432,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, ) ) diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index cd1db39..23798bf 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -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,15 @@ 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]) + if affine[d] is not None: + origin, step = affine[d] + 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 +255,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, + (:class:`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. + :meth:`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 + :class:`~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,10 +288,10 @@ 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 @@ -254,17 +304,23 @@ class SQLBackendArray(xr.backends.BackendArray): 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 +347,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. """ 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 +381,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 +395,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 +409,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 +508,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 +516,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 +628,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 +648,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`. 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 :func:`_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 +679,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, @@ -666,13 +796,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 = { 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..872134d --- /dev/null +++ b/xarray_sql/lazyscan.py @@ -0,0 +1,362 @@ +"""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: + +* :meth:`~LazyResultHandle.schema` — result column names/types, without + executing the query; +* :meth:`~LazyResultHandle.distinct` — one column's distinct values + (coordinate discovery; the caller sorts); +* :meth:`~LazyResultHandle.fetch` — the result narrowed by per-dimension + windows and projected to the requested columns, as Arrow batches. + +Windows are passed as :data:`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 +:class:`~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, 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[str, 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).""" + + +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 re-executes, which is + exactly the re-executable contract. 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 :meth:`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/roundtrip.py b/xarray_sql/roundtrip.py new file mode 100644 index 0000000..62d3a98 --- /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 +:mod:`xarray_sql.lazyscan`. DuckDB relations are re-executable but +refuse the chunked path (a thread-safety limitation noted on +:class:`~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 :func:`_open_stream` recognizes, then objects + with a ``to_arrow_table()`` method (DataFusion DataFrames and the + :class:`~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 + :meth:`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/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"