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

[Mmoncadaisla]

What

Adds xarray_sql/proj.py, an optional pyproj extension (pip install xarray-sql[geo]) that registers an ST_Transform-style scalar UDF:

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 UDF registration time (which is how benchmarks 07/09 previously hard-coded it, each with its own copy of the UDF). Any CRS spelling pyproj.CRS accepts works (authority codes, WKT, PROJ strings), and because the arguments are ordinary SQL expressions the CRS may vary per row — e.g. a CASE expression selecting the UTM zone from the longitude.

XarrayContext auto-registers reproject() when pyproj is installed (mirroring the cftime() UDF); proj.register(ctx, name=...) is the explicit hook for plain SessionContexts or custom names. Benchmarks 07 and 09 now use the extension and their local UDF definitions are removed.

The threading discovery (worth flagging)

The docstring in 07 said "PROJ's context is not thread-safe," and the benchmark serialized the UDF by forcing a single chunk. While porting this I found the real failure mode is narrower and more interesting:

• Plain Python threads run concurrent Transformer.from_crs + transform() with no locking and no crashes (pyproj 3.7.2 / PROJ 9.5.1 — verified with a standalone ThreadPoolExecutor stress repro).
• Constructing a Transformer on a DataFusion runtime (tokio) thread segfaults inside PROJ's CRS machinery whenever it overlaps other PROJ work. Neither per-thread transformer caching nor a full readers-writer lock around all pyproj calls fixed it — the crash is specific to running PROJ construction on the runtime's threads (whose stacks are much smaller than Python threads'; that is the likely mechanism).

So the extension never calls pyproj on the calling thread: all PROJ work runs on a small dedicated pool of Python-owned worker threads, each caching one transformer per CRS pair. pyproj releases the GIL during transforms, so concurrent partitions still transform in parallel — the single-chunk/serial-UDF workaround in the benchmarks is gone rather than relocated.

Why the CRS is a query argument (design context)

Passing the CRS pair as SQL arguments isn't just ergonomics — given today's DataFusion, it's the only honest place to put it. DataFusion's type system cannot yet attach metadata like a CRS to a column type: that's the extension-types gap tracked in apache/datafusion#12644, and the reason spatial support overall (apache/datafusion#7859) was blocked for a year before the GeoArrow dense-union workaround. Until columns can carry their CRS, any "implicit CRS" design would have to smuggle it through UDF registration state — which is exactly the hard-coding this PR removes.

This also means the UDF stays complementary to where the ecosystem is heading: when extension types land and GeoArrow columns know their CRS (e.g. via datafusion-contrib/datafusion-geo), a geometry-typed ST_Transform(geom, dst) can arrive alongside reproject() rather than replace it — raster coordinates enter xarray-sql as plain float columns either way.

Semantics

• Returns a {x, y} struct (always_xy order) from one call — one PROJ transform per row.
• NULL/non-finite coordinates and NULL CRS arguments → NaN; PROJ's inf for out-of-domain points is normalized to NaN so results round-trip to xarray like any other missing value.
• An invalid CRS raises (pyproj.exceptions.CRSError) and fails the query loudly rather than returning wrong coordinates.

Testing

tests/test_proj.py (pyproj added to the test extra): equivalence with direct pyproj on a multi-partition grid, round-trip UTM↔lon/lat, per-row CRS via CASE, NULL/out-of-domain handling, invalid-CRS failure, and custom-name registration on a plain SessionContext. The suite passed 20/20 consecutive runs locally (the round-trip test is the one that segfaulted reliably before the worker-pool design — it mixes two CRS pairs in one query).

Caveats: benchmarks 07/09 need Earth Engine credentials I don't have wired up here — they compile and the SQL shape is covered by the tests, but a maintainer EE run would be a good final check.

Docs

docs/geospatial.md (section 6 narrative + SQL snippets), both READMEs, and the module docstring double as the extension's documentation.

🤖 Generated with Claude Code

[Mmoncadaisla]

Closing for now — opened prematurely; will resubmit once it's ready for review.

[alxmrs]

xref: wrt the tokio issue, check out #145. We may want to tackle that soon(ish) since it might be useful for proj, too.

[alxmrs]

pip install xarray-sql[proj]

Maybe we should make this pip install xarray-sql[geo] -- wdyt?

[alxmrs]

LGTM! A few technical fixes, then I think this should be ready to land. Thanks for the amazing contribution! :)

[alxmrs]

Should we modify this to include the geo extra? (is that a thing?)

[Mmoncadaisla]

Yes! "xarray-sql[geo]" works there and uv resolves the extra against the [tool.uv.sources] path source (I verified it with a quick uv run --script probe: the geo extra pulls pyproj into the script's environment). Adding it

[alxmrs]

A few random nits. Thanks for this significant contribution. LGTM.