diff --git a/README.md b/README.md index a254d29..fb9f0f8 100644 --- a/README.md +++ b/README.md @@ -210,8 +210,10 @@ against an xarray/array reference** to floating-point tolerance: reproduces the published result that GraphCast beats Pangu at every lead. * **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a table of regions. -* **Reprojection and regridding** — a scalar PROJ UDF (validated against Earth - Engine's own geodesy via [Xee](https://github.com/google/Xee)) and a +* **Reprojection and regridding** — a `reproject(x, y, src_crs, dst_crs)` + scalar PROJ UDF, shipped as the optional pyproj extension + (`pip install xarray-sql[geo]`, validated against Earth Engine's own + geodesy via [Xee](https://github.com/google/Xee)) and a sparse-weight-table `JOIN` (regridding real SRTM terrain). Every case matches its array reference. The headline finding: these operations diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py index 75b4d68..02815e3 100644 --- a/benchmarks/geospatial/07_reproject_udf.py +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -24,9 +24,13 @@ world already does it — PostGIS ``ST_Transform`` and DuckDB-spatial ``ST_Transform`` are scalar PROJ wrappers. -So we register a PROJ-backed scalar UDF and reproject in SQL:: +xarray-sql ships that UDF as its pyproj extension (``xarray_sql.proj``): +with pyproj installed, every ``XarrayContext`` speaks CRS out of the box, +and the CRS pair is part of the query rather than baked into the UDF:: - SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat + 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 reference is Earth Engine itself.** There is *one* dataset: a single UTM @@ -39,9 +43,9 @@ for the *same* pixels. The reference is a different geodesy engine, not PROJ again, and they agree to sub-metre precision. -PROJ's context is not thread-safe and DataFusion evaluates projection -expressions concurrently, so we return *both* coordinates from one -struct-returning UDF and keep the source in a single chunk (one serial UDF). +The extension returns *both* coordinates from one struct-returning call +(one PROJ transform per row) and runs all PROJ work on its own worker +pool, so the query parallelizes across partitions safely. Requires Earth Engine access: ``earthengine authenticate`` once, then an initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise. @@ -49,11 +53,7 @@ from __future__ import annotations -import numpy as np -import pyarrow as pa -import pyproj import xarray as xr -from datafusion import udf import xarray_sql as xql @@ -73,46 +73,6 @@ _SCALE_M = 2_000 # 2 km pixels → a ~50×60 grid -def register_reproject_udf( - ctx, src_crs: str, dst_crs: str, name: str = "reproject" -) -> None: - """Register a ``reproject(x, y) -> {lon, lat}`` PROJ scalar UDF. - - Mirrors ``xarray_sql.cftime.make_cftime_udf``: a vectorized scalar UDF over - Arrow arrays. ``always_xy=True`` keeps argument order (easting, northing) → - (lon, lat) regardless of CRS axis conventions. Like PostGIS/DuckDB - ``ST_Transform``, it returns *both* output coordinates from one call — here - as an Arrow struct, so callers write ``reproject(x, y)['lon']``. - - Returning a struct (rather than two separate UDFs) is deliberate: PROJ's - context is not thread-safe, and DataFusion evaluates independent projection - expressions concurrently — two PROJ UDFs in one SELECT race and crash. One - struct-returning UDF does the transform exactly once per row, on one thread. - """ - ret = pa.struct([("lon", pa.float64()), ("lat", pa.float64())]) - - def _fn(x: pa.Array, y: pa.Array) -> pa.Array: - # Build the Transformer inside the call so it lives on the worker - # thread that uses it (PROJ contexts are thread-bound). - transformer = pyproj.Transformer.from_crs( - src_crs, dst_crs, always_xy=True - ) - xs = np.asarray(x.to_numpy(zero_copy_only=False), dtype="float64") - ys = np.asarray(y.to_numpy(zero_copy_only=False), dtype="float64") - lon, lat = transformer.transform(xs, ys) - return pa.StructArray.from_arrays( - [ - pa.array(np.asarray(lon, "float64")), - pa.array(np.asarray(lat, "float64")), - ], - names=["lon", "lat"], - ) - - ctx.register_udf( - udf(_fn, [pa.float64(), pa.float64()], ret, "immutable", name) - ) - - def _open_ee_lonlat_grid() -> xr.Dataset: """Open ``ee.Image.pixelLonLat()`` on a UTM grid via Xee. @@ -153,17 +113,19 @@ def main() -> None: f"{_SRC_CRS} → {_DST_CRS}" ) + # XarrayContext registers reproject() automatically (the pyproj + # extension). The chunking deliberately splits the ~60-row grid into + # 15-row slabs → 4 partitions, forcing DataFusion to evaluate the UDF + # concurrently: the extension runs PROJ on its own worker pool, so + # parallel partitions are safe (previously this required one chunk → + # one partition → a serial UDF). ctx = xql.XarrayContext() - # Single chunk → single partition → serial UDF (PROJ is not thread-safe). - ctx.from_dataset( - "grid", ds, chunks={"y": ds.sizes["y"], "x": ds.sizes["x"]} - ) - register_reproject_udf(ctx, _SRC_CRS, _DST_CRS) + ctx.from_dataset("grid", ds, chunks={"y": 15, "x": ds.sizes["x"]}) - sql = """ + sql = f""" SELECT x, y, - reproject(x, y)['lon'] AS lon, - reproject(x, y)['lat'] AS lat + reproject(x, y, '{_SRC_CRS}', '{_DST_CRS}')['x'] AS lon, + reproject(x, y, '{_SRC_CRS}', '{_DST_CRS}')['y'] AS lat FROM grid ORDER BY y, x """ diff --git a/benchmarks/geospatial/09_warp.py b/benchmarks/geospatial/09_warp.py index 33d3b5f..4516e2b 100644 --- a/benchmarks/geospatial/09_warp.py +++ b/benchmarks/geospatial/09_warp.py @@ -55,11 +55,9 @@ from __future__ import annotations import numpy as np -import pyarrow as pa import pyproj import shapely.geometry as sgeom import xarray as xr -from datafusion import udf import xarray_sql as xql @@ -81,35 +79,6 @@ _DST_SCALE_DEG = 0.02 # ~2 km target cells -def _register_reproject_udf(ctx, src_crs, dst_crs, name="reproject"): - """Register ``reproject(a, b) -> {x, y}`` — case 07's PROJ scalar UDF. - - Vectorized over each Arrow batch; ``always_xy=True`` keeps (easting, northing) - /(lon, lat) order. Returns both output coordinates from one struct-returning - call (PROJ contexts are not thread-safe, so one UDF, evaluated serially). - """ - ret = pa.struct([("x", pa.float64()), ("y", pa.float64())]) - - def _fn(a: pa.Array, b: pa.Array) -> pa.Array: - transformer = pyproj.Transformer.from_crs( - src_crs, dst_crs, always_xy=True - ) - xs = np.asarray(a.to_numpy(zero_copy_only=False), dtype="float64") - ys = np.asarray(b.to_numpy(zero_copy_only=False), dtype="float64") - ox, oy = transformer.transform(xs, ys) - return pa.StructArray.from_arrays( - [ - pa.array(np.asarray(ox, "float64")), - pa.array(np.asarray(oy, "float64")), - ], - names=["x", "y"], - ) - - ctx.register_udf( - udf(_fn, [pa.float64(), pa.float64()], ret, "immutable", name) - ) - - def _open_srtm( grid_crs: str, scale: tuple[float, float], xy_names ) -> xr.DataArray: @@ -212,8 +181,9 @@ def main() -> None: f"{len(tlat)}×{len(tlon)} ({_SRC_CRS} → {_DST_CRS})" ) + # XarrayContext registers reproject() automatically (the pyproj + # extension) — the direction is spelled in the query itself. ctx = xql.XarrayContext() - _register_reproject_udf(ctx, _DST_CRS, _SRC_CRS) # The target grid as a (dst_lat, dst_lon) table. LON, LAT = np.meshgrid(tlon, tlat) @@ -227,10 +197,10 @@ def main() -> None: ctx.from_dataset("target", target, chunks={"cell": LON.size}) # 1) SQL reprojects the target grid into the source CRS (case 07's UDF). - reproj_sql = """ + reproj_sql = f""" SELECT dst_lat, dst_lon, - reproject(dst_lon, dst_lat)['x'] AS sx, - reproject(dst_lon, dst_lat)['y'] AS sy + reproject(dst_lon, dst_lat, '{_DST_CRS}', '{_SRC_CRS}')['x'] AS sx, + reproject(dst_lon, dst_lat, '{_DST_CRS}', '{_SRC_CRS}')['y'] AS sy FROM target """ show_sql(reproj_sql, label="SQL — reproject target grid (PROJ UDF)") diff --git a/benchmarks/geospatial/README.md b/benchmarks/geospatial/README.md index 9daf102..44f8f6b 100644 --- a/benchmarks/geospatial/README.md +++ b/benchmarks/geospatial/README.md @@ -21,7 +21,7 @@ plain-English definition of the operation, and computes the same numbers. | 04 | `04_anomaly.py` | climatology broadcast-subtract | climatology CTE self-`JOIN` | | 05 | `05_forecast_skill.py` | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` + aggregate | | 06 | `06_zonal_vector.py` | rasterize + mask per region | range `JOIN` raster↔regions | -| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()`), à la PostGIS `ST_Transform` | +| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()` from the pyproj extension), à la PostGIS `ST_Transform` | | 08 | `08_regrid_weights.py` | interpolation to a new grid | sparse-weight table `JOIN` + weighted `GROUP BY` | | 09 | `09_warp.py` | reproject **and** resample (warp) | reproject **UDF** (07) → weight table `JOIN` (08) | diff --git a/docs/geospatial.md b/docs/geospatial.md index 085a1b0..2c0b317 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -212,12 +212,16 @@ paradigm. They split cleanly along one line: **is the operation row-independent? **Reprojection is.** Moving a coordinate from one CRS to another depends only on that coordinate, so it is a *scalar function* — exactly what PostGIS and -DuckDB-spatial already ship as `ST_Transform`. We register a PROJ-backed scalar -UDF (mirroring the `cftime()` UDF already in `xarray_sql/cftime.py`) and -reproject in SQL: +DuckDB-spatial already ship as `ST_Transform`. xarray-sql ships it as an +optional pyproj extension (`pip install xarray-sql[geo]`): with pyproj +installed, every `XarrayContext` registers a PROJ-backed +`reproject(x, y, src_crs, dst_crs)` scalar UDF, so the CRS pair — any CRS +pyproj understands — is part of the query rather than baked into the function: ```sql -SELECT x, y, reproject(x, y)['lon'] AS lon, reproject(x, y)['lat'] AS lat +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 ``` @@ -226,8 +230,10 @@ 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* reprojection reference, not PROJ-vs-PROJ. The SQL UDF and EE agree to sub-metre -precision. The script flags one practical gotcha (PROJ is not thread-safe, so the -UDF runs serially), but the caveat that matters here is conceptual: reprojection +precision. There is one practical gotcha — PROJ objects must not be shared +across threads, so the extension runs all PROJ work on its own worker pool, +keeping the UDF safe (and parallel) under DataFusion's concurrent partitions — +but the caveat that matters here is conceptual: reprojection moves the coordinates without resampling the data onto a new grid — and *that* is the next operation. @@ -257,8 +263,9 @@ into bilinear weights, and the 08 `JOIN` applies them: ```sql -- 1. reproject the target grid into source coordinates (the 07 UDF) -SELECT dst_lat, dst_lon, reproject(dst_lon, dst_lat)['x'] AS sx, - reproject(dst_lon, dst_lat)['y'] AS sy +SELECT dst_lat, dst_lon, + reproject(dst_lon, dst_lat, 'EPSG:4326', 'EPSG:32610')['x'] AS sx, + reproject(dst_lon, dst_lat, 'EPSG:4326', 'EPSG:32610')['y'] AS sy FROM target -- 2. apply the bilinear weights built from those points (the 08 JOIN) SELECT w.dst_lat AS lat, w.dst_lon AS lon, SUM(s.value * w.weight) AS warped diff --git a/pyproject.toml b/pyproject.toml index 4273d31..e4bc3a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,8 +36,12 @@ dependencies = [ ] [project.optional-dependencies] +geo = [ + "pyproj", +] test = [ "cftime", + "pyproj", "pytest", "xarray[io]", "gcsfs", diff --git a/tests/test_proj.py b/tests/test_proj.py new file mode 100644 index 0000000..b10c9df --- /dev/null +++ b/tests/test_proj.py @@ -0,0 +1,158 @@ +"""Tests for the pyproj CRS-transform extension (`xarray_sql.proj`).""" + +import numpy as np +import pyproj +import pytest +import xarray as xr +from datafusion import SessionContext + +from xarray_sql import XarrayContext, proj + +UTM10 = "EPSG:32610" # UTM zone 10N (metres) +UTM11 = "EPSG:32611" # UTM zone 11N (metres) +WGS84 = "EPSG:4326" # lon/lat degrees +WEBMERC = "EPSG:3857" # Web Mercator (metres) + + +@pytest.fixture +def utm_grid(): + """A 60x50 UTM zone 10N grid over the San Francisco Bay Area. + + Chunked into several partitions so DataFusion evaluates the UDF + concurrently — exercising the per-thread transformer cache. + """ + x = np.linspace(530_000.0, 630_000.0, 50) + y = np.linspace(4_140_000.0, 4_250_000.0, 60) + xx, yy = np.meshgrid(x, y) + return xr.Dataset( + {"value": (["y", "x"], np.hypot(xx, yy))}, + coords={"y": y, "x": x}, + ).chunk({"y": 15, "x": 50}) + + +def test_reproject_matches_pyproj(utm_grid): + ctx = XarrayContext() + ctx.from_dataset("grid", utm_grid) + result = ctx.sql( + f""" + SELECT x, y, + reproject(x, y, '{UTM10}', '{WGS84}')['x'] AS lon, + reproject(x, y, '{UTM10}', '{WGS84}')['y'] AS lat + FROM grid + ORDER BY y, x + """ + ).to_pandas() + + transformer = pyproj.Transformer.from_crs(UTM10, WGS84, always_xy=True) + ref_lon, ref_lat = transformer.transform( + result["x"].to_numpy(), result["y"].to_numpy() + ) + np.testing.assert_allclose(result["lon"], ref_lon, rtol=0, atol=1e-9) + np.testing.assert_allclose(result["lat"], ref_lat, rtol=0, atol=1e-9) + + +def test_reproject_roundtrip(utm_grid): + ctx = XarrayContext() + ctx.from_dataset("grid", utm_grid) + result = ctx.sql( + f""" + WITH lonlat AS ( + SELECT x, y, + reproject(x, y, '{UTM10}', '{WGS84}')['x'] AS lon, + reproject(x, y, '{UTM10}', '{WGS84}')['y'] AS lat + FROM grid + ) + SELECT x, y, + reproject(lon, lat, '{WGS84}', '{UTM10}')['x'] AS rx, + reproject(lon, lat, '{WGS84}', '{UTM10}')['y'] AS ry + FROM lonlat + ORDER BY y, x + """ + ).to_pandas() + # A metre-based CRS round-trips to well under a millimetre. + np.testing.assert_allclose(result["rx"], result["x"], rtol=0, atol=1e-4) + np.testing.assert_allclose(result["ry"], result["y"], rtol=0, atol=1e-4) + + +def test_per_row_crs(): + """The CRS arguments are expressions, so they may vary row by row.""" + lon = np.linspace(-125.9, -114.1, 24) # spans UTM zones 10N and 11N + lat = np.linspace(32.5, 41.5, 10) + LON, LAT = np.meshgrid(lon, lat) + pts = xr.Dataset( + { + "lon": (["i"], LON.ravel()), + "lat": (["i"], LAT.ravel()), + }, + coords={"i": np.arange(LON.size)}, + ).chunk({"i": LON.size}) + + ctx = XarrayContext() + ctx.from_dataset("pts", pts) + result = ctx.sql( + f""" + SELECT lon, lat, + reproject(lon, lat, '{WGS84}', + CASE WHEN lon < -120.0 + THEN '{UTM10}' ELSE '{UTM11}' END)['x'] AS e, + reproject(lon, lat, '{WGS84}', + CASE WHEN lon < -120.0 + THEN '{UTM10}' ELSE '{UTM11}' END)['y'] AS n + FROM pts + ORDER BY i + """ + ).to_pandas() + + for zone, mask in [ + (UTM10, result["lon"] < -120.0), + (UTM11, result["lon"] >= -120.0), + ]: + transformer = pyproj.Transformer.from_crs(WGS84, zone, always_xy=True) + ref_e, ref_n = transformer.transform( + result.loc[mask, "lon"].to_numpy(), + result.loc[mask, "lat"].to_numpy(), + ) + np.testing.assert_allclose( + result.loc[mask, "e"], ref_e, rtol=0, atol=1e-6 + ) + np.testing.assert_allclose( + result.loc[mask, "n"], ref_n, rtol=0, atol=1e-6 + ) + + +def test_null_and_out_of_domain_yield_nan(): + ctx = XarrayContext() + result = ctx.sql( + f""" + SELECT + reproject(CAST(NULL AS DOUBLE), 45.0, + '{WGS84}', '{WEBMERC}')['x'] AS null_coord, + reproject(0.0, 100.0, '{WGS84}', '{WEBMERC}')['y'] AS bad_lat, + reproject(0.0, 45.0, CAST(NULL AS VARCHAR), + '{WEBMERC}')['x'] AS null_crs + """ + ).to_pandas() + assert np.isnan(result["null_coord"].iloc[0]) + assert np.isnan(result["bad_lat"].iloc[0]) + assert np.isnan(result["null_crs"].iloc[0]) + + +def test_invalid_crs_raises(): + ctx = XarrayContext() + with pytest.raises(Exception): + ctx.sql( + "SELECT reproject(0.0, 0.0, 'EPSG:999999', 'EPSG:4326')" + ).to_pandas() + + +def test_register_on_plain_session_context_with_custom_name(): + ctx = SessionContext() + proj.register(ctx, name="st_transform") + result = ctx.sql( + f""" + SELECT st_transform(-122.0, 37.0, '{WGS84}', '{WEBMERC}')['x'] AS gx + """ + ).to_pandas() + transformer = pyproj.Transformer.from_crs(WGS84, WEBMERC, always_xy=True) + ref_x, _ = transformer.transform(-122.0, 37.0) + np.testing.assert_allclose(result["gx"].iloc[0], ref_x, rtol=0, atol=1e-6) diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py new file mode 100644 index 0000000..88f0a04 --- /dev/null +++ b/xarray_sql/proj.py @@ -0,0 +1,223 @@ +"""PROJ-backed CRS transforms for SQL — an optional pyproj 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.** + Constructing a transformer on a DataFusion runtime thread segfaults + inside PROJ, while the identical work on Python-owned threads is + stable — so the UDF hands each batch to the pool rather than calling + pyproj in place. Each pool thread caches one transformer per CRS pair + (transformers must not be shared across threads), which also + amortizes the expensive PROJ database lookups of construction 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, :class:`xarray_sql.XarrayContext` registers ``reproject()`` +automatically; :func:`register` is the explicit hook for plain +DataFusion ``SessionContext`` objects or custom UDF names. +""" + +from __future__ import annotations + +import os +import threading +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyproj +from datafusion import udf + +__all__ = ["register"] + +#: Arrow type returned by ``reproject()``: destination coordinates in +#: ``always_xy`` order — ``x`` is easting/longitude, ``y`` is +#: northing/latitude. +RETURN_TYPE = pa.struct([("x", pa.float64()), ("y", pa.float64())]) + + +# --------------------------------------------------------------------------- +# The PROJ worker pool +# --------------------------------------------------------------------------- +# +# DataFusion evaluates UDFs on its runtime's worker threads. Plain Python +# threads run pyproj construction and transforms concurrently without +# incident (pyproj keeps its PROJ contexts thread-local), but constructing +# a ``Transformer`` *on a DataFusion runtime thread* segfaults inside +# PROJ's CRS machinery — Rust runtime threads are provisioned differently +# from Python threads (notably a much smaller stack). So the UDF never +# calls pyproj in place: every batch is handed to a small pool of +# Python-owned worker threads. pyproj releases the GIL during the +# transform loop, so concurrent partitions still run in parallel across +# the pool. + +_local = threading.local() +_pool_lock = threading.Lock() +_pool: ThreadPoolExecutor | None = None + + +def _proj_pool() -> ThreadPoolExecutor: + """Return the process-wide pool that runs all pyproj work.""" + global _pool + if _pool is None: + with _pool_lock: + if _pool is None: + _pool = ThreadPoolExecutor( + max_workers=os.cpu_count() or 4, + thread_name_prefix="xarray-sql-proj", + ) + return _pool + + +def _transformer(src_crs: str, dst_crs: str) -> pyproj.Transformer: + """Return a cached ``Transformer`` owned by the calling pool thread. + + PROJ transformers are not safe to share across threads, so each pool + thread keeps its own transformer per ``(src, dst)`` pair; the cache + also amortizes construction (expensive PROJ database lookups) across + record batches. ``always_xy=True`` fixes the argument order to + (easting/longitude, northing/latitude) regardless of the CRS's + declared axis order. + """ + cache = getattr(_local, "transformers", None) + if cache is None: + cache = _local.transformers = {} + key = (src_crs, dst_crs) + transformer = cache.get(key) + if transformer is None: + transformer = cache[key] = pyproj.Transformer.from_crs( + src_crs, dst_crs, always_xy=True + ) + return transformer + + +def _transform_chunk( + src_crs: str, dst_crs: str, xs: np.ndarray, ys: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Transform one coordinate chunk; runs on a PROJ pool thread.""" + return _transformer(src_crs, dst_crs).transform(xs, ys) + + +# --------------------------------------------------------------------------- +# The UDF +# --------------------------------------------------------------------------- + + +def _reproject( + x: pa.Array, y: pa.Array, src_crs: pa.Array, dst_crs: pa.Array +) -> pa.Array: + """Vectorized ``reproject`` kernel over one Arrow record batch. + + DataFusion broadcasts scalar arguments (the usual literal CRS + strings) to full-length arrays before calling in, so all four + arguments arrive with one value per row. The common case — one CRS + pair for the whole batch — never touches the strings row by row: + uniqueness is established with a vectorized Arrow kernel and the + batch becomes a single PROJ call. (Materializing the CRS columns + as Python strings costs two object allocations per row, which at + billions of rows dwarfs the transform itself.) Only when the CRS + genuinely varies within the batch are rows grouped by pair and + transformed per group. + """ + xs = np.asarray(x.to_numpy(zero_copy_only=False), dtype="float64") + ys = np.asarray(y.to_numpy(zero_copy_only=False), dtype="float64") + + out_x = np.full(xs.shape, np.nan) + out_y = np.full(ys.shape, np.nan) + valid = np.isfinite(xs) & np.isfinite(ys) + src_unique = pc.unique(src_crs) + dst_unique = pc.unique(dst_crs) + + if len(src_unique) == 1 and len(dst_unique) == 1: + src, dst = src_unique[0].as_py(), dst_unique[0].as_py() + if src is not None and dst is not None and valid.any(): + tx, ty = ( + _proj_pool() + .submit(_transform_chunk, src, dst, xs[valid], ys[valid]) + .result() + ) + out_x[valid] = tx + out_y[valid] = ty + else: + pairs = list(zip(src_crs.to_pylist(), dst_crs.to_pylist())) + for src, dst in set(pairs): + if src is None or dst is None: + continue + mask = valid & np.fromiter( + (p == (src, dst) for p in pairs), dtype=bool, count=len(pairs) + ) + if not mask.any(): + continue + tx, ty = ( + _proj_pool() + .submit(_transform_chunk, src, dst, xs[mask], ys[mask]) + .result() + ) + out_x[mask] = tx + out_y[mask] = ty + + # PROJ signals out-of-domain points with inf; normalize to NaN so + # the result round-trips to xarray like any other missing value. + invalid = ~(np.isfinite(out_x) & np.isfinite(out_y)) + out_x[invalid] = np.nan + out_y[invalid] = np.nan + + return pa.StructArray.from_arrays( + [pa.array(out_x), pa.array(out_y)], names=["x", "y"] + ) + + +def register(ctx, name: str = "reproject") -> None: + """Register the ``reproject(x, y, src_crs, dst_crs)`` scalar UDF. + + Works on any DataFusion ``SessionContext`` (``XarrayContext`` + registers it automatically when pyproj is installed). The UDF + returns a ``{x, y}`` struct of destination coordinates, so a query + selects components with subscripts:: + + SELECT reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon + FROM grid + + Args: + ctx: The DataFusion session context to register the UDF on. + name: SQL name for the function (default ``"reproject"``). + """ + ctx.register_udf( + udf( + _reproject, + [pa.float64(), pa.float64(), pa.utf8(), pa.utf8()], + RETURN_TYPE, + "immutable", + name, + ) + ) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 0ec60ad..226a9fb 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -8,6 +8,11 @@ from .ds import XarrayDataFrame from .reader import read_xarray_table +try: # pyproj is an optional dependency (`pip install xarray-sql[geo]`). + from . import proj as _proj +except ImportError: # pragma: no cover - depends on the environment + _proj = None + class XarrayContext(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" @@ -21,6 +26,10 @@ def __init__(self, *args, **kwargs): # in SQL (e.g. ``"air"`` for a uniform-dim Dataset, or # ``"era5.surface"`` for one entry from a multi-dim-group split). self._registered_datasets: dict[str, xr.Dataset] = {} + # With pyproj installed, every context speaks CRS out of the box: + # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform. + if _proj is not None: + _proj.register(self) def from_dataset( self,