Skip to content

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

Open
Mmoncadaisla wants to merge 4 commits into
xqlsystems:mainfrom
Mmoncadaisla:feat/proj-udf-extension
Open

feat: pyproj CRS extension — generic reproject(x, y, src_crs, dst_crs) UDF#225
Mmoncadaisla wants to merge 4 commits into
xqlsystems:mainfrom
Mmoncadaisla:feat/proj-udf-extension

Conversation

@Mmoncadaisla

@Mmoncadaisla Mmoncadaisla commented Jul 9, 2026

Copy link
Copy Markdown

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

…) UDF

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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
@Mmoncadaisla

Copy link
Copy Markdown
Author

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

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 <noreply@anthropic.com>
@alxmrs

alxmrs commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

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

alxmrs commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

pip install xarray-sql[proj]

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

@Mmoncadaisla Mmoncadaisla reopened this Jul 13, 2026
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 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants