Add SpatialIndex for fast spatial queries

[espg]

Addresses #72

Summary

Adds spherely.SpatialIndex, a spatial index over a collection of Geography objects for fast candidate retrieval — conceptually similar to shapely's STRtree. Without it, spatial joins are O(N×M) brute-force over the vectorized predicates; this lets callers cheaply pre-filter candidate pairs.

It's backed by s2geography::GeographyIndex (a MutableS2ShapeIndex), which is already available in the pinned s2geography 0.2.0 — no upstream or dependency changes needed.

API

tree = spherely.SpatialIndex(geographies)   # array-like of Geography
len(tree)                                   # number of indexed geographies
tree.geometries                             # object-ndarray of the inputs (input order)

# scalar query -> sorted np.intp array of tree indices
tree.query(geom)                            # coarse: cells overlap (candidate set)
tree.query(geom, predicate="intersects")    # refined

# array query -> (2, K): row 0 = input index, row 1 = tree index (shapely-compatible)
tree.query(geom_array, predicate="intersects")

• predicate=None returns the coarse cell-overlap candidate set (a superset of true intersections).
• Supported predicates: "intersects", "within", "contains", "covers", "covered_by", "touches", "equals" — each a subset of the candidate set, with semantics matching the existing scalar predicates (a tree geometry t matches when predicate(query_geom, t) is True).
• "disjoint" is rejected with ValueError (it's the complement of the candidate set, so it can't be answered from the index alone); unknown predicate names also raise ValueError.
• Empty geographies are indexed but never returned by queries.

Implementation notes

• src/spatial_index.cpp — the SpatialIndex binding. Builds the index with GeographyIndex::Add(geog, i); queries compute a covering via S2RegionCoverer over Geography::Region() and walk GeographyIndex::Iterator::Query. The input geographies are held as a numpy object-dtype array, which both keeps the underlying C++ Geography objects alive (the index only borrows their S2Shapes) and backs the geometries property.
• src/predicates.hpp + src/predicates_common.cpp — a small shared get_predicate(name) factory returning a closure over two ShapeIndexGeography, reusing the same s2geog::s2_* calls / S2BooleanOperation::Options as predicates.cpp (which remains the reference for the vectorized predicates).
• Wired into src/spherely.cpp and CMakeLists.txt; type stubs in src/spherely.pyi; docs in docs/api.rst and docs/api_hidden.rst.

Out of scope (possible follow-ups)

• query_nearest / distance queries via S2ClosestEdgeQuery.
• Index serialization (S2's encoded shape index enables mmap'd, shippable indexes — noted as high-value in Querying geographies (spatial index) #72).
• Configurable coverer (max_cells) and a dedicated point-only fast path (S2PointIndex).

Testing

• New tests/test_spatial_index.py (10 tests): build / len / geometries, coarse vs. predicate-refined queries, array (2, K) layout, empty-geography handling, and ValueError cases.
• Full suite passes (162 passed, 1 skipped); mypy and pre-commit (black + clang-format) clean.

[espg]

Some additional background on this--

I maintain two libraries: zagg which aggregates input geospatial data to zarr grids, and mortie , which implements healpix / morton spatial indexing. The main usecase for zagg right now is to query the NASA CMR , and then map spatially coincident data granuals to worker process that can aggregate to a grid in parallel. It started off by using mortie to map spatially coincident data files together, but we've re-architected it to work on arbitrary grids beyond healpix.

To do the 'map CMR geometries' to worker step, there's a few options:

1. Morton indexing, which is what we've typically used
2. STRtree against Shapely
3. Use Spherely

Option 1 was previously fast, but buggy and had omission errors; fixing these has made it less fast. It also requires a parameter to be set on the size of the grid cells for the intersection, which impacts run speed and also the commission error rate (i.e., how many false positives due to over-sized grid cells).

Option 2 is fast, but doesn't work globally-- euclidean geometry breaks down at the poles. We have to reproject to a stereographic projection to build Spatial Range Tree in that metric space... this is a hassle. It's either automatic and brittle, or, differed to the user to select and subject to user error (and less usable overall).

Option 3 works today, is correct (no omission or commission), and doesn't require any parameters to be set-- but it doesn't scale well. Individual polygon checks are fairly fast given that the routines are compiled with low constants, but when doing a NASA CMR catalog subset we're hitting 10's or 100's of thousands of geometries, so the actual algorithmic scaling starts to matter.

Hence this PR, which implement a spatial index with things like intersect and contains predicates.

The test case that I benchmarked this against is below, and builds a search across 76,575 polygons:

Backend comparison — cycle-22 ATL06 (76,575 pairs)

mortie 0.7.2 MOC sweep (the espg/mortie#33 fix → zero omission at every order):

backend	order	omission	commission	time
morton_coverage_moc	6	0	15	5.4 s
morton_coverage_moc	8	0	2	11.8 s
morton_coverage_moc	10	0	0 (exact)	44.5 s
spherely + SpatialIndex (prototype)	—	0	0	~2 s
spherely brute (exact ref)	—	—	—	~56 s

As you can see above, having the spatial index drops the intersection time from about a minute, to 2 seconds (and unlike mortie doesn't require any parameter settings).

[espg]

Note that the three failing checks (ci/environment-dev.yml, Python 3.14, ubuntu/macos/windows) are unrelated to this PR — they fail in an upstream dependency build before any spherely code is compiled.

Those legs build s2geography main from source against an unpinned s2geometry>=0.11.1. conda-forge published s2geometry 0.14.0 on 2026-06-06, and s2geography main doesn't compile against it — 0.14.0's header layout breaks the s2/base/port.h include chain (via s2coords_internal.h), failing while compiling s2geography's own accessors-geog.cc:

  .../envs/spherely-dev/include/s2/s2coords_internal.h:22:10:
     fatal error: s2/base/port.h: No such file or directory

This is independent of the change here: the dev legs passed as recently as the June 3 dependabot run (which resolved a pre-0.14 s2geometry); they broke the moment 0.14.0 landed, and this PR was simply the first CI run afterward. Any PR run now hits the same wall. The ci/environment.yml legs stay green across 3.10–3.14 because they use the conda s2geography 0.2.0 package, whose metadata keeps s2geometry on a compatible version.

[espg]

@benbovy is there anything I can do to move this along towards a merge?

I was thinking that query_nearest / distance queries via S2ClosestEdgeQuery would make sense as a follow up to this, but if adding them here helps nudge towards a merge I'm happy to tackle them as part of this.

[benbovy]

Thanks @espg, this is a great addition!

Sorry for the slow review, I'm not very active on this repository and I missed your initial submission a few weeks ago.

I left a few comments, overall this looks great! @jorisvandenbossche do you want to take a look at it?

I was thinking that query_nearest / distance queries via S2ClosestEdgeQuery would make sense as a follow up to this.

Yes it is perfectly fine to keep that for a follow-up PR.

Note that the three failing checks (ci/environment-dev.yml, Python 3.14, ubuntu/macos/windows) are unrelated to this PR — they fail in an upstream dependency build before any spherely code is compiled.

Those have been fixed in #119. Could you merge the main branch here please?

I maintain two libraries: zagg which aggregates input geospatial data to zarr grids, and mortie , which implements healpix / morton spatial indexing.

Thanks for providing additional context! Actually I've been working recently on a similar project that converts other data sources to HEALPix. It is not yet publicly released but I'd be happy to discuss it once it is.

[benbovy]

Suggested change

	** The input geographies are held as a numpy object-dtype array, which both
	** keeps the underlying C++ Geography objects (whose S2Shapes are borrowed by
	** the index) alive and backs the ``geometries`` property.
	*/
	*/

Not very informative.

[espg]

trimmed

[benbovy]

Could use function overloading instead?

py::array_t<py::ssize_t> query(PyObjectGeography geography, std::optional<std::string> predicate) const

and

py::array_t<py::ssize_t> query(py::array_t<PyObjectGeography> geographies, std::optional<std::string> predicate) const

[espg]

done

[benbovy]

Note: a shallow copy of the input Geography objects is safe as long as we don't support in-place coordinate replacement. Currently we don't support it but if we eventually do so - rather unlikely I guess - we'll probably need deep-copy here (i.e., clone the underlying S2 objects).

[espg]

Agreed — the index borrows the S2Shapes from the input geographies, so in-place coordinate replacement would invalidate it. If that ever lands, cloning the underlying S2 objects here (and rebuilding on mutation) would be needed.

Since Geography is currently move-only/immutable from Python, the shallow copy stays, but in the same spirit the constructor now stores a shallow copy() of the input array itself: previously, replacing an element of the caller's array after building the index could drop the last reference to a Geography whose S2Shapes the index still borrows. (shapely's STRtree has the same exposure via its geometries attribute, so this only closes the input-array half, but it's a one-liner.) Added a test for it.

[benbovy]

Suggested change

	std::vector<int> query_one(Geography* query_geog, const PredicateFunc* pred) const {
	std::vector<int> query_one(const Geography* query_geog, const PredicateFunc* pred) const {

(or use const references)

[espg]

Done — it now takes const Geography&. This required making Geography::geog_index()
const, since it lazily constructs the ShapeIndexGeography on first use; the cached
pointer in geography.hpp is now mutable (pure caching, no observable state change).

[benbovy]

I guess we could simply use np.testing.assert_array_equal instead of sets. IIUC query results are sorted.

[espg]

query results are indeed sorted (scalar queries) and in deterministic
input-major order (array queries), so the exact expected arrays are asserted directly now

[benbovy]

Could use switch instead?

[espg]

C++ can't switch on std::string, so this now goes through a small
unordered_map<std::string, PredicateId> lookup followed by a switch on the enum —
which also gets us -Wswitch coverage if a predicate is ever added to the enum but not
handled.

If you'd rather keep the simpler if/else chain, happy to revert.

[benbovy]

It's backed by s2geography::GeographyIndex (a MutableS2ShapeIndex), which is already available in the pinned s2geography 0.2.0 — no upstream or dependency changes needed.

s2geography 0.2.0 is indeed pinned in CI for the release, but currently there's no minimal version constraint in the CMake configuration. We should probably do it at some point.

In the meantime, for consistency could you move src/predicates_common.cpp and src/spatial_index.cpp under the if clause below alongside geoarrow and projections, please?

[espg]

Done — src/predicates_common.cpp and src/spatial_index.cpp now sit alongside
geoarrow.cpp and projections.cpp in the VERSION_GREATER_EQUAL "0.2.0" clause.
Agreed a minimum version constraint in the CMake config would be good at some point.

[espg]

Thanks for the review @benbovy! My turn to apologize for the slow response; got sucked into the summer holidays earlier in the month.

Actually I've been working on a similar project that converts other data sources to
HEALPix. It is not yet publicly released but I'd be happy to discuss it once it is.

Would enjoy that! If a call is ever easier than issue threads (and doable with the time difference), feel free to grab a slot on my calendly

On the HEALPix side, I just finalized the 1.0 morton indexing spec — if you happen to take a look, I'd be really interested in your thoughts. Like I mentioned, zagg is the writer side of all this, but I'm starting to tackle the reader side now with a small library, moczarr, which is a sparse-DGGS xarray reader for zarr stores of HEALPix cells under mortie's packed z-order ("morton") encoding, with MOC-declared coverage. It plugs into xdggs through the public registry (register_dggs("morton"), shipped as a moczarr[xdggs] extra), and its MOC/interval-set-backed lazy index is essentially a working instance of what was discussed in xarray-contrib/xdggs#143 for larger-than-memory cell_ids: selection and alignment run in coverage/interval arithmetic, and the cell-id coordinate is only fabricated on demand. Longer term I'd like to bring the morton grid kind into xdggs proper — the concrete asks are collected in englacial/zagg#72 — would certainly appreciate your read on whether/how that fits the xdggs roadmap.

This SpatialIndex PR is upstream plumbing for the same pipeline: spherely does the geometry→granule candidate mapping that feeds the gridding. Eventually, we can hopefully view the output using gridlook client side or read into xarray for pytorch ingest. Anyways, on the review pass for this PR, Claude Fable did most of the code updates this run, and seemed to do a good job-- here's it's summary report:

---

All comments are addressed in the latest push, and the branch is rebased on main to pick up the CI fix from #119.

Most suggestions were applied as-is; a few were adapted slightly for C++/pybind11 reasons
(details in the inline replies):

• query is now split into two proper overloads as suggested. The scalar overload takes
  the registered Geography type directly rather than PyObjectGeography, since
  PyObjectGeography doesn't have a pybind11 type caster outside of py::vectorize
  (it's only wired up through the vectorize_arg specialization in pybind11.hpp).
  pybind11's two-pass overload resolution then does the scalar-vs-array dispatch that was
  previously hand-rolled.
• To make query_one(const Geography&, ...) possible, Geography::geog_index() is now
  const with the lazily-built index pointer marked mutable (caching only — it doesn't
  affect the observable state of the Geography).
• get_predicate now uses a switch — via a small name→enum map first, since C++ can't
  switch on strings directly. Happy to revert to the plain if/else chain if you prefer.
• Following up on your shallow-copy note: the constructor now stores a shallow copy()
  of the input object array, so replacing elements of the caller's array after building
  can no longer drop Geography objects whose S2Shapes the index still borrows
  (with a test).

The refinement test is also reworked so the coarse candidate set is a strict superset of the predicate-refined result (a point outside a triangle but inside its cell covering), so the test now actually exercises the predicate filtering.

[espg]

Pushed one more change to this branch: a SpatialIndex.build() method.

This came up while running on larger spatial collections and running into memory bounded slow downs... the type of case where you'd want to build and persist the index. I'll post another follow PR shortly for spatial index persistence and go over the timings in more detail there.

What was added

MutableS2ShapeIndex builds lazily — the constructor only queues the shapes, and the first query pays for the whole build. build() forces that build to happen when you ask for it, and optionally lets you raise s2geometry's temporary memory budget for the duration:

index = spherely.SpatialIndex(geographies)
assert not index.is_built                    # the constructor only queues the shapes
index.build()                                # build now, under the default budget
assert index.is_built

big = spherely.SpatialIndex(lots_of_geographies)
big.build(tmp_memory_budget=4 * 1024**3)     # ...or under a 4 GiB budget

There is also a read-only SpatialIndex.is_built property (a thin wrapper over MutableS2ShapeIndex::is_fresh()) so that "has this cost been paid yet?" is answerable without timing a query. An index over an empty collection has nothing queued and is is_built == True from the start.

The argument is in bytes, matching the unit of the underlying s2geometry setting (FLAGS_s2shape_index_tmp_memory_budget, declared in s2/mutable_s2shape_index.h, default 104857600, i.e. 100 MB). It is set through an RAII guard and restored when the call returns, including on the exception path, so a single tuned build never leaks the override into anything else in the process. None does not mean "no budget" — it leaves s2geometry's setting untouched, so the build still batches and costs what the first query would have; it just happens where you asked for it. Pass a raised value to change the build's cost; the build itself goes through s2's own MutableS2ShapeIndex::ForceBuild(), which is documented for exactly this ("to ensure that the first subsequent query is as fast as possible") and is already a no-op once the index is fresh — so build() is idempotent and free to call after queries have already forced the build.

The knob is only useful raised. A budget below the one in effect splits the build into more batches and makes it slower — 1 MB measured about ten times slower than the default on a 634k-edge collection — and the unit is easy to get wrong in the expensive direction: tmp_memory_budget=100 is 100 bytes, not 100 MB, and is accepted. Zero and negatives raise ValueError naming both the offending value and the s2 default; values that do not fit in an int64 raise TypeError from the binding.

Why it's added here

Any user who indexes a large collection hits this today, and hits it in a confusing place: the constructor looks fast and then an unrelated-looking first query hangs.

The default temp memory is bad at scale-- each batch re-absorbs the index cells built so far, so many batches therefore cost superlinearly rather than additively. (Query time is linear and has nothing to do with any of this.)

Measured on 555,867 real ICESat-2 ATL03 footprint polygons (17.68M vertices), single-threaded, M1 Max —

	build (s)	2,721-query loop (s)	pairs returned	peak RSS (MB)
build() (default 100 MB)	196.51	5.27	190,625	2,614
build(tmp_memory_budget=4 * 1024**3)	10.07	5.30	190,625	4,336

19.5× on the build, and the fitted complexity drops from N^1.79 to N^1.08 — i.e. it goes from superlinear to essentially linear in edge count.

The two arms return byte-identical results (same 190,625 (input, tree) pairs, same SHA-256 over the pair array), and query time is unchanged, which is the point: this is purely a build-scheduling knob, not a change to what the index contains.

edit: Both arms return the same predicate-refined results — the 190,625 pairs above are from predicate="intersects", byte-identical, same SHA-256. The unrefined candidate set is not stable across budgets: s2 batches differently, and candidate lists at different budgets are mutually incomparable (candidates can appear and disappear). What is invariant is the set of true matches, so any query with a predicate is unaffected. A second build() call returns in 5 µs.

The cost is transient memory: peak RSS goes from 2.6 GB to 4.3 GB at this N, since the larger budget is exactly a licence to hold more scratch at once. Peak RSS can therefore rise by roughly the budget granted, and s2 needs on the order of 226 bytes of temporary memory per edge — so n_edges * 226 is the number to pass to build in a single batch (17.68M vertices here ≈ 4GB, which is why 4 GiB was enough). That rule of thumb is in the docstring. One thing worth stating plainly: build() holds the GIL for its whole duration, so no other Python thread runs until it returns.