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]

Future work: should cftime be a standard UDF (bc it's standard in Xarray), or be part of the geo extension? CC: @ghostiee-11.

[Mmoncadaisla]

IMHO it doesn't feel like it belongs in the geo package so I'd say it's better in the standard lib as it is right now with lazy/guarded imports via _maybe_register_cftime_udf

[alxmrs]

Should we use the C-based proj dependency and call this from our native backend? (This is a good question for Claude).

https://github.com/georust/proj

Could that approach address this concurrent/threading issue that tends to segfault?

[alxmrs]

On the other hand, I think this would make the [geo] extra setting impossible, we would have to bundle proj. Much to consider.

[Mmoncadaisla]

After researching a little bit more, it seems that on threads not created by Python (DataFusion's tokio workers), the Python thread state is created and destroyed around every UDF call, so the context dies each time and the next call on that OS thread picks up a dangling pointer.

I verified the fix from pyproj4/pyproj#1541 empirically with the help of Claude so this should be resolved upstream in the next pyproj release, but it's been a while since the last one, and no released version has it today.

Moreover I'd still keep the pool even after that release, for performance rather than safety. Because everything pyproj keeps per thread (context, transformer internals) dies between calls, in-place execution rebuilds them per batch even on fixed pyproj.

I measured it as you suggested, the pool round-trip costs around 10 µs/batch, while rebuilding a transformer costs 0.07 ms (UTM→WGS84, warm context) up to 10-12 ms for pairs like EPSG:3035 or EPSG:27700, vs 0.25 ms to transform a whole 2,000-row batch.

See the benchmark snippet + results (click to expand)

import statistics, time
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pyproj


def med_ms(fn, n=20):
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        ts.append((time.perf_counter() - t0) * 1e3)
    return statistics.median(ts)


# 1. the pool "tax": one submit round-trip per batch
pool = ThreadPoolExecutor(4)
pool.submit(lambda: None).result()  # warm
print(f"pool submit round-trip: {med_ms(lambda: pool.submit(lambda: None).result(), 200) * 1e3:.0f} us")

# 2. what the cache avoids: rebuilding the transformer (warm PROJ context)
for src in ("EPSG:32610", "EPSG:3035", "EPSG:27700"):
    pyproj.Transformer.from_crs(src, "EPSG:4326", always_xy=True)  # cold first
    ms = med_ms(lambda: pyproj.Transformer.from_crs(src, "EPSG:4326", always_xy=True))
    print(f"rebuild {src}->EPSG:4326: {ms:6.2f} ms")

# 3. the useful work, for scale: transforming one 2,000-row batch
t = pyproj.Transformer.from_crs("EPSG:32610", "EPSG:4326", always_xy=True)
xs = np.random.uniform(4e5, 6e5, 2_000)
ys = np.random.uniform(4e6, 5e6, 2_000)
t.transform(xs, ys)  # warm
print(f"transform 2,000-row batch: {med_ms(lambda: t.transform(xs, ys)):.2f} ms")

# 4. end-to-end: the shipped UDF, per-thread cache vs rebuild-per-batch
import xarray as xr
import xarray_sql as xql
from xarray_sql import proj as xproj

ds = xr.Dataset(
    {"v": (("y", "x"), np.zeros((400, 400)))},
    coords={"x": np.linspace(4e5, 6e5, 400), "y": np.linspace(4e6, 5e6, 400)},
)
ctx = xql.XarrayContext()
ctx.from_dataset("grid", ds, chunks={"y": 5})  # 160k rows -> 80 batches/query

print("end-to-end (160k rows / 80 batches per query):")
cached = xproj._transformer
for crs in ("EPSG:32610", "EPSG:3035"):
    q = f"SELECT AVG(reproject(x, y, '{crs}', 'EPSG:4326')['x']) FROM grid"
    for label, fn in [
        ("cached ", cached),
        ("rebuild", lambda s, d: pyproj.Transformer.from_crs(s, d, always_xy=True)),
    ]:
        xproj._transformer = fn  # rebuild-per-batch still runs on the pool: safe on 3.7.2
        ctx.sql(q).collect()  # warm-up
        print(f"  {crs:>10} {label}: {med_ms(lambda: ctx.sql(q).collect(), 5):7.1f} ms")
xproj._transformer = cached

Output on my machine (macOS arm64, pyproj 3.7.2 / PROJ 9.5.1):

pool submit round-trip: 7 us
rebuild EPSG:32610->EPSG:4326:   0.07 ms
rebuild EPSG:3035->EPSG:4326:  11.99 ms
rebuild EPSG:27700->EPSG:4326:  10.21 ms
transform 2,000-row batch: 0.24 ms
end-to-end (160k rows / 80 batches per query):
  EPSG:32610 cached :    54.9 ms
  EPSG:32610 rebuild:    60.8 ms
   EPSG:3035 cached :    56.5 ms
   EPSG:3035 rebuild:  1027.7 ms

[Mmoncadaisla]

Re: https://github.com/georust/proj

I think it'd work because it'd give PROJ one persistent context per thread, which is what the current approach via the Python worker pool does too. I have not measured performance using the rust approach, but with the numbers above and given what you mentioned:

On the other hand, I think this would make the [geo] extra setting impossible, we would have to bundle proj. Much to consider.

I don't think it's worth it primarily because of this, but we could dig deeper and get some numbers on the georust/proj path to compare against the python work pool to see what we'd be giving up, WDYT?

[alxmrs]

I like the idea of having a possible future pathway (via rust/native code) for high performance for this feature . For now, I think prioritizing stability/correctness first is right, with decent performance as a secondary concern. Let's establish users of this feature (I think there could be many) before optimizing to this level.

Given this, I trust your judgement if we should use a threadpool or recreate the options in pure scope each time. Than you for measuring.

[alxmrs]

How expensive is PROJ database lookups s.t. it is worth it to share state? Is that overhead more than the cost of Python process pooling? I think it's best to measure empirically instead of assume. It would be far simpler to reconstruct these objects and not share state.

[Mmoncadaisla]

See the comment above

[alxmrs]

Is this copy intentional? If we moved this to pyarrow, could it be cheaper (or, would that break the pyproj contract?)?

[alxmrs]

Could we get away with a numpy view?

[Mmoncadaisla]

It's actually already a view whenever the batch has no NULLs: to_numpy(zero_copy_only=False) returns a zero-copy, read-only window onto Arrow's buffer (same memory address), and np.asarray(..., dtype="float64") is a no-op since the UDF signature guarantees float64.

The zero_copy_only=False flag is just the fallback permission for batches with NULLs, which is needed since Arrow and Numpy have different representations of NULLs.

Repro: same memory address, no copy without NULLs

import numpy as np
import pyarrow as pa

a = pa.array(np.linspace(0.0, 1.0, 1_000_000))
v = a.to_numpy(zero_copy_only=False)

# 1. The numpy array points at Arrow's own buffer, byte for byte
arrow_addr = a.buffers()[1].address   # Arrow's data buffer ([0] is the validity bitmap)
numpy_addr = v.ctypes.data            # where numpy thinks its data lives
print(f"arrow buffer address : {hex(arrow_addr)}")
print(f"numpy data address   : {hex(numpy_addr)}")
print(f"same memory          : {arrow_addr == numpy_addr}")

# 2. two 'independent' conversions are windows onto the same buffer
print(f"shares_memory        : {np.shares_memory(v, a.to_numpy(zero_copy_only=False))}")

# 3. view arithmetic: a sliced Arrow array lands at buffer + offset * 8 bytes
vs = a.slice(1000, 10).to_numpy(zero_copy_only=False)
print(f"slice offset check   : {vs.ctypes.data == arrow_addr + 1000 * 8}")

# 4. the flags agree: borrowed buffer, read-only (Arrow buffers are immutable)
print(f"owndata={v.flags.owndata}  writeable={v.flags.writeable}")
print(f"asarray(float64) is a no-op: {np.asarray(v, dtype='float64') is v}")

# 5. the contrast: with nulls, to_numpy returns a DIFFERENT address (fresh
#    buffer, validity bitmap materialized as NaN - the UDF's NULL->NaN contract)
b = pa.array([1.0, None, 3.0])
nb = b.to_numpy(zero_copy_only=False)
print(f"with nulls, same addr: {nb.ctypes.data == b.buffers()[1].address}  values: {nb}")

Output on my machine:

arrow buffer address : 0x140080000
numpy data address   : 0x140080000
same memory          : True
shares_memory        : True
slice offset check   : True
owndata=False  writeable=False
asarray(float64) is a no-op: True
with nulls, same addr: False  values: [ 1. nan  3.]

[alxmrs]

I think we could simplify this if..else to one block that operated on a list, but sometimes operated on a list of one item. Is there a performance cost to doing that way, or an advantage to separating out the two cases like this? I'm not sure I understand the unique logic, but in theory, couldn't we have a small conditional to select the first/only unique item as a list and then apply it to the second block? I think that would be less repetition.

[Mmoncadaisla]

You're right that the submit + scatter block is duplicated — I've unified it in fbe5f2c: both branches now just build a list of (src, dst, mask) groups (a single one with mask = valid when the CRS is constant), and one loop does the pool submit + scatter.

I kept the detection of the two cases separate though. In the common case (constant CRS — SQL literals that DataFusion broadcasts to full-length columns) uniqueness is checked with a vectorized Arrow kernel, while the general path has to materialize the CRS columns as Python strings (two allocations per row). I measured it and pc.unique is ~45x cheaper than to_pylist, so I think that small conditional is worth keeping.

Performance details on pc.unique vs to_pylist approach

import statistics
import time

import pyarrow as pa
import pyarrow.compute as pc


def med_ms(fn, n=10):
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        ts.append((time.perf_counter() - t0) * 1e3)
    return statistics.median(ts)


# DataFusion broadcasts the (usually literal) CRS args to full-length
# string arrays, so this is what the kernel actually receives.
for n in (8_192, 1_000_000):
    crs = pa.array(["EPSG:32610"] * n)
    t_unique = med_ms(lambda: pc.unique(crs))
    t_pylist = med_ms(lambda: crs.to_pylist())
    print(f"n={n:>9,}: pc.unique {t_unique:8.3f} ms | to_pylist {t_pylist:8.2f} ms  ({t_pylist / t_unique:,.0f}x)")

Output on my machine (pyarrow, macOS arm64):

n=    8,192: pc.unique    0.054 ms | to_pylist     2.32 ms  (43x)
n=1,000,000: pc.unique    6.578 ms | to_pylist   306.81 ms  (47x)

[alxmrs]

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

[alxmrs]

Is it the pyproj or geo extension?

[alxmrs]

I really dislike Claude's use of the term "slab"

[alxmrs]

Is it a "pyproj" or "geo" extension?

[alxmrs]

is it a "pyproj" or "geo" extension?

[alxmrs]

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