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

	SpatialIndex(py::object geographies) {
	auto arr = py::array_t<PyObjectGeography>::ensure(geographies);
	if (!arr) {
	throw py::type_error("geographies must be an array-like of Geography objects");
	}
	SpatialIndex(py::array_t<PyObjectGeography> geographies) {

This is cleaner I think (or maybe const py::array_t<PyObjectGeography>& geographies).

I wasn't aware of py::array_t<>::ensure (it is not in pybind11's API reference docs?), but I don't think it will check that array elements are Geography objects (it will just check that the array has the numpy.object dtype). The Geography instance check is done below element by element when calling data[i].as_geog_ptr().

Using py::array_t<PyObjectGeography> directly should handle array-like input arguments as well.

[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.

[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

[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).

[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)

[benbovy]

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

[benbovy]

Use np.testing.assert_array_equal.

[benbovy]

Suggested change

	version = "0.1.1+spatialindex"
	version = "0.1.1"

[benbovy]

Could use switch instead?

[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?