Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers - #227
Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers#227Mmoncadaisla wants to merge 83 commits into
Conversation
Prototype of the "xarray-sql as the Xarray <-> engine translator" idea: the library owns two seams and nothing else. Seam 1 (register) attaches a lazy Dataset as a table on an engine's own connection; seam 2 (round-trip) turns any engine's Arrow result plus a template Dataset back into a labeled xr.Dataset. Dialects, geometry, H3, and optimizers stay with each engine and its extension ecosystem. - xarray_sql/backends/: adapter protocol + dispatch on connection type; DataFusion adapter delegates to the existing table provider, DuckDB adapter registers a re-scannable Arrow C-stream view (fresh lazy reader per scan: lazy AND re-queryable, no pushdown yet). - xarray_sql/roundtrip.py: engine-agnostic to_dataset() accepting DuckDB relations, pyarrow Tables/readers, or any __arrow_c_stream__ object; reuses the batches->Dataset core extracted from ds._materialize. - xql.register / xql.to_dataset exported at top level; [duckdb] extra. - docs/engines.md: the engine model, DuckDB usage, relation to duckdb-zarr (engine-native Zarr path; this adapter covers the rest of the xarray reader surface plus the labeled round-trip). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the default DuckDB registration object with XarrayPushdownDataset, a pyarrow.dataset.Dataset subclass (the pattern Lance uses for LanceDataset). DuckDB classifies it by isinstance and calls scanner(columns=..., filter=...) once per query, enabling: - projection pushdown: only the data variables a query mentions are loaded from storage; - chunk pruning: per-dimension shadow FileSystemDataset fragments carry each chunk's coordinate range as a partition_expression, so Arrow's guarantee simplification decides satisfiability for any predicate shape with no expression parsing on our side. One shadow per dimension keeps fragment counts at sum(n_d) instead of prod(n_d); - parallel production: surviving chunks are loaded by a bounded prefetch thread pool. DuckDB deletes pushed comparison conjuncts from its plan and never re-applies them, so the scanner always applies the exact expression via pyarrow Scanner; pruning is only an optimization. Filter-only columns absent from the projection are discovered by probing the expression against an empty table and widening on the miss. 10M-row benchmark vs the v1 stream: full scan 0.52s -> 0.035s, 1% time-filtered scan 0.24s -> 0.006s. A native-resolution bounding-box GROUP BY over a 9.13B-pixel cloud GeoTIFF answers in 0.8s on a plain duckdb connection (previously minutes). XarrayArrowStream stays as the dependency-light no-pushdown fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Bucketed two-level shadow pruning: axes with more chunks than the 1024-fragment fanout get a coarse bucket shadow refined lazily per surviving bucket, bounding pruning cost for finely partitioned datasets (e.g. hourly-chunked reanalysis time axes with hundreds of thousands of chunks). Refinement is skipped when a predicate keeps most buckets, where it cannot pay for itself. - Shadow datasets carry the full table schema so compound predicates referencing several columns bind and prune correctly. - Mixed-dimension datasets split into one table per dimension group (<name>_<dims>), sharing a single read of the dimension coordinates across sub-tables. - register() forwards adapter-specific kwargs (batch_size, prefetch on DuckDB; table_names on DataFusion). - Edge cases covered by tests: fully-pruned empty scans, LIMIT early termination, descending coordinate axes, uint8/string variables, 5000-chunk bucketed pruning loading exactly one chunk, kwarg forwarding. - duckdb extra pinned to >=1.4 (semantics verified against 1.5.4); engines doc gains production notes; pushdown benchmark added to benchmarks/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Coordinate columns of a C-ordered partition are exact repeat/tile patterns of the dimension coords: dim k's flat column is its values each repeated prod(shape[k+1:]) times, tiled prod(shape[:k]) times. Building each column once per partition with those two sequential-write kernels and emitting batches as zero-copy Arrow slices replaces the per-batch division/modulo plus gather, making the pivot ~2.9x faster (53M -> 154M rows/s on a 10M-row, 3-dim dataset) with bitwise-identical output. Both engines benefit: iter_record_batches feeds the DataFusion table provider and the DuckDB pushdown scanner alike. The fast path holds the partition's full coordinate columns in memory (rows x 8 bytes x n_dims), so it is gated at 8M rows; larger partitions (e.g. single-time-step reanalysis) keep the O(batch_size) streaming path. The memory-profile test bands move accordingly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
rioxarray serializes GDAL reads behind a lock by default, capping any scan at single-stream speed regardless of the adapter's prefetch pool. lock=False measured 6x on full scans of a 9-billion-pixel cloud GeoTIFF (277s -> 43s) and makes remote reads as fast as a local copy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Move XarrayPushdownDataset and XarrayArrowStream to
backends/pyarrow.py: the pushdown dataset is a real
pyarrow.dataset.Dataset, so it serves every consumer of that protocol,
not just DuckDB. The new public constructor xql.arrow_dataset(ds)
makes that explicit:
pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) # Polars, lazy
con.register("t", xql.arrow_dataset(ds)) # DuckDB
xql.arrow_dataset(ds).to_table(columns=..., filter=...) # pyarrow
Verified with Polars 1.42: scans are lazy, predicates and projections
push into the dataset (a filtered group-by read 1 of 20 chunks and 3
of 5 columns), results match xarray exactly, and Polars frames
round-trip through xql.to_dataset via the Arrow PyCapsule protocol.
The DuckDB adapter is now a thin registration shim over the shared
module. polars added to the test extra.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_fragments(filter=...) yields one fragment per source chunk, pruned by the same per-dimension shadow index the scanner uses. This is the protocol datafusion-python's register_dataset consumes (one DataFusion partition per fragment, filter pushdown marked Exact — safe because fragment scanners apply the expression row-exactly) and enables the Dask pattern from_map(lambda f: f.to_table().to_pandas(), ds.get_fragments()). Fragments carry a __dask_tokenize__ hook since the parent dataset is deliberately unpicklable. One arrow_dataset() object now serves DuckDB, Polars, DataFusion, Dask, and pyarrow itself. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two fast paths in the batches->Dataset core (shared by every engine's round-trip): - Uniformly spaced axes (rasters, regular time steps, ascending or descending) locate each row with rint((value - origin) / step) instead of a per-row binary search plus argsort remap. 2.2x on a 25M-row chunk-ordered window (0.96s -> 0.43s); axes that are not affine within a quarter step keep the searchsorted path. - Results that form the complete grid in C order (ORDER BY'd scans, single-chunk windows) skip positional scatter entirely and reshape the value column. Both are detected, never assumed; sparse and arbitrarily ordered results fall back to the general scatter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Registered xarray tables are virtual: every query re-streams the source. For statistics asked repeatedly, xql.materialize(con, name, query, order_by=...) pays the scan once into a native engine table (sorted so coordinate columns compress and zone maps prune), and xql.pyramid(con, name, table, aggs=..., base_cell=..., levels=...) builds a CARTO-tileset-style multi-resolution pre-aggregated cube: level 0 bins the source in a single pass, coarser levels roll up from the level below, so any zoom/extent query is a range scan over a small table. Aggregates are restricted to decomposable kinds (sum/count/min/max) so roll-ups stay exact; averages derive from sum + count at query time. Both helpers dispatch through a new EngineAdapter.run_sql seam and are tested identically on DuckDB connections and DataFusion contexts (the SQL they emit is restricted to the shared dialect subset; DataFusion INSERT is lazy, so run_sql collects). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Measured guidance in one place: parallel source reads (LIBERTIFF / lock=False, 11x on full raster scans; zarr async.concurrency), chunk sizing for the scan, adapter knobs (prefetch, batch_size), what pushdown does and does not cover, materialize/pyramid for repeated statistics, and the round-trip fast paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- NaN/NaT in a chunk's dimension coordinate poisoned the shadow guarantee into (dim >= NaN), which Arrow simplifies every predicate against as false — chunks with matching rows were silently pruned. Such spans now carry an always-true guarantee (unprunable). - _affine_axis accepted NaN-poisoned axes (NaN > tol is false), letting the affine scatter place values in wrong cells when a result's dim column contains NULLs. The acceptance test is now a <= .all() so NaN rejects and the positional searchsorted path handles it. - Projected scans omit dimension columns from the scan schema; the cftime coordinate-conversion loop assumed every dim had a schema field and raised KeyError on e.g. SELECT SUM(v) over a 360_day dataset. Dims absent from the schema are now skipped. - pyramid() rebinned float cell origins level-over-level, occasionally aliasing boundary points into a neighboring cell. Cells are now tracked as integer indices that halve exactly per level, with float origins kept as query labels. Also: docs notes from stress testing (Polars is_in float-literal pushdown caveat, per-thread DuckDB cursor re-registration pattern). Regression tests for all four fixes; battle-test matrix passed on duckdb 1.4.5 LTS and 1.5.4, 4-thread concurrency, dtype zoo, and a 600-query memory soak. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Polars (>= 1.40) passes batch_size through Dataset.to_batches to size its streaming-engine morsels; the scanner accepted the kwarg and dropped it, so consumers had no control over batch granularity. The emitted batch size matters beyond Polars: engines parallelize per record batch (DuckDB slices each batch into 2048-row vectors across threads, Acero schedules one filter task per batch), so it is the downstream-parallelism granule. Also pins two consumer-contract properties with tests: batch_size reaches the emitted batches through both scanner() and to_batches(), and the table schema never contains view types — one view-typed column disables DuckDB filter pushdown for the entire table (duckdb-python#227). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
count_rows() previously routed through a full scanner() scan, pivoting every surviving chunk to count rows the chunk grid already knows. Three-way split instead: pruned chunks contribute nothing, chunks whose coordinate ranges PROVE the filter true contribute their size arithmetically (no I/O), and only undecided boundary chunks are scanned — reading just the columns the filter references. Strictness is decided by Arrow itself, no expression parsing: a chunk with conjunctive guarantee G satisfies filter F everywhere iff G AND NOT F is unsatisfiable, which get_fragments(filter=~F) over per-chunk guarantee fragments answers (the Iceberg inclusive/strict evaluator pattern). Everything undecidable — NaN spans, data-variable predicates, >4096 survivors, unsupported expression shapes — falls back conservatively to an exact boundary scan. Also fixes columns=[] being treated as "all columns" (falsy-list bug): an explicitly empty projection is now a real projection, served by zero-column batches whose row counts are chunk arithmetic. Measured (10M-row synthetic, time chunked by 10): - unfiltered count: full scan -> 0.01 ms, zero chunks read - 5-day mid-chunk time range: exact count in 19 ms reading only the 2 boundary chunks (11 interior chunks proven arithmetically) - data-variable filter (no guarantee): still row-exact, all chunks scanned as before Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Finely chunked sources (an hourly-stepped time axis is one chunk per hour) pay one store round-trip per surviving chunk. With coalesce_rows=N, runs of consecutive surviving chunks along the most finely chunked dimension are merged into single isel reads of at most N rows, after pruning and per-run (no gaps are ever read). On Zarr sources the merged read fetches its member chunks through the store's own concurrent batching instead of one request per chunk through the prefetch pool. Scanner path only: get_fragments() keeps one fragment per source chunk so fragment consumers (DataFusion, dask) retain their parallelism granularity. Measured on ARCO-ERA5 over anonymous GCS (1.32M hourly time chunks, prefetch=16, coalesce_rows=8M), identical results both ways: - 1 day x Iberia bbox: 2.23s -> 1.18s (24 reads -> 4) - 1 week x full globe (174M rows): 9.82s -> 6.62s (168 reads -> 24) Peak RSS grows with prefetch x merged-block size as documented (0.7 GB -> 1.2 GB here); size coalesce_rows/prefetch together. Off by default: memory scales with the merged block, and local or coarsely chunked sources gain nothing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Each scan previously created (and tore down) its own ThreadPoolExecutor. Beyond the per-scan setup cost, spawning OS threads from inside an engine's scan callback is exactly the embedded-engine hazard the ecosystem keeps rediscovering (pg_duckdb serializes all host calls; ParadeDB routes engine work through dedicated pools): thread startup contends with concurrent.futures' process-global shutdown lock and with the consumer's own pool management, observed as intermittent deadlocks when scans are driven from dask worker threads. The pool now lives on the dataset, its threads started at construction time — never inside an engine callback. Scans that stop early (LIMIT) cancel their queued loads instead of shutting the pool down. Single-block scans (a lazy round-trip window that maps onto one source chunk) skip the pool entirely. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
xql.to_dataset gains chunks= and coords=: data variables are reconstructed window by window on access, each window re-executing the engine's query narrowed to its coordinate range with the engine's own typed expression API (never rendered SQL text). Over a table registered through xarray-sql, the window's range predicate flows back into chunk pruning at the source, so accessing one output chunk reads only the source chunks it maps onto. The engine-specific surface of the existing DataFusion lazy path (expression building, per-dim distinct discovery, schema access) is extracted into LazyResultHandle implementations in the new xarray_sql.lazyscan module; SQLBackendArray and _build_lazy_scan are now engine-neutral, and the DataFusion wrapper path is byte-identical in behavior (39 pre-existing round-trip tests unchanged and green). Contiguous window requests become two-literal range predicates (engines can prune on them); stepped/fancy indexers fall back to explicit value lists, exact by construction. Engine support: - Polars LazyFrame/DataFrame: full chunked support; per-window fetches run on the streaming engine. Verified deadlock-free under threaded dask (6/6 stress runs) and correct on descending coordinates, stepped indexers, filtered and aggregated queries. - DataFusion DataFrame: unchanged, now also reachable through the engine-agnostic xql.to_dataset. - DuckDB relations: eager round-trip fully supported through a dedicated single engine thread (concurrent materialization of derived relations corrupts shared pending-query state — verified — so all handle calls are funnelled through one thread). Chunked reconstruction FAILS FAST with guidance instead of hanging: re-executing a relation that scans a Python-backed table while other threads start/stop deadlocks intermittently inside duckdb-python/CPython (~50% of runs, macOS/CPython 3.12; persists with SET threads=1, connection serialization, and pool pre-warming; Polars is clean under the identical topology). Reproducer kept in the working notes for an upstream report. coords="template" skips per-dim DISTINCT discovery when the result spans the template's full extent, making construction free of source reads for unfiltered scans on any engine. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Documents which engines support chunked reconstruction and why DuckDB relations fail fast (upstream duckdb-python deadlock, reproduced and isolated); states the memory bound (prefetch x pivoted-block-size) with the ARCO-ERA5 measurements that back it, and the coalesce_rows memory/latency tradeoff. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Asserts the SHAPE of the work — exactly which source chunks each query reads (via the scanner's iteration callback) and exact row counts — not just answers or timings, so a pruning/coalescing/fast-path regression that silently falls back to scanning everything fails loudly (the plan-shape-assertion pattern; benchmark-suite tripwires caught regressions plain tests missed in comparable projects). Measured this run (anonymous GCS, 1.32M-chunk hourly time axis): - day+bbox: 2.8s / exactly 24 chunk reads (4 reads coalesced, 1.6s) - week globe (174M rows): 16.6s / exactly 168 reads - count(*) over January (772M rows): 0.09s / ZERO reads (arithmetic) - Polars lazy round-trip: construction 0 reads; 1-day window compute reads exactly its 4 coalesced blocks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
register(..., geometry=(x_dim, y_dim)) appends a geometry point column
derived from the coordinate dims, synthesized per batch at scan time —
the pivot's coordinate columns already carry the values, so the column
costs an annotation plus (for WKB) a vectorized 21-byte encode of the
rows actually scanned.
Two encodings, chosen per destination:
- "wkb" (default): geoarrow.wkb extension metadata + CRS. DuckDB >=1.2
with spatial loaded ingests the column as GEOMETRY('OGC:CRS84'), so
ST_Within(geometry, ...) works with no ST_Point construction in SQL.
- "point": GeoArrow-native separated coordinates; the struct children
ARE the coordinate arrays (no copy, no parse) for consumers that
execute on native layouts — verified with GeoPandas 1.x
GeoDataFrame.from_arrow, CRS carried through.
Documented sharp edge, measured: engines do not push ST_* functions
into the scan, so a geometry-only predicate defeats chunk pruning and
encodes every row (~29x slower than the paired form on a 10M-row
grid: 101ms bbox vs 2.9s ST_Within-only vs 118ms bbox+ST_Within).
The geospatial docs now state the idiom: bbox conjuncts for pruning,
geometry for exactness.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
bcba93a to
f9919b6
Compare
Polars translates float is_in literals imprecisely (reproduced on 1.42: is_in([<stored coordinate>]) matches zero rows). The handle now renders float value lists as OR-chains of degenerate is_between ranges, which compare exactly. Non-float dims keep is_in. A stepped lazy-round-trip window over non-representable float coordinates (linspace(-45, 45, 19)) previously scattered nothing and returned garbage; now row-exact. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The eager path materialized unconditionally: a billion-row result (or a sparse one whose dense coordinate-product grid dwarfs its Arrow payload) exhausted memory rather than erroring. max_result_bytes= now raises a clean ValueError with the running size at both danger points — while collecting the Arrow stream, and before allocating dense arrays (checked against the coordinate product, which is where sparse diagonals blow up). Error before OOM, never truncate; opt-in and unlimited by default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…pruning The strictness analysis behind count_rows built one guarantee fragment per surviving chunk, capped at 4096 survivors; broader filters fell back to scanning every survivor (a January count on an hourly axis scanned 744 network chunks; step >= 100 over 1M chunks scanned ~1M). Replaced with hierarchical classification: each level buckets the surviving index lists into at most 4096 span-products, decides whole buckets at once with two guarantee-simplification passes (G AND NOT F unsat => proven; G AND F unsat => pruned), and recurses only into mixed cells. Per-chunk coordinate bounds are vectorized (np.minimum.reduceat, cached), so million-chunk axes classify in milliseconds. The prune side also refines the per-dimension pruning with cross-dimension information: paired ranges across dims ((A AND B) OR (C AND D)) no longer read the cross combinations. Measured: count over 1M single-row chunks with a near-universal filter 999,900 rows in 0.21s with ZERO chunk reads (previously scanned every survivor); cross-dim paired ranges read 2 boundary chunks instead of 4; 200 randomized differential checks against numpy ground truth all exact. NaN spans, non-numeric dims and simplifier-rejected expressions still land conservatively in the exact boundary scan. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The previous version used the single-lat-chunk fixture (no cross to refine) and hand-computed the wrong expected count; now differential against numpy on a grid chunked in both dims. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Count-based admission (prefetch=N blocks) under-uses the network with small chunks and overshoots memory with coalesced ones, because block sizes vary. prefetch_bytes= gates admission on estimated pivoted bytes in flight (rows x schema row width), the Lance io_buffer_size semantics; prefetch keeps bounding concurrency (thread count). With a byte budget, raising coalesce_rows no longer multiplies peak RSS. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Engines never push ST_* functions into the scan, so a geometry-only WHERE reads every chunk (documented, measured 29x). This helper renders the bbox range conjuncts from a geometry's envelope (a (xmin, ymin, xmax, ymax) tuple or anything with .bounds — shapely geometries qualify), making the bbox+geometry idiom one f-string instead of hand-copied bounds. Optional pad= margin for ST_DWithin-style predicates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mple) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ot streams The two results the re-execution path could not serve now work: spill=True streams the result exactly ONCE with bounded memory into a temporary Parquet file — through the engine handle where one exists (DuckDB spills on its dedicated engine thread, dodging the duckdb-under-worker-threads deadlock entirely; Polars uses its native streaming sink; DataFusion streams batch by batch) or straight from the Arrow stream for one-shot tables/readers — and the ordinary lazy reconstruction then runs over a Polars scan of that file, whose per-window predicates get Parquet row-group pruning. The temp file is deleted when the returned Dataset is garbage collected. Trade-off vs re-execution, by design: one full pass plus temporary disk instead of pay-per-window — the right shape when most of the result will be touched; Polars/DataFusion re-execution remains the default for window-at-a-time access. Verified: the previously deadlocking DuckDB chunked compute topology now passes 6/6 stress runs with repeated full computes; the source is read exactly once (chunk counter), and windows read only the spill file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
alxmrs
left a comment
There was a problem hiding this comment.
First pass at a review. Really excited about this.
| * ``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). |
There was a problem hiding this comment.
Ideally, it would be good to compare the "arrow-datafusion" vs the "table-provider-datafusion" paths deliberately.
There was a problem hiding this comment.
Agreed, I changed it in e3f95de. so the two paths are now explicit engines datafusion runs the native table provider and raises at startup if xarray_sql._native is missing (no silent fallback anymore) and datafusion-arrow which corresponds to the pure-Python pyarrow-dataset path. I'll re-run the full benchmark with all 4 engines and post the results here as well as update the docs.
There was a problem hiding this comment.
Results are in. Full matrix (portable cases x 4 engines x 5 cold reps, fresh process per rep, every answer asserted against the xarray reference) on e2-standard-8/16/32 in us-central1; each VM built the native module from source before its first datafusion cell (313s to 516s depending on the machine). Docs updated in e7c9438.
Colors rank engines within each row (🟩 fastest to 🟥 slowest); mind that on 05 the whole spread is a few milliseconds.
e2-standard-8, medians of 5:
| Case | DataFusion (native) | DataFusion (pyarrow) | DuckDB | Polars | xarray reference |
|---|---|---|---|---|---|
| 01 NDVI | 🟩 4.308 s | 🟨 4.380 s | 🟧 5.325 s | 🟥 5.725 s | 0.444 s |
| 02 Climatology | 🟧 7.235 s | 🟥 8.746 s | 🟩 4.622 s | 🟨 4.936 s | 2.367 s |
| 03 Zonal mean | 🟥 4.763 s | 🟨 3.614 s | 🟩 2.958 s | 🟧 3.705 s | 0.829 s |
| 04 Anomaly | 🟨 10.416 s | 🟥 16.262 s | 🟩 9.817 s | 🟧 13.837 s | 4.410 s |
| 05 Forecast skill | 🟩 1.791 s | 🟧 1.814 s | 🟨 1.797 s | 🟥 2.405 s | 0.247 s |
| 06 Zonal stats | 🟩 2.390 s | 🟨 6.131 s | 🟥 9.112 s | unsupported | 1.813 s |
What I take from the native vs pyarrow comparison: the native table provider wins where rows are consumed in bulk (06 range JOIN: 2.4 vs 6.1 s; 04 self JOIN: 10.4 vs 16.3 s at a third of the peak memory; ~1.2x on 02), the two are at parity on the read-bound cases (01, 05), and on the plain zonal mean 03 the pyarrow path is actually faster. DuckDB stays the fastest consumer of the shared scan on the plain group-bys. The 16 and 32 vCPU sizes show no practical scaling difference, same as before (shared-core e2 plus a single-stream cold read dominate).
e2-standard-16 and e2-standard-32 tables (click to expand)
e2-standard-16, medians of 5:
| Case | DataFusion (native) | DataFusion (pyarrow) | DuckDB | Polars | xarray reference |
|---|---|---|---|---|---|
| 01 NDVI | 🟩 4.181 s | 🟨 4.813 s | 🟧 5.915 s | 🟥 6.566 s | 0.607 s |
| 02 Climatology | 🟧 9.190 s | 🟥 11.174 s | 🟩 6.888 s | 🟨 7.519 s | 2.326 s |
| 03 Zonal mean | 🟧 5.042 s | 🟨 4.993 s | 🟩 4.109 s | 🟥 5.470 s | 1.102 s |
| 04 Anomaly | 🟨 13.105 s | 🟥 19.615 s | 🟩 11.711 s | 🟧 14.192 s | 4.202 s |
| 05 Forecast skill | 🟩 2.828 s | 🟧 2.966 s | 🟨 2.859 s | 🟥 3.532 s | 0.338 s |
| 06 Zonal stats | 🟩 2.914 s | 🟥 8.633 s | 🟨 6.791 s | unsupported | 2.353 s |
e2-standard-32, medians of 5:
| Case | DataFusion (native) | DataFusion (pyarrow) | DuckDB | Polars | xarray reference |
|---|---|---|---|---|---|
| 01 NDVI | 🟩 3.648 s | 🟨 4.241 s | 🟧 5.179 s | 🟥 6.038 s | 0.489 s |
| 02 Climatology | 🟧 6.029 s | 🟥 7.860 s | 🟩 3.950 s | 🟨 4.365 s | 1.882 s |
| 03 Zonal mean | 🟥 3.879 s | 🟨 3.388 s | 🟩 2.799 s | 🟧 3.449 s | 0.726 s |
| 04 Anomaly | 🟨 9.445 s | 🟥 14.057 s | 🟩 5.956 s | 🟧 11.250 s | 4.066 s |
| 05 Forecast skill | 🟨 1.708 s | 🟧 1.842 s | 🟩 1.672 s | 🟥 2.172 s | 0.237 s |
| 06 Zonal stats | 🟩 2.200 s | 🟥 5.830 s | 🟨 5.371 s | unsupported | 1.593 s |
| "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"] |
There was a problem hiding this comment.
Nice, it's a good idea to vary across VM sizes.
There was a problem hiding this comment.
I was unsure whether to use standard VMs or compute optimized VMs. WDYT?
There was a problem hiding this comment.
I used standard ones because for me it was more about comparing the different engines under identical resources regardless of what those were, and seeing whether the performance scaled across engines.
However, it is true that analytics engines are often compared on compute optimised VMs (e.g. ClickBench) and perhaps that'd be what users pick when running xarray-sql workloads too. Do you have a suggestion for a particular family and sizes?
|
|
||
| === "DuckDB" | ||
|
|
||
| **Re-executing relations from worker threads deadlocks.** |
There was a problem hiding this comment.
Yep 😅 at least the library fails fast there (chunks= on a DuckDB relation raises instead of hanging) and the guard is pinned by a test, so it cannot silently regress into a hang.
I think it's worth dedicating a little bit more of time on this one, but perhaps we can do so in the future, WDYT?
| assert len(reads) <= 2 | ||
|
|
||
|
|
||
| def test_prefetch_bytes_scan_reads_every_block_once(): |
There was a problem hiding this comment.
Can we make this test file more succinct? Let's try to make them test the core contract.
Further, I find it's good to have claude "hunt" for possible bugs through tests, and then keep those conditions as tests.
I also prefer tests that are more on the integration side of things that unit.
| (including extensions), and ``xql.to_dataset`` rebuilds a labeled | ||
| Dataset from the Arrow result. | ||
| """ | ||
|
|
There was a problem hiding this comment.
I could see us introducing a "testing" submodule where we include the benchmark engine harness. That utility would be useful for all tests and benchmarks. WDYT?
There was a problem hiding this comment.
In such a module, we could include utilities for fuzzing / simulation tests (e.g. maybe with an opt dep if hypothesis).
|
When this PR lands, I think we can say it fixes #4 (!!)! |
alxmrs
left a comment
There was a problem hiding this comment.
Ok, I think I'll have to make another round of review (3/4). But, what I'm reading so far is really excellent work. Bravo, and thank you so much for your contribution.
| 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. |
There was a problem hiding this comment.
I kind of like the "seam" concepts here, but I also find them to be claudish (context from the conversation, not self contained).
| @@ -0,0 +1,111 @@ | |||
| """Engine-adapter dispatch for :func:`xarray_sql.register`. | |||
|
|
|||
| An *engine adapter* implements one seam: given an engine's native | |||
There was a problem hiding this comment.
See note about seams. Otherwise, this looks good.
| *, | ||
| chunks: Chunks = None, | ||
| **kwargs: Any, | ||
| ) -> Any: |
There was a problem hiding this comment.
optional: Let's use generic typing for Con to make this protocol more type safe.
| from ..df import Chunks | ||
|
|
||
|
|
||
| @runtime_checkable |
There was a problem hiding this comment.
If we used generic types, then maybe we wouldn't need this.
|
|
||
|
|
||
| def register( | ||
| con: Any, |
|
|
||
| def distinct(self, column: str) -> np.ndarray: ... | ||
|
|
||
| def fetch( |
There was a problem hiding this comment.
Do we want this to return an iterator of recordbatches? Or, is a light light weight?
There was a problem hiding this comment.
e.g. it could also return a record batch iterator IIUC.
|
|
||
| def spill_parquet(self, path: str) -> None: ... | ||
|
|
||
| # Handles may additionally offer ``stream(columns)``, yielding the |
There was a problem hiding this comment.
IIUC, we may want to make the oppose the default: we should always iterate through the dataset if we can (to scale to very large data), but sometimes we are OK with materializing the dataset.
Do I have the wrong mental model?
There was a problem hiding this comment.
I very well could simply not "grok" how the handle system works yet. Would love your perspective.
|
|
||
| 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 |
There was a problem hiding this comment.
What is "the re-executable contract"?
| """ | ||
|
|
||
| supports_chunked = False | ||
| """Chunked (lazy) reconstruction is disabled for DuckDB relations. |
There was a problem hiding this comment.
Would be nice to track this in a GH issue, if we ever aim to address this. Very cool analysis here.
| 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 |
alxmrs
left a comment
There was a problem hiding this comment.
Ok, given the last outstanding issues are addressed from the previous review (we can also just file follow up issues), then I think this will be ready to merge. Excellent work, Miguel.
GEOBENCH_ENGINE=datafusion previously fell back to the pyarrow-dataset path when the native module was missing, recording which path ran only in the flavor field. Split them: datafusion requires the compiled native module and raises when it is absent; datafusion-arrow selects the pure-Python pyarrow path. Drop the GEOBENCH_NO_NATIVE hook the fallback needed and add datafusion-arrow to the suite matrix. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The suite shipped only the pure-Python sources, so remote (and local tarball-tree) datafusion cells could never run the native table provider. Ship the Rust crate alongside and provision xarray_sql._native in the first datafusion cell per VM: copy the driver's build when it imports on that platform, otherwise install rustup and build once with the project's maturin backend via pip wheel. The tarball is byte-deterministic so the digest-keyed source root — and with it the build cache — survives warm-VM reruns, and the provisioning outcome is recorded in the cell result so a failed build surfaces as that cell's error. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
EngineContext branched on the engine name in __init__, from_dataset, and sql_to_dataset. Each engine is now a subclass implementing _connect/_register/_execute (the native path overrides from_dataset and sql_to_dataset, keeping its byte-for-byte XarrayContext behavior), selected once through the _IMPLS registry that also defines the valid engine names. No behavior change; flavors and error messages are identical. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Three reliability fixes in the VM driver: the _meta jsonl line is written from a local record instead of results[-1], which another VM thread could replace between the append and the write; the probe retry loop no longer sleeps 90s after its final failed attempt; and a run where any VM never produced results now exits nonzero instead of passing for a complete one. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
PEP 723 inline metadata with the driver dependencies and the editable xarray-sql source, matching the case scripts, so 'uv run benchmarks/geospatial/engine_suite.py' works without a prepared environment. Documents the one operational constraint: remote runs launch from outside the repo, since Coiled's package sync resolves the repo's uv.lock when the cwd contains one. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
v1/v2 referred to iterations of a working session, not to anything in the tree. The two paths are now labeled by behavior: the re-scannable stream (no pushdown) and the default register() pushdown dataset. Also fixes the stale usage line, which pointed at a file that no longer exists. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
benchmarks/era5_out_of_core.py was assertions wearing a benchmark coat: it verified exactly which source chunks each query reads and exact row counts, with timing prints on top. Those assertions now live in tests/test_era5_integration.py under the integration marker (first use; registered in pyproject), which the CI unit run already excludes via -m 'not integration'. Adds a re-scannability test the benchmark only exercised implicitly. Timing across engines stays with benchmarks/geospatial/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Assert that the stream, pushdown, and in-memory-ceiling paths return the same answers, report medians with min/max spread instead of best-of-n, and drop the DataFusion section: cross-engine timing lives in benchmarks/geospatial/, while the three DuckDB adapter paths are what this file uniquely measures. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
engines.md links LanceDB where it names the shared pattern, and XarrayPushdownDataset's docstring gains a references section pointing at LanceDataset, the production subclass built the same way. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The suite now provisions the native module on its VMs, so the cross-engine table gains the DataFusion-native column measured on the same hardware as the other engines. Numbers from the fresh 3-VM run (duckdb 1.5.5), the scope paragraph rewritten for the four-column reality, and the ChunkedArray bug note now cites its regression test instead of describing the fix relative to the writing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DataFusion's section now carries the same relation note DuckDB's has: zarr-datafusion is the engine-native path for plain Zarr, the adapter covers everything else xarray opens plus the round-trip. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A fictional adapter modeled on the DuckDB one: matches by type inspection, register through XarrayPushdownDataset, dispatch via the register_adapter decorator, and the dataset-protocol vs plain-stream trade-off stated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
duckdb and polars floors (and their version-floor comment) were duplicated between the engine extras and the test extra; the test extra now references xarray_sql[duckdb,polars], so a floor bump has one home. Lockfile regenerated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ckdb # Conflicts: # README.md # pyproject.toml # tests/test_df.py # xarray_sql/df.py
Four were introduced by the recent benchmark-harness work (untyped result dict, GzipFile into tarfile.open); the rest predate it on the branch (Optional narrowing mypy cannot see through subscripts and closures, Any-typed returns from pyarrow calls). No behavior change; narrowing is expressed with locals and the pyarrow returns are wrapped in their declared types. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A zero batch_size never advances the zero-column scan's row loop (IndexError today, an infinite loop without it); negative values break batch synthesis. Fail with ValueError when the dataset is built instead of mid-scan. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sections follow the dataset's contract (protocol surface, consumer integrations, projection, pruning/counting, scheduling knobs, re-scannability, lifecycle). The four structurally identical count_rows tests become one parametrized contract table. New pinned conditions from adversarial probing: NaN/NaT coordinate chunks are scanned rather than pruned, filter-only columns are read but not returned, an empty projection with a filter reads only the filter column, concurrent and alternating scans stay exact, degenerate tuning values degrade the schedule but never the answer, and non-positive batch_size fails at construction. Every previous regression pin is kept. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
test_era5_integration.py named the instrument, not the subject; it is now test_arrow_dataset_integration.py, pairing with the contract file the way xee pairs ext_test with ext_integration_test. The suite is a dataset x engine matrix: StoreCase entries with expectations computed from declared cadence plus the store's own coordinates, and an engine registry (DuckDB SQL, Polars expressions, DataFusion SQL) that five of the seven tests parametrize over. New coverage: per-engine windowed exactness against a direct xarray read, projection isolation, and concurrent disjoint queries. Fragment consumers (DataFusion) are flagged in the registry since scanner-level coalescing does not apply to them, and the Polars lazy round-trip skips on polars >= 1.43, whose streaming re-execution regressed from 7s to over 10 minutes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
From adversarial review: a run whose cells all failed still exited 0 (only missing _meta was checked); the SKIPPED marker outranked a nonzero exit status, recording crashed cases as skips; and an interrupted source extraction left a corrupt tree every retry reused. The driver now exits nonzero on any error/timeout cell, trusts exit status over the skip marker, and marks extraction complete only after it finishes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Summary
This branch prototypes xarray-sql as the Xarray ↔ engine translator: the library owns exactly two seams — register (a lazy
xarray.Datasetbecomes a table on an engine's own connection) and round-trip (any engine's Arrow result becomes a labeledxr.Datasetagain). 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"]One object serves every engine:
xql.arrow_dataset(ds)is a realpyarrow.dataset.Dataset(the pattern Lance uses forLanceDataset), so DuckDB registers it, Polars scans it viascan_pyarrow_dataset, DataFusion consumes it viaregister_dataset(through the fragments API), and Dask maps overget_fragments(). Ibis works throughibis.duckdb.from_connectionwith 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"]Three properties worth calling out:
FileSystemDatasetfragments carry each chunk's coordinate range as apartition_expression(their paths are never opened). Sound for every predicate shape — equality,OR,IN,NOT.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, lazyxr.Dataset— each window re-executes the engine's query narrowed to its coordinate range, which flows back into chunk pruning at the source.spill=Trueprovides 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"]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):GROUP BY, native res, full table registereddocs/performance.md)AVGscan (vs v1 stream)to_dataset)Peak scan memory is a contract, not an accident: bounded by
prefetch × pivoted-block-sizeregardless 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, anomalyJOIN, 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.pyrepeats the portable cases across DataFusion (both the native table provider and the pure-Python pyarrow path, as separatedatafusion/datafusion-arrowengines), DuckDB, and Polars on GCE VMs, provisioning the compiled native module on each VM. The write-up isdocs/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) andXarrayArrowStream(re-scannable C-stream fallback).xarray_sql/roundtrip.py+xarray_sql/lazyscan.py— engine-agnosticxql.to_datasetaccepting DuckDB relations, Polars frames, DataFusion DataFrames,pyarrowtables/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 pointgeometrycolumn (WKB for DuckDB-nativeGEOMETRY, or GeoArrow-separated for GeoPandas/lonboard), CRS tagged;xql.bbox_conjunctsrenders prunable bbox predicates from any geometry's envelope.df.py/ds.py.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), reworkeddocs/examples.mdand README.[duckdb](duckdb>=1.4) and[polars](polars>=1.33) extras.Testing
275 tests, all green. Beyond unit coverage:
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.max_result_bytesenforcement on eager collection and dense allocation.ChunkedArrayslipping into the pivot's fast path — all fixed with regression tests.Known limitations / follow-ups
chunks=raises immediately): re-execution from worker threads deadlocks inside duckdb-python 1.4–1.5 on CPython 3.12, sospill=Trueis the chunked path there. Documented with the full story indocs/limitations.md.XarrayPushdownDatasetsubclasses pyarrow's cythonDatasetwithout initializing the native base (as Lance does); dangerous inherited members are stubbed, and the contract is pinned by tests — re-verify on pyarrow upgrades.is_infloat literals imprecisely (reproducible without xarray-sql); documented, use range predicates. Polars also has no geometry types — thegeometrycolumn arrives as plain binary there.🤖 Generated with Claude Code