feat: pyproj CRS extension — generic reproject(x, y, src_crs, dst_crs) UDF

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

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,21 +43,17 @@
 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.
 """
 
 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
     """

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

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

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

pyproject.toml

@@ -36,8 +36,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+geo = [
+  "pyproj",
+]
 test = [
   "cftime",
+  "pyproj",
   "pytest",
   "xarray[io]",
   "gcsfs",