SpatialIndex follow-ups: query_nearest, index persistence, and a constructor fix

[espg]

Stacked on #118 — that should merge first.

Follow-up to the #118 review thread (the query_nearest split from that discussion) plus the serialization piece from #72.

Motivation

Same pipeline as #118: mapping NASA CMR granule footprints to spatial workers in zagg. The index over the full 555,867-granule catalog is for the full ICESat-2 mission to date, so identical for every query against that mission, and the whole build cost is per-process waste after the first invoke.

555,867 polygons (17.68M vertices) against 2,721 query polygons, single-threaded, M1 Max. The mortie column is the HEALPix/MOC path the same pipeline uses, for scale:

	spherely	mortie
query without an index	>20 min ¹	64.0 s
build the index	19.5 s ²	37.7 s
index size	632 MB	814 MB ³
load the index	0.65 s	2.30 s
query with the index	5.05 s	18.9 s
load + query	5.70 s	21.2 s

¹ elementwise spherely.intersects; killed at 20 min, never completed.
² 8.35 s building geographies + 0.05 s constructor (14.97 before this branch's fix) + 10.07 s build(tmp_memory_budget=4 GiB) (198.34 at the default — see #118) + ~1 s encode().
³ parquet defaults; ~350 MB with zstd and tuned row groups.

So ~19.5 s once and 0.65 s per process after, against ~18.5 s in every process that wants to query. query_nearest is the distance-query half requested in #72; the encode work is the "serialize the index / memory-map and query" idea from the same thread. Baselines in espg#2.

Right now, we default to mortie and pay around a minute; that could get dropped by ~2x from optimizations in that library. Building an index in mortie has high memory usage, and the index improvement is marginal (mainly a result of the input geometries, which are thin and cross many cells). #118 plus the constructor fix below gets ~25 s — a 2.6× speedup over the current pipeline — even building and discarding the index every time. But having this merged drops the query time by more than 10x after the first index build is persisted.

The other benefit is separation of concerns; larger machine builds the index and pays the memory overhead, follow on users that hit the index do it memory efficient at small runtimes.

---

This diff includes #118's commits; the new work is the last 4 5 commits (query_nearest, encode, from_encoded, constructor, bugfix).

1. query_nearest — shapely STRtree.query_nearest parity

tree.query_nearest(geog)                        # sorted indices of nearest tree geographies
tree.query_nearest(geog_array)                  # (2, K): input index / tree index pairs
tree.query_nearest(geog, return_distance=True)  # (indices, distances)
tree.query_nearest(geog, max_distance=..., exclusive=..., all_matches=...)

S2ClosestEdgeQuery over the existing index. Distances in units of a radius keyword defaulting to EARTH_RADIUS_METERS, matching spherely.distance; interiors count on both sides, also matching it.

Two behaviours worth flagging:

• all_matches=False returns the lowest tied index rather than an arbitrary one — deterministic, and a valid instance of shapely's documented "arbitrary".
• max_distance is genuinely inclusive: the converted bound is widened a few ULP, because the meters→S1ChordAngle round trip otherwise disagrees with the reported distance by ~2 ULP, so re-querying with a distance the call just returned could come back empty. A single fixture passes by luck ~72% of the time; pinned by a randomized boundary test.

Correctness is pinned against the independent spherely.distance path on randomized trees, including antimeridian-crossing and pole-adjacent geometry and >2-way exact ties.

2. Index persistence — build once, ship the index, open ~instantly

blob = tree.encode()                              # bytes
tree2 = spherely.SpatialIndex.from_encoded(blob)  # lazy view; near-instant open

• encode() (default include_geographies=True) writes a full-fidelity blob: candidate and predicate queries, query_nearest, exclusive and .geometries all work on the decoded index, through the same code paths as a built index — no second implementation of predicate semantics to drift. Geographies are reconstructed lazily per entry on first touch and cached; opening a 5k-entry blob is ~0.1 ms, and candidate-only queries reconstruct zero geographies.
• include_geographies=False writes a structure-only blob: candidate queries and query_nearest work; predicate refinement, exclusive and .geometries raise, naming the full-fidelity option as the remedy.
• Single-copy layout — the index structure's shapes are served out of the per-geography blocks via a custom ShapeFactory, so shape data is stored once. Full-fidelity blobs run ≈1.1× (polygon-heavy) to ≈2× (small point-only) the size of structure-only ones.
• A small versioned header plus s2's own encoders (CompactEncodeTaggedShapes, MutableS2ShapeIndex::Encode, Geography::EncodeTagged). EncodeTagged/DecodeTagged are experimental upstream, so the docs say plainly this is a build-coupled cache/ship format tied to the bundled s2geometry/s2geography versions — not archival.
• Defensive decode — header/version/flags checks, count and offset-table validation before any allocation, a cell↔value-table cross-check after Init, per-block kind verification. Corrupt or truncated input raises ValueError, fuzzed over every truncation prefix and a set of crafted mutations. A corrupt block first touched from inside a distance query surfaces cleanly and leaves the index usable.

3. Constructor: bypassing s2geography::GeographyIndex::Add

The last commit adds shapes to the wrapped MutableShapeIndex() directly, keeping the shape id → tree index mapping in spherely.

Add grows that mapping with values_.reserve(values_.size() + num_shapes) then resize(new_shape_id + 1) per shape. libc++ reserves exactly and the resize leaves capacity == size, so every subsequent Add reallocates and copies the whole vector — quadratic in geography count. On a 555,867-polygon catalog that is 14.974 s of pure vector growth; doing it here takes 0.048 s (312×), with a byte-identical encode() blob — same index, same fill rule, including multi-shape and empty geographies.

The cost: the wrapped GeographyIndex's own values_ is now permanently empty, so its value() and Iterator must not be used. spherely uses neither (it already walks the shape index itself) and there is a comment saying so, but it is a sharp edge. Three options here — first one implemented since it fixes it now (question 4 below):

1. As written — reverts in one commit once s2geography grows that vector geometrically.
2. Drop GeographyIndex here for a plain MutableS2ShapeIndex. Nothing else in spherely uses the wrapper, so the hazard disappears — but so does the clean revert.
3. Fix s2geography instead — grow values_ geometrically, take it as a version floor. Two-line change on their side; happy to open that PR.

Testing

293 passing (181 at the #118 head). The one touch this branch makes to #118-reviewed code — the candidate cell walk now uses a generic S2ShapeIndex::Iterator so built and decoded indexes share one query path — is pinned by a differential test, byte-identical to the #118 branch across 1,644 comparisons spanning all geometry kinds, predicates, and antimeridian/pole cases. Built-vs-decoded parity is tested exhaustively (candidates, all predicates, nearest indices and distances, multi-shape collections, empties), plus laziness pins via an internal decode counter, concurrent-query stress, and a GIL-release test on encode(). mypy, black and clang-format clean; warning-free under -Wconversion -Wsign-conversion.

Questions for review

1. from_encoded accepts bytes/bytearray/memoryview and copies once. A zero-copy/mmap path — the "memory map it and query it" idea from Querying geographies (spatial index) #72 — is a natural follow-up if there's appetite; EncodedS2ShapeIndex is built for it.
2. The persistence container is fork-defined framing around s2's encoders. If serialization should ultimately live at the s2geography level per GeographyIndex improvement ideas paleolimbot/s2geography#45, this is structured to migrate — the header/id-map framing is the only spherely-specific part.
3. The block decoder verifies geography kind on decode; the same check would harden Geography.__setstate__ against crafted state tuples. Left out since pickle is pre-existing surface — separate PR if wanted.
4. Which of the three constructor options above do you want? Can we go with option 1 until option 3 lands upstream?