feat: pyproj CRS extension — generic reproject(x, y, src_crs, dst_crs) UDF#225
Open
Mmoncadaisla wants to merge 4 commits into
Open
feat: pyproj CRS extension — generic reproject(x, y, src_crs, dst_crs) UDF#225Mmoncadaisla wants to merge 4 commits into
Mmoncadaisla wants to merge 4 commits into
Conversation
…) 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>
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>
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. |
Collaborator
Maybe we should make this |
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds
xarray_sql/proj.py, an optional pyproj extension (pip install xarray-sql[geo]) that registers anST_Transform-style scalar UDF: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.CRSaccepts works (authority codes, WKT, PROJ strings), and because the arguments are ordinary SQL expressions the CRS may vary per row — e.g. aCASEexpression selecting the UTM zone from the longitude.XarrayContextauto-registersreproject()when pyproj is installed (mirroring thecftime()UDF);proj.register(ctx, name=...)is the explicit hook for plainSessionContexts 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:
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).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 alongsidereproject()rather than replace it — raster coordinates enter xarray-sql as plain float columns either way.Semantics
{x, y}struct (always_xyorder) from one call — one PROJ transform per row.inffor out-of-domain points is normalized to NaN so results round-trip to xarray like any other missing value.pyproj.exceptions.CRSError) and fails the query loudly rather than returning wrong coordinates.Testing
tests/test_proj.py(pyproj added to thetestextra): equivalence with direct pyproj on a multi-partition grid, round-trip UTM↔lon/lat, per-row CRS viaCASE, NULL/out-of-domain handling, invalid-CRS failure, and custom-name registration on a plainSessionContext. 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