xarray-sql¶

Engine adapters — the register seam of xarray-sql.

xarray-sql translates data, not queries, across two seams — the two boundaries between xarray and a query engine that neither side builds for itself: register (a lazy xarray.Dataset becomes a table on the engine's own connection; this package) and round-trip (an Arrow result becomes a labeled Dataset again; xarray_sql.to_dataset). SQL dialects, geometry, H3, and optimizers belong to each engine and its extension ecosystem.

Adapters register themselves on import via register_adapter; register dispatches on the connection type.

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=...)

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)

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)

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=...)

Bridge between cftime calendars and Arrow/DataFusion types.

cftime (https://unidata.github.io/cftime/) provides datetime objects for calendars used in climate science — noleap, 360-day, all-leap, julian, etc. Arrow and DataFusion have no native concept of non-Gregorian calendars, so this module handles the conversion in two tiers:

• Gregorian-like calendars (standard, gregorian, proleptic_gregorian, noleap/365_day, all_leap/366_day): mapped to pa.timestamp('us') so that string-based SQL filters like WHERE time > '1980-01-01' work naturally. Microsecond resolution avoids the 1678–2262 overflow of nanoseconds while preserving sub-second precision.

• Non-Gregorian calendars (360_day, julian): mapped to pa.int64() with xarray:units and xarray:calendar metadata on the Arrow field. This preserves the original CF-convention encoding losslessly. A cftime() DataFusion UDF (registered automatically by XarrayContext.from_dataset) provides ergonomic SQL filtering.

SELECT * FROM ds360 WHERE time > cftime('0500-01-01')

Reconstruct xarray Datasets from SQL query results.

The inverse of the forward Dataset-to-table pivot done by xarray_sql.df.pivot. Internally defines an XarrayDataFrame wrapper around the DataFusion DataFrame returned by XarrayContext.sql, with a XarrayDataFrame.to_dataset method that round-trips a query result back to xr.Dataset.

Reconstruction is controlled by the chunks argument to XarrayDataFrame.to_dataset -- the xarray idiom for tuning how a result is partitioned -- rather than by reflecting on the query plan:

• Eager (chunks=None, or the default "inherit" when the result keeps no multi-chunk source dimension): the plan executes exactly once via execute_stream and the result is scattered into a dense in-memory Dataset. This is the right default for reductions (aggregations), whose results are small, and it never re-executes.
• Lazy / chunked (chunks is a mapping, "auto", or "inherit" over a multi-chunk source dimension): data variables are backed by SQLBackendArray wrapped in xarray.core.indexing.LazilyIndexedArray and chunked via xarray's configured chunk manager (dask, cubed, ...). Each chunk maps onto the source partitions and reads its coordinate range on access by translating the indexer into a DataFusion filter expression, so only the requested partitions are materialized as Arrow RecordBatch es and scattered into numpy.

.compute() materializes the whole Dataset in memory.

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<x: double, y: double> 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.

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}'))"
)

Re-executable engine handles behind the lazy chunked round-trip.

The lazy path of to_dataset(chunks=...) re-executes the engine's query per accessed chunk, narrowed to that chunk's coordinate window and columns. That requires the engine result to be re-executable — a handle onto the query, not a one-shot stream of its rows. Each handle here adapts one engine's native lazy surface to the three operations the reconstruction needs:

• schema — result column names/types, without executing the query;
• distinct — one column's distinct values (coordinate discovery; the caller sorts);
• fetch — the result narrowed by per-dimension windows and projected to the requested columns, as Arrow batches.

Windows are passed as DimSpec values instead of rendered SQL so each engine can express them with its own typed expression API — strings would re-open every literal-formatting pitfall (timestamps, floats, quoting) per dialect.

Handles compose with the registration seam: when the wrapped query scans a Dataset registered through xarray-sql's pushdown machinery, the per-chunk range filter flows back through the engine into XarrayPushdownDataset, so each output chunk's access reads only the source chunks it maps onto.

PROJ-backed CRS transforms for SQL — the optional geo extension.

Geospatial SQL dialects expose coordinate reference system (CRS) transforms as a scalar function — PostGIS and DuckDB-spatial both call it ST_Transform — because a CRS transform is row-independent: each point's new coordinate depends only on its own old coordinate. This module brings the same capability to xarray-sql as a vectorized scalar UDF over Arrow arrays::

SELECT x, y,
       reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon,
       reproject(x, y, 'EPSG:32610', 'EPSG:4326')['y'] AS lat
FROM grid

The CRS pair is part of the query, not baked in at registration time, so one registered UDF serves any transform — and, because the arguments are ordinary SQL expressions, the CRS may even vary per row (e.g. a CASE expression selecting the UTM zone from the longitude).

Design notes:

• Both output coordinates come from one call, returned as an Arrow struct {x, y} (in always_xy order: easting/longitude first). Splitting the transform into two scalar UDFs would run PROJ twice per row and, worse, evaluate the two projections concurrently on separate expression trees.
• All pyproj work runs on a dedicated pool of Python threads. DataFusion's runtime workers are not Python-created threads, and pyproj (< 3.8, see pyproj#1541) leaves a dangling PJ_CONTEXT behind when their ephemeral Python thread states are torn down, so calling pyproj in place segfaults — the UDF hands each batch to the pool instead. Pool threads are long-lived, so each caches one transformer per CRS pair (transformers must not be shared across threads), amortizing the expensive construction — PROJ database lookups and candidate-operation selection — across record batches. Concurrent partitions still transform in parallel across the pool.
• Any CRS spelling pyproj.CRS accepts works: authority codes (EPSG:4326), WKT, PROJ strings (+proj=utm +zone=10), etc. An unknown CRS raises pyproj.exceptions.CRSError and fails the query loudly rather than returning wrong coordinates.
• Non-finite or NULL input coordinates yield NaN output (PROJ itself would return inf); NULL CRS arguments yield NaN as well.

Requires pyproj (pip install xarray-sql[geo]). When pyproj is installed, xarray_sql.XarrayContext registers reproject() automatically; register is the explicit hook for plain DataFusion SessionContext objects or custom UDF names.

SELECT reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon
FROM grid

Lazy Arrow stream reader for xarray Datasets.

This module provides XarrayRecordBatchReader, which implements the Arrow PyCapsule Interface (arrow_c_stream) to enable zero-copy, lazy streaming of xarray data to DataFusion and other Arrow consumers.

The implementation delegates to PyArrow's RecordBatchReader for the actual stream implementation, wrapping xarray block iteration in a generator.

# This query will skip loading partitions with time < '2020-02-01'
result = ctx.sql('SELECT * FROM air WHERE time > "2020-02-01"').collect()

Engine-agnostic round-trip: Arrow query results → labeled xr.Dataset.

The second seam of xarray-sql. Any engine's result — a DuckDB relation, a pyarrow.Table, a pyarrow.RecordBatchReader, or any object implementing the Arrow PyCapsule stream protocol — plus the registered Dataset as a template is enough to rebuild a labeled, metadata-carrying Dataset. Nothing here is engine-specific: results arrive as Arrow record batches regardless of which engine executed the SQL.

Reconstruction is eager by default (the result is materialized once into a dense in-memory Dataset). Passing chunks= selects the lazy/chunked path instead: data variables are reconstructed on access, window by window, by re-executing the engine's query narrowed to each chunk's coordinate range. That requires the result to be re-executable — a Polars LazyFrame (or eager DataFrame) or a DataFusion DataFrame — not a one-shot Arrow stream; see xarray_sql.lazyscan. DuckDB relations are re-executable but refuse the chunked path (a thread-safety limitation noted on DuckDBHandle); pair them with spill=True instead.

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)