From bf9d069063e2a61cc314d51ad386529f8a42bf1c Mon Sep 17 00:00:00 2001 From: Miguel Moncada Date: Thu, 9 Jul 2026 12:48:57 +0200 Subject: [PATCH 1/4] =?UTF-8?q?feat:=20pyproj=20CRS=20extension=20?= =?UTF-8?q?=E2=80=94=20generic=20reproject(x,=20y,=20src=5Fcrs,=20dst=5Fcr?= =?UTF-8?q?s)=20UDF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds xarray_sql/proj.py, an optional pyproj extension that registers an ST_Transform-style scalar UDF. The CRS pair is part of the query (any spelling pyproj.CRS accepts, and it may vary per row via ordinary SQL expressions) instead of being baked in at UDF registration time, which is how the geospatial benchmarks previously hard-coded it. All PROJ work runs on a dedicated pool of Python-owned worker threads: constructing a Transformer on a DataFusion runtime thread segfaults inside PROJ, while identical concurrent work on Python threads is stable (pyproj 3.7 / PROJ 9.5). Each pool thread caches one transformer per CRS pair, and pyproj releases the GIL during transforms, so the UDF is parallel across partitions — the previous single-chunk/serial-UDF workaround in the benchmarks is no longer needed. XarrayContext auto-registers reproject() when pyproj is installed (pip install xarray-sql[proj]); proj.register() is the explicit hook for plain SessionContexts or custom names. Benchmarks 07 and 09 now use the extension instead of their local hard-coded UDFs. Co-Authored-By: Claude Fable 5 --- README.md | 6 +- benchmarks/geospatial/07_reproject_udf.py | 71 ++------ benchmarks/geospatial/09_warp.py | 40 +---- benchmarks/geospatial/README.md | 2 +- docs/geospatial.md | 23 ++- pyproject.toml | 4 + tests/test_proj.py | 158 +++++++++++++++++ xarray_sql/proj.py | 205 ++++++++++++++++++++++ xarray_sql/sql.py | 9 + 9 files changed, 417 insertions(+), 101 deletions(-) create mode 100644 tests/test_proj.py create mode 100644 xarray_sql/proj.py diff --git a/README.md b/README.md index a254d29..e132b35 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[proj]`, 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..94cb32d 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,18 @@ def main() -> None: f"{_SRC_CRS} → {_DST_CRS}" ) + # XarrayContext registers reproject() automatically (the pyproj + # extension). Several partitions: the extension runs PROJ on its own + # worker pool, so the UDF is safe under DataFusion's parallelism. 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"]} + "grid", ds, chunks={"y": max(1, ds.sizes["y"] // 4), "x": ds.sizes["x"]} ) - register_reproject_udf(ctx, _SRC_CRS, _DST_CRS) - 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..8b10f22 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[proj]`): 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..6e29701 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,8 +36,12 @@ dependencies = [ ] [project.optional-dependencies] +proj = [ + "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..3aadda5 --- /dev/null +++ b/xarray_sql/proj.py @@ -0,0 +1,205 @@ +"""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[proj]``). 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 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 — is a single vectorized PROJ call; when + the CRS varies by row, rows are 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") + pairs = list(zip(src_crs.to_pylist(), dst_crs.to_pylist())) + + out_x = np.full(xs.shape, np.nan) + out_y = np.full(ys.shape, np.nan) + valid = np.isfinite(xs) & np.isfinite(ys) + + 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..9b0f1f5 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[proj]`). + 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, From de833c6290c5fa8648eb44e18127c47c5761853d Mon Sep 17 00:00:00 2001 From: Miguel Moncada Date: Thu, 9 Jul 2026 13:20:13 +0200 Subject: [PATCH 2/4] docs: spell out the deliberate multi-partition chunking in benchmark 07 The y-chunking exists to force several DataFusion partitions so the reproject() UDF demonstrably runs in parallel; say so explicitly instead of hiding it behind a computed chunk size. Co-Authored-By: Claude Fable 5 --- benchmarks/geospatial/07_reproject_udf.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py index 94cb32d..02815e3 100644 --- a/benchmarks/geospatial/07_reproject_udf.py +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -114,12 +114,13 @@ def main() -> None: ) # XarrayContext registers reproject() automatically (the pyproj - # extension). Several partitions: the extension runs PROJ on its own - # worker pool, so the UDF is safe under DataFusion's parallelism. + # 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() - ctx.from_dataset( - "grid", ds, chunks={"y": max(1, ds.sizes["y"] // 4), "x": ds.sizes["x"]} - ) + ctx.from_dataset("grid", ds, chunks={"y": 15, "x": ds.sizes["x"]}) sql = f""" SELECT x, y, From 925c7a6fba24ec8699f961eb39e721688af712fa Mon Sep 17 00:00:00 2001 From: Miguel Moncada Date: Thu, 9 Jul 2026 14:18:59 +0200 Subject: [PATCH 3/4] perf: vectorized constant-CRS fast path in the reproject kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel materialized the two broadcast CRS string columns with to_pylist() — two Python object allocations per row before PROJ did any work. At benchmark scale that dominated everything: a 606M-pixel aggregate-only reprojection took 496s while the same scan without the UDF took 4.5s. Establish CRS uniqueness with pyarrow.compute.unique (vectorized C) instead. A batch with one CRS pair — the overwhelmingly common case — becomes a single vectorized PROJ call with no per-row Python. The per-row grouping path remains for genuinely varying CRS (e.g. a CASE expression picking the UTM zone). Same 606M-pixel run after: 16.0s (37.9M px/s), a 31x speedup, with bit-identical results and flat (~2GB) streaming memory. Co-Authored-By: Claude Fable 5 --- xarray_sql/proj.py | 58 ++++++++++++++++++++++++++++++---------------- 1 file changed, 38 insertions(+), 20 deletions(-) diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py index 3aadda5..2f84190 100644 --- a/xarray_sql/proj.py +++ b/xarray_sql/proj.py @@ -54,6 +54,7 @@ import numpy as np import pyarrow as pa +import pyarrow.compute as pc import pyproj from datafusion import udf @@ -140,33 +141,50 @@ def _reproject( 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 — is a single vectorized PROJ call; when - the CRS varies by row, rows are grouped by pair and transformed - per group. + 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") - pairs = list(zip(src_crs.to_pylist(), dst_crs.to_pylist())) out_x = np.full(xs.shape, np.nan) out_y = np.full(ys.shape, np.nan) valid = np.isfinite(xs) & np.isfinite(ys) - - 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 + 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. From da926478745dd2783a1da6b398a115748fed6cba Mon Sep 17 00:00:00 2001 From: Miguel Moncada Date: Mon, 13 Jul 2026 08:36:37 +0200 Subject: [PATCH 4/4] rename pip extra from [proj] to [geo] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aligns with the base + add-ons direction: geo is the umbrella extra for CRS reprojection today and polygon/vector support later, rather than being named after the first feature it happens to ship. The module stays proj.py — it does CRS reprojection specifically. Co-Authored-By: Claude Fable 5 --- README.md | 2 +- docs/geospatial.md | 2 +- pyproject.toml | 2 +- xarray_sql/proj.py | 2 +- xarray_sql/sql.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index e132b35..fb9f0f8 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,7 @@ against an xarray/array reference** to floating-point tolerance: table of regions. * **Reprojection and regridding** — a `reproject(x, y, src_crs, dst_crs)` scalar PROJ UDF, shipped as the optional pyproj extension - (`pip install xarray-sql[proj]`, validated against Earth Engine's own + (`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). diff --git a/docs/geospatial.md b/docs/geospatial.md index 8b10f22..2c0b317 100644 --- a/docs/geospatial.md +++ b/docs/geospatial.md @@ -213,7 +213,7 @@ 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`. xarray-sql ships it as an -optional pyproj extension (`pip install xarray-sql[proj]`): with pyproj +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: diff --git a/pyproject.toml b/pyproject.toml index 6e29701..e4bc3a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ ] [project.optional-dependencies] -proj = [ +geo = [ "pyproj", ] test = [ diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py index 2f84190..88f0a04 100644 --- a/xarray_sql/proj.py +++ b/xarray_sql/proj.py @@ -40,7 +40,7 @@ * 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[proj]``). When pyproj is +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. diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 9b0f1f5..226a9fb 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -8,7 +8,7 @@ from .ds import XarrayDataFrame from .reader import read_xarray_table -try: # pyproj is an optional dependency (`pip install xarray-sql[proj]`). +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