Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers

[Mmoncadaisla]

Summary

This branch prototypes xarray-sql as the Xarray ↔ engine translator: the library owns exactly two seams — register (a lazy xarray.Dataset becomes a table on an engine's own connection) and round-trip (any engine's Arrow result becomes a labeled xr.Dataset again). SQL dialects, geometry functions, H3, and optimizers stay with each engine and its extension ecosystem. No transpiler, no unified dialect.

flowchart TB
    A["Zarr · NetCDF · GRIB · GeoTIFF · Earth Engine"] --> B["lazy xarray.Dataset"]
    B --> C["seam 1 — register<br/>xql.register(con, name, ds) / xql.arrow_dataset(ds)"]
    C --> D["DuckDB<br/>+ spatial, h3, …"]
    C --> E["Polars"]
    C --> F["DataFusion"]
    C --> G["Dask"]
    D --> H["seam 2 — round-trip<br/>xql.to_dataset(result, template=ds)"]
    E --> H
    F --> H
    H --> I["labeled xr.Dataset — SQL in, array out"]

Loading

One object serves every engine: xql.arrow_dataset(ds) is a real pyarrow.dataset.Dataset (the pattern Lance uses for LanceDataset), so DuckDB registers it, Polars scans it via scan_pyarrow_dataset, DataFusion consumes it via register_dataset (through the fragments API), and Dask maps over get_fragments(). Ibis works through ibis.duckdb.from_connection with zero code.

How the pushdown scan works

flowchart TB
    Q["engine calls scanner(columns, filter)"] --> P["prune chunks<br/>per-dimension shadow fragments;<br/>Arrow guarantee simplification decides satisfiability"]
    Q --> J["project<br/>only referenced variables are read"]
    P --> C["coalesce (opt-in)<br/>merge consecutive surviving chunks<br/>into single reads"]
    C --> L["prefetch pool<br/>bounded by prefetch (threads)<br/>and prefetch_bytes (memory)"]
    J --> L
    L --> X["exact filter<br/>pyarrow applies the pushed expression row-exactly"]
    X --> R["Arrow batches → engine"]

Loading

Three properties worth calling out:

• Correctness never depends on pruning. DuckDB deletes pushed comparison conjuncts from its plan and does not re-apply them (verified in duckdb-python source); the scanner therefore always applies the exact expression via pyarrow, and pruning is only an optimization.
• No expression parsing. Chunk pruning delegates satisfiability to Arrow's own guarantee simplification: shadow FileSystemDataset fragments carry each chunk's coordinate range as a partition_expression (their paths are never opened). Sound for every predicate shape — equality, OR, IN, NOT.
• Bounded for finely chunked axes. One shadow per dimension (Σ nᵈ fragments, not Π nᵈ), bucketed two-level for axes beyond 1024 chunks — a 745k-chunk hourly time axis prunes in milliseconds. count(*) never scans: unfiltered counts are chunk arithmetic, and coordinate-range counts read at most the boundary chunks.

The chunked round-trip is engine-generic

xql.to_dataset(result, chunks=...) reconstructs a query result as a chunked, lazy xr.Dataset — each window re-executes the engine's query narrowed to its coordinate range, which flows back into chunk pruning at the source. spill=True provides the alternative one-pass shape: stream once (bounded memory) into a temporary Parquet file that windows re-execute against.

flowchart TB
    R["xql.to_dataset(result, ...)"] --> K{"chunks=?"}
    K -- "None (default)" --> E["eager: materialize once<br/>(max_result_bytes= guards the stream<br/>and the dense grid)"]
    K -- "mapping / auto / inherit" --> SP{"spill=?"}
    SP -- "False (default)" --> RX["re-execution<br/>Polars & DataFusion results"]
    SP -- "True / directory" --> SPL["one-pass spill → temp Parquet<br/>the chunked path for DuckDB relations<br/>and one-shot Arrow streams"]

Loading

coords="template" skips coordinate discovery entirely for full-extent results: on ARCO-ERA5 (1.32M hourly chunks) it 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 it.

Performance

Measured on a public 9-billion-pixel cloud-optimized GeoTIFF and a 10M-row synthetic benchmark (benchmarks/duckdb_pushdown.py):

Scenario	before	after
Bounding-box GROUP BY, native res, full table registered	minutes+	~0.8 s
Full-table aggregation (with reader guidance from docs/performance.md)	277 s	24 s
10M rows, full AVG scan (vs v1 stream)	0.52 s	0.027 s
10M rows, 1% time filter	0.24 s	0.006 s
xarray→Arrow pivot throughput	53M rows/s	154M rows/s
25M-row window round-trip (to_dataset)	0.96 s	0.43 s

Peak scan memory is a contract, not an accident: bounded by prefetch × pivoted-block-size regardless of data scanned (a 772M-row month-scale ARCO-ERA5 aggregation peaks at the same ~0.75 GB RSS as the week-scale scan). The pivot fast path (whole-partition repeat/tile coordinate columns, zero-copy batch slices) and the round-trip fast paths (affine-axis scatter, grid reshape) live in shared code, so the existing DataFusion engine benefits equally.

Beyond microbenchmarks, benchmarks/geospatial/ runs nine staples of geospatial/climate analysis (NDVI, climatology, anomaly JOIN, forecast skill vs WeatherBench 2, raster×vector zonal stats, PROJ-UDF reprojection, weight-table regridding, warp) in SQL against real cloud datasets, each asserted against an xarray reference; engine_suite.py repeats the portable cases across DataFusion (both the native table provider and the pure-Python pyarrow path, as separate datafusion / datafusion-arrow engines), DuckDB, and Polars on GCE VMs, provisioning the compiled native module on each VM. The write-up is docs/geospatial.md.

What's in the branch

• xarray_sql/backends/ — adapter protocol with dispatch on connection type; DataFusion adapter delegates to the existing table provider; DuckDB adapter registers the pushdown dataset (mixed-dimension datasets split into one table per dim group, sharing coordinate reads).
• xarray_sql/backends/pyarrow.py — XarrayPushdownDataset (projection pushdown, shadow pruning, coalescing, prefetch, fragments API) and XarrayArrowStream (re-scannable C-stream fallback).
• xarray_sql/roundtrip.py + xarray_sql/lazyscan.py — engine-agnostic xql.to_dataset accepting DuckDB relations, Polars frames, DataFusion DataFrames, pyarrow tables/readers, or any __arrow_c_stream__ object; eager, re-executing, and spill-backed chunked reconstruction; metadata recovery from a template Dataset.
• xarray_sql/geometry.py — register(..., geometry=("x", "y")) derives a GeoArrow point geometry column (WKB for DuckDB-native GEOMETRY, or GeoArrow-separated for GeoPandas/lonboard), CRS tagged; xql.bbox_conjuncts renders prunable bbox predicates from any geometry's envelope.
• Pivot and round-trip fast paths in df.py / ds.py.
• Docs: docs/engines.md (the engine model, per-engine usage, support matrix), docs/performance.md (measured tuning guide), docs/limitations.md (known issues, pinned by tests), docs/geospatial.md (the relational-operations write-up with the benchmark results), reworked docs/examples.md and README.
• Packaging: [duckdb] (duckdb>=1.4) and [polars] (polars>=1.33) extras.

Testing

275 tests, all green. Beyond unit coverage:

• Compatibility matrix: the full pushdown battery (pruning verified via read callbacks, mid-chunk equality, OR/IN, filter-column-outside-projection, LIMIT, empty results, NULL semantics) passes identically on duckdb 1.4.5 LTS and 1.5.4, plus the Polars battery on 1.42.
• Concurrency: 4 threads × 50 iterations over a shared dataset; joins of two registered datasets; thread-safety of the lazy shadow cache.
• Memory: 600-query soak, RSS plateaus (no leak); streamed max_result_bytes enforcement on eager collection and dense allocation.
• End-to-end: every geospatial benchmark case asserts SQL == xarray reference to floating-point tolerance before its timing counts.
• Adversarial review found real bugs along the way — NaN-poisoned pruning guarantees (silent row loss), NaN acceptance in the affine scatter, a cftime crash under projection, a ChunkedArray slipping into the pivot's fast path — all fixed with regression tests.

Known limitations / follow-ups

• DuckDB relations refuse the re-executing chunked round-trip (chunks= raises immediately): re-execution from worker threads deadlocks inside duckdb-python 1.4–1.5 on CPython 3.12, so spill=True is the chunked path there. Documented with the full story in docs/limitations.md.
• XarrayPushdownDataset subclasses pyarrow's cython Dataset without initializing the native base (as Lance does); dangerous inherited members are stubbed, and the contract is pinned by tests — re-verify on pyarrow upgrades.
• Upstream Polars translates is_in float literals imprecisely (reproducible without xarray-sql); documented, use range predicates. Polars also has no geometry types — the geometry column arrives as plain binary there.
• Roadmap candidates from ecosystem research: an optional H3/quadbin cell column emitted at registration (integer join key for cross-raster joins), an obstore-backed Zarr path, and a RaQuet writer for warehouse interchange.

🤖 Generated with Claude Code