From 0e96b1a3fb70b2e39fce15407f0174192e08b9a0 Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Sat, 6 Jun 2026 00:59:20 -0700 Subject: [PATCH 01/12] prototype for spatial indexing for fast polygon contains operation --- CMakeLists.txt | 2 + docs/api.rst | 12 +++ docs/api_hidden.rst | 3 + spatial_index_plan.md | 178 +++++++++++++++++++++++++++++++ src/predicates.hpp | 30 ++++++ src/predicates_common.cpp | 70 +++++++++++++ src/spatial_index.cpp | 203 ++++++++++++++++++++++++++++++++++++ src/spherely.cpp | 2 + src/spherely.pyi | 16 +++ tests/test_spatial_index.py | 112 ++++++++++++++++++++ 10 files changed, 628 insertions(+) create mode 100644 spatial_index_plan.md create mode 100644 src/predicates.hpp create mode 100644 src/predicates_common.cpp create mode 100644 src/spatial_index.cpp create mode 100644 tests/test_spatial_index.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 3cd1d11..91df11b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -77,6 +77,8 @@ set(CPP_SOURCES src/geography.cpp src/io.cpp src/predicates.cpp + src/predicates_common.cpp + src/spatial_index.cpp src/spherely.cpp) if(${s2geography_VERSION} VERSION_GREATER_EQUAL "0.2.0") diff --git a/docs/api.rst b/docs/api.rst index 1c906da..b2e4e93 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -99,6 +99,18 @@ between two geographies. covers covered_by +.. _api_spatial_index: + +Spatial indexing +---------------- + +Index a collection of geographies for fast spatial queries. + +.. autosummary:: + :toctree: _api_generated/ + + SpatialIndex + .. _api_overlays: Overlays (boolean operations) diff --git a/docs/api_hidden.rst b/docs/api_hidden.rst index 7862012..9241a7d 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -6,6 +6,9 @@ :toctree: _api_generated/ Geography + SpatialIndex + SpatialIndex.geometries + SpatialIndex.query Projection Projection.lnglat Projection.pseudo_mercator diff --git a/spatial_index_plan.md b/spatial_index_plan.md new file mode 100644 index 0000000..6d89622 --- /dev/null +++ b/spatial_index_plan.md @@ -0,0 +1,178 @@ +# Plan: `SpatialIndex` for spherely (issue #72) + +## Context + +spherely (S2-backed analog of shapely) currently has no way to spatially index a +large collection of `Geography` objects. Every pairwise predicate (`intersects`, +`contains`, …) is computed brute-force via `py::vectorize`, which is O(N·M) for a +spatial join. Issue #72 asks for an STRtree-like index so candidate pairs can be +pre-filtered cheaply, as shapely's `STRtree` does. + +**Feasibility: confirmed — all primitives already exist in the pinned +`s2geography 0.2.0` (conda-forge), so no upstream changes are needed.** Key facts +from exploration: + +- `s2geography::GeographyIndex` (header-only, `s2geography/index.h`) wraps a + `MutableS2ShapeIndex` and is purpose-built as "a GEOSSTRTree index" over a + *vector* of geographies. Public API (verified at the `0.2.0` tag): + - `void Add(const Geography& geog, int value)` — `value` is the caller's array index. + - `int value(int shape_id) const` — maps an internal shape id back to that array index. + - `const MutableS2ShapeIndex& ShapeIndex() const`. + - nested `Iterator` with `void Query(const std::vector& covering, std::unordered_set* out)` + which fills `out` with matching array indices (cell-overlap candidates). +- `s2geography::Geography::Region() -> std::unique_ptr` and + `GetCellUnionBound(std::vector*)` exist, so we can compute a query + covering via `S2RegionCoverer`. +- The R `s2` package implements exactly this pattern (`s2-matrix.cpp`): build + `GeographyIndex`, then per query feature `coverer.GetCovering(*Geog().Region(), &cells)` + → `iterator.Query(cells, &indices)`. +- spherely's existing predicate functions (`s2geog::s2_intersects`, `s2_contains`, + etc., used in `src/predicates.cpp`) operate on `ShapeIndexGeography` and can be + reused directly for the refinement step, via the lazy + `Geography::geog_index()` accessor (`src/geography.hpp:82`). + +**Decided scope (per user):** ship the index + `query()` with optional predicate +refinement. Distance/nearest (`S2ClosestEdgeQuery`) and serialization are +explicitly deferred to a follow-up PR. Class name: **`SpatialIndex`**. + +## Target Python API + +```python +tree = spherely.SpatialIndex(geographies) # iterable / object-ndarray of Geography +len(tree) # number of indexed geographies +tree.geometries # object-ndarray of the inputs (same order) + +# scalar query -> sorted np.intp array of tree indices +idx = tree.query(geom) # coarse: cells overlap (candidate set) +idx = tree.query(geom, predicate="intersects") # refined + +# array query -> shape (2, K): row 0 = input index, row 1 = tree index (shapely-compatible) +pairs = tree.query(np.array([...]), predicate="intersects") +``` + +- `predicate=None` → coarse cell-overlap candidates (superset of true intersections). +- Supported `predicate` values (all are subsets of the intersecting-candidate set, + so refinement is correct): `"intersects"`, `"within"`, `"contains"`, `"covers"`, + `"covered_by"`, `"touches"`, `"equals"`. Semantics match + shapely: a tree geom `t` matches when `predicate(query_geom, t)` is True. +- `"disjoint"` is **rejected** in v1 (it is *not* a subset of the overlap candidate + set; computing it via the index needs the complement — defer). Raise `ValueError`. + +## Implementation + +### 1. New C++ source: `src/spatial_index.cpp` +Mirror the structure of `src/predicates.cpp` (a `void init_spatial_index(py::module&)` +entry point, `namespace s2geog = s2geography`, reuse `PyObjectGeography`). + +Includes: ``, ``, ``, +``, `"geography.hpp"`, `"pybind11.hpp"`. + +Define `class SpatialIndex`: +- **Members:** `std::unique_ptr m_index;` and + `py::array m_geographies;` (an object-dtype ndarray held to (a) keep the indexed + `Geography` C++ objects alive — `GeographyIndex` only borrows their `S2Shape`s — + and (b) back the `geometries` property). +- **Constructor `SpatialIndex(py::array_t geographies)`:** + store the array; `m_index = std::make_unique();` then loop + `i` over the flat array, `auto* g = geographies.data()[i]->as_geog_ptr();` + `m_index->Add(g->geog(), static_cast(i));`. Empty geographies contribute no + shapes (documented: they are never returned). Validate 1-D input; reject `None` + entries with a clear `TypeError` (reuse `PyObjectGeography::check_type`). +- **`__len__`** → number of input geographies. +- **`geometries` property** → return `m_geographies`. +- **Private helper `query_one(Geography* query_geog, predicate) -> std::vector`:** + 1. `auto region = query_geog->geog().Region();` + 2. `S2RegionCoverer coverer;` (default options for v1) → + `std::vector covering; coverer.GetCovering(*region, &covering);` + 3. `std::unordered_set hits; GeographyIndex::Iterator it(m_index.get); it.Query(covering, &hits);` + 4. If `predicate != None`: keep only `c` where + `predfn(query_geog->geog_index(), candidate->geog_index(), options)` is true, + where `candidate = m_geographies.data()[c]->as_geog_ptr()`. Map predicate name → + callable over two `ShapeIndexGeography` (see step 2 reuse note below). + 5. copy set → `std::vector`, `std::sort`. +- **`query(scalar Geography)`** → `py::array_t` (sorted). +- **`query(array of Geography)`** → build two `std::vector` (input idx, tree idx) + and return a `(2, K)` `py::array_t` (shapely's `query` bulk layout). + +### 2. Reuse predicate functions (small shared helper) +The predicate lambdas currently live *inside* `init_predicates` in +`src/predicates.cpp` and aren't reachable. Add a tiny shared mapping so both files +use one source of truth. Lightest option: + +- Add `src/predicates.hpp` declaring a factory, e.g. + `using PredicateFunc = std::function;` + and `PredicateFunc get_predicate(const std::string& name);` returning a closure + over the right `s2geog::s2_intersects`/`s2_contains` call + the appropriate + `S2BooleanOperation::Options` (e.g. closed model for `covers`/`covered_by`, + the dual-options trick for `touches`, argument-swap for `within`/`covered_by`). +- Implement it in `predicates.cpp` (or a new `predicates_common.cpp`) and have + `init_predicates` register the public vectorized predicates by reusing the same + closures, avoiding duplication. `SpatialIndex::query_one` calls `get_predicate`. +- Unknown / unsupported name (incl. `"disjoint"`) → throw `py::value_error`. + +### 3. Wire-up +- `src/spherely.cpp`: declare `void init_spatial_index(py::module&);` and call it in + `PYBIND11_MODULE` (next to `init_predicates`). +- `CMakeLists.txt`: add `src/spatial_index.cpp` (and `predicates_common.cpp` if + split out) to `CPP_SOURCES`. No new link deps — `GeographyIndex` is header-only and + `s2geography`/`s2::s2` are already linked. Gate on s2geography ≥ 0.2.0 if needed + (already the project floor: `s2geography >=0.2.0,<0.3`). + +### 4. Type stubs — `src/spherely.pyi` +Add (near the predicates section): +```python +class SpatialIndex: + def __init__(self, geographies: Iterable[Geography]) -> None: ... + def __len__(self) -> int: ... + @property + def geometries(self) -> T_NDArray_Geography: ... + @overload + def query(self, geography: Geography, predicate: str | None = None) -> npt.NDArray[np.intp]: ... + @overload + def query(self, geography: Iterable[Geography], predicate: str | None = None) -> npt.NDArray[np.intp]: ... +``` + +### 5. Docs +- `docs/api.rst`: new "Spatial indexing" section listing `SpatialIndex`. +- `docs/api_hidden.rst`: add the class' members so autosummary picks them up + (match the existing hidden-entry pattern). + +### 6. Tests — `tests/test_spatial_index.py` +Mirror `tests/test_predicates.py` conventions (typed, numpy assertions). Cover: +- build + `len()` + `geometries` round-trips and dtype/order. +- coarse `query(scalar)` returns candidate indices (sorted, np.intp). +- `query(scalar, predicate="intersects")` refines vs coarse (strict subset case). +- each supported predicate gives shapely-equivalent results on a small hand-checked + fixture (points/lines/polygons spanning a couple of cells). +- array query returns `(2, K)` `(input_idx, tree_idx)` layout. +- empty geography in the tree is never returned; empty query returns empty. +- `predicate="disjoint"` and unknown predicate raise; `None` entry raises `TypeError`. + +## Verification + +```bash +pixi run -e test compile # build the extension (configure+compile tasks) +pixi run -e test tests # run pytest, incl. new test_spatial_index.py +pixi run -e lint mypy # type-check stubs against tests +``` +Manual smoke test in `pixi shell -e test`: +```python +import numpy as np, spherely +geoms = spherely.points([0,1,2,50], [0,1,2,50]) +t = spherely.SpatialIndex(geoms) +assert len(t) == 4 +q = spherely.create_polygon([(-1,-1),(3,-1),(3,3),(-1,3)]) +print(t.query(q)) # -> [0 1 2] +print(t.query(q, predicate="contains")) # refined subset +``` +Cross-check a handful of results against shapely's `STRtree` on the planar +equivalents to confirm the join semantics line up. + +## Deferred to follow-up (out of scope here) +- `query_nearest` / `nearest` / distance filtering via `S2ClosestEdgeQuery` + (+ `S2ClosestEdgeQuery::ShapeIndexTarget`, header ``). +- Index serialization / pickling (S2's `EncodedS2ShapeIndex` enables mmap'd, + shippable indexes — noted as high value in issue #72). +- Configurable coverer (`max_cells`) and a dedicated point-only fast path + (`S2PointIndex`) for large point sets. +``` diff --git a/src/predicates.hpp b/src/predicates.hpp new file mode 100644 index 0000000..81bc480 --- /dev/null +++ b/src/predicates.hpp @@ -0,0 +1,30 @@ +#ifndef SPHERELY_PREDICATES_H_ +#define SPHERELY_PREDICATES_H_ + +#include + +#include +#include + +namespace spherely { + +namespace s2geog = s2geography; + +// A binary spatial predicate operating on two indexed geographies, i.e. +// ``predicate(a, b)`` where ``a`` and ``b`` are ``ShapeIndexGeography``. +using PredicateFunc = + std::function; + +// Returns the predicate closure matching ``name`` (one of "intersects", +// "within", "contains", "covers", "covered_by", "touches", "equals"). +// +// Throws ``pybind11::value_error`` for unknown names and for "disjoint" (which +// cannot be evaluated as a refinement of a spatial-index candidate set). +// +// The closures mirror the semantics of the vectorized predicates registered in +// ``predicates.cpp``; that file remains the reference implementation. +PredicateFunc get_predicate(const std::string& name); + +} // namespace spherely + +#endif // SPHERELY_PREDICATES_H_ diff --git a/src/predicates_common.cpp b/src/predicates_common.cpp new file mode 100644 index 0000000..f457577 --- /dev/null +++ b/src/predicates_common.cpp @@ -0,0 +1,70 @@ +#include +#include +#include + +#include "predicates.hpp" + +namespace py = pybind11; +namespace s2geog = s2geography; + +namespace spherely { + +namespace { + +S2BooleanOperation::Options closed_options() { + S2BooleanOperation::Options options; + options.set_polyline_model(S2BooleanOperation::PolylineModel::CLOSED); + options.set_polygon_model(S2BooleanOperation::PolygonModel::CLOSED); + return options; +} + +S2BooleanOperation::Options open_options() { + S2BooleanOperation::Options options; + options.set_polyline_model(S2BooleanOperation::PolylineModel::OPEN); + options.set_polygon_model(S2BooleanOperation::PolygonModel::OPEN); + return options; +} + +} // namespace + +PredicateFunc get_predicate(const std::string& name) { + using Index = s2geog::ShapeIndexGeography; + + if (name == "intersects") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_intersects(a, b, options); + }; + } else if (name == "within") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_contains(b, a, options); + }; + } else if (name == "contains") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_contains(a, b, options); + }; + } else if (name == "equals") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_equals(a, b, options); + }; + } else if (name == "covers") { + return [options = closed_options()](const Index& a, const Index& b) { + return s2geog::s2_contains(a, b, options); + }; + } else if (name == "covered_by") { + return [options = closed_options()](const Index& a, const Index& b) { + return s2geog::s2_contains(b, a, options); + }; + } else if (name == "touches") { + return [closed = closed_options(), open = open_options()](const Index& a, const Index& b) { + return s2geog::s2_intersects(a, b, closed) && !s2geog::s2_intersects(a, b, open); + }; + } else if (name == "disjoint") { + throw py::value_error( + "the 'disjoint' predicate is not supported by SpatialIndex.query: it is not a " + "refinement of the spatial-index candidate set"); + } else { + throw py::value_error("invalid predicate: '" + name + "'"); + } +} + +} // namespace spherely diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp new file mode 100644 index 0000000..ef75fe2 --- /dev/null +++ b/src/spatial_index.cpp @@ -0,0 +1,203 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "geography.hpp" +#include "predicates.hpp" +#include "pybind11.hpp" + +namespace py = pybind11; +namespace s2geog = s2geography; +using namespace spherely; + +/* +** A spatial index over a collection of Geography objects, backed by +** s2geography::GeographyIndex (a MutableS2ShapeIndex). Provides fast candidate +** lookup similar to shapely's STRtree. +** +** 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. +*/ +class SpatialIndex { +public: + SpatialIndex(py::object geographies) { + auto arr = py::array_t::ensure(geographies); + if (!arr) { + throw py::type_error("geographies must be an array-like of Geography objects"); + } + if (arr.ndim() != 1) { + throw py::type_error("geographies must be a 1-dimensional array"); + } + + m_geographies = arr; + m_index = std::make_unique(); + + auto n = arr.size(); + auto* data = static_cast(arr.request().ptr); + for (py::ssize_t i = 0; i < n; i++) { + auto* geog_ptr = data[i].as_geog_ptr(); + m_index->Add(geog_ptr->geog(), static_cast(i)); + } + } + + std::size_t size() const { + return static_cast(m_geographies.size()); + } + + py::array geometries() const { + return m_geographies; + } + + // Dispatch between scalar (single Geography -> 1-d index array) and + // array-like (-> (2, K) array of (input index, tree index) pairs) queries. + py::object query(py::object geography, std::optional predicate) const { + PredicateFunc pred; + bool has_pred = predicate.has_value(); + if (has_pred) { + pred = get_predicate(*predicate); + } + + // PyObjectGeography is layout-compatible with py::object (see + // PyObjectGeography::from_geog), so a reference cast is safe here. + auto& maybe_geog = static_cast(geography); + if (maybe_geog.is_geog_ptr()) { + auto results = query_one(maybe_geog.as_geog_ptr(), has_pred ? &pred : nullptr); + return to_index_array(results); + } + + auto arr = py::array_t::ensure(geography); + if (!arr || arr.ndim() != 1) { + throw py::type_error( + "query geography must be a Geography or a 1-dimensional array of Geography"); + } + + auto n = arr.size(); + auto* data = static_cast(arr.request().ptr); + std::vector input_idx; + std::vector tree_idx; + for (py::ssize_t i = 0; i < n; i++) { + auto results = query_one(data[i].as_geog_ptr(), has_pred ? &pred : nullptr); + for (int t : results) { + input_idx.push_back(i); + tree_idx.push_back(static_cast(t)); + } + } + + auto k = input_idx.size(); + py::array_t out({static_cast(2), static_cast(k)}); + auto out_data = out.mutable_unchecked<2>(); + for (std::size_t j = 0; j < k; j++) { + out_data(0, static_cast(j)) = input_idx[j]; + out_data(1, static_cast(j)) = tree_idx[j]; + } + return out; + } + +private: + std::unique_ptr m_index; + py::array_t m_geographies; + + // Return the sorted tree indices whose cells overlap the query geography, + // optionally refined by ``pred`` (predicate(query, candidate)). + std::vector query_one(Geography* query_geog, const PredicateFunc* pred) const { + std::unordered_set hits; + + auto region = query_geog->geog().Region(); + S2RegionCoverer coverer; + std::vector covering; + coverer.GetCovering(*region, &covering); + + s2geog::GeographyIndex::Iterator iter(m_index.get()); + iter.Query(covering, &hits); + + std::vector results; + if (pred == nullptr) { + results.assign(hits.begin(), hits.end()); + } else { + const auto& query_index = query_geog->geog_index(); + auto* data = static_cast(m_geographies.request().ptr); + for (int candidate : hits) { + auto* cand_geog = data[candidate].as_geog_ptr(); + if ((*pred)(query_index, cand_geog->geog_index())) { + results.push_back(candidate); + } + } + } + + std::sort(results.begin(), results.end()); + return results; + } + + static py::array_t to_index_array(const std::vector& indices) { + auto n = indices.size(); + py::array_t out(static_cast(n)); + auto out_data = out.mutable_unchecked<1>(); + for (std::size_t i = 0; i < n; i++) { + out_data(static_cast(i)) = static_cast(indices[i]); + } + return out; + } +}; + +void init_spatial_index(py::module& m) { + py::class_(m, "SpatialIndex", R"pbdoc( + A spatial index for a collection of geographies. + + The index allows fast retrieval of the geographies that may interact + with a given geography, e.g., for accelerating spatial joins. It is + built on top of s2geometry's shape index and is conceptually similar to + shapely's ``STRtree``. + + Parameters + ---------- + geographies : array_like + An array (or other sequence) of :py:class:`Geography` objects. Empty + geographies are indexed but never returned by queries. + + )pbdoc") + .def(py::init(), py::arg("geographies")) + .def("__len__", &SpatialIndex::size) + .def_property_readonly("geometries", + &SpatialIndex::geometries, + "The array of geographies in the index (in input order).") + .def("query", + &SpatialIndex::query, + py::arg("geography"), + py::arg("predicate") = py::none(), + R"pbdoc(query(geography, predicate=None) + + Return the integer indices of all geographies in the index whose cells + overlap the given geography (a coarse candidate set), optionally refined + by a spatial predicate. + + Parameters + ---------- + geography : :py:class:`Geography` or array_like + The geography or geographies to query. + predicate : str, optional + If provided, only return candidates for which + ``predicate(geography, indexed_geography)`` is True. One of + "intersects", "within", "contains", "covers", "covered_by", + "touches" or "equals". + + Returns + ------- + ndarray + If ``geography`` is a scalar, a 1-d array of (sorted) indices into + the index. If ``geography`` is an array, a ``(2, N)`` array where the + first row holds the input geography indices and the second row the + matching index (tree) indices. + + )pbdoc"); +} diff --git a/src/spherely.cpp b/src/spherely.cpp index 1f4dce7..42571db 100644 --- a/src/spherely.cpp +++ b/src/spherely.cpp @@ -8,6 +8,7 @@ namespace py = pybind11; void init_geography(py::module&); void init_creation(py::module&); void init_predicates(py::module&); +void init_spatial_index(py::module&); void init_boolean_operations(py::module&); void init_accessors(py::module&); void init_io(py::module&); @@ -29,6 +30,7 @@ PYBIND11_MODULE(spherely, m) { init_geography(m); init_creation(m); init_predicates(m); + init_spatial_index(m); init_boolean_operations(m); init_accessors(m); init_io(m); diff --git a/src/spherely.pyi b/src/spherely.pyi index e9f03d0..146a43c 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -232,6 +232,22 @@ touches: _VFunc_Nin2_Nout1[Literal["touches"], bool, bool] covers: _VFunc_Nin2_Nout1[Literal["covers"], bool, bool] covered_by: _VFunc_Nin2_Nout1[Literal["covered_by"], bool, bool] +# spatial index + +class SpatialIndex: + def __init__(self, geographies: Iterable[Geography]) -> None: ... + def __len__(self) -> int: ... + @property + def geometries(self) -> T_NDArray_Geography: ... + @overload + def query( + self, geography: Geography, predicate: str | None = None + ) -> npt.NDArray[np.intp]: ... + @overload + def query( + self, geography: Iterable[Geography], predicate: str | None = None + ) -> npt.NDArray[np.intp]: ... + # boolean operations union: _VFunc_Nin2_Nout1[Literal["union"], Geography, Geography] diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py new file mode 100644 index 0000000..94cb484 --- /dev/null +++ b/tests/test_spatial_index.py @@ -0,0 +1,112 @@ +from typing import Any + +import numpy as np +import numpy.typing as npt +import pytest + +import spherely + + +@pytest.fixture +def geographies() -> npt.NDArray[Any]: + # three points near (0, 0)..(2, 2) and one far-away point + return np.array( + [ + spherely.create_point(0, 0), + spherely.create_point(1, 1), + spherely.create_point(2, 2), + spherely.create_point(50, 50), + ] + ) + + +def test_spatial_index_len(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + assert len(tree) == 4 + + +def test_spatial_index_geometries(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + geoms = tree.geometries + assert isinstance(geoms, np.ndarray) + assert geoms.shape == (4,) + assert all(spherely.equals(a, b) for a, b in zip(geoms, geographies)) + + +def test_spatial_index_from_list() -> None: + tree = spherely.SpatialIndex( + [spherely.create_point(0, 0), spherely.create_point(1, 1)] + ) + assert len(tree) == 2 + + +def test_query_scalar(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + result = tree.query(poly) + assert result.dtype == np.intp + # the three nearby points are candidates; the far point is not + np.testing.assert_array_equal(result, [0, 1, 2]) + + +def test_query_predicate_refines(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + coarse = tree.query(poly) + refined = tree.query(poly, predicate="contains") + # refinement is a subset of the coarse candidate set + assert {int(x) for x in refined}.issubset({int(x) for x in coarse}) + # all three nearby points are actually inside the polygon + np.testing.assert_array_equal(refined, [0, 1, 2]) + + +def test_query_predicate_intersects(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + result = tree.query(point, predicate="intersects") + np.testing.assert_array_equal(result, [1]) + + +def test_query_array(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + queries = np.array( + [ + spherely.create_point(1, 1), + spherely.create_point(50, 50), + ] + ) + result = tree.query(queries, predicate="intersects") + assert result.shape[0] == 2 + # (input_index, tree_index) pairs + pairs = {(int(a), int(b)) for a, b in zip(result[0], result[1])} + assert pairs == {(0, 1), (1, 3)} + + +def test_query_empty_geography_never_returned() -> None: + geoms = np.array( + [ + spherely.create_point(0, 0), + spherely.create_polygon(None), # empty + ] + ) + tree = spherely.SpatialIndex(geoms) + assert len(tree) == 2 + poly = spherely.create_polygon([(-1, -1), (1, -1), (1, 1), (-1, 1), (-1, -1)]) + result = tree.query(poly) + assert 1 not in {int(x) for x in result} + + +def test_query_disjoint_rejected(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + with pytest.raises(ValueError, match="disjoint"): + tree.query(point, predicate="disjoint") + + +def test_query_invalid_predicate(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + with pytest.raises(ValueError, match="invalid predicate"): + tree.query(point, predicate="not_a_predicate") From a05b71f03ccea610c0c78a2cd2be294492a92787 Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Sat, 6 Jun 2026 13:13:20 -0700 Subject: [PATCH 02/12] culling planning doc --- spatial_index_plan.md | 178 ------------------------------------------ 1 file changed, 178 deletions(-) delete mode 100644 spatial_index_plan.md diff --git a/spatial_index_plan.md b/spatial_index_plan.md deleted file mode 100644 index 6d89622..0000000 --- a/spatial_index_plan.md +++ /dev/null @@ -1,178 +0,0 @@ -# Plan: `SpatialIndex` for spherely (issue #72) - -## Context - -spherely (S2-backed analog of shapely) currently has no way to spatially index a -large collection of `Geography` objects. Every pairwise predicate (`intersects`, -`contains`, …) is computed brute-force via `py::vectorize`, which is O(N·M) for a -spatial join. Issue #72 asks for an STRtree-like index so candidate pairs can be -pre-filtered cheaply, as shapely's `STRtree` does. - -**Feasibility: confirmed — all primitives already exist in the pinned -`s2geography 0.2.0` (conda-forge), so no upstream changes are needed.** Key facts -from exploration: - -- `s2geography::GeographyIndex` (header-only, `s2geography/index.h`) wraps a - `MutableS2ShapeIndex` and is purpose-built as "a GEOSSTRTree index" over a - *vector* of geographies. Public API (verified at the `0.2.0` tag): - - `void Add(const Geography& geog, int value)` — `value` is the caller's array index. - - `int value(int shape_id) const` — maps an internal shape id back to that array index. - - `const MutableS2ShapeIndex& ShapeIndex() const`. - - nested `Iterator` with `void Query(const std::vector& covering, std::unordered_set* out)` - which fills `out` with matching array indices (cell-overlap candidates). -- `s2geography::Geography::Region() -> std::unique_ptr` and - `GetCellUnionBound(std::vector*)` exist, so we can compute a query - covering via `S2RegionCoverer`. -- The R `s2` package implements exactly this pattern (`s2-matrix.cpp`): build - `GeographyIndex`, then per query feature `coverer.GetCovering(*Geog().Region(), &cells)` - → `iterator.Query(cells, &indices)`. -- spherely's existing predicate functions (`s2geog::s2_intersects`, `s2_contains`, - etc., used in `src/predicates.cpp`) operate on `ShapeIndexGeography` and can be - reused directly for the refinement step, via the lazy - `Geography::geog_index()` accessor (`src/geography.hpp:82`). - -**Decided scope (per user):** ship the index + `query()` with optional predicate -refinement. Distance/nearest (`S2ClosestEdgeQuery`) and serialization are -explicitly deferred to a follow-up PR. Class name: **`SpatialIndex`**. - -## Target Python API - -```python -tree = spherely.SpatialIndex(geographies) # iterable / object-ndarray of Geography -len(tree) # number of indexed geographies -tree.geometries # object-ndarray of the inputs (same order) - -# scalar query -> sorted np.intp array of tree indices -idx = tree.query(geom) # coarse: cells overlap (candidate set) -idx = tree.query(geom, predicate="intersects") # refined - -# array query -> shape (2, K): row 0 = input index, row 1 = tree index (shapely-compatible) -pairs = tree.query(np.array([...]), predicate="intersects") -``` - -- `predicate=None` → coarse cell-overlap candidates (superset of true intersections). -- Supported `predicate` values (all are subsets of the intersecting-candidate set, - so refinement is correct): `"intersects"`, `"within"`, `"contains"`, `"covers"`, - `"covered_by"`, `"touches"`, `"equals"`. Semantics match - shapely: a tree geom `t` matches when `predicate(query_geom, t)` is True. -- `"disjoint"` is **rejected** in v1 (it is *not* a subset of the overlap candidate - set; computing it via the index needs the complement — defer). Raise `ValueError`. - -## Implementation - -### 1. New C++ source: `src/spatial_index.cpp` -Mirror the structure of `src/predicates.cpp` (a `void init_spatial_index(py::module&)` -entry point, `namespace s2geog = s2geography`, reuse `PyObjectGeography`). - -Includes: ``, ``, ``, -``, `"geography.hpp"`, `"pybind11.hpp"`. - -Define `class SpatialIndex`: -- **Members:** `std::unique_ptr m_index;` and - `py::array m_geographies;` (an object-dtype ndarray held to (a) keep the indexed - `Geography` C++ objects alive — `GeographyIndex` only borrows their `S2Shape`s — - and (b) back the `geometries` property). -- **Constructor `SpatialIndex(py::array_t geographies)`:** - store the array; `m_index = std::make_unique();` then loop - `i` over the flat array, `auto* g = geographies.data()[i]->as_geog_ptr();` - `m_index->Add(g->geog(), static_cast(i));`. Empty geographies contribute no - shapes (documented: they are never returned). Validate 1-D input; reject `None` - entries with a clear `TypeError` (reuse `PyObjectGeography::check_type`). -- **`__len__`** → number of input geographies. -- **`geometries` property** → return `m_geographies`. -- **Private helper `query_one(Geography* query_geog, predicate) -> std::vector`:** - 1. `auto region = query_geog->geog().Region();` - 2. `S2RegionCoverer coverer;` (default options for v1) → - `std::vector covering; coverer.GetCovering(*region, &covering);` - 3. `std::unordered_set hits; GeographyIndex::Iterator it(m_index.get); it.Query(covering, &hits);` - 4. If `predicate != None`: keep only `c` where - `predfn(query_geog->geog_index(), candidate->geog_index(), options)` is true, - where `candidate = m_geographies.data()[c]->as_geog_ptr()`. Map predicate name → - callable over two `ShapeIndexGeography` (see step 2 reuse note below). - 5. copy set → `std::vector`, `std::sort`. -- **`query(scalar Geography)`** → `py::array_t` (sorted). -- **`query(array of Geography)`** → build two `std::vector` (input idx, tree idx) - and return a `(2, K)` `py::array_t` (shapely's `query` bulk layout). - -### 2. Reuse predicate functions (small shared helper) -The predicate lambdas currently live *inside* `init_predicates` in -`src/predicates.cpp` and aren't reachable. Add a tiny shared mapping so both files -use one source of truth. Lightest option: - -- Add `src/predicates.hpp` declaring a factory, e.g. - `using PredicateFunc = std::function;` - and `PredicateFunc get_predicate(const std::string& name);` returning a closure - over the right `s2geog::s2_intersects`/`s2_contains` call + the appropriate - `S2BooleanOperation::Options` (e.g. closed model for `covers`/`covered_by`, - the dual-options trick for `touches`, argument-swap for `within`/`covered_by`). -- Implement it in `predicates.cpp` (or a new `predicates_common.cpp`) and have - `init_predicates` register the public vectorized predicates by reusing the same - closures, avoiding duplication. `SpatialIndex::query_one` calls `get_predicate`. -- Unknown / unsupported name (incl. `"disjoint"`) → throw `py::value_error`. - -### 3. Wire-up -- `src/spherely.cpp`: declare `void init_spatial_index(py::module&);` and call it in - `PYBIND11_MODULE` (next to `init_predicates`). -- `CMakeLists.txt`: add `src/spatial_index.cpp` (and `predicates_common.cpp` if - split out) to `CPP_SOURCES`. No new link deps — `GeographyIndex` is header-only and - `s2geography`/`s2::s2` are already linked. Gate on s2geography ≥ 0.2.0 if needed - (already the project floor: `s2geography >=0.2.0,<0.3`). - -### 4. Type stubs — `src/spherely.pyi` -Add (near the predicates section): -```python -class SpatialIndex: - def __init__(self, geographies: Iterable[Geography]) -> None: ... - def __len__(self) -> int: ... - @property - def geometries(self) -> T_NDArray_Geography: ... - @overload - def query(self, geography: Geography, predicate: str | None = None) -> npt.NDArray[np.intp]: ... - @overload - def query(self, geography: Iterable[Geography], predicate: str | None = None) -> npt.NDArray[np.intp]: ... -``` - -### 5. Docs -- `docs/api.rst`: new "Spatial indexing" section listing `SpatialIndex`. -- `docs/api_hidden.rst`: add the class' members so autosummary picks them up - (match the existing hidden-entry pattern). - -### 6. Tests — `tests/test_spatial_index.py` -Mirror `tests/test_predicates.py` conventions (typed, numpy assertions). Cover: -- build + `len()` + `geometries` round-trips and dtype/order. -- coarse `query(scalar)` returns candidate indices (sorted, np.intp). -- `query(scalar, predicate="intersects")` refines vs coarse (strict subset case). -- each supported predicate gives shapely-equivalent results on a small hand-checked - fixture (points/lines/polygons spanning a couple of cells). -- array query returns `(2, K)` `(input_idx, tree_idx)` layout. -- empty geography in the tree is never returned; empty query returns empty. -- `predicate="disjoint"` and unknown predicate raise; `None` entry raises `TypeError`. - -## Verification - -```bash -pixi run -e test compile # build the extension (configure+compile tasks) -pixi run -e test tests # run pytest, incl. new test_spatial_index.py -pixi run -e lint mypy # type-check stubs against tests -``` -Manual smoke test in `pixi shell -e test`: -```python -import numpy as np, spherely -geoms = spherely.points([0,1,2,50], [0,1,2,50]) -t = spherely.SpatialIndex(geoms) -assert len(t) == 4 -q = spherely.create_polygon([(-1,-1),(3,-1),(3,3),(-1,3)]) -print(t.query(q)) # -> [0 1 2] -print(t.query(q, predicate="contains")) # refined subset -``` -Cross-check a handful of results against shapely's `STRtree` on the planar -equivalents to confirm the join semantics line up. - -## Deferred to follow-up (out of scope here) -- `query_nearest` / `nearest` / distance filtering via `S2ClosestEdgeQuery` - (+ `S2ClosestEdgeQuery::ShapeIndexTarget`, header ``). -- Index serialization / pickling (S2's `EncodedS2ShapeIndex` enables mmap'd, - shippable indexes — noted as high value in issue #72). -- Configurable coverer (`max_cells`) and a dedicated point-only fast path - (`S2PointIndex`) for large point sets. -``` From 0381dcd99fc6f829af47d04ee89d3a008f8a2fab Mon Sep 17 00:00:00 2001 From: Shane Date: Tue, 9 Jun 2026 18:15:52 -0700 Subject: [PATCH 03/12] Bump version to 0.1.1+spatialindex for fork wheel builds --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 1606a70..164a31b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build" [project] name = "spherely" -version = "0.1.1" +version = "0.1.1+spatialindex" description = "Manipulation and analysis of geometric objects on the sphere" keywords = ["gis", "geometry", "s2geometry", "shapely"] readme = "README.md" From 7b2604830cc81ea0ef34895ef1fe1f1e7fd4efbc Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Tue, 21 Jul 2026 16:37:16 -0700 Subject: [PATCH 04/12] addressing #118 review --- CMakeLists.txt | 8 +-- pyproject.toml | 2 +- src/geography.hpp | 6 ++- src/predicates_common.cpp | 97 ++++++++++++++++++++++++------------- src/spatial_index.cpp | 90 ++++++++++++++++++---------------- src/spherely.pyi | 8 ++- tests/test_spatial_index.py | 47 ++++++++++++------ 7 files changed, 161 insertions(+), 97 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 91df11b..bd6f1d3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -77,12 +77,14 @@ set(CPP_SOURCES src/geography.cpp src/io.cpp src/predicates.cpp - src/predicates_common.cpp - src/spatial_index.cpp src/spherely.cpp) if(${s2geography_VERSION} VERSION_GREATER_EQUAL "0.2.0") - set(CPP_SOURCES ${CPP_SOURCES} src/geoarrow.cpp src/projections.cpp) + set(CPP_SOURCES ${CPP_SOURCES} + src/geoarrow.cpp + src/predicates_common.cpp + src/projections.cpp + src/spatial_index.cpp) endif() add_library(spherely MODULE ${CPP_SOURCES}) diff --git a/pyproject.toml b/pyproject.toml index 164a31b..1606a70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build" [project] name = "spherely" -version = "0.1.1+spatialindex" +version = "0.1.1" description = "Manipulation and analysis of geometric objects on the sphere" keywords = ["gis", "geometry", "s2geometry", "shapely"] readme = "README.md" diff --git a/src/geography.hpp b/src/geography.hpp index ca0c89f..3f61ab8 100644 --- a/src/geography.hpp +++ b/src/geography.hpp @@ -79,7 +79,7 @@ class Geography { return reinterpret_cast(&geog()); } - inline const s2geog::ShapeIndexGeography& geog_index() { + inline const s2geog::ShapeIndexGeography& geog_index() const { if (!m_s2geog_index_ptr) { m_s2geog_index_ptr = std::make_unique(geog()); } @@ -111,7 +111,9 @@ class Geography { private: S2GeographyPtr m_s2geog_ptr; - S2GeographyIndexPtr m_s2geog_index_ptr; + // lazily built by geog_index() (mutable: caching only, doesn't affect + // the observable state of the Geography) + mutable S2GeographyIndexPtr m_s2geog_index_ptr; bool m_is_empty = false; GeographyType m_geog_type; diff --git a/src/predicates_common.cpp b/src/predicates_common.cpp index f457577..2d6701b 100644 --- a/src/predicates_common.cpp +++ b/src/predicates_common.cpp @@ -2,6 +2,8 @@ #include #include +#include + #include "predicates.hpp" namespace py = pybind11; @@ -25,46 +27,75 @@ S2BooleanOperation::Options open_options() { return options; } +enum class PredicateId { + intersects, + within, + contains, + equals, + covers, + covered_by, + touches, + disjoint +}; + +const std::unordered_map predicate_ids = { + {"intersects", PredicateId::intersects}, + {"within", PredicateId::within}, + {"contains", PredicateId::contains}, + {"equals", PredicateId::equals}, + {"covers", PredicateId::covers}, + {"covered_by", PredicateId::covered_by}, + {"touches", PredicateId::touches}, + {"disjoint", PredicateId::disjoint}, +}; + } // namespace PredicateFunc get_predicate(const std::string& name) { using Index = s2geog::ShapeIndexGeography; - if (name == "intersects") { - return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { - return s2geog::s2_intersects(a, b, options); - }; - } else if (name == "within") { - return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { - return s2geog::s2_contains(b, a, options); - }; - } else if (name == "contains") { - return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { - return s2geog::s2_contains(a, b, options); - }; - } else if (name == "equals") { - return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { - return s2geog::s2_equals(a, b, options); - }; - } else if (name == "covers") { - return [options = closed_options()](const Index& a, const Index& b) { - return s2geog::s2_contains(a, b, options); - }; - } else if (name == "covered_by") { - return [options = closed_options()](const Index& a, const Index& b) { - return s2geog::s2_contains(b, a, options); - }; - } else if (name == "touches") { - return [closed = closed_options(), open = open_options()](const Index& a, const Index& b) { - return s2geog::s2_intersects(a, b, closed) && !s2geog::s2_intersects(a, b, open); - }; - } else if (name == "disjoint") { - throw py::value_error( - "the 'disjoint' predicate is not supported by SpatialIndex.query: it is not a " - "refinement of the spatial-index candidate set"); - } else { + auto it = predicate_ids.find(name); + if (it == predicate_ids.end()) { throw py::value_error("invalid predicate: '" + name + "'"); } + + switch (it->second) { + case PredicateId::intersects: + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_intersects(a, b, options); + }; + case PredicateId::within: + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_contains(b, a, options); + }; + case PredicateId::contains: + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_contains(a, b, options); + }; + case PredicateId::equals: + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_equals(a, b, options); + }; + case PredicateId::covers: + return [options = closed_options()](const Index& a, const Index& b) { + return s2geog::s2_contains(a, b, options); + }; + case PredicateId::covered_by: + return [options = closed_options()](const Index& a, const Index& b) { + return s2geog::s2_contains(b, a, options); + }; + case PredicateId::touches: + return [closed = closed_options(), open = open_options()](const Index& a, + const Index& b) { + return s2geog::s2_intersects(a, b, closed) && !s2geog::s2_intersects(a, b, open); + }; + case PredicateId::disjoint: + throw py::value_error( + "the 'disjoint' predicate is not supported by SpatialIndex.query: it is not a " + "refinement of the spatial-index candidate set"); + } + + throw py::value_error("invalid predicate: '" + name + "'"); // unreachable } } // namespace spherely diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index ef75fe2..42169b1 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -24,30 +24,24 @@ using namespace spherely; ** A spatial index over a collection of Geography objects, backed by ** s2geography::GeographyIndex (a MutableS2ShapeIndex). Provides fast candidate ** lookup similar to shapely's STRtree. -** -** 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. */ class SpatialIndex { public: - SpatialIndex(py::object geographies) { - auto arr = py::array_t::ensure(geographies); - if (!arr) { - throw py::type_error("geographies must be an array-like of Geography objects"); - } - if (arr.ndim() != 1) { + SpatialIndex(const py::array_t& geographies) { + if (geographies.ndim() != 1) { throw py::type_error("geographies must be a 1-dimensional array"); } - m_geographies = arr; + // store a shallow copy of the input array so that replacing its + // elements afterwards cannot drop Geography objects whose S2Shapes + // are still borrowed by the index + m_geographies = py::array_t::ensure(geographies.attr("copy")()); m_index = std::make_unique(); - auto n = arr.size(); - auto* data = static_cast(arr.request().ptr); + auto n = m_geographies.size(); + auto* data = static_cast(m_geographies.request().ptr); for (py::ssize_t i = 0; i < n; i++) { - auto* geog_ptr = data[i].as_geog_ptr(); - m_index->Add(geog_ptr->geog(), static_cast(i)); + m_index->Add(data[i].as_geog_ptr()->geog(), static_cast(i)); } } @@ -59,35 +53,31 @@ class SpatialIndex { return m_geographies; } - // Dispatch between scalar (single Geography -> 1-d index array) and - // array-like (-> (2, K) array of (input index, tree index) pairs) queries. - py::object query(py::object geography, std::optional predicate) const { - PredicateFunc pred; - bool has_pred = predicate.has_value(); - if (has_pred) { - pred = get_predicate(*predicate); - } - - // PyObjectGeography is layout-compatible with py::object (see - // PyObjectGeography::from_geog), so a reference cast is safe here. - auto& maybe_geog = static_cast(geography); - if (maybe_geog.is_geog_ptr()) { - auto results = query_one(maybe_geog.as_geog_ptr(), has_pred ? &pred : nullptr); - return to_index_array(results); - } + // Scalar query: single Geography -> 1-d array of sorted tree indices. + py::array_t query(const Geography& geography, + std::optional predicate) const { + auto pred = make_predicate(predicate); + auto results = query_one(geography, pred ? &*pred : nullptr); + return to_index_array(results); + } - auto arr = py::array_t::ensure(geography); - if (!arr || arr.ndim() != 1) { + // Array query: 1-d array of Geography -> (2, K) array of + // (input index, tree index) pairs. + py::array_t query(const py::array_t& geographies, + std::optional predicate) const { + if (geographies.ndim() != 1) { throw py::type_error( "query geography must be a Geography or a 1-dimensional array of Geography"); } - auto n = arr.size(); - auto* data = static_cast(arr.request().ptr); + auto pred = make_predicate(predicate); + + auto n = geographies.size(); + auto* data = static_cast(geographies.request().ptr); std::vector input_idx; std::vector tree_idx; for (py::ssize_t i = 0; i < n; i++) { - auto results = query_one(data[i].as_geog_ptr(), has_pred ? &pred : nullptr); + auto results = query_one(*data[i].as_geog_ptr(), pred ? &*pred : nullptr); for (int t : results) { input_idx.push_back(i); tree_idx.push_back(static_cast(t)); @@ -106,14 +96,22 @@ class SpatialIndex { private: std::unique_ptr m_index; + // also keeps the Geography objects alive (the index borrows their S2Shapes) py::array_t m_geographies; + static std::optional make_predicate(const std::optional& name) { + if (!name.has_value()) { + return std::nullopt; + } + return get_predicate(*name); + } + // Return the sorted tree indices whose cells overlap the query geography, // optionally refined by ``pred`` (predicate(query, candidate)). - std::vector query_one(Geography* query_geog, const PredicateFunc* pred) const { + std::vector query_one(const Geography& query_geog, const PredicateFunc* pred) const { std::unordered_set hits; - auto region = query_geog->geog().Region(); + auto region = query_geog.geog().Region(); S2RegionCoverer coverer; std::vector covering; coverer.GetCovering(*region, &covering); @@ -125,7 +123,7 @@ class SpatialIndex { if (pred == nullptr) { results.assign(hits.begin(), hits.end()); } else { - const auto& query_index = query_geog->geog_index(); + const auto& query_index = query_geog.geog_index(); auto* data = static_cast(m_geographies.request().ptr); for (int candidate : hits) { auto* cand_geog = data[candidate].as_geog_ptr(); @@ -166,13 +164,16 @@ void init_spatial_index(py::module& m) { geographies are indexed but never returned by queries. )pbdoc") - .def(py::init(), py::arg("geographies")) + .def(py::init&>(), + py::arg("geographies"), + "__init__(self, geographies)") .def("__len__", &SpatialIndex::size) .def_property_readonly("geometries", &SpatialIndex::geometries, "The array of geographies in the index (in input order).") .def("query", - &SpatialIndex::query, + py::overload_cast>(&SpatialIndex::query, + py::const_), py::arg("geography"), py::arg("predicate") = py::none(), R"pbdoc(query(geography, predicate=None) @@ -199,5 +200,10 @@ void init_spatial_index(py::module& m) { first row holds the input geography indices and the second row the matching index (tree) indices. - )pbdoc"); + )pbdoc") + .def("query", + py::overload_cast&, std::optional>( + &SpatialIndex::query, py::const_), + py::arg("geography"), + py::arg("predicate") = py::none()); } diff --git a/src/spherely.pyi b/src/spherely.pyi index 146a43c..13410d0 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -234,6 +234,10 @@ covered_by: _VFunc_Nin2_Nout1[Literal["covered_by"], bool, bool] # spatial index +_IndexPredicate = Literal[ + "intersects", "within", "contains", "covers", "covered_by", "touches", "equals" +] + class SpatialIndex: def __init__(self, geographies: Iterable[Geography]) -> None: ... def __len__(self) -> int: ... @@ -241,11 +245,11 @@ class SpatialIndex: def geometries(self) -> T_NDArray_Geography: ... @overload def query( - self, geography: Geography, predicate: str | None = None + self, geography: Geography, predicate: _IndexPredicate | None = None ) -> npt.NDArray[np.intp]: ... @overload def query( - self, geography: Iterable[Geography], predicate: str | None = None + self, geography: Iterable[Geography], predicate: _IndexPredicate | None = None ) -> npt.NDArray[np.intp]: ... # boolean operations diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 94cb484..2892ba7 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -40,6 +40,16 @@ def test_spatial_index_from_list() -> None: assert len(tree) == 2 +def test_spatial_index_copies_input_array() -> None: + geoms = np.array([spherely.create_point(0, 0), spherely.create_point(1, 1)]) + tree = spherely.SpatialIndex(geoms) + # replacing an element of the input array does not affect the index + geoms[0] = spherely.create_point(40, 40) + result = tree.query(spherely.create_point(0, 0), predicate="intersects") + np.testing.assert_array_equal(result, [0]) + assert spherely.equals(tree.geometries[0], spherely.create_point(0, 0)) + + def test_query_scalar(geographies: npt.NDArray[Any]) -> None: tree = spherely.SpatialIndex(geographies) poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) @@ -50,16 +60,27 @@ def test_query_scalar(geographies: npt.NDArray[Any]) -> None: np.testing.assert_array_equal(result, [0, 1, 2]) -def test_query_predicate_refines(geographies: npt.NDArray[Any]) -> None: +def test_query_predicate_refines() -> None: + # point 1 is outside the triangle (beyond the hypotenuse) but close enough + # that it falls within the triangle's coarse cell covering; point 2 is far + # away and is not a candidate at all + geographies = np.array( + [ + spherely.create_point(1, 1), + spherely.create_point(4, 4), + spherely.create_point(50, 50), + ] + ) tree = spherely.SpatialIndex(geographies) - poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + triangle = spherely.create_polygon([(0, 0), (5, 0), (0, 5), (0, 0)]) + + coarse = tree.query(triangle) + refined = tree.query(triangle, predicate="contains") - coarse = tree.query(poly) - refined = tree.query(poly, predicate="contains") - # refinement is a subset of the coarse candidate set - assert {int(x) for x in refined}.issubset({int(x) for x in coarse}) - # all three nearby points are actually inside the polygon - np.testing.assert_array_equal(refined, [0, 1, 2]) + # the coarse candidate set includes the false positive, the refined + # (exact) result does not + np.testing.assert_array_equal(coarse, [0, 1]) + np.testing.assert_array_equal(refined, [0]) def test_query_predicate_intersects(geographies: npt.NDArray[Any]) -> None: @@ -78,10 +99,8 @@ def test_query_array(geographies: npt.NDArray[Any]) -> None: ] ) result = tree.query(queries, predicate="intersects") - assert result.shape[0] == 2 # (input_index, tree_index) pairs - pairs = {(int(a), int(b)) for a, b in zip(result[0], result[1])} - assert pairs == {(0, 1), (1, 3)} + np.testing.assert_array_equal(result, [[0, 1], [1, 3]]) def test_query_empty_geography_never_returned() -> None: @@ -95,18 +114,18 @@ def test_query_empty_geography_never_returned() -> None: assert len(tree) == 2 poly = spherely.create_polygon([(-1, -1), (1, -1), (1, 1), (-1, 1), (-1, -1)]) result = tree.query(poly) - assert 1 not in {int(x) for x in result} + np.testing.assert_array_equal(result, [0]) def test_query_disjoint_rejected(geographies: npt.NDArray[Any]) -> None: tree = spherely.SpatialIndex(geographies) point = spherely.create_point(1, 1) with pytest.raises(ValueError, match="disjoint"): - tree.query(point, predicate="disjoint") + tree.query(point, predicate="disjoint") # type: ignore[call-overload] def test_query_invalid_predicate(geographies: npt.NDArray[Any]) -> None: tree = spherely.SpatialIndex(geographies) point = spherely.create_point(1, 1) with pytest.raises(ValueError, match="invalid predicate"): - tree.query(point, predicate="not_a_predicate") + tree.query(point, predicate="not_a_predicate") # type: ignore[call-overload] From f44e66400110ba047cec6ed31ece5a5000fce7dd Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 15:48:49 -0700 Subject: [PATCH 05/12] exposing temporary memory budget to caller on build spatial index --- docs/api_hidden.rst | 2 + src/spatial_index.cpp | 159 ++++++++++++++++++++++++++++++++++++ src/spherely.pyi | 5 ++ tests/test_spatial_index.py | 147 +++++++++++++++++++++++++++++++++ 4 files changed, 313 insertions(+) diff --git a/docs/api_hidden.rst b/docs/api_hidden.rst index 9241a7d..d2d538a 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -7,7 +7,9 @@ Geography SpatialIndex + SpatialIndex.build SpatialIndex.geometries + SpatialIndex.is_built SpatialIndex.query Projection Projection.lnglat diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index 42169b1..e20f28b 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -1,4 +1,6 @@ #include +#include +#include #include #include #include @@ -6,6 +8,7 @@ #include #include +#include #include #include #include @@ -20,6 +23,40 @@ namespace py = pybind11; namespace s2geog = s2geography; using namespace spherely; +namespace { + +/* +** Sets s2geometry's index-build temporary memory budget +** (``FLAGS_s2shape_index_tmp_memory_budget``, declared in +** s2/mutable_s2shape_index.h) for the lifetime of the guard and puts the +** previous value back on destruction, including when the build throws. +** +** The flag is process-wide, so the guard exists to keep the override from +** leaking out of the single build it was requested for. It is read and written +** with absl::GetFlag / absl::SetFlag; s2/base/commandlineflags.h is included +** for those (it pulls in ) so that the absl flags +** dependency is expressed through s2's own header, the way s2 does it. +*/ +class TmpMemoryBudgetGuard { +public: + explicit TmpMemoryBudgetGuard(std::int64_t bytes) + : m_previous(absl::GetFlag(FLAGS_s2shape_index_tmp_memory_budget)) { + absl::SetFlag(&FLAGS_s2shape_index_tmp_memory_budget, bytes); + } + + ~TmpMemoryBudgetGuard() { + absl::SetFlag(&FLAGS_s2shape_index_tmp_memory_budget, m_previous); + } + + TmpMemoryBudgetGuard(const TmpMemoryBudgetGuard&) = delete; + TmpMemoryBudgetGuard& operator=(const TmpMemoryBudgetGuard&) = delete; + +private: + std::int64_t m_previous; +}; + +} // namespace + /* ** A spatial index over a collection of Geography objects, backed by ** s2geography::GeographyIndex (a MutableS2ShapeIndex). Provides fast candidate @@ -49,6 +86,36 @@ class SpatialIndex { return static_cast(m_geographies.size()); } + // True once the queued updates have been applied, i.e. the index can be + // queried without doing any build work. + bool is_built() const { + return m_index->ShapeIndex().is_fresh(); + } + + // Apply the updates queued by the constructor, optionally under a raised + // s2 temporary memory budget. A no-op once the index is built. + void build(std::optional tmp_memory_budget) const { + if (tmp_memory_budget.has_value() && *tmp_memory_budget <= 0) { + throw py::value_error("tmp_memory_budget must be a positive number of bytes, got " + + std::to_string(*tmp_memory_budget) + + " (s2geometry's default budget is 104857600 bytes)"); + } + + // nothing is pending, so there is no build to tune: return without + // touching the process-wide budget flag + if (is_built()) { + return; + } + + const auto& shape_index = m_index->ShapeIndex(); + if (tmp_memory_budget.has_value()) { + TmpMemoryBudgetGuard guard(*tmp_memory_budget); + shape_index.ForceBuild(); + } else { + shape_index.ForceBuild(); + } + } + py::array geometries() const { return m_geographies; } @@ -157,6 +224,11 @@ void init_spatial_index(py::module& m) { built on top of s2geometry's shape index and is conceptually similar to shapely's ``STRtree``. + The underlying shape index is built lazily: the constructor only queues + the geographies, and the first query pays for the build. Use + :py:meth:`SpatialIndex.build` to do that work up front, and to tune it + for large collections. + Parameters ---------- geographies : array_like @@ -167,7 +239,87 @@ void init_spatial_index(py::module& m) { .def(py::init&>(), py::arg("geographies"), "__init__(self, geographies)") + .def("build", + &SpatialIndex::build, + py::arg("tmp_memory_budget") = py::none(), + R"pbdoc(build(tmp_memory_budget=None) + + Force the deferred index build, optionally under a raised temporary + memory budget. + + The index is built lazily, so the constructor is cheap and the first + query pays the whole build cost. This method makes that cost explicit + (and keeps it out of the timing of the first query), and lets it be + tuned for large collections. + + s2geometry builds the index in batches sized to fit a temporary memory + budget of 104857600 bytes (100 MB) by default, and each batch + re-absorbs the cells built so far. Many batches therefore cost + noticeably more than one, and collections of millions of edges can + spend minutes in the first query. Raising the budget so that the index + is built in fewer (ideally one) batches makes the build close to linear + in the number of edges. + + Parameters + ---------- + tmp_memory_budget : int, optional + Temporary memory budget for this build, **in bytes** -- the unit of + the underlying s2geometry setting + (``FLAGS_s2shape_index_tmp_memory_budget``, default 104857600, i.e. + 100 MB). Must be strictly positive, and must fit in a signed 64-bit + integer (larger values raise ``TypeError``). The previous value is + put back when the call returns, so the override does not leak into + any later build. ``None`` does *not* mean "no budget": it leaves + s2geometry's setting untouched (100 MB unless something else in the + process changed it), so the build still batches and costs what the + first query would have. Pass a raised value to avoid that. + + This 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. Mind the unit: ``tmp_memory_budget=100`` is 100 + *bytes*, not 100 MB, and is accepted. + + Notes + ----- + Memory: the budget is a licence to hold that much scratch at once, so + peak resident memory can rise by roughly the amount granted. s2geometry + needs on the order of 140 bytes of temporary memory per edge, so + ``n_edges * 140`` is a reasonable budget for building in a single + batch. Nothing is granted unless a budget is passed, so the extra + memory is always opt-in. + + Calling this method is optional and never changes query results. It is + idempotent: a second call, or a call after the index has already been + built by a query, does nothing (and leaves the budget untouched). See + :py:attr:`SpatialIndex.is_built`. + + Threads: this call holds the GIL for its whole duration, so no other + Python thread runs until the build finishes -- minutes, for a large + collection. That is not a regression (a first query does the same build + under the same GIL), but it is worth knowing before calling it from a + thread. The budget is a process-wide s2geometry setting for the + duration of the call, so an index being built concurrently in a + non-Python thread would see the override too. + + Examples + -------- + >>> import spherely + >>> index = spherely.SpatialIndex(geographies) + >>> index.build(tmp_memory_budget=4 * 1024**3) # 4 GiB + + )pbdoc") .def("__len__", &SpatialIndex::size) + .def_property_readonly("is_built", + &SpatialIndex::is_built, + R"pbdoc(Whether the underlying shape index has been built. + + ``False`` between the constructor and the first + :py:meth:`SpatialIndex.build` or query, ``True`` afterwards. An index + over an empty collection has nothing to queue and is ``True`` from the + start. + + )pbdoc") .def_property_readonly("geometries", &SpatialIndex::geometries, "The array of geographies in the index (in input order).") @@ -206,4 +358,11 @@ void init_spatial_index(py::module& m) { &SpatialIndex::query, py::const_), py::arg("geography"), py::arg("predicate") = py::none()); + + // Not public API: exposed so that the test suite can assert that + // SpatialIndex.build() puts the process-wide budget back. + m.def( + "_s2_tmp_memory_budget", + []() { return absl::GetFlag(FLAGS_s2shape_index_tmp_memory_budget); }, + "Return s2geometry's current index-build temporary memory budget, in bytes."); } diff --git a/src/spherely.pyi b/src/spherely.pyi index 13410d0..51a8462 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -240,8 +240,11 @@ _IndexPredicate = Literal[ class SpatialIndex: def __init__(self, geographies: Iterable[Geography]) -> None: ... + def build(self, tmp_memory_budget: int | None = None) -> None: ... def __len__(self) -> int: ... @property + def is_built(self) -> bool: ... + @property def geometries(self) -> T_NDArray_Geography: ... @overload def query( @@ -252,6 +255,8 @@ class SpatialIndex: self, geography: Iterable[Geography], predicate: _IndexPredicate | None = None ) -> npt.NDArray[np.intp]: ... +def _s2_tmp_memory_budget() -> int: ... + # boolean operations union: _VFunc_Nin2_Nout1[Literal["union"], Geography, Geography] diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 2892ba7..65de511 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -50,6 +50,153 @@ def test_spatial_index_copies_input_array() -> None: assert spherely.equals(tree.geometries[0], spherely.create_point(0, 0)) +def test_is_built(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + assert tree.is_built is False + tree.build() + assert tree.is_built is True + + +def test_is_built_after_query(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + assert tree.is_built is False + # the first query builds the index + tree.query(spherely.create_point(1, 1)) + assert tree.is_built is True + + +def test_build_is_idempotent(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + assert not tree.is_built + tree.build() + assert tree.is_built + tree.build() + # a budget override on an already-built index is accepted and ignored + tree.build(tmp_memory_budget=1) + assert tree.is_built + + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + np.testing.assert_array_equal(tree.query(poly), [0, 1, 2]) + + +def test_build_after_query(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + expected = tree.query(point, predicate="intersects") + # the query already built the index; building again changes nothing + assert tree.is_built + tree.build() + assert tree.is_built + np.testing.assert_array_equal(tree.query(point, predicate="intersects"), expected) + + +def test_build_does_not_change_results(geographies: npt.NDArray[Any]) -> None: + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + lazy = spherely.SpatialIndex(geographies) + eager = spherely.SpatialIndex(geographies) + eager.build() + assert eager.is_built + assert not lazy.is_built + + np.testing.assert_array_equal(eager.query(poly), lazy.query(poly)) + np.testing.assert_array_equal( + eager.query(poly, predicate="contains"), lazy.query(poly, predicate="contains") + ) + + +# the s2 default and two raised budgets -- raising is the only useful +# direction, see the SpatialIndex.build docstring +@pytest.mark.parametrize("budget", [104_857_600, 1024**3, 4 * 1024**3]) +def test_build_budget_does_not_change_results( + geographies: npt.NDArray[Any], budget: int +) -> None: + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + default = spherely.SpatialIndex(geographies) + default.build() + tuned = spherely.SpatialIndex(geographies) + tuned.build(tmp_memory_budget=budget) + assert tuned.is_built + + np.testing.assert_array_equal(tuned.query(poly), default.query(poly)) + np.testing.assert_array_equal( + tuned.query(poly, predicate="contains"), + default.query(poly, predicate="contains"), + ) + + +def test_build_budget_below_default_is_legal(geographies: npt.NDArray[Any]) -> None: + # a budget far below the default is accepted and still correct; it only + # makes the build slower (more batches), so it is not a value to copy + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + default = spherely.SpatialIndex(geographies) + default.build() + tiny = spherely.SpatialIndex(geographies) + tiny.build(tmp_memory_budget=1) + assert tiny.is_built + + np.testing.assert_array_equal(tiny.query(poly), default.query(poly)) + + +@pytest.mark.parametrize("budget", [0, -1, -(1024**3)]) +def test_build_invalid_budget(geographies: npt.NDArray[Any], budget: int) -> None: + tree = spherely.SpatialIndex(geographies) + with pytest.raises(ValueError, match="positive"): + tree.build(tmp_memory_budget=budget) + assert not tree.is_built + + # the failed call left the index buildable + tree.build() + assert tree.is_built + np.testing.assert_array_equal( + tree.query(spherely.create_point(1, 1), predicate="intersects"), [1] + ) + + # the budget is validated before the "already built" early return, so an + # invalid value is rejected even when there is nothing left to build + with pytest.raises(ValueError, match="positive"): + tree.build(tmp_memory_budget=budget) + + +def test_build_budget_out_of_range(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + # the budget is a C int64; anything larger is rejected by the binding + with pytest.raises(TypeError): + tree.build(tmp_memory_budget=2**63) + assert not tree.is_built + + +def test_build_budget_is_restored(geographies: npt.NDArray[Any]) -> None: + # the override is scoped to a single build: the process-wide s2 setting is + # back to its previous value once the call returns + before = spherely._s2_tmp_memory_budget() + + tuned = spherely.SpatialIndex(geographies) + # deliberately different from the budget in effect, so that a failure to + # put the previous value back is visible here + tuned.build(tmp_memory_budget=before + 4096) + assert spherely._s2_tmp_memory_budget() == before + + # ...so an index built afterwards without one behaves exactly as if the + # override had never happened + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + plain = spherely.SpatialIndex(geographies) + plain.build() + assert spherely._s2_tmp_memory_budget() == before + np.testing.assert_array_equal(plain.query(poly), tuned.query(poly)) + + +def test_build_empty_index() -> None: + tree = spherely.SpatialIndex([]) + # nothing was queued, so there is nothing to build + assert tree.is_built + tree.build(tmp_memory_budget=1024) + assert len(tree) == 0 + np.testing.assert_array_equal(tree.query(spherely.create_point(0, 0)), []) + + def test_query_scalar(geographies: npt.NDArray[Any]) -> None: tree = spherely.SpatialIndex(geographies) poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) From dbfceb3402c63cd3ae615a8a7a73f32bed4265fd Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 16:16:44 -0700 Subject: [PATCH 06/12] fixing windows build failures, condensing docstrings and comments --- src/spatial_index.cpp | 137 ++++++++++++++++++++---------------------- 1 file changed, 66 insertions(+), 71 deletions(-) diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index e20f28b..abe255a 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -1,5 +1,6 @@ +#include +#include #include -#include #include #include #include @@ -25,34 +26,52 @@ using namespace spherely; namespace { +constexpr const char* kTmpMemoryBudgetFlag = "s2shape_index_tmp_memory_budget"; + /* -** Sets s2geometry's index-build temporary memory budget -** (``FLAGS_s2shape_index_tmp_memory_budget``, declared in -** s2/mutable_s2shape_index.h) for the lifetime of the guard and puts the -** previous value back on destruction, including when the build throws. +** Sets s2geometry's process-wide index-build temporary memory budget for the +** lifetime of the guard and puts the previous value back on destruction, +** including when the build throws, so an override cannot leak out of the build +** it was requested for. ** -** The flag is process-wide, so the guard exists to keep the override from -** leaking out of the single build it was requested for. It is read and written -** with absl::GetFlag / absl::SetFlag; s2/base/commandlineflags.h is included -** for those (it pulls in ) so that the absl flags -** dependency is expressed through s2's own header, the way s2 does it. +** The flag is reached by name through absl's reflection API rather than as +** ``FLAGS_s2shape_index_tmp_memory_budget``: ABSL_DECLARE_FLAG expands to a +** plain ``extern`` with no ``dllimport``, so taking that data symbol's address +** does not link against a shared s2 on Windows. The cost is that the value +** round-trips through a string and the parse can fail; a lookup that finds no +** such flag skips the override. */ class TmpMemoryBudgetGuard { public: explicit TmpMemoryBudgetGuard(std::int64_t bytes) - : m_previous(absl::GetFlag(FLAGS_s2shape_index_tmp_memory_budget)) { - absl::SetFlag(&FLAGS_s2shape_index_tmp_memory_budget, bytes); + : m_flag(absl::FindCommandLineFlag(kTmpMemoryBudgetFlag)) { + if (m_flag == nullptr) { + return; + } + m_previous = m_flag->CurrentValue(); + std::string error; + if (!m_flag->ParseFrom(std::to_string(bytes), &error)) { + m_flag = nullptr; // nothing was changed, so nothing to restore + throw py::value_error("could not set tmp_memory_budget: " + error); + } } ~TmpMemoryBudgetGuard() { - absl::SetFlag(&FLAGS_s2shape_index_tmp_memory_budget, m_previous); + if (m_flag == nullptr) { + return; + } + // a destructor cannot report a failed restore, and the value it is + // putting back came from the flag itself, so swallow the result + std::string error; + m_flag->ParseFrom(m_previous, &error); } TmpMemoryBudgetGuard(const TmpMemoryBudgetGuard&) = delete; TmpMemoryBudgetGuard& operator=(const TmpMemoryBudgetGuard&) = delete; private: - std::int64_t m_previous; + absl::CommandLineFlag* m_flag; + std::string m_previous; }; } // namespace @@ -86,8 +105,7 @@ class SpatialIndex { return static_cast(m_geographies.size()); } - // True once the queued updates have been applied, i.e. the index can be - // queried without doing any build work. + // True once the queued updates have been applied (no build work pending). bool is_built() const { return m_index->ShapeIndex().is_fresh(); } @@ -101,8 +119,7 @@ class SpatialIndex { " (s2geometry's default budget is 104857600 bytes)"); } - // nothing is pending, so there is no build to tune: return without - // touching the process-wide budget flag + // nothing pending: no build to tune, so leave the process-wide flag alone if (is_built()) { return; } @@ -225,9 +242,8 @@ void init_spatial_index(py::module& m) { shapely's ``STRtree``. The underlying shape index is built lazily: the constructor only queues - the geographies, and the first query pays for the build. Use - :py:meth:`SpatialIndex.build` to do that work up front, and to tune it - for large collections. + the geographies and the first query pays for the build. Use + :py:meth:`SpatialIndex.build` to do that work up front instead. Parameters ---------- @@ -247,60 +263,34 @@ void init_spatial_index(py::module& m) { Force the deferred index build, optionally under a raised temporary memory budget. - The index is built lazily, so the constructor is cheap and the first - query pays the whole build cost. This method makes that cost explicit - (and keeps it out of the timing of the first query), and lets it be - tuned for large collections. - - s2geometry builds the index in batches sized to fit a temporary memory - budget of 104857600 bytes (100 MB) by default, and each batch - re-absorbs the cells built so far. Many batches therefore cost - noticeably more than one, and collections of millions of edges can - spend minutes in the first query. Raising the budget so that the index - is built in fewer (ideally one) batches makes the build close to linear - in the number of edges. + The index is built lazily, so without this call the first query pays + the whole build cost. Calling it is optional, never changes query + results, and does nothing once the index is built (see + :py:attr:`SpatialIndex.is_built`). Parameters ---------- tmp_memory_budget : int, optional - Temporary memory budget for this build, **in bytes** -- the unit of - the underlying s2geometry setting - (``FLAGS_s2shape_index_tmp_memory_budget``, default 104857600, i.e. - 100 MB). Must be strictly positive, and must fit in a signed 64-bit - integer (larger values raise ``TypeError``). The previous value is - put back when the call returns, so the override does not leak into - any later build. ``None`` does *not* mean "no budget": it leaves - s2geometry's setting untouched (100 MB unless something else in the - process changed it), so the build still batches and costs what the - first query would have. Pass a raised value to avoid that. - - This 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. Mind the unit: ``tmp_memory_budget=100`` is 100 - *bytes*, not 100 MB, and is accepted. + Temporary memory budget for this build, in bytes. s2geometry builds + the index in batches sized to fit it, and many batches cost more + than one, so the knob is only useful *raised*; about 140 bytes per + edge is enough to build in a single batch. Must be strictly + positive, and must fit in a signed 64-bit integer (larger values + raise ``TypeError``). The previous value is put back when the call + returns, so the override does not leak into any later build. + + ``None`` does *not* mean "no budget": it leaves s2geometry's + setting untouched (104857600 bytes, i.e. 100 MB, unless something + else in the process changed it), so the build still batches. Notes ----- - Memory: the budget is a licence to hold that much scratch at once, so - peak resident memory can rise by roughly the amount granted. s2geometry - needs on the order of 140 bytes of temporary memory per edge, so - ``n_edges * 140`` is a reasonable budget for building in a single - batch. Nothing is granted unless a budget is passed, so the extra - memory is always opt-in. - - Calling this method is optional and never changes query results. It is - idempotent: a second call, or a call after the index has already been - built by a query, does nothing (and leaves the budget untouched). See - :py:attr:`SpatialIndex.is_built`. - - Threads: this call holds the GIL for its whole duration, so no other - Python thread runs until the build finishes -- minutes, for a large - collection. That is not a regression (a first query does the same build - under the same GIL), but it is worth knowing before calling it from a - thread. The budget is a process-wide s2geometry setting for the - duration of the call, so an index being built concurrently in a - non-Python thread would see the override too. + Peak resident memory can rise by roughly the budget granted, and + nothing is granted unless a budget is passed. + + This call holds the GIL for its whole duration -- minutes, for a large + collection -- and the budget is a process-wide s2geometry setting for + that duration. Examples -------- @@ -316,8 +306,7 @@ void init_spatial_index(py::module& m) { ``False`` between the constructor and the first :py:meth:`SpatialIndex.build` or query, ``True`` afterwards. An index - over an empty collection has nothing to queue and is ``True`` from the - start. + over an empty collection is ``True`` from the start. )pbdoc") .def_property_readonly("geometries", @@ -363,6 +352,12 @@ void init_spatial_index(py::module& m) { // SpatialIndex.build() puts the process-wide budget back. m.def( "_s2_tmp_memory_budget", - []() { return absl::GetFlag(FLAGS_s2shape_index_tmp_memory_budget); }, + []() -> std::int64_t { + const auto* flag = absl::FindCommandLineFlag(kTmpMemoryBudgetFlag); + if (flag == nullptr) { + throw py::value_error(std::string("no such s2 flag: ") + kTmpMemoryBudgetFlag); + } + return std::stoll(flag->CurrentValue()); + }, "Return s2geometry's current index-build temporary memory budget, in bytes."); } From aa0fa23550cd13490feb3fba22f5d02a4e951448 Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 17:52:20 -0700 Subject: [PATCH 07/12] windows fixes for symbol mapping --- CMakeLists.txt | 29 +++++++++- src/s2_tmp_memory_budget.cpp | 52 +++++++++++++++++ src/s2_tmp_memory_budget.hpp | 21 +++++++ src/spatial_index.cpp | 83 +++++++++++---------------- src/spherely.pyi | 1 + tests/test_spatial_index.py | 106 ++++++++++++++++++++++++++++++++--- 6 files changed, 233 insertions(+), 59 deletions(-) create mode 100644 src/s2_tmp_memory_budget.cpp create mode 100644 src/s2_tmp_memory_budget.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index bd6f1d3..5618167 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -47,6 +47,31 @@ endif() # Compile definitions and flags +# How src/s2_tmp_memory_budget.cpp declares +# FLAGS_s2shape_index_tmp_memory_budget. A shared s2 exports it as *data*, and +# an import library defines only __imp_ for data, so the reference has +# to be routed through that entry with __declspec(dllimport); a static s2 +# defines the bare symbol and wants the plain extern that s2's own header +# already provides. Both are in use -- conda-forge ships a DLL, vcpkg builds +# s2geometry static on Windows -- so take the answer from the imported target. +# +# The fallback is the dllimport form because the two mistakes are not +# symmetric: a plain extern against a DLL cannot link at all (LNK2001), while +# dllimport against a static s2 links as long as the object defining the flag +# is in the image -- which it is here, since it also defines +# MutableS2ShapeIndex -- at the cost of warning LNK4217. +if(WIN32) + get_target_property(_s2_type s2::s2 TYPE) + if(_s2_type STREQUAL "STATIC_LIBRARY") + set(SPHERELY_S2_DLLIMPORT 0) + else() + set(SPHERELY_S2_DLLIMPORT 1) + endif() + message(STATUS "s2::s2 is ${_s2_type}: SPHERELY_S2_DLLIMPORT=${SPHERELY_S2_DLLIMPORT}") +else() + set(SPHERELY_S2_DLLIMPORT 0) +endif() + if (MSVC) # used in s2geometry's CMakeLists.txt but not defined in target # TODO: move this in FindS2.cmake? @@ -84,6 +109,7 @@ if(${s2geography_VERSION} VERSION_GREATER_EQUAL "0.2.0") src/geoarrow.cpp src/predicates_common.cpp src/projections.cpp + src/s2_tmp_memory_budget.cpp src/spatial_index.cpp) endif() @@ -95,7 +121,8 @@ target_compile_definitions( VERSION_INFO=${PROJECT_VERSION} S2GEOGRAPHY_VERSION=${s2geography_VERSION} S2GEOGRAPHY_VERSION_MAJOR=${s2geography_VERSION_MAJOR} - S2GEOGRAPHY_VERSION_MINOR=${s2geography_VERSION_MINOR}) + S2GEOGRAPHY_VERSION_MINOR=${s2geography_VERSION_MINOR} + SPHERELY_S2_DLLIMPORT=${SPHERELY_S2_DLLIMPORT}) target_link_libraries(spherely PRIVATE pybind11::module pybind11::lto pybind11::windows_extras diff --git a/src/s2_tmp_memory_budget.cpp b/src/s2_tmp_memory_budget.cpp new file mode 100644 index 0000000..320fb80 --- /dev/null +++ b/src/s2_tmp_memory_budget.cpp @@ -0,0 +1,52 @@ +#include "s2_tmp_memory_budget.hpp" + +#include + +/* +** ``FLAGS_s2shape_index_tmp_memory_budget`` is defined by s2geometry and +** declared by its headers with ABSL_DECLARE_FLAG, which expands to a plain +** ``extern``. +** +** That is enough on ELF and Mach-O, and on Windows whenever s2 is linked +** statically. It is not enough against a *shared* s2, where the flag is a data +** export: CMake's WINDOWS_EXPORT_ALL_SYMBOLS does put data symbols in the +** export table -- s2.dll exports +** ``?FLAGS_s2shape_index_tmp_memory_budget@@3V?$Flag@_J@flags_internal@lts_*@absl@@A`` +** from .data, and s2.lib carries the matching ``__imp_`` entry -- but for data +** an import library provides *only* that entry, so the importing side has to +** route the reference through it with ``__declspec(dllimport)``. A plain +** ``extern`` asks the linker for the bare symbol and does not link (LNK2001). +** +** Which form is right therefore depends on how s2 was built, and both happen: +** conda-forge ships a DLL, vcpkg builds s2geometry static on Windows. CMake +** decides and passes SPHERELY_S2_DLLIMPORT; see CMakeLists.txt for why the +** fallback here is the dllimport form. +** +** The two declarations cannot both be visible in one translation unit, so the +** dllimport one is made here, where no s2 header is included. +*/ +#ifndef SPHERELY_S2_DLLIMPORT +#ifdef _WIN32 +#define SPHERELY_S2_DLLIMPORT 1 +#else +#define SPHERELY_S2_DLLIMPORT 0 +#endif +#endif + +#if SPHERELY_S2_DLLIMPORT +extern __declspec(dllimport) absl::Flag FLAGS_s2shape_index_tmp_memory_budget; +#else +#include +#endif + +namespace spherely { + +std::int64_t get_s2_tmp_memory_budget() { + return static_cast(absl::GetFlag(FLAGS_s2shape_index_tmp_memory_budget)); +} + +void set_s2_tmp_memory_budget(std::int64_t bytes) { + absl::SetFlag(&FLAGS_s2shape_index_tmp_memory_budget, bytes); +} + +} // namespace spherely diff --git a/src/s2_tmp_memory_budget.hpp b/src/s2_tmp_memory_budget.hpp new file mode 100644 index 0000000..bd0ae65 --- /dev/null +++ b/src/s2_tmp_memory_budget.hpp @@ -0,0 +1,21 @@ +#ifndef SPHERELY_S2_TMP_MEMORY_BUDGET_H_ +#define SPHERELY_S2_TMP_MEMORY_BUDGET_H_ + +#include + +namespace spherely { + +/* +** Accessors for s2geometry's process-wide index-build temporary memory budget +** (``FLAGS_s2shape_index_tmp_memory_budget``, declared in +** s2/mutable_s2shape_index.h), in bytes. +** +** They are defined out of line because reaching that flag needs a declaration +** that cannot coexist with s2's own in one translation unit -- see the .cpp. +*/ +std::int64_t get_s2_tmp_memory_budget(); +void set_s2_tmp_memory_budget(std::int64_t bytes); + +} // namespace spherely + +#endif // SPHERELY_S2_TMP_MEMORY_BUDGET_H_ diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index abe255a..ab92fda 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -1,5 +1,3 @@ -#include -#include #include #include #include @@ -19,6 +17,7 @@ #include "geography.hpp" #include "predicates.hpp" #include "pybind11.hpp" +#include "s2_tmp_memory_budget.hpp" namespace py = pybind11; namespace s2geog = s2geography; @@ -26,52 +25,31 @@ using namespace spherely; namespace { -constexpr const char* kTmpMemoryBudgetFlag = "s2shape_index_tmp_memory_budget"; - /* -** Sets s2geometry's process-wide index-build temporary memory budget for the -** lifetime of the guard and puts the previous value back on destruction, -** including when the build throws, so an override cannot leak out of the build -** it was requested for. +** Sets s2geometry's process-wide index-build temporary memory budget +** (``FLAGS_s2shape_index_tmp_memory_budget``) for the lifetime of the guard and +** puts the previous value back when it goes out of scope, so an override cannot +** leak out of the build it was requested for. ** -** The flag is reached by name through absl's reflection API rather than as -** ``FLAGS_s2shape_index_tmp_memory_budget``: ABSL_DECLARE_FLAG expands to a -** plain ``extern`` with no ``dllimport``, so taking that data symbol's address -** does not link against a shared s2 on Windows. The cost is that the value -** round-trips through a string and the parse can fail; a lookup that finds no -** such flag skips the override. +** The flag is read and written through spherely::get/set_s2_tmp_memory_budget, +** which own the declaration it needs (see s2_tmp_memory_budget.cpp) -- reaching +** it directly from this file would not link against a shared s2 on Windows. */ class TmpMemoryBudgetGuard { public: - explicit TmpMemoryBudgetGuard(std::int64_t bytes) - : m_flag(absl::FindCommandLineFlag(kTmpMemoryBudgetFlag)) { - if (m_flag == nullptr) { - return; - } - m_previous = m_flag->CurrentValue(); - std::string error; - if (!m_flag->ParseFrom(std::to_string(bytes), &error)) { - m_flag = nullptr; // nothing was changed, so nothing to restore - throw py::value_error("could not set tmp_memory_budget: " + error); - } + explicit TmpMemoryBudgetGuard(std::int64_t bytes) : m_previous(get_s2_tmp_memory_budget()) { + set_s2_tmp_memory_budget(bytes); } ~TmpMemoryBudgetGuard() { - if (m_flag == nullptr) { - return; - } - // a destructor cannot report a failed restore, and the value it is - // putting back came from the flag itself, so swallow the result - std::string error; - m_flag->ParseFrom(m_previous, &error); + set_s2_tmp_memory_budget(m_previous); } TmpMemoryBudgetGuard(const TmpMemoryBudgetGuard&) = delete; TmpMemoryBudgetGuard& operator=(const TmpMemoryBudgetGuard&) = delete; private: - absl::CommandLineFlag* m_flag; - std::string m_previous; + std::int64_t m_previous; }; } // namespace @@ -264,16 +242,16 @@ void init_spatial_index(py::module& m) { memory budget. The index is built lazily, so without this call the first query pays - the whole build cost. Calling it is optional, never changes query - results, and does nothing once the index is built (see - :py:attr:`SpatialIndex.is_built`). + the whole build cost. Calling it is optional, never changes the result + of a query made with a ``predicate``, and does nothing once the index + is built (see :py:attr:`SpatialIndex.is_built`). Parameters ---------- tmp_memory_budget : int, optional Temporary memory budget for this build, in bytes. s2geometry builds the index in batches sized to fit it, and many batches cost more - than one, so the knob is only useful *raised*; about 140 bytes per + than one, so the knob is only useful *raised*; about 226 bytes per edge is enough to build in a single batch. Must be strictly positive, and must fit in a signed 64-bit integer (larger values raise ``TypeError``). The previous value is put back when the call @@ -288,6 +266,14 @@ void init_spatial_index(py::module& m) { Peak resident memory can rise by roughly the budget granted, and nothing is granted unless a budget is passed. + The budget changes where s2geometry places its batch boundaries, and so + how it subdivides the index. A query made without a ``predicate`` can + therefore return a different candidate set under a different budget -- + candidates may be added *or* dropped, and the sets for two budgets need + not be nested. What does not change is that the candidate set is always + a superset of the true matches, so any ``predicate`` refines it to the + same answer. + This call holds the GIL for its whole duration -- minutes, for a large collection -- and the budget is a process-wide s2geometry setting for that duration. @@ -348,16 +334,13 @@ void init_spatial_index(py::module& m) { py::arg("geography"), py::arg("predicate") = py::none()); - // Not public API: exposed so that the test suite can assert that - // SpatialIndex.build() puts the process-wide budget back. - m.def( - "_s2_tmp_memory_budget", - []() -> std::int64_t { - const auto* flag = absl::FindCommandLineFlag(kTmpMemoryBudgetFlag); - if (flag == nullptr) { - throw py::value_error(std::string("no such s2 flag: ") + kTmpMemoryBudgetFlag); - } - return std::stoll(flag->CurrentValue()); - }, - "Return s2geometry's current index-build temporary memory budget, in bytes."); + // Not public API: exposed so that the test suite can move the process-wide + // budget off its default and assert what SpatialIndex.build() puts back. + m.def("_s2_tmp_memory_budget", + &get_s2_tmp_memory_budget, + "Return s2geometry's current index-build temporary memory budget, in bytes."); + m.def("_set_s2_tmp_memory_budget", + &set_s2_tmp_memory_budget, + py::arg("bytes"), + "Set s2geometry's index-build temporary memory budget, in bytes."); } diff --git a/src/spherely.pyi b/src/spherely.pyi index 51a8462..5a53a2e 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -256,6 +256,7 @@ class SpatialIndex: ) -> npt.NDArray[np.intp]: ... def _s2_tmp_memory_budget() -> int: ... +def _set_s2_tmp_memory_budget(bytes: int) -> None: ... # boolean operations diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 65de511..7fa64ce 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -1,3 +1,4 @@ +from collections.abc import Iterator from typing import Any import numpy as np @@ -7,6 +8,11 @@ import spherely +def _pairs(result: npt.NDArray[np.intp]) -> set[tuple[Any, ...]]: + # the (2, N) result of an array query as a set of (input, tree) pairs + return set(map(tuple, result.T)) + + @pytest.fixture def geographies() -> npt.NDArray[Any]: # three points near (0, 0)..(2, 2) and one far-away point @@ -108,7 +114,7 @@ def test_build_does_not_change_results(geographies: npt.NDArray[Any]) -> None: # the s2 default and two raised budgets -- raising is the only useful # direction, see the SpatialIndex.build docstring @pytest.mark.parametrize("budget", [104_857_600, 1024**3, 4 * 1024**3]) -def test_build_budget_does_not_change_results( +def test_build_budget_does_not_change_matches( geographies: npt.NDArray[Any], budget: int ) -> None: poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) @@ -119,11 +125,12 @@ def test_build_budget_does_not_change_results( tuned.build(tmp_memory_budget=budget) assert tuned.is_built - np.testing.assert_array_equal(tuned.query(poly), default.query(poly)) - np.testing.assert_array_equal( - tuned.query(poly, predicate="contains"), - default.query(poly, predicate="contains"), - ) + # the budget moves s2's batch boundaries and with them the subdivision of + # the index, so the unrefined candidate set is only guaranteed to stay a + # superset of the matches -- what is invariant is the refined answer + matches = tuned.query(poly, predicate="contains") + np.testing.assert_array_equal(matches, default.query(poly, predicate="contains")) + assert set(matches) <= set(tuned.query(poly)) def test_build_budget_below_default_is_legal(geographies: npt.NDArray[Any]) -> None: @@ -137,7 +144,57 @@ def test_build_budget_below_default_is_legal(geographies: npt.NDArray[Any]) -> N tiny.build(tmp_memory_budget=1) assert tiny.is_built - np.testing.assert_array_equal(tiny.query(poly), default.query(poly)) + matches = tiny.query(poly, predicate="contains") + np.testing.assert_array_equal(matches, default.query(poly, predicate="contains")) + assert set(matches) <= set(tiny.query(poly)) + + +def _overlapping_geographies() -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: + # 600 overlapping 30-gons and 50 polygon queries over the same patch: big + # enough that the coarse candidate set holds real false positives, and that + # a budget below the default splits the build into several batches (18000 + # edges at s2geometry's 226 temporary bytes per edge is about 4 MB) + rng = np.random.default_rng(7) + + def polygon( + cx: float, cy: float, radius: float, nvertices: int + ) -> spherely.Geography: + angles = np.linspace(0, 2 * np.pi, nvertices, endpoint=False) + coords = [ + (cx + radius * np.cos(a), cy + radius * np.sin(a)) for a in angles.tolist() + ] + return spherely.create_polygon(coords + [coords[0]]) + + indexed = np.array( + [polygon(rng.uniform(-1, 1), rng.uniform(-1, 1), 0.3, 30) for _ in range(600)] + ) + queries = np.array( + [polygon(rng.uniform(-1, 1), rng.uniform(-1, 1), 0.2, 8) for _ in range(50)] + ) + return indexed, queries + + +def test_build_budget_does_not_change_matches_on_larger_collection() -> None: + indexed, queries = _overlapping_geographies() + + default = spherely.SpatialIndex(indexed) + default.build() + expected = default.query(queries, predicate="intersects") + + # unlike the four-geography fixture above, here the coarse candidate set is + # strictly larger than the matches, so "superset of the matches" is a real + # constraint rather than an accident of the collection's size + assert _pairs(expected) < _pairs(default.query(queries)) + + # 1 MiB forces the build into several batches (this collection needs about + # 4 MB of temporary space) while 4 GiB builds it in one, as the default + # does here: the candidate sets need not agree, the matches do + for budget in (1024**2, 4 * 1024**3): + tuned = spherely.SpatialIndex(indexed) + tuned.build(tmp_memory_budget=budget) + matches = tuned.query(queries, predicate="intersects") + np.testing.assert_array_equal(matches, expected) + assert _pairs(matches) <= _pairs(tuned.query(queries)) @pytest.mark.parametrize("budget", [0, -1, -(1024**3)]) @@ -185,7 +242,40 @@ def test_build_budget_is_restored(geographies: npt.NDArray[Any]) -> None: plain = spherely.SpatialIndex(geographies) plain.build() assert spherely._s2_tmp_memory_budget() == before - np.testing.assert_array_equal(plain.query(poly), tuned.query(poly)) + np.testing.assert_array_equal( + plain.query(poly, predicate="contains"), + tuned.query(poly, predicate="contains"), + ) + + +@pytest.fixture +def restore_s2_budget() -> Iterator[None]: + # for the tests below, which move the process-wide s2 setting themselves + before = spherely._s2_tmp_memory_budget() + yield + spherely._set_s2_tmp_memory_budget(before) + + +def test_s2_tmp_memory_budget_round_trip(restore_s2_budget: None) -> None: + # the accessors reach the real s2 flag: what is set is what is read back + spherely._set_s2_tmp_memory_budget(4096) + assert spherely._s2_tmp_memory_budget() == 4096 + spherely._set_s2_tmp_memory_budget(8192) + assert spherely._s2_tmp_memory_budget() == 8192 + + +def test_build_budget_restores_previous_value( + geographies: npt.NDArray[Any], restore_s2_budget: None +) -> None: + # test_build_budget_is_restored runs at s2geometry's default, where putting + # the previous value back and resetting to the default are the same thing; + # start from a non-default value so that only the former passes + spherely._set_s2_tmp_memory_budget(4096) + + tree = spherely.SpatialIndex(geographies) + tree.build(tmp_memory_budget=8192) + assert tree.is_built + assert spherely._s2_tmp_memory_budget() == 4096 def test_build_empty_index() -> None: From 3aaad1e1faaba388e7837139309f79f6d8d5ac6e Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 18:50:27 -0700 Subject: [PATCH 08/12] add SpatialIndex.query_nearest --- docs/api_hidden.rst | 1 + src/spatial_index.cpp | 313 ++++++++++++++++++++++++++++++- src/spherely.pyi | 67 +++++++ tests/test_spatial_index.py | 356 ++++++++++++++++++++++++++++++++++++ 4 files changed, 730 insertions(+), 7 deletions(-) diff --git a/docs/api_hidden.rst b/docs/api_hidden.rst index d2d538a..43615fd 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -11,6 +11,7 @@ SpatialIndex.geometries SpatialIndex.is_built SpatialIndex.query + SpatialIndex.query_nearest Projection Projection.lnglat Projection.pseudo_mercator diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index ab92fda..f79b1ea 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -1,19 +1,25 @@ #include #include +#include +#include #include #include +#include #include #include #include #include +#include #include +#include #include #include #include #include #include +#include "constants.hpp" #include "geography.hpp" #include "predicates.hpp" #include "pybind11.hpp" @@ -146,14 +152,77 @@ class SpatialIndex { } } - auto k = input_idx.size(); - py::array_t out({static_cast(2), static_cast(k)}); - auto out_data = out.mutable_unchecked<2>(); + return to_pairs_array(input_idx, tree_idx); + } + + // Scalar nearest query: single Geography -> 1-d array of sorted tree + // indices of the nearest geographies (optionally with their distance). + py::object query_nearest(const Geography& geography, + std::optional max_distance, + bool return_distance, + bool exclusive, + bool all_matches, + double radius) const { + S1ChordAngle distance; + auto results = query_nearest_one( + geography, to_chord_angle(max_distance, radius), exclusive, all_matches, &distance); + auto indices = to_index_array(results); + if (!return_distance) { + return indices; + } + + auto n = results.size(); + py::array_t distances(static_cast(n)); + auto dist_data = distances.mutable_unchecked<1>(); + auto dist = distance.ToAngle().radians() * radius; + for (std::size_t i = 0; i < n; i++) { + dist_data(static_cast(i)) = dist; + } + return py::make_tuple(indices, distances); + } + + // Array nearest query: 1-d array of Geography -> (2, K) array of + // (input index, tree index) pairs (optionally with their distances). + py::object query_nearest(const py::array_t& geographies, + std::optional max_distance, + bool return_distance, + bool exclusive, + bool all_matches, + double radius) const { + if (geographies.ndim() != 1) { + throw py::type_error( + "query geography must be a Geography or a 1-dimensional array of Geography"); + } + + auto max_dist = to_chord_angle(max_distance, radius); + auto n = geographies.size(); + auto* data = static_cast(geographies.request().ptr); + std::vector input_idx; + std::vector tree_idx; + std::vector dists; + for (py::ssize_t i = 0; i < n; i++) { + S1ChordAngle distance; + auto results = query_nearest_one( + *data[i].as_geog_ptr(), max_dist, exclusive, all_matches, &distance); + for (int t : results) { + input_idx.push_back(i); + tree_idx.push_back(static_cast(t)); + dists.push_back(distance.ToAngle().radians() * radius); + } + } + + auto pairs = to_pairs_array(input_idx, tree_idx); + if (!return_distance) { + return pairs; + } + + auto k = dists.size(); + py::array_t distances(static_cast(k)); + auto dist_data = distances.mutable_unchecked<1>(); for (std::size_t j = 0; j < k; j++) { - out_data(0, static_cast(j)) = input_idx[j]; - out_data(1, static_cast(j)) = tree_idx[j]; + dist_data(static_cast(j)) = dists[j]; } - return out; + return py::make_tuple(pairs, distances); } private: @@ -199,6 +268,156 @@ class SpatialIndex { return results; } + // Convert a distance in meters (given the radius) to a chord angle. + static std::optional to_chord_angle(std::optional distance, + double radius) { + // the radius scales both the returned distances and ``max_distance``, + // so it has to be checked even when there is no distance bound + if (!std::isfinite(radius) || radius <= 0) { + throw py::value_error("radius must be a finite value greater than 0"); + } + if (!distance.has_value()) { + return std::nullopt; + } + // NaN would compare false here and silently mean "unbounded" + if (!std::isfinite(*distance) || *distance <= 0) { + throw py::value_error("max_distance must be a finite value greater than 0"); + } + auto bound = S1ChordAngle(S1Angle::Radians(*distance / radius)); + // the meters -> radians -> chord angle conversion here and the chord + // angle -> radians -> meters conversion of the returned distances do + // not round trip exactly (they disagree by a couple of ULPs), so a + // distance reported by ``query_nearest`` would otherwise be rejected + // as a max_distance about a quarter of the time. Widen the bound by a + // few ULPs to keep it inclusive (S1ChordAngle::set_inclusive_max_distance + // only bumps it by one). FromLength2 clamps to a straight angle. + return S1ChordAngle::FromLength2(bound.length2() * + (1 + 8 * std::numeric_limits::epsilon())); + } + + // Return the sorted tree indices of every geography whose distance to the + // target is exactly ``dist`` (``dist`` must be the minimum distance over + // the geographies not skipped, so that no geography has an edge in + // ``(0, dist)``). + std::vector collect_at_distance(S2ClosestEdgeQuery& query, + S2ClosestEdgeQuery::ShapeIndexTarget& target, + S1ChordAngle dist) const { + query.mutable_options()->set_max_results(S2ClosestEdgeQuery::Options::kMaxMaxResults); + query.mutable_options()->set_inclusive_max_distance(dist); + std::unordered_set hits; + for (const auto& result : query.FindClosestEdges(&target)) { + if (result.distance() == dist) { + hits.insert(m_index->value(result.shape_id())); + } + } + std::vector results(hits.begin(), hits.end()); + std::sort(results.begin(), results.end()); + return results; + } + + // ``collect_at_distance``, reduced to the lowest tree index if + // ``all_matches`` is false: which tie is returned is unspecified by the + // API but is picked deterministically (the results are sorted). + std::vector nearest_at_distance(S2ClosestEdgeQuery& query, + S2ClosestEdgeQuery::ShapeIndexTarget& target, + S1ChordAngle dist, + bool all_matches) const { + auto results = collect_at_distance(query, target, dist); + if (!all_matches && results.size() > 1) { + results.resize(1); + } + return results; + } + + // Return the sorted tree indices of the geographies nearest to + // ``query_geog`` (every tie at the minimum distance if ``all_matches``, + // otherwise the lowest of them), storing the minimum distance in + // ``*distance`` (Infinity if there is no result). Geographies farther + // than ``max_distance`` (inclusive) are never returned; if ``exclusive``, + // neither are geographies equal to ``query_geog``. + std::vector query_nearest_one(const Geography& query_geog, + std::optional max_distance, + bool exclusive, + bool all_matches, + S1ChordAngle* distance) const { + *distance = S1ChordAngle::Infinity(); + if (query_geog.is_empty()) { + return {}; + } + + S2ClosestEdgeQuery query(&m_index->ShapeIndex()); + // count the interior of indexed polygons (resp. of the query + // geography) as distance zero, consistent with spherely.distance + query.mutable_options()->set_include_interiors(true); + if (max_distance.has_value()) { + query.mutable_options()->set_inclusive_max_distance(*max_distance); + } + S2ClosestEdgeQuery::ShapeIndexTarget target(&query_geog.geog_index().ShapeIndex()); + target.set_include_interiors(true); + + auto closest = query.FindClosestEdge(&target); + if (closest.is_empty()) { + return {}; + } + auto min_dist = closest.distance(); + + if (!exclusive || min_dist > S1ChordAngle::Zero()) { + // geographies equal to the query would be at distance zero, so + // nothing needs to be excluded here + *distance = min_dist; + return nearest_at_distance(query, target, min_dist, all_matches); + } + + // exclusive with matches at distance zero: split those into + // geographies equal to the query (excluded) and others (returned) + auto zero_hits = collect_at_distance(query, target, S1ChordAngle::Zero()); + auto equals_pred = get_predicate("equals"); + const auto& query_index = query_geog.geog_index(); + auto* data = static_cast(m_geographies.request().ptr); + std::vector non_equal; + for (int t : zero_hits) { + if (!equals_pred(query_index, data[t].as_geog_ptr()->geog_index())) { + non_equal.push_back(t); + } + } + if (!non_equal.empty()) { + *distance = S1ChordAngle::Zero(); + if (!all_matches) { + return {non_equal.front()}; // sorted: the lowest tree index + } + return non_equal; + } + + // every geography at distance zero equals the query: look for the + // nearest one at a distance > 0. Every edge of an equal geography is + // at distance zero (equal geographies cover each other), so the first + // result at a positive distance -- results are sorted by distance -- + // gives the minimum distance over the non-equal geographies. + if (max_distance.has_value()) { + query.mutable_options()->set_inclusive_max_distance(*max_distance); + } else { + query.mutable_options()->set_max_distance(S1ChordAngle::Infinity()); + } + constexpr int kMaxResults = S2ClosestEdgeQuery::Options::kMaxMaxResults; + for (int max_results = 16;;) { + query.mutable_options()->set_max_results(max_results); + auto results = query.FindClosestEdges(&target); + for (const auto& result : results) { + if (result.distance() > S1ChordAngle::Zero()) { + *distance = result.distance(); + return nearest_at_distance(query, target, result.distance(), all_matches); + } + } + if (results.size() < static_cast(max_results) || + max_results == kMaxResults) { + // all edges within max_distance seen, none at distance > 0 + return {}; + } + // saturate instead of overflowing (which would be undefined) + max_results = max_results > kMaxResults / 2 ? kMaxResults : max_results * 2; + } + } + static py::array_t to_index_array(const std::vector& indices) { auto n = indices.size(); py::array_t out(static_cast(n)); @@ -208,6 +427,18 @@ class SpatialIndex { } return out; } + + static py::array_t to_pairs_array(const std::vector& input_idx, + const std::vector& tree_idx) { + auto k = input_idx.size(); + py::array_t out({static_cast(2), static_cast(k)}); + auto out_data = out.mutable_unchecked<2>(); + for (std::size_t j = 0; j < k; j++) { + out_data(0, static_cast(j)) = input_idx[j]; + out_data(1, static_cast(j)) = tree_idx[j]; + } + return out; + } }; void init_spatial_index(py::module& m) { @@ -332,7 +563,75 @@ void init_spatial_index(py::module& m) { py::overload_cast&, std::optional>( &SpatialIndex::query, py::const_), py::arg("geography"), - py::arg("predicate") = py::none()); + py::arg("predicate") = py::none()) + .def( + "query_nearest", + py::overload_cast, bool, bool, bool, double>( + &SpatialIndex::query_nearest, py::const_), + py::arg("geography"), + py::arg("max_distance") = py::none(), + py::arg("return_distance") = false, + py::arg("exclusive") = false, + py::arg("all_matches") = true, + py::arg("radius") = numeric_constants::EARTH_RADIUS_METERS, + R"pbdoc(query_nearest(geography, max_distance=None, return_distance=False, exclusive=False, all_matches=True, radius=spherely.EARTH_RADIUS_METERS) + + Return the integer indices of the geographies in the index that are + nearest to the given geography, similar to shapely's + ``STRtree.query_nearest``. + + The distance to a geography is zero if the query geography intersects + it (including polygon interiors), consistent with + :py:func:`spherely.distance`. Empty geographies (in the index or as + query) are never matched. + + Parameters + ---------- + geography : :py:class:`Geography` or array_like + The geography or geographies to query. + max_distance : float, optional + The maximum distance (in meters, inclusive) within which to query + for the nearest geographies. Must be a finite value greater than 0 + (pass None, the default, for no bound). + return_distance : bool, default False + If True, also return the distance(s) to the nearest geographies. + exclusive : bool, default False + If True, geographies in the index that are equal to the query + geography are not returned. + all_matches : bool, default True + If True, return all geographies tied at the minimum distance. If + False, return only one of them: the one with the lowest index + (shapely leaves the choice arbitrary, spherely makes it + deterministic). + radius : float, optional + Radius of Earth in meters, default 6,371,010. Used to scale the + returned distances and to interpret ``max_distance``. Must be a + finite value greater than 0. + + Returns + ------- + ndarray or tuple of ndarray + If ``geography`` is a scalar, a 1-d array of (sorted) indices into + the index. If ``geography`` is an array, a ``(2, K)`` array where + the first row holds the input geography indices and the second row + the matching index (tree) indices. If ``return_distance`` is True, + a ``(indices, distances)`` tuple where ``distances`` holds the + distance in meters for each match. + + )pbdoc") + .def("query_nearest", + py::overload_cast&, + std::optional, + bool, + bool, + bool, + double>(&SpatialIndex::query_nearest, py::const_), + py::arg("geography"), + py::arg("max_distance") = py::none(), + py::arg("return_distance") = false, + py::arg("exclusive") = false, + py::arg("all_matches") = true, + py::arg("radius") = numeric_constants::EARTH_RADIUS_METERS); // Not public API: exposed so that the test suite can move the process-wide // budget off its default and assert what SpatialIndex.build() puts back. diff --git a/src/spherely.pyi b/src/spherely.pyi index 5a53a2e..58a8245 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -254,6 +254,73 @@ class SpatialIndex: def query( self, geography: Iterable[Geography], predicate: _IndexPredicate | None = None ) -> npt.NDArray[np.intp]: ... + # every argument may also be passed positionally (as in shapely); the + # trailing overloads cover a ``return_distance`` value only known at + # runtime + @overload + def query_nearest( + self, + geography: Geography, + max_distance: float | None = None, + return_distance: Literal[False] = False, + exclusive: bool = False, + all_matches: bool = True, + radius: float = 6371010.0, + ) -> npt.NDArray[np.intp]: ... + @overload + def query_nearest( + self, + geography: Iterable[Geography], + max_distance: float | None = None, + return_distance: Literal[False] = False, + exclusive: bool = False, + all_matches: bool = True, + radius: float = 6371010.0, + ) -> npt.NDArray[np.intp]: ... + @overload + def query_nearest( + self, + geography: Geography, + max_distance: float | None = None, + return_distance: Literal[True] = ..., + exclusive: bool = False, + all_matches: bool = True, + radius: float = 6371010.0, + ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.float64]]: ... + @overload + def query_nearest( + self, + geography: Iterable[Geography], + max_distance: float | None = None, + return_distance: Literal[True] = ..., + exclusive: bool = False, + all_matches: bool = True, + radius: float = 6371010.0, + ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.float64]]: ... + @overload + def query_nearest( + self, + geography: Geography, + max_distance: float | None = None, + return_distance: bool = False, + exclusive: bool = False, + all_matches: bool = True, + radius: float = 6371010.0, + ) -> ( + npt.NDArray[np.intp] | tuple[npt.NDArray[np.intp], npt.NDArray[np.float64]] + ): ... + @overload + def query_nearest( + self, + geography: Iterable[Geography], + max_distance: float | None = None, + return_distance: bool = False, + exclusive: bool = False, + all_matches: bool = True, + radius: float = 6371010.0, + ) -> ( + npt.NDArray[np.intp] | tuple[npt.NDArray[np.intp], npt.NDArray[np.float64]] + ): ... def _s2_tmp_memory_budget() -> int: ... def _set_s2_tmp_memory_budget(bytes: int) -> None: ... diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 7fa64ce..52e6d3d 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -366,3 +366,359 @@ def test_query_invalid_predicate(geographies: npt.NDArray[Any]) -> None: point = spherely.create_point(1, 1) with pytest.raises(ValueError, match="invalid predicate"): tree.query(point, predicate="not_a_predicate") # type: ignore[call-overload] + + +def test_query_nearest_scalar(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + result = tree.query_nearest(spherely.create_point(1.2, 1.2)) + assert result.dtype == np.intp + np.testing.assert_array_equal(result, [1]) + + +def test_query_nearest_ties() -> None: + # two points exactly symmetric about the query point on the equator + tree = spherely.SpatialIndex( + [ + spherely.create_point(-2, 0), + spherely.create_point(2, 0), + spherely.create_point(50, 0), + ] + ) + query = spherely.create_point(0, 0) + np.testing.assert_array_equal(tree.query_nearest(query), [0, 1]) + + single = tree.query_nearest(query, all_matches=False) + np.testing.assert_array_equal(single, [0]) + + +def test_query_nearest_ties_lowest_index() -> None: + # more than two geographies tied at the minimum distance: all_matches=False + # returns the lowest index, not an arbitrary member of the tie + tree = spherely.SpatialIndex( + [ + spherely.create_point(50, 0), # far away + spherely.create_point(2, 0), + spherely.create_point(-2, 0), + spherely.create_point(2, 0), # duplicate of index 1 + ] + ) + query = spherely.create_point(0, 0) + + np.testing.assert_array_equal(tree.query_nearest(query), [1, 2, 3]) + np.testing.assert_array_equal(tree.query_nearest(query, all_matches=False), [1]) + + # array form (and with the distances) + queries = np.array([query, spherely.create_point(49, 0)]) + single, distances = tree.query_nearest( + queries, all_matches=False, return_distance=True + ) + np.testing.assert_array_equal(single, [[0, 1], [1, 0]]) + np.testing.assert_allclose( + distances, + [ + spherely.distance(query, spherely.create_point(2, 0)), + spherely.distance(queries[1], spherely.create_point(50, 0)), + ], + ) + + +def test_query_nearest_exclusive_ties_lowest_index() -> None: + # same, on the exclusive path: the geographies equal to the query are + # dropped and the lowest of the remaining tie is returned + tree = spherely.SpatialIndex( + [ + spherely.create_point(0, 0), # equal to the query + spherely.create_point(2, 0), + spherely.create_point(-2, 0), + spherely.create_point(2, 0), + ] + ) + query = spherely.create_point(0, 0) + np.testing.assert_array_equal(tree.query_nearest(query, exclusive=True), [1, 2, 3]) + np.testing.assert_array_equal( + tree.query_nearest(query, exclusive=True, all_matches=False), [1] + ) + + # ... and on the "everything at distance zero is equal" fallback path + tree = spherely.SpatialIndex( + [ + spherely.create_point(0, 0), + spherely.create_point(0, 0), + spherely.create_point(2, 0), + spherely.create_point(-2, 0), + ] + ) + np.testing.assert_array_equal(tree.query_nearest(query, exclusive=True), [2, 3]) + np.testing.assert_array_equal( + tree.query_nearest(query, exclusive=True, all_matches=False), [2] + ) + + +def test_query_nearest_interior_distance_zero() -> None: + # a point inside an indexed polygon is at distance zero from it, + # consistent with spherely.distance + poly = spherely.create_polygon([(0, 0), (10, 0), (10, 10), (0, 10), (0, 0)]) + tree = spherely.SpatialIndex([poly, spherely.create_point(5.1, 5.1)]) + inside = spherely.create_point(5, 5) + assert spherely.distance(inside, poly) == 0 + np.testing.assert_array_equal(tree.query_nearest(inside), [0]) + + # and symmetrically: an indexed point inside a query polygon + tree = spherely.SpatialIndex( + [spherely.create_point(5, 5), spherely.create_point(20, 20)] + ) + query_poly = spherely.create_polygon([(0, 0), (10, 0), (10, 10), (0, 10), (0, 0)]) + np.testing.assert_array_equal(tree.query_nearest(query_poly), [0]) + + +def test_query_nearest_empty_geographies() -> None: + tree = spherely.SpatialIndex( + [ + spherely.create_polygon(None), # empty, never returned + spherely.create_point(10, 10), + ] + ) + np.testing.assert_array_equal(tree.query_nearest(spherely.create_point(0, 0)), [1]) + + # empty query geography -> no result + result = tree.query_nearest(spherely.create_polygon(None)) + assert result.shape == (0,) + + # empty tree -> no result + tree = spherely.SpatialIndex(np.array([], dtype=object)) + result = tree.query_nearest(spherely.create_point(0, 0)) + assert result.shape == (0,) + + +def test_query_nearest_return_distance(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + query = spherely.create_point(1.2, 1.2) + + indices, distances = tree.query_nearest(query, return_distance=True) + np.testing.assert_array_equal(indices, [1]) + assert distances.dtype == np.float64 + np.testing.assert_allclose(distances, spherely.distance(query, geographies[1])) + + # the arguments may also be passed positionally, as in shapely + positional = tree.query_nearest(query, None, True) + np.testing.assert_array_equal(positional[0], indices) + np.testing.assert_array_equal(positional[1], distances) + + # a return_distance only known at runtime is typed as the union of both + flag = bool(len(indices)) + runtime = tree.query_nearest(query, return_distance=flag) + assert isinstance(runtime, tuple) + + # radius scales the returned distances + _, unscaled = tree.query_nearest(query, return_distance=True, radius=1) + np.testing.assert_allclose( + unscaled * spherely.EARTH_RADIUS_METERS, distances, rtol=1e-15 + ) + + +def test_query_nearest_array(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + queries = np.array( + [ + spherely.create_point(1.2, 1.2), + spherely.create_point(49, 49), + ] + ) + + result = tree.query_nearest(queries) + assert result.dtype == np.intp + # (input_index, tree_index) pairs + np.testing.assert_array_equal(result, [[0, 1], [1, 3]]) + + pairs, distances = tree.query_nearest(queries, return_distance=True) + np.testing.assert_array_equal(pairs, result) + expected = [ + spherely.distance(queries[0], geographies[1]), + spherely.distance(queries[1], geographies[3]), + ] + np.testing.assert_allclose(distances, expected) + + +def test_query_nearest_array_ties_and_empty() -> None: + tree = spherely.SpatialIndex( + [ + spherely.create_point(-2, 0), + spherely.create_point(2, 0), + spherely.create_point(50, 0), + ] + ) + queries = np.array( + [ + spherely.create_point(0, 0), # tied between 0 and 1 + spherely.create_polygon(None), # empty: contributes no pairs + spherely.create_point(49, 0), + ] + ) + result = tree.query_nearest(queries) + np.testing.assert_array_equal(result, [[0, 0, 2], [0, 1, 2]]) + + single = tree.query_nearest(queries, all_matches=False) + # ties are broken deterministically on the lowest tree index + np.testing.assert_array_equal(single, [[0, 2], [0, 2]]) + + +def test_query_nearest_max_distance(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + query = spherely.create_point(1.2, 1.2) + dist = spherely.distance(query, geographies[1]) + + np.testing.assert_array_equal( + tree.query_nearest(query, max_distance=dist * 1.01), [1] + ) + # inclusive bound + np.testing.assert_array_equal(tree.query_nearest(query, max_distance=dist), [1]) + + # no geography within max_distance -> empty result (shapely semantics) + result = tree.query_nearest(query, max_distance=1.0) + assert result.shape == (0,) + indices, distances = tree.query_nearest( + query, max_distance=1.0, return_distance=True + ) + assert indices.shape == (0,) + assert distances.shape == (0,) + + # array form + queries = np.array([query, spherely.create_point(49, 49)]) + result = tree.query_nearest(queries, max_distance=dist * 1.01) + np.testing.assert_array_equal(result, [[0], [1]]) + + # array form with no match at all: (2, 0) pairs and 0 distances + pairs, distances = tree.query_nearest( + queries, max_distance=1.0, return_distance=True + ) + assert pairs.shape == (2, 0) + assert pairs.dtype == np.intp + assert distances.shape == (0,) + assert distances.dtype == np.float64 + + +def test_query_nearest_max_distance_inclusive_boundary() -> None: + # the distance returned by query_nearest must itself be accepted as an + # (inclusive) max_distance: a fixture or two would not catch the ULP-level + # round trip error between meters and chord angles, so sweep many pairs + rng = np.random.default_rng(0) + lon = rng.uniform(-180, 180, 50) + lat = rng.uniform(-90, 90, 50) + tree = spherely.SpatialIndex(np.asarray(spherely.points(lon, lat))) + + qlon = rng.uniform(-180, 180, 200) + qlat = rng.uniform(-90, 90, 200) + for query in np.asarray(spherely.points(qlon, qlat)): + expected, distances = tree.query_nearest(query, return_distance=True) + result = tree.query_nearest(query, max_distance=distances[0]) + np.testing.assert_array_equal(result, expected) + + +@pytest.mark.parametrize("max_distance", [0, -1.0, np.nan, np.inf, -np.inf]) +def test_query_nearest_invalid_max_distance( + geographies: npt.NDArray[Any], max_distance: float +) -> None: + # NaN in particular used to be silently accepted as "no bound" + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + match = "max_distance must be a finite value greater than 0" + with pytest.raises(ValueError, match=match): + tree.query_nearest(point, max_distance=max_distance) + with pytest.raises(ValueError, match=match): + tree.query_nearest(np.array([point]), max_distance=max_distance) + + +@pytest.mark.parametrize("radius", [0, -1.0, np.nan, np.inf]) +def test_query_nearest_invalid_radius( + geographies: npt.NDArray[Any], radius: float +) -> None: + # the radius filters here (it scales max_distance): a zero radius would + # make every distance zero, a negative one would flip the chord angle + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + match = "radius must be a finite value greater than 0" + with pytest.raises(ValueError, match=match): + tree.query_nearest(point, radius=radius) + with pytest.raises(ValueError, match=match): + tree.query_nearest(point, max_distance=1000.0, radius=radius) + with pytest.raises(ValueError, match=match): + tree.query_nearest(np.array([point]), radius=radius) + + +def test_query_nearest_exclusive() -> None: + tree = spherely.SpatialIndex( + [ + spherely.create_point(0, 0), + spherely.create_point(-2, 0), + spherely.create_point(2, 0), + ] + ) + query = spherely.create_point(0, 0) + np.testing.assert_array_equal(tree.query_nearest(query), [0]) + result = tree.query_nearest(query, exclusive=True) + # the equal geography is skipped, the two symmetric neighbors are tied + np.testing.assert_array_equal(result, [1, 2]) + + # a non-equal geography at distance zero is still returned + query = spherely.create_point(1, 1) + poly = spherely.create_polygon([(0, 0), (2, 0), (2, 2), (0, 2), (0, 0)]) + tree = spherely.SpatialIndex([spherely.create_point(1, 1), poly]) + result, distances = tree.query_nearest(query, exclusive=True, return_distance=True) + np.testing.assert_array_equal(result, [1]) + np.testing.assert_array_equal(distances, [0.0]) + + +def test_query_nearest_exclusive_all_equal() -> None: + # all geographies at distance zero are equal to the query + tree = spherely.SpatialIndex( + [spherely.create_point(1, 1), spherely.create_point(1, 1)] + ) + query = spherely.create_point(1, 1) + result = tree.query_nearest(query, exclusive=True) + assert result.shape == (0,) + + # with a farther non-equal geography, that one is returned + far = spherely.create_point(3, 3) + tree = spherely.SpatialIndex( + [spherely.create_point(1, 1), spherely.create_point(1, 1), far] + ) + indices, distances = tree.query_nearest(query, exclusive=True, return_distance=True) + np.testing.assert_array_equal(indices, [2]) + np.testing.assert_allclose(distances, spherely.distance(query, far)) + + # unless it is beyond max_distance + result = tree.query_nearest(query, exclusive=True, max_distance=1.0) + assert result.shape == (0,) + + +def test_query_nearest_exclusive_all_equal_beyond_batch() -> None: + # more geographies equal to the query than the initial max_results batch + # of the fallback path (16), so the query has to be re-run with a larger + # batch until the first result at a distance > 0 shows up + query = spherely.create_point(1, 1) + equal = [spherely.create_point(1, 1) for _ in range(40)] + far = spherely.create_point(3, 3) + tree = spherely.SpatialIndex(equal + [far]) + + indices, distances = tree.query_nearest(query, exclusive=True, return_distance=True) + np.testing.assert_array_equal(indices, [len(equal)]) + np.testing.assert_allclose(distances, spherely.distance(query, far)) + + # ... and with nothing else in the index, the loop still terminates + tree = spherely.SpatialIndex(equal) + assert tree.query_nearest(query, exclusive=True).shape == (0,) + + +def test_query_nearest_against_distance_oracle() -> None: + rng = np.random.default_rng(42) + lon = rng.uniform(-180, 180, 20) + lat = rng.uniform(-90, 90, 20) + geographies = np.asarray(spherely.points(lon, lat)) + tree = spherely.SpatialIndex(geographies) + + for qlon, qlat in [(0, 0), (179.5, 0.5), (-179.5, -0.5), (0, 89.9), (0, -89.9)]: + query = spherely.create_point(qlon, qlat) + result = tree.query_nearest(query) + distances = np.asarray(spherely.distance(query, geographies)) + expected = np.flatnonzero(distances == distances.min()) + np.testing.assert_array_equal(result, expected) From fd47d1506af015a2caf0122ae4f7db21909a710b Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 18:52:53 -0700 Subject: [PATCH 09/12] add SpatialIndex.encode --- docs/api_hidden.rst | 1 + src/spatial_index.cpp | 156 ++++++++++++++++++++++++++++++++++++ src/spherely.pyi | 1 + tests/test_spatial_index.py | 29 +++++++ 4 files changed, 187 insertions(+) diff --git a/docs/api_hidden.rst b/docs/api_hidden.rst index 43615fd..9143301 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -8,6 +8,7 @@ Geography SpatialIndex SpatialIndex.build + SpatialIndex.encode SpatialIndex.geometries SpatialIndex.is_built SpatialIndex.query diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index f79b1ea..c1b80d5 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include #include @@ -16,6 +18,7 @@ #include #include #include +#include #include #include @@ -121,6 +124,65 @@ class SpatialIndex { return m_geographies; } + // Serialize the index to bytes: a small header, the shape id -> tree + // index mapping, then either the Geography objects themselves (a fixed + // width offset/length table followed by one block per tree entry) or a + // compact S2 shape stream, and finally the index structure. + // + // The shape data is stored once either way: with the geographies included + // the blocks hold it and no separate shape stream is written. + // + // Everything past the header is s2geometry's / s2geography's own + // encoding, so the format is tied to the library versions spherely is + // built against (only the header is versioned here). + py::bytes encode(bool include_geographies) const { + // gather the wrapped geographies under the GIL; their S2 data is + // owned by the Python objects this index keeps alive + std::vector geogs; + if (include_geographies) { + auto* data = static_cast(m_geographies.request().ptr); + geogs.reserve(static_cast(m_geographies.size())); + for (py::ssize_t i = 0; i < m_geographies.size(); i++) { + geogs.push_back(data[i].as_geog_ptr()); + } + } + + Encoder encoder; + + const auto& index = m_index->ShapeIndex(); + auto num_shapes = index.num_shape_ids(); + auto num_geographies = static_cast(m_geographies.size()); + + { + // encoding is pure C++ work on objects this index keeps alive and + // can take a while for a large index: let other threads run + py::gil_scoped_release release; + + // widened to size_t: the header plus one varint per shape value + // and one for each of the two counts (int would overflow past + // ~2e8 shapes) + encoder.Ensure(static_cast(kHeaderBytes) + + static_cast(Encoder::kVarintMax64) * + (static_cast(num_shapes) + 2)); + encoder.put32(kEncodingMagic); + encoder.put8(kEncodingVersion); + encoder.put8(include_geographies ? kFlagHasGeographies : 0); + encoder.put_varint64(num_geographies); + encoder.put_varint64(static_cast(num_shapes)); + for (int i = 0; i < num_shapes; i++) { + encoder.put_varint64(static_cast(m_index->value(i))); + } + + if (include_geographies) { + encode_geography_blocks(encoder, geogs); + } else { + s2shapeutil::CompactEncodeTaggedShapes(index, &encoder); + } + index.Encode(&encoder); + } + return py::bytes(encoder.base(), encoder.length()); + } + // Scalar query: single Geography -> 1-d array of sorted tree indices. py::array_t query(const Geography& geography, std::optional predicate) const { @@ -226,10 +288,63 @@ class SpatialIndex { } private: + // encoding header: magic (the bytes "SPIX", which put32 writes least + // significant byte first), format version, flags + static constexpr std::uint32_t kEncodingMagic = 0x58495053; + // format version of the header below; only this exact value is accepted + // when decoding, so any future format change is a bump here plus a + // decoder branch + static constexpr std::uint8_t kEncodingVersion = 1; + // flags: the blob carries a per-geography block section + static constexpr std::uint8_t kFlagHasGeographies = 1; + static constexpr int kHeaderBytes = + sizeof(kEncodingMagic) + sizeof(kEncodingVersion) + sizeof(std::uint8_t); + // one fixed-width (offset, length) pair per geography, so that block j is + // addressable without scanning the section + static constexpr std::size_t kBlockTableEntryBytes = 2 * sizeof(std::uint64_t); + std::unique_ptr m_index; // also keeps the Geography objects alive (the index borrows their S2Shapes) py::array_t m_geographies; + // Append the per-geography section to ``encoder``: a fixed-width + // (offset, length) table (offsets relative to the first block, blocks + // contiguous in tree order) followed by the blocks themselves. Each block + // mirrors the Geography pickle tuple: the geography type (int8), the + // empty flag (uint8) and the geography serialized with s2geography's + // ``EncodeTagged``, so decoding goes through exactly the same code path + // as unpickling. + static void encode_geography_blocks(Encoder& encoder, + const std::vector& geogs) { + std::vector blocks; + blocks.reserve(geogs.size()); + std::size_t blocks_bytes = 0; + for (const Geography* geog : geogs) { + Encoder block; + block.Ensure(2); + using IntType = std::underlying_type_t; + block.put8(static_cast(static_cast(geog->geog_type()))); + block.put8(geog->is_empty() ? 1 : 0); + // default options, as in Geography::encode (the pickle support): + // exact vertices (FAST hint), no covering, no lazy decoding + s2geog::EncodeOptions encode_opts; + geog->geog().EncodeTagged(&block, encode_opts); + blocks_bytes += block.length(); + blocks.push_back(std::move(block)); + } + + encoder.Ensure(kBlockTableEntryBytes * blocks.size() + blocks_bytes); + std::uint64_t offset = 0; + for (const auto& block : blocks) { + encoder.put64(offset); + encoder.put64(static_cast(block.length())); + offset += static_cast(block.length()); + } + for (const auto& block : blocks) { + encoder.putn(block.base(), block.length()); + } + } + static std::optional make_predicate(const std::optional& name) { if (!name.has_value()) { return std::nullopt; @@ -529,6 +644,47 @@ void init_spatial_index(py::module& m) { .def_property_readonly("geometries", &SpatialIndex::geometries, "The array of geographies in the index (in input order).") + .def("encode", + &SpatialIndex::encode, + py::arg("include_geographies") = true, + R"pbdoc(encode(include_geographies=True) + + Serialize the index to bytes. + + By default the Geography objects themselves are serialized along with + the index structure. The shape data is stored only once (inside the + serialized geographies). With ``include_geographies=False`` an + index-only blob is written instead. It is the smaller of the two, by + an amount that depends on the geometry mix: in measurements the + default blob ran from about 1.1x the index-only one (polygon-heavy + indexes) to about 2x (small point-only ones). + + .. warning:: + The encoded bytes are not a portable interchange format: past a + small spherely header they are s2geometry's (and s2geography's) + own encoding, which is tied to the library versions spherely was + built against. The per-geography blocks in particular round-trip + through ``s2geography::Geography::EncodeTagged`` / + ``DecodeTagged``, which s2geography labels EXPERIMENTAL. A blob + written by one build is not guaranteed to be readable by a build + linked against a different s2geometry (loading it raises + ``ValueError`` rather than returning wrong results). Encode and + decode with the same spherely build, and treat the bytes as a + cache, not as an archival format. + + Parameters + ---------- + include_geographies : bool, default True + If True, serialize the Geography objects along with the index so + that the decoded index supports every operation. If False, write + the smaller index-only blob. + + Returns + ------- + bytes + The encoded index. + + )pbdoc") .def("query", py::overload_cast>(&SpatialIndex::query, py::const_), diff --git a/src/spherely.pyi b/src/spherely.pyi index 58a8245..47012b9 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -246,6 +246,7 @@ class SpatialIndex: def is_built(self) -> bool: ... @property def geometries(self) -> T_NDArray_Geography: ... + def encode(self, include_geographies: bool = True) -> bytes: ... @overload def query( self, geography: Geography, predicate: _IndexPredicate | None = None diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 52e6d3d..cd23465 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -1,4 +1,5 @@ from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from typing import Any import numpy as np @@ -722,3 +723,31 @@ def test_query_nearest_against_distance_oracle() -> None: distances = np.asarray(spherely.distance(query, geographies)) expected = np.flatnonzero(distances == distances.min()) np.testing.assert_array_equal(result, expected) + + +def test_encode(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + encoded = tree.encode() + assert isinstance(encoded, bytes) + assert len(encoded) > 0 + # magic, format version and flags (geographies included by default) + assert encoded[:6] == b"SPIX\x01\x01" + # encoding is deterministic + assert tree.encode() == encoded + + thin = tree.encode(include_geographies=False) + assert thin[:6] == b"SPIX\x01\x00" + # the thin blob drops the per-geography blocks (and their offset table) + assert len(thin) < len(encoded) + assert tree.encode(include_geographies=False) == thin + + +def test_encode_concurrent() -> None: + # encode() releases the GIL around the (pure C++) encoding work + geoms = [spherely.create_point(x, x % 80) for x in range(200)] + tree = spherely.SpatialIndex(geoms) + expected = tree.encode() + + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(lambda _: tree.encode(), range(8))) + assert all(result == expected for result in results) From 96313f019ac59bb26d5c186951e2e7b7b516b8f3 Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 18:53:39 -0700 Subject: [PATCH 10/12] add SpatialIndex.from_encoded with a lazy decoded index --- docs/api_hidden.rst | 1 + src/spatial_index.cpp | 650 ++++++++++++++++++++++++-- src/spherely.pyi | 6 + tests/test_spatial_index.py | 897 +++++++++++++++++++++++++++++++++--- 4 files changed, 1466 insertions(+), 88 deletions(-) diff --git a/docs/api_hidden.rst b/docs/api_hidden.rst index 9143301..ac37458 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -9,6 +9,7 @@ SpatialIndex SpatialIndex.build SpatialIndex.encode + SpatialIndex.from_encoded SpatialIndex.geometries SpatialIndex.is_built SpatialIndex.query diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index c1b80d5..e7a8fc0 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -6,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -20,6 +22,7 @@ #include #include #include +#include #include #include "constants.hpp" @@ -88,12 +91,26 @@ class SpatialIndex { } } + // in encoded mode the EncodedS2ShapeIndex points into ``m_encoded``, + // whose buffer may live inside the string object itself (small string + // optimization): moving or copying the index would leave it dangling + SpatialIndex(SpatialIndex&&) = delete; + SpatialIndex& operator=(SpatialIndex&&) = delete; + std::size_t size() const { + if (is_encoded()) { + return static_cast(m_num_geographies); + } return static_cast(m_geographies.size()); } // True once the queued updates have been applied (no build work pending). + // An index loaded from encoded bytes never has any: its cells are decoded + // on demand, so there is nothing to build. bool is_built() const { + if (is_encoded()) { + return true; + } return m_index->ShapeIndex().is_fresh(); } @@ -106,22 +123,55 @@ class SpatialIndex { " (s2geometry's default budget is 104857600 bytes)"); } + // an encoded index has no queued updates and no mutable shape index to + // build: the budget check above is all that applies to it + if (is_encoded()) { + return; + } + // nothing pending: no build to tune, so leave the process-wide flag alone if (is_built()) { return; } - const auto& shape_index = m_index->ShapeIndex(); + const auto& index = m_index->ShapeIndex(); if (tmp_memory_budget.has_value()) { TmpMemoryBudgetGuard guard(*tmp_memory_budget); - shape_index.ForceBuild(); + index.ForceBuild(); } else { - shape_index.ForceBuild(); + index.ForceBuild(); } } + // Number of tree entries whose Geography has been reconstructed from the + // encoded blob so far (always 0 for a built index). Exposed as the + // ``_decoded_count`` property so that tests can pin the laziness + // contract deterministically. + std::size_t decoded_count() const { + return m_decoded_count; + } + py::array geometries() const { - return m_geographies; + check_has_geographies("geometries"); + if (!is_encoded()) { + return m_geographies; + } + // reconstruct every remaining entry (insertion order is the tree + // order) and cache the resulting array + if (!m_decoded_geometries) { + auto result = py::array_t(m_num_geographies); + py::buffer_info buf = result.request(); + py::object* data = static_cast(buf.ptr); + for (py::ssize_t i = 0; i < m_num_geographies; i++) { + decoded_geog(i); + data[i] = m_decoded_slots[static_cast(i)]; + } + m_decoded_geometries = std::move(result); + } + // m_decoded_geometries is stored as a py::object, so borrow a new + // reference to it as a py::array rather than downcasting the + // reference itself + return py::reinterpret_borrow(m_decoded_geometries); } // Serialize the index to bytes: a small header, the shape id -> tree @@ -129,13 +179,33 @@ class SpatialIndex { // width offset/length table followed by one block per tree entry) or a // compact S2 shape stream, and finally the index structure. // - // The shape data is stored once either way: with the geographies included - // the blocks hold it and no separate shape stream is written. + // The shape data is stored once either way. Without the geographies the + // stream written by CompactEncodeTaggedShapes backs the decoded index + // directly (via LazyDecodeShapeFactory, which decodes each shape on + // first access only). With the geographies there is no separate shape + // stream: the decoded index pulls shape k out of the reconstruction of + // the geography that owns it (see materialize_shape), which defers all + // shape data to first touch but makes that touch reconstruct the whole + // owning geography. // // Everything past the header is s2geometry's / s2geography's own // encoding, so the format is tied to the library versions spherely is // built against (only the header is versioned here). py::bytes encode(bool include_geographies) const { + if (is_encoded()) { + // the original blob is all this index has: hand it back when it + // matches the request rather than fabricating the other flavor + if (include_geographies != m_has_geographies) { + throw py::value_error( + m_has_geographies + ? "encode(include_geographies=False) is not supported for an index " + "loaded from encoded bytes that include the geographies" + : "encode(include_geographies=True) is not supported for an index " + "loaded from bytes encoded with include_geographies=False"); + } + return py::bytes(m_encoded); + } + // gather the wrapped geographies under the GIL; their S2 data is // owned by the Python objects this index keeps alive std::vector geogs; @@ -174,6 +244,8 @@ class SpatialIndex { } if (include_geographies) { + // single copy of the shape data: the geography blocks stand + // in for the shape stream (see materialize_shape) encode_geography_blocks(encoder, geogs); } else { s2shapeutil::CompactEncodeTaggedShapes(index, &encoder); @@ -183,6 +255,120 @@ class SpatialIndex { return py::bytes(encoder.base(), encoder.length()); } + // Load an index serialized with ``encode()``, wrapping the bytes in a + // lazy EncodedS2ShapeIndex: opening is near-instant and index cells / + // shapes are only decoded when queries touch them. + static std::unique_ptr from_encoded(const py::buffer& encoded) { + // any contiguous bytes-like object (bytes, bytearray, memoryview...); + // the bytes are copied, so the buffer is not kept alive afterwards + py::buffer_info info = encoded.request(); + if (info.ndim != 1 || info.itemsize != 1 || info.strides[0] != 1) { + throw py::type_error("encoded must be a contiguous bytes-like object"); + } + + auto self = std::unique_ptr(new SpatialIndex()); + self->m_encoded.assign(static_cast(info.ptr), + static_cast(info.size)); + + Decoder decoder(self->m_encoded.data(), self->m_encoded.size()); + auto fail = []() { + throw py::value_error("invalid encoded SpatialIndex"); + }; + if (decoder.avail() < kHeaderBytes || decoder.get32() != kEncodingMagic) { + fail(); + } + if (decoder.get8() != kEncodingVersion) { + throw py::value_error("unsupported encoded SpatialIndex version"); + } + auto flags = decoder.get8(); + if ((flags & ~kFlagHasGeographies) != 0) { + throw py::value_error("unsupported encoded SpatialIndex flags"); + } + self->m_has_geographies = (flags & kFlagHasGeographies) != 0; + + std::uint64_t num_geographies; + std::uint64_t num_shapes; + if (!decoder.get_varint64(&num_geographies) || !decoder.get_varint64(&num_shapes)) { + fail(); + } + // both counts come straight from the (untrusted) input: reject values + // that cannot describe this buffer before allocating anything. Tree + // indices and shape ids are both stored as int (as in the built index), + // and every encoded value takes at least one byte. + constexpr auto kMaxInt = static_cast(std::numeric_limits::max()); + if (num_geographies > kMaxInt || num_shapes > kMaxInt || num_shapes > decoder.avail()) { + fail(); + } + // a fat blob spends at least one block table entry plus one (minimal) + // block per geography: a count that cannot fit in what is left of the + // buffer is rejected here, before the per-geography vectors are sized + // on it. Without this a 12 byte blob claiming 2^31-1 geographies asks + // for tens of gigabytes before ``init_geography_blocks`` ever looks at + // ``avail()``. + if (self->m_has_geographies && + num_geographies > decoder.avail() / (kBlockTableEntryBytes + kBlockMinBytes)) { + fail(); + } + self->m_num_geographies = static_cast(num_geographies); + self->m_values.reserve(num_shapes); + for (std::uint64_t i = 0; i < num_shapes; i++) { + std::uint64_t value; + if (!decoder.get_varint64(&value) || value >= num_geographies) { + fail(); + } + self->m_values.push_back(static_cast(value)); + } + + if (self->m_has_geographies) { + self->m_shape_counts.assign(static_cast(num_geographies), 0); + self->m_first_shape.assign(static_cast(num_geographies), -1); + int previous = -1; + for (std::size_t shape_id = 0; shape_id < self->m_values.size(); shape_id++) { + int value = self->m_values[shape_id]; + // materialize_shape addresses shape k as (tree entry, k - + // first shape of the entry), which requires each entry's + // shapes to form one contiguous ascending run, as written by + // encode: a non-monotonic value table cannot come from it + if (value < previous) { + fail(); + } + if (value > previous) { + self->m_first_shape[static_cast(value)] = + static_cast(shape_id); + } + self->m_shape_counts[static_cast(value)]++; + previous = value; + } + self->m_decoded_slots.resize(static_cast(num_geographies)); + if (!self->init_geography_blocks(decoder, num_geographies)) { + fail(); + } + } + + self->m_encoded_index = std::make_unique(); + // a fat blob carries no shape stream of its own: the index pulls + // shapes out of the geography blocks on demand. A thin blob's shapes + // are decoded lazily from its compact shape stream. + bool init_ok = + self->m_has_geographies + ? self->m_encoded_index->Init(&decoder, GeographyBlockShapeFactory(self.get())) + : self->m_encoded_index->Init(&decoder, + s2shapeutil::LazyDecodeShapeFactory(&decoder)); + if (!init_ok) { + fail(); + } + // the value table is indexed by shape id: a table that does not cover + // every decoded shape would be read out of bounds by queries + if (self->m_values.size() != + static_cast(self->m_encoded_index->num_shape_ids())) { + fail(); + } + if (self->m_has_geographies && !self->check_index_shape_ids()) { + fail(); + } + return self; + } + // Scalar query: single Geography -> 1-d array of sorted tree indices. py::array_t query(const Geography& geography, std::optional predicate) const { @@ -256,6 +442,11 @@ class SpatialIndex { "query geography must be a Geography or a 1-dimensional array of Geography"); } + // checked up front (as the predicate is in the array ``query``) so + // that an empty input array raises too rather than silently returning + if (exclusive) { + check_has_geographies("query_nearest with exclusive=True"); + } auto max_dist = to_chord_angle(max_distance, radius); auto n = geographies.size(); auto* data = static_cast(geographies.request().ptr); @@ -302,10 +493,48 @@ class SpatialIndex { // one fixed-width (offset, length) pair per geography, so that block j is // addressable without scanning the section static constexpr std::size_t kBlockTableEntryBytes = 2 * sizeof(std::uint64_t); - - std::unique_ptr m_index; - // also keeps the Geography objects alive (the index borrows their S2Shapes) + // a block holds at least the geography type and empty bytes plus the + // 4-byte tag written by s2geography::Geography::EncodeTagged + static constexpr std::uint64_t kBlockMinBytes = 2 + 4; + + // Both indexes below borrow from the members declared above them and are + // therefore declared last, so that they are destroyed first: members are + // destroyed in reverse declaration order, and an S2ShapeIndex touches its + // shapes while being torn down. + + // built mode (constructed from an array of geographies) + // m_geographies keeps the Geography objects alive (the index borrows + // their S2Shapes) py::array_t m_geographies; + std::unique_ptr m_index; + + // encoded mode (loaded with from_encoded); the encoded bytes are kept + // alive here as EncodedS2ShapeIndex works directly on them + std::string m_encoded; + std::vector m_values; + py::ssize_t m_num_geographies = 0; + // whether the blob carries per-geography blocks, and where each block + // lives inside ``m_encoded`` (absolute offset, length) + bool m_has_geographies = false; + std::vector> m_geog_blocks; + // number of shapes per tree entry according to the id map (used to + // cross-check reconstructed geographies against the encoded index), and + // the first shape id of each entry's run (-1 for entries with no shape) + std::vector m_shape_counts; + std::vector m_first_shape; + // lazily reconstructed Geography objects, one slot per tree entry: a slot + // is filled on first touch and cached for the lifetime of the index + // (mutable: caching only, doesn't affect the observable query results) + mutable std::vector m_decoded_slots; + mutable std::size_t m_decoded_count = 0; + // cached ``geometries`` array of the reconstructed objects + mutable py::object m_decoded_geometries; + // holds S2Shapes materialized from m_decoded_slots (fat) or decoded out + // of m_encoded (thin): declared after both + std::unique_ptr m_encoded_index; + + // only used internally by ``from_encoded`` (encoded mode) + SpatialIndex() = default; // Append the per-geography section to ``encoder``: a fixed-width // (offset, length) table (offsets relative to the first block, blocks @@ -345,13 +574,312 @@ class SpatialIndex { } } - static std::optional make_predicate(const std::optional& name) { + // Parse and validate the per-geography section (see + // ``encode_geography_blocks``) at the decoder's current position, filling + // ``m_geog_blocks``, and skip the decoder past it. Returns false on any + // inconsistency. + bool init_geography_blocks(Decoder& decoder, std::uint64_t num_geographies) { + // ``from_encoded`` bounds num_geographies by what is left of the + // buffer, so this product cannot overflow + auto table_bytes = static_cast(kBlockTableEntryBytes * num_geographies); + if (decoder.avail() < table_bytes) { + return false; + } + auto blocks_begin = decoder.pos() + table_bytes; + // what is left for the blocks themselves once the table is consumed + auto blocks_avail = static_cast(decoder.avail() - table_bytes); + m_geog_blocks.reserve(num_geographies); + std::uint64_t expected_offset = 0; + for (std::uint64_t j = 0; j < num_geographies; j++) { + std::uint64_t offset = decoder.get64(); + std::uint64_t length = decoder.get64(); + // blocks are contiguous and in order: reject gaps, overlaps and + // runt blocks + if (offset != expected_offset || length < kBlockMinBytes || length > blocks_avail) { + return false; + } + m_geog_blocks.emplace_back(blocks_begin + static_cast(offset), + static_cast(length)); + expected_offset = offset + length; + // bounded on every entry rather than once after the loop: the + // running sum is free to wrap past 2^64 before a single check at + // the end ever looks at it, and each block would then have been + // recorded at an in-range offset with an out-of-range length + if (expected_offset > blocks_avail) { + return false; + } + } + decoder.skip(static_cast(expected_offset)); + return true; + } + + // Check that every shape id clipped by the decoded index cells is covered + // by the value table (fat blobs only). + // + // A thin blob carries its own shape stream, so ``num_shape_ids()`` is an + // independent count and the table size check in ``from_encoded`` is a real + // cross-check. A fat blob's index is sized from + // ``GeographyBlockShapeFactory::size()``, i.e. from the value table + // itself, so that check compares the table with itself and nothing + // reconciles the shape ids stored in the index cells with it. A cell + // clipping a shape id past the end of the table makes + // ``EncodedS2ShapeIndex::shape()`` index ``shapes_`` out of bounds + // (unchecked), which segfaults inside the distance queries. + // + // This is one linear pass over the index cells at load, of the same order + // as the value table and block table passes ``from_encoded`` already + // makes. It decodes cells but never shapes, so no geography is + // reconstructed here. + bool check_index_shape_ids() const { + auto num_shape_ids = static_cast(m_values.size()); + int max_shape_id = -1; + for (S2ShapeIndex::Iterator iter(&shape_index(), S2ShapeIndex::BEGIN); !iter.done(); + iter.Next()) { + const S2ShapeIndexCell& cell = iter.cell(); + for (int k = 0; k < cell.num_clipped(); k++) { + int shape_id = cell.clipped(k).shape_id(); + if (shape_id < 0 || shape_id >= num_shape_ids) { + return false; + } + max_shape_id = std::max(max_shape_id, shape_id); + } + } + // exact, not just in range: ``encode`` writes one value per shape of + // the index it serializes and every such shape is clipped into at + // least one cell, so a table with entries past the last shape the + // cells reference did not come from it (and would make + // ``num_shape_ids()`` -- which is this table's size for a fat blob -- + // over-report the shapes the index actually has) + return max_shape_id + 1 == num_shape_ids; + } + + bool is_encoded() const { + return m_encoded_index != nullptr; + } + + // The operations that need the Geography objects work on a built index + // and on one decoded from a blob written with the geographies included; + // they raise on an index decoded from an index-only blob. + void check_has_geographies(const std::string& what) const { + if (is_encoded() && !m_has_geographies) { + throw py::value_error(what + + " is not supported for an index loaded from bytes encoded " + "with include_geographies=False (encode with " + "include_geographies=True to keep the Geography objects)"); + } + } + + // Whether a block's geography type byte can describe what its tagged + // encoding actually holds. ``Geography::decode`` takes the type byte at + // face value -- as it does when unpickling, where it comes from a trusted + // ``encode`` -- so nothing else reconciles the two: a relabelled block + // would decode into a geography whose ``get_type_id`` disagrees with the + // geography itself (a Polygon byte on a point block, say). + // + // A singular type and its MULTI* form share one kind, so this narrows the + // byte to a pair rather than to a single value: swapping the two labels + // stays undetectable here (it would take counting the parts, which is + // ``extract_geog_properties``' job and is not reachable from a decoded + // Geography). That is a wrong ``get_type_id`` label on a corrupt blob, + // with no effect on the geometry or on any query -- the same latitude + // unpickling a hand-made state tuple already has. + static bool type_matches_kind(GeographyType geog_type, s2geog::GeographyKind kind) { + switch (kind) { + case s2geog::GeographyKind::POINT: + case s2geog::GeographyKind::CELL_CENTER: + return geog_type == GeographyType::Point || geog_type == GeographyType::MultiPoint; + case s2geog::GeographyKind::POLYLINE: + return geog_type == GeographyType::LineString || + geog_type == GeographyType::MultiLineString; + case s2geog::GeographyKind::POLYGON: + return geog_type == GeographyType::Polygon || + geog_type == GeographyType::MultiPolygon; + case s2geog::GeographyKind::GEOGRAPHY_COLLECTION: + return geog_type == GeographyType::GeometryCollection; + default: + // the kinds rejected just above; GeographyType::None belongs + // with those, so it never matches either + return false; + } + } + + // Reconstruct the Geography of tree entry ``j`` from its block, or return + // the cached reconstruction. Decoding goes through Geography::decode (the + // pickle path), so the objects are indistinguishable from unpickled ones. + // + // The caller must hold the GIL. The slot cache is a plain vector and the + // counter a plain size_t, so the cache-hit path below is only safe + // because every entry point reaches it with the GIL held (no query path + // releases it) -- the ``gil_scoped_acquire`` guard sits on the decode + // path and does not, and cannot, cover the slot read above it. It is + // there so that the Python API calls of a decode still work should a + // caller ever release the GIL around the C++ query work; making that + // caller safe would take more than moving the guard up. + Geography* decoded_geog(py::ssize_t i) const { + auto j = static_cast(i); + py::object& slot = m_decoded_slots[j]; + if (!slot) { + py::gil_scoped_acquire acquire; + const auto& [offset, length] = m_geog_blocks[j]; + const char* block = m_encoded.data() + offset; + auto geog_type = static_cast(static_cast(block[0])); + auto empty = static_cast(block[1]); + if (geog_type < -1 || geog_type > 6 || empty > 1) { + throw py::value_error("invalid encoded SpatialIndex geography block"); + } + auto encoded_geog = py::bytes(block + 2, length - 2); + std::unique_ptr geog; + try { + auto tuple = py::make_tuple(geog_type, empty != 0, encoded_geog); + geog = std::make_unique(Geography::decode(tuple)); + } catch (const py::error_already_set&) { + throw; + } catch (const std::exception&) { + throw py::value_error("invalid encoded SpatialIndex geography block"); + } + // ``encode`` only ever writes the kinds a spherely Geography can + // wrap: the remaining ones are not just unexpected here, they are + // unsound. UNINITIALIZED and SHAPE_INDEX have no WKT + // representation (accessors raise "Unsupported Geography + // subclass"), and an ENCODED_SHAPE_INDEX decodes lazily over the + // buffer it was handed -- the temporary string inside + // Geography::decode -- so it would outlive its own data. + switch (geog->geog().kind()) { + case s2geog::GeographyKind::UNINITIALIZED: + case s2geog::GeographyKind::SHAPE_INDEX: + case s2geog::GeographyKind::ENCODED_SHAPE_INDEX: + throw py::value_error("invalid encoded SpatialIndex geography block"); + default: + break; + } + if (!type_matches_kind(geog->geog_type(), geog->geog().kind())) { + throw py::value_error("invalid encoded SpatialIndex geography block"); + } + // the id map and the block must agree on the shape count, or the + // reconstructed geography is not the one the index was built on + if (geog->num_shapes() != m_shape_counts[j]) { + throw py::value_error("invalid encoded SpatialIndex geography block"); + } + auto decoded = PyObjectGeography::from_geog(std::move(geog)); + // everything between the ``!slot`` test and here runs + // Python-visible work (py::bytes, Geography::decode, from_geog), + // any of which can trigger a collection whose __del__ re-enters + // this index and fills the same slot. The index may already have + // cached an S2Shape borrowing that inner object, so overwriting + // the slot would leave the cached shape pointing at a dropped + // geography: keep the entry that got there first. + if (!slot) { + slot = std::move(decoded); + m_decoded_count++; + } + } + return static_cast(slot).as_geog_ptr(); + } + + // Geography of tree entry ``i``: from the built array, or lazily + // reconstructed from its encoded block. Predicate refinement and equality + // run on these objects through the same code paths either way. + Geography* candidate_geog(py::ssize_t i) const { + if (is_encoded()) { + return decoded_geog(i); + } + auto* data = static_cast(m_geographies.request().ptr); + return data[i].as_geog_ptr(); + } + + // Produce shape ``shape_id`` of a fat blob's index by reconstructing (or + // reusing) the geography that owns it. The returned shape borrows the + // geography's data, which stays alive in the slot cache for the lifetime + // of this index. Called from GeographyBlockShapeFactory inside query + // execution, always with the GIL held (see decoded_geog). + std::unique_ptr materialize_shape(int shape_id) const { + auto j = m_values.at(static_cast(shape_id)); + Geography* geog = decoded_geog(j); + // decoded_geog checked the block against the id map's shape count, + // so the run arithmetic below cannot leave the geography's range + int sub = shape_id - m_first_shape[static_cast(j)]; + return geog->geog().Shape(sub); + } + + // ShapeFactory backing the EncodedS2ShapeIndex of a fat blob: there is + // no second copy of the shape data in the blob, shapes are pulled from + // the per-geography blocks instead. The owner pointer is stable + // (SpatialIndex is immovable) and outlives the encoded index holding the + // cloned factory, which is declared last among the members precisely so + // that it is torn down before everything it borrows. + class GeographyBlockShapeFactory : public S2ShapeIndex::ShapeFactory { + public: + explicit GeographyBlockShapeFactory(const SpatialIndex* owner) : m_owner(owner) {} + + int size() const override { + return static_cast(m_owner->m_values.size()); + } + + std::unique_ptr operator[](int shape_id) const override { + return m_owner->materialize_shape(shape_id); + } + + std::unique_ptr Clone() const override { + return std::make_unique(*this); + } + + private: + const SpatialIndex* m_owner; + }; + + const S2ShapeIndex& shape_index() const { + if (is_encoded()) { + return *m_encoded_index; + } + return m_index->ShapeIndex(); + } + + // tree index of the geography that owns the given shape + int tree_value(int shape_id) const { + if (is_encoded()) { + // ``from_encoded`` checks the table covers every shape id; ``at`` + // keeps a mismatch a Python exception rather than a bad read + return m_values.at(static_cast(shape_id)); + } + return m_index->value(shape_id); + } + + std::optional make_predicate(const std::optional& name) const { if (!name.has_value()) { return std::nullopt; } + check_has_geographies("query with a predicate"); return get_predicate(*name); } + // Collect the tree indices of the shapes in the index cells overlapping + // ``cell_id`` (mirrors s2geography::GeographyIndex::Iterator::Query, + // which is bound to the mutable index of the built mode). + void query_cell(S2ShapeIndex::Iterator& iter, + const S2CellId& cell_id, + std::unordered_set& hits) const { + auto add_cell_shapes = [&]() { + const S2ShapeIndexCell& index_cell = iter.cell(); + for (int k = 0; k < index_cell.num_clipped(); k++) { + hits.insert(tree_value(index_cell.clipped(k).shape_id())); + } + }; + + S2CellRelation relation = iter.Locate(cell_id); + if (relation == S2CellRelation::INDEXED) { + // the index has this cell (or an ancestor of it) + add_cell_shapes(); + } else if (relation == S2CellRelation::SUBDIVIDED) { + // the index has child cells of ``cell_id`` (the iterator is + // positioned at the first one): visit them all + while (!iter.done() && cell_id.contains(iter.id())) { + add_cell_shapes(); + iter.Next(); + } + } + // else: DISJOINT (do nothing) + } + // Return the sorted tree indices whose cells overlap the query geography, // optionally refined by ``pred`` (predicate(query, candidate)). std::vector query_one(const Geography& query_geog, const PredicateFunc* pred) const { @@ -362,17 +890,18 @@ class SpatialIndex { std::vector covering; coverer.GetCovering(*region, &covering); - s2geog::GeographyIndex::Iterator iter(m_index.get()); - iter.Query(covering, &hits); + S2ShapeIndex::Iterator iter(&shape_index()); + for (const S2CellId& cell_id : covering) { + query_cell(iter, cell_id, hits); + } std::vector results; if (pred == nullptr) { results.assign(hits.begin(), hits.end()); } else { const auto& query_index = query_geog.geog_index(); - auto* data = static_cast(m_geographies.request().ptr); for (int candidate : hits) { - auto* cand_geog = data[candidate].as_geog_ptr(); + auto* cand_geog = candidate_geog(candidate); if ((*pred)(query_index, cand_geog->geog_index())) { results.push_back(candidate); } @@ -422,7 +951,7 @@ class SpatialIndex { std::unordered_set hits; for (const auto& result : query.FindClosestEdges(&target)) { if (result.distance() == dist) { - hits.insert(m_index->value(result.shape_id())); + hits.insert(tree_value(result.shape_id())); } } std::vector results(hits.begin(), hits.end()); @@ -456,11 +985,14 @@ class SpatialIndex { bool all_matches, S1ChordAngle* distance) const { *distance = S1ChordAngle::Infinity(); + if (exclusive) { + check_has_geographies("query_nearest with exclusive=True"); + } if (query_geog.is_empty()) { return {}; } - S2ClosestEdgeQuery query(&m_index->ShapeIndex()); + S2ClosestEdgeQuery query(&shape_index()); // count the interior of indexed polygons (resp. of the query // geography) as distance zero, consistent with spherely.distance query.mutable_options()->set_include_interiors(true); @@ -488,10 +1020,9 @@ class SpatialIndex { auto zero_hits = collect_at_distance(query, target, S1ChordAngle::Zero()); auto equals_pred = get_predicate("equals"); const auto& query_index = query_geog.geog_index(); - auto* data = static_cast(m_geographies.request().ptr); std::vector non_equal; for (int t : zero_hits) { - if (!equals_pred(query_index, data[t].as_geog_ptr()->geog_index())) { + if (!equals_pred(query_index, candidate_geog(t)->geog_index())) { non_equal.push_back(t); } } @@ -590,7 +1121,9 @@ void init_spatial_index(py::module& m) { The index is built lazily, so without this call the first query pays the whole build cost. Calling it is optional, never changes the result of a query made with a ``predicate``, and does nothing once the index - is built (see :py:attr:`SpatialIndex.is_built`). + is built (see :py:attr:`SpatialIndex.is_built`) -- including on an + index loaded with :py:meth:`SpatialIndex.from_encoded`, which has no + build to force. Parameters ---------- @@ -638,12 +1171,18 @@ void init_spatial_index(py::module& m) { ``False`` between the constructor and the first :py:meth:`SpatialIndex.build` or query, ``True`` afterwards. An index - over an empty collection is ``True`` from the start. + over an empty collection is ``True`` from the start, and so is one + loaded with :py:meth:`SpatialIndex.from_encoded` (which decodes on + demand and has nothing to build). )pbdoc") .def_property_readonly("geometries", &SpatialIndex::geometries, "The array of geographies in the index (in input order).") + .def_property_readonly("_decoded_count", + &SpatialIndex::decoded_count, + "Number of geographies reconstructed so far by an index loaded " + "from encoded bytes (testing/introspection only).") .def("encode", &SpatialIndex::encode, py::arg("include_geographies") = true, @@ -652,12 +1191,17 @@ void init_spatial_index(py::module& m) { Serialize the index to bytes. By default the Geography objects themselves are serialized along with - the index structure. The shape data is stored only once (inside the - serialized geographies). With ``include_geographies=False`` an - index-only blob is written instead. It is the smaller of the two, by - an amount that depends on the geometry mix: in measurements the - default blob ran from about 1.1x the index-only one (polygon-heavy - indexes) to about 2x (small point-only ones). + the index structure, so that :py:meth:`SpatialIndex.from_encoded` + restores a fully functional index. The shape data is stored only once + (inside the serialized geographies, which the decoded index reads its + shapes from). With ``include_geographies=False`` an index-only blob + is written instead. It is the smaller of the two, by an amount that + depends on the geometry mix: in measurements the default blob ran + from about 1.1x the index-only one (polygon-heavy indexes) to about + 2x (small point-only ones). The loaded index does not support the + operations that need the Geography objects, and its distance queries + decode individual shapes rather than whole geographies (see + :py:meth:`SpatialIndex.from_encoded`). .. warning:: The encoded bytes are not a portable interchange format: past a @@ -684,6 +1228,60 @@ void init_spatial_index(py::module& m) { bytes The encoded index. + )pbdoc") + .def_static("from_encoded", + &SpatialIndex::from_encoded, + py::arg("encoded"), + R"pbdoc(from_encoded(encoded) + + Load an index serialized with :py:meth:`SpatialIndex.encode`. + + The returned index is a lazy view on the encoded bytes: loading it is + near-instant regardless of the index size, and the index cells are + only decoded on demand by queries. For bytes written with + ``include_geographies=True`` (the default) the Geography of a tree + entry is reconstructed the first time an operation touches that + entry: by predicate refinement of a candidate, by a distance query + (``query_nearest``) examining one of the entry's shapes, or by the + ``geometries`` property (which reconstructs all entries). Candidate + queries (``query`` without a predicate) only walk the index cells + and reconstruct nothing. Reconstruction of an entry decodes that + geography completely (through the same code path as unpickling a + Geography) and the result is cached, so each entry pays this cost at + most once and untouched entries never pay it. + + ``query_nearest`` is the costly one: it reconstructs every entry + whose shapes the distance search examines, and it reconstructs them + whole even though it only needs one shape of each. That is a few + tens of geographies per call -- in measurements on point indexes of + 100 to 20,000 entries, between about 10 and 80, so near-constant in + the size of the index rather than proportional to it -- against the + single shapes an index-only blob decodes for the same query. + + The decoded index supports every operation with the same results -- + including the same integer indices -- as the index it was encoded + from. For bytes written with ``include_geographies=False`` the + Geography objects are not available, so the ``geometries`` property, + ``query`` with a predicate and ``query_nearest`` with + ``exclusive=True`` raise ``ValueError``. + + Only bytes written by :py:meth:`SpatialIndex.encode` from the same + spherely build are supported -- see the warning there about the + format's coupling to s2geometry. Corrupt or foreign bytes raise + ``ValueError``. + + Parameters + ---------- + encoded : bytes-like + An index serialized with :py:meth:`SpatialIndex.encode`, as any + contiguous bytes-like object (``bytes``, ``bytearray``, + ``memoryview``, ...). The bytes are copied. + + Returns + ------- + :py:class:`SpatialIndex` + The loaded index. + )pbdoc") .def("query", py::overload_cast>(&SpatialIndex::query, diff --git a/src/spherely.pyi b/src/spherely.pyi index 47012b9..1a78c6f 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -246,7 +246,13 @@ class SpatialIndex: def is_built(self) -> bool: ... @property def geometries(self) -> T_NDArray_Geography: ... + # testing / introspection only: pins the laziness contract of an index + # loaded from encoded bytes, not part of the public API + @property + def _decoded_count(self) -> int: ... def encode(self, include_geographies: bool = True) -> bytes: ... + @staticmethod + def from_encoded(encoded: bytes | bytearray | memoryview) -> SpatialIndex: ... @overload def query( self, geography: Geography, predicate: _IndexPredicate | None = None diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index cd23465..4681f5f 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -1,6 +1,7 @@ +import struct from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor -from typing import Any +from typing import Any, Callable import numpy as np import numpy.typing as npt @@ -8,6 +9,8 @@ import spherely +MakeTree = Callable[[Any], spherely.SpatialIndex] + def _pairs(result: npt.NDArray[np.intp]) -> set[tuple[Any, ...]]: # the (2, N) result of an array query as a set of (input, tree) pairs @@ -27,13 +30,47 @@ def geographies() -> npt.NDArray[Any]: ) -def test_spatial_index_len(geographies: npt.NDArray[Any]) -> None: +def _make_tree(param: str, geographies: Any) -> spherely.SpatialIndex: + # build a SpatialIndex directly or round-trip it through + # encode / from_encoded, with ("fat") or without ("thin") the Geography + # objects: queries must behave identically on all of them tree = spherely.SpatialIndex(geographies) + if param == "fat": + tree = spherely.SpatialIndex.from_encoded(tree.encode()) + elif param == "thin": + tree = spherely.SpatialIndex.from_encoded( + tree.encode(include_geographies=False) + ) + return tree + + +@pytest.fixture(params=["built", "fat", "thin"]) +def make_tree(request: pytest.FixtureRequest) -> MakeTree: + def make(geographies: Any) -> spherely.SpatialIndex: + return _make_tree(request.param, geographies) + + return make + + +@pytest.fixture(params=["built", "fat"]) +def make_full_tree(request: pytest.FixtureRequest) -> MakeTree: + # trees that support the full API: predicates, exclusive nearest queries + # and the geometries property (a thin decoded index raises on those) + def make(geographies: Any) -> spherely.SpatialIndex: + return _make_tree(request.param, geographies) + + return make + + +def test_spatial_index_len(geographies: npt.NDArray[Any], make_tree: MakeTree) -> None: + tree = make_tree(geographies) assert len(tree) == 4 -def test_spatial_index_geometries(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_spatial_index_geometries( + geographies: npt.NDArray[Any], make_full_tree: MakeTree +) -> None: + tree = make_full_tree(geographies) geoms = tree.geometries assert isinstance(geoms, np.ndarray) assert geoms.shape == (4,) @@ -288,8 +325,8 @@ def test_build_empty_index() -> None: np.testing.assert_array_equal(tree.query(spherely.create_point(0, 0)), []) -def test_query_scalar(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_scalar(geographies: npt.NDArray[Any], make_tree: MakeTree) -> None: + tree = make_tree(geographies) poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) result = tree.query(poly) @@ -298,7 +335,7 @@ def test_query_scalar(geographies: npt.NDArray[Any]) -> None: np.testing.assert_array_equal(result, [0, 1, 2]) -def test_query_predicate_refines() -> None: +def test_query_predicate_refines(make_full_tree: MakeTree) -> None: # point 1 is outside the triangle (beyond the hypotenuse) but close enough # that it falls within the triangle's coarse cell covering; point 2 is far # away and is not a candidate at all @@ -309,7 +346,7 @@ def test_query_predicate_refines() -> None: spherely.create_point(50, 50), ] ) - tree = spherely.SpatialIndex(geographies) + tree = make_full_tree(geographies) triangle = spherely.create_polygon([(0, 0), (5, 0), (0, 5), (0, 0)]) coarse = tree.query(triangle) @@ -321,15 +358,17 @@ def test_query_predicate_refines() -> None: np.testing.assert_array_equal(refined, [0]) -def test_query_predicate_intersects(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_predicate_intersects( + geographies: npt.NDArray[Any], make_full_tree: MakeTree +) -> None: + tree = make_full_tree(geographies) point = spherely.create_point(1, 1) result = tree.query(point, predicate="intersects") np.testing.assert_array_equal(result, [1]) -def test_query_array(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_array(geographies: npt.NDArray[Any], make_full_tree: MakeTree) -> None: + tree = make_full_tree(geographies) queries = np.array( [ spherely.create_point(1, 1), @@ -341,44 +380,50 @@ def test_query_array(geographies: npt.NDArray[Any]) -> None: np.testing.assert_array_equal(result, [[0, 1], [1, 3]]) -def test_query_empty_geography_never_returned() -> None: +def test_query_empty_geography_never_returned(make_tree: MakeTree) -> None: geoms = np.array( [ spherely.create_point(0, 0), spherely.create_polygon(None), # empty ] ) - tree = spherely.SpatialIndex(geoms) + tree = make_tree(geoms) assert len(tree) == 2 poly = spherely.create_polygon([(-1, -1), (1, -1), (1, 1), (-1, 1), (-1, -1)]) result = tree.query(poly) np.testing.assert_array_equal(result, [0]) -def test_query_disjoint_rejected(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_disjoint_rejected( + geographies: npt.NDArray[Any], make_full_tree: MakeTree +) -> None: + tree = make_full_tree(geographies) point = spherely.create_point(1, 1) with pytest.raises(ValueError, match="disjoint"): tree.query(point, predicate="disjoint") # type: ignore[call-overload] -def test_query_invalid_predicate(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_invalid_predicate( + geographies: npt.NDArray[Any], make_full_tree: MakeTree +) -> None: + tree = make_full_tree(geographies) point = spherely.create_point(1, 1) with pytest.raises(ValueError, match="invalid predicate"): tree.query(point, predicate="not_a_predicate") # type: ignore[call-overload] -def test_query_nearest_scalar(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_nearest_scalar( + geographies: npt.NDArray[Any], make_tree: MakeTree +) -> None: + tree = make_tree(geographies) result = tree.query_nearest(spherely.create_point(1.2, 1.2)) assert result.dtype == np.intp np.testing.assert_array_equal(result, [1]) -def test_query_nearest_ties() -> None: +def test_query_nearest_ties(make_tree: MakeTree) -> None: # two points exactly symmetric about the query point on the equator - tree = spherely.SpatialIndex( + tree = make_tree( [ spherely.create_point(-2, 0), spherely.create_point(2, 0), @@ -392,10 +437,10 @@ def test_query_nearest_ties() -> None: np.testing.assert_array_equal(single, [0]) -def test_query_nearest_ties_lowest_index() -> None: +def test_query_nearest_ties_lowest_index(make_tree: MakeTree) -> None: # more than two geographies tied at the minimum distance: all_matches=False # returns the lowest index, not an arbitrary member of the tie - tree = spherely.SpatialIndex( + tree = make_tree( [ spherely.create_point(50, 0), # far away spherely.create_point(2, 0), @@ -423,10 +468,10 @@ def test_query_nearest_ties_lowest_index() -> None: ) -def test_query_nearest_exclusive_ties_lowest_index() -> None: +def test_query_nearest_exclusive_ties_lowest_index(make_full_tree: MakeTree) -> None: # same, on the exclusive path: the geographies equal to the query are # dropped and the lowest of the remaining tie is returned - tree = spherely.SpatialIndex( + tree = make_full_tree( [ spherely.create_point(0, 0), # equal to the query spherely.create_point(2, 0), @@ -441,7 +486,7 @@ def test_query_nearest_exclusive_ties_lowest_index() -> None: ) # ... and on the "everything at distance zero is equal" fallback path - tree = spherely.SpatialIndex( + tree = make_full_tree( [ spherely.create_point(0, 0), spherely.create_point(0, 0), @@ -455,25 +500,23 @@ def test_query_nearest_exclusive_ties_lowest_index() -> None: ) -def test_query_nearest_interior_distance_zero() -> None: +def test_query_nearest_interior_distance_zero(make_tree: MakeTree) -> None: # a point inside an indexed polygon is at distance zero from it, # consistent with spherely.distance poly = spherely.create_polygon([(0, 0), (10, 0), (10, 10), (0, 10), (0, 0)]) - tree = spherely.SpatialIndex([poly, spherely.create_point(5.1, 5.1)]) + tree = make_tree([poly, spherely.create_point(5.1, 5.1)]) inside = spherely.create_point(5, 5) assert spherely.distance(inside, poly) == 0 np.testing.assert_array_equal(tree.query_nearest(inside), [0]) # and symmetrically: an indexed point inside a query polygon - tree = spherely.SpatialIndex( - [spherely.create_point(5, 5), spherely.create_point(20, 20)] - ) + tree = make_tree([spherely.create_point(5, 5), spherely.create_point(20, 20)]) query_poly = spherely.create_polygon([(0, 0), (10, 0), (10, 10), (0, 10), (0, 0)]) np.testing.assert_array_equal(tree.query_nearest(query_poly), [0]) -def test_query_nearest_empty_geographies() -> None: - tree = spherely.SpatialIndex( +def test_query_nearest_empty_geographies(make_tree: MakeTree) -> None: + tree = make_tree( [ spherely.create_polygon(None), # empty, never returned spherely.create_point(10, 10), @@ -486,13 +529,15 @@ def test_query_nearest_empty_geographies() -> None: assert result.shape == (0,) # empty tree -> no result - tree = spherely.SpatialIndex(np.array([], dtype=object)) + tree = make_tree(np.array([], dtype=object)) result = tree.query_nearest(spherely.create_point(0, 0)) assert result.shape == (0,) -def test_query_nearest_return_distance(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_nearest_return_distance( + geographies: npt.NDArray[Any], make_tree: MakeTree +) -> None: + tree = make_tree(geographies) query = spherely.create_point(1.2, 1.2) indices, distances = tree.query_nearest(query, return_distance=True) @@ -517,8 +562,10 @@ def test_query_nearest_return_distance(geographies: npt.NDArray[Any]) -> None: ) -def test_query_nearest_array(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_nearest_array( + geographies: npt.NDArray[Any], make_tree: MakeTree +) -> None: + tree = make_tree(geographies) queries = np.array( [ spherely.create_point(1.2, 1.2), @@ -540,8 +587,8 @@ def test_query_nearest_array(geographies: npt.NDArray[Any]) -> None: np.testing.assert_allclose(distances, expected) -def test_query_nearest_array_ties_and_empty() -> None: - tree = spherely.SpatialIndex( +def test_query_nearest_array_ties_and_empty(make_tree: MakeTree) -> None: + tree = make_tree( [ spherely.create_point(-2, 0), spherely.create_point(2, 0), @@ -563,8 +610,10 @@ def test_query_nearest_array_ties_and_empty() -> None: np.testing.assert_array_equal(single, [[0, 2], [0, 2]]) -def test_query_nearest_max_distance(geographies: npt.NDArray[Any]) -> None: - tree = spherely.SpatialIndex(geographies) +def test_query_nearest_max_distance( + geographies: npt.NDArray[Any], make_tree: MakeTree +) -> None: + tree = make_tree(geographies) query = spherely.create_point(1.2, 1.2) dist = spherely.distance(query, geographies[1]) @@ -598,14 +647,14 @@ def test_query_nearest_max_distance(geographies: npt.NDArray[Any]) -> None: assert distances.dtype == np.float64 -def test_query_nearest_max_distance_inclusive_boundary() -> None: +def test_query_nearest_max_distance_inclusive_boundary(make_tree: MakeTree) -> None: # the distance returned by query_nearest must itself be accepted as an # (inclusive) max_distance: a fixture or two would not catch the ULP-level # round trip error between meters and chord angles, so sweep many pairs rng = np.random.default_rng(0) lon = rng.uniform(-180, 180, 50) lat = rng.uniform(-90, 90, 50) - tree = spherely.SpatialIndex(np.asarray(spherely.points(lon, lat))) + tree = make_tree(np.asarray(spherely.points(lon, lat))) qlon = rng.uniform(-180, 180, 200) qlat = rng.uniform(-90, 90, 200) @@ -617,10 +666,10 @@ def test_query_nearest_max_distance_inclusive_boundary() -> None: @pytest.mark.parametrize("max_distance", [0, -1.0, np.nan, np.inf, -np.inf]) def test_query_nearest_invalid_max_distance( - geographies: npt.NDArray[Any], max_distance: float + geographies: npt.NDArray[Any], make_tree: MakeTree, max_distance: float ) -> None: # NaN in particular used to be silently accepted as "no bound" - tree = spherely.SpatialIndex(geographies) + tree = make_tree(geographies) point = spherely.create_point(1, 1) match = "max_distance must be a finite value greater than 0" with pytest.raises(ValueError, match=match): @@ -631,11 +680,11 @@ def test_query_nearest_invalid_max_distance( @pytest.mark.parametrize("radius", [0, -1.0, np.nan, np.inf]) def test_query_nearest_invalid_radius( - geographies: npt.NDArray[Any], radius: float + geographies: npt.NDArray[Any], make_tree: MakeTree, radius: float ) -> None: # the radius filters here (it scales max_distance): a zero radius would # make every distance zero, a negative one would flip the chord angle - tree = spherely.SpatialIndex(geographies) + tree = make_tree(geographies) point = spherely.create_point(1, 1) match = "radius must be a finite value greater than 0" with pytest.raises(ValueError, match=match): @@ -646,8 +695,8 @@ def test_query_nearest_invalid_radius( tree.query_nearest(np.array([point]), radius=radius) -def test_query_nearest_exclusive() -> None: - tree = spherely.SpatialIndex( +def test_query_nearest_exclusive(make_full_tree: MakeTree) -> None: + tree = make_full_tree( [ spherely.create_point(0, 0), spherely.create_point(-2, 0), @@ -663,24 +712,22 @@ def test_query_nearest_exclusive() -> None: # a non-equal geography at distance zero is still returned query = spherely.create_point(1, 1) poly = spherely.create_polygon([(0, 0), (2, 0), (2, 2), (0, 2), (0, 0)]) - tree = spherely.SpatialIndex([spherely.create_point(1, 1), poly]) + tree = make_full_tree([spherely.create_point(1, 1), poly]) result, distances = tree.query_nearest(query, exclusive=True, return_distance=True) np.testing.assert_array_equal(result, [1]) np.testing.assert_array_equal(distances, [0.0]) -def test_query_nearest_exclusive_all_equal() -> None: +def test_query_nearest_exclusive_all_equal(make_full_tree: MakeTree) -> None: # all geographies at distance zero are equal to the query - tree = spherely.SpatialIndex( - [spherely.create_point(1, 1), spherely.create_point(1, 1)] - ) + tree = make_full_tree([spherely.create_point(1, 1), spherely.create_point(1, 1)]) query = spherely.create_point(1, 1) result = tree.query_nearest(query, exclusive=True) assert result.shape == (0,) # with a farther non-equal geography, that one is returned far = spherely.create_point(3, 3) - tree = spherely.SpatialIndex( + tree = make_full_tree( [spherely.create_point(1, 1), spherely.create_point(1, 1), far] ) indices, distances = tree.query_nearest(query, exclusive=True, return_distance=True) @@ -692,30 +739,32 @@ def test_query_nearest_exclusive_all_equal() -> None: assert result.shape == (0,) -def test_query_nearest_exclusive_all_equal_beyond_batch() -> None: +def test_query_nearest_exclusive_all_equal_beyond_batch( + make_full_tree: MakeTree, +) -> None: # more geographies equal to the query than the initial max_results batch # of the fallback path (16), so the query has to be re-run with a larger # batch until the first result at a distance > 0 shows up query = spherely.create_point(1, 1) equal = [spherely.create_point(1, 1) for _ in range(40)] far = spherely.create_point(3, 3) - tree = spherely.SpatialIndex(equal + [far]) + tree = make_full_tree(equal + [far]) indices, distances = tree.query_nearest(query, exclusive=True, return_distance=True) np.testing.assert_array_equal(indices, [len(equal)]) np.testing.assert_allclose(distances, spherely.distance(query, far)) # ... and with nothing else in the index, the loop still terminates - tree = spherely.SpatialIndex(equal) + tree = make_full_tree(equal) assert tree.query_nearest(query, exclusive=True).shape == (0,) -def test_query_nearest_against_distance_oracle() -> None: +def test_query_nearest_against_distance_oracle(make_tree: MakeTree) -> None: rng = np.random.default_rng(42) lon = rng.uniform(-180, 180, 20) lat = rng.uniform(-90, 90, 20) geographies = np.asarray(spherely.points(lon, lat)) - tree = spherely.SpatialIndex(geographies) + tree = make_tree(geographies) for qlon, qlat in [(0, 0), (179.5, 0.5), (-179.5, -0.5), (0, 89.9), (0, -89.9)]: query = spherely.create_point(qlon, qlat) @@ -751,3 +800,727 @@ def test_encode_concurrent() -> None: with ThreadPoolExecutor(max_workers=4) as pool: results = list(pool.map(lambda _: tree.encode(), range(8))) assert all(result == expected for result in results) + + +def test_from_encoded_preserves_indices() -> None: + # multi-shape and empty geographies make the shape id <-> tree index + # mapping non-trivial: it must survive the encode / from_encoded round + # trip so that both indexes return the same integer indices + geoms = [ + spherely.create_point(0, 0), + spherely.create_collection( # two shapes + [ + spherely.create_point(30, 30), + spherely.create_linestring([(31, 31), (32, 32)]), + ] + ), + spherely.create_polygon(None), # empty: no shape + spherely.create_point(-60, -60), + ] + tree = spherely.SpatialIndex(geoms) + decoded = spherely.SpatialIndex.from_encoded(tree.encode()) + + assert len(decoded) == len(tree) == 4 + for query in geoms[:2] + geoms[3:] + [spherely.create_point(30.5, 30.5)]: + np.testing.assert_array_equal( + decoded.query_nearest(query), tree.query_nearest(query) + ) + np.testing.assert_array_equal(decoded.query(query), tree.query(query)) + + # a decoded index has no shapes of its own to re-encode: encode() hands + # back the blob it was loaded from, so comparing those bytes with the + # original would be circular. Check instead that the blob survives another + # round trip and still answers queries identically. + redecoded = spherely.SpatialIndex.from_encoded(decoded.encode()) + assert len(redecoded) == len(tree) + for query in geoms[:2] + geoms[3:] + [spherely.create_point(30.5, 30.5)]: + np.testing.assert_array_equal( + redecoded.query_nearest(query), tree.query_nearest(query) + ) + np.testing.assert_array_equal(redecoded.query(query), tree.query(query)) + + +def test_from_encoded_bytes_like() -> None: + tree = spherely.SpatialIndex( + [spherely.create_point(0, 0), spherely.create_point(1, 1)] + ) + encoded = tree.encode() + query = spherely.create_point(0.9, 0.9) + for buffer in [encoded, bytearray(encoded), memoryview(encoded)]: + decoded = spherely.SpatialIndex.from_encoded(buffer) + np.testing.assert_array_equal(decoded.query_nearest(query), [1]) + + with pytest.raises(TypeError): + spherely.SpatialIndex.from_encoded(memoryview(encoded)[::2]) + with pytest.raises(TypeError): + spherely.SpatialIndex.from_encoded("not bytes") # type: ignore[arg-type] + + +def test_from_encoded_empty_tree() -> None: + tree = spherely.SpatialIndex(np.array([], dtype=object)) + decoded = spherely.SpatialIndex.from_encoded(tree.encode()) + assert len(decoded) == 0 + result = decoded.query_nearest(spherely.create_point(0, 0)) + assert result.shape == (0,) + assert decoded.geometries.shape == (0,) + + +@pytest.mark.parametrize("include_geographies", [True, False]) +def test_from_encoded_is_built( + geographies: npt.NDArray[Any], include_geographies: bool +) -> None: + # a decoded index decodes its cells on demand rather than building them, + # so there is never anything left for build() to force + tree = spherely.SpatialIndex(geographies) + decoded = spherely.SpatialIndex.from_encoded( + tree.encode(include_geographies=include_geographies) + ) + assert decoded.is_built + decoded.build() + decoded.build(tmp_memory_budget=1024) + assert decoded.is_built + assert 1 in decoded.query(spherely.create_point(1, 1)) + + # the budget is validated before the "nothing to build" early return + with pytest.raises(ValueError, match="positive"): + decoded.build(tmp_memory_budget=0) + + +def test_thin_encoded_unsupported_operations(geographies: npt.NDArray[Any]) -> None: + # an index loaded from a blob without the geographies cannot support the + # operations that need the Geography objects; the error points at the + # include_geographies=True remedy + tree = spherely.SpatialIndex.from_encoded( + spherely.SpatialIndex(geographies).encode(include_geographies=False) + ) + point = spherely.create_point(1, 1) + match = "not supported for an index loaded.*include_geographies=True" + with pytest.raises(ValueError, match=match): + tree.geometries # noqa: B018 + with pytest.raises(ValueError, match=match): + tree.query(point, predicate="intersects") + with pytest.raises(ValueError, match=match): + tree.query_nearest(point, exclusive=True) + # including for an (empty) array input, which never reaches the scalar path + with pytest.raises(ValueError, match=match): + tree.query_nearest(np.array([], dtype=object), exclusive=True) + with pytest.raises(ValueError, match=match): + tree.query_nearest(np.array([point]), exclusive=True) + + +def test_fat_encoded_full_parity(geographies: npt.NDArray[Any]) -> None: + # an index loaded from a default (include_geographies=True) blob supports + # the operations a thin one cannot, with the same results as the built one + built = spherely.SpatialIndex(geographies) + tree = spherely.SpatialIndex.from_encoded(built.encode()) + point = spherely.create_point(1, 1) + + np.testing.assert_array_equal( + tree.query(point, predicate="intersects"), + built.query(point, predicate="intersects"), + ) + np.testing.assert_array_equal( + tree.query_nearest(point, exclusive=True), + built.query_nearest(point, exclusive=True), + ) + geoms = tree.geometries + assert geoms.shape == (4,) + assert all(spherely.equals(a, b) for a, b in zip(geoms, geographies)) + + +def test_fat_encoded_lazy_reconstruction(geographies: npt.NDArray[Any]) -> None: + built = spherely.SpatialIndex(geographies) + assert built._decoded_count == 0 + + tree = spherely.SpatialIndex.from_encoded(built.encode()) + # opening and candidate-only queries walk the index cells and + # reconstruct nothing + assert tree._decoded_count == 0 + tree.query(spherely.create_point(1, 1)) + assert tree._decoded_count == 0 + + # predicate refinement reconstructs (only) the touched candidates, and + # each entry at most once: the far-away point (50, 50) is never a + # candidate, so it is never reconstructed + tree.query(spherely.create_point(1, 1), predicate="intersects") + touched = tree._decoded_count + assert 1 <= touched < len(geographies) + tree.query(spherely.create_point(1, 1), predicate="intersects") + assert tree._decoded_count == touched + + # geometries materializes every entry + tree.geometries # noqa: B018 + assert tree._decoded_count == len(geographies) + + +def test_from_encoded_invalid_bytes() -> None: + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(b"not a spatial index") + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(b"") + + encoded = spherely.SpatialIndex( + [spherely.create_point(0, 0), spherely.create_point(1, 1)] + ).encode() + with pytest.raises(ValueError, match="unsupported encoded SpatialIndex version"): + spherely.SpatialIndex.from_encoded(encoded[:4] + b"\xff" + encoded[5:]) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(encoded[: len(encoded) // 2]) + + +def _varint(value: int) -> bytes: + # the (LEB128) varint encoding used by s2's Encoder::put_varint64 + out = bytearray() + while value >= 0x80: + out.append((value & 0x7F) | 0x80) + value >>= 7 + out.append(value) + return bytes(out) + + +def test_from_encoded_corrupt_value_table() -> None: + # the header is followed by the geography count, the shape count and one + # varint value per shape; all three counts are small enough here to be + # encoded on a single byte each. A thin blob is used so that the spliced + # bytes below keep targeting the value table (in a blob with geographies + # the block table would fail validation first). + geoms = [spherely.create_point(i, 0) for i in range(4)] + encoded = spherely.SpatialIndex(geoms).encode(include_geographies=False) + header = 6 + assert encoded[header] == 4 # num_geographies + assert encoded[header + 1] == 4 # num_shapes (one per point) + values = header + 2 + + # a shape count smaller than the number of encoded shapes leaves the + # value table too short: queries would read past its end + truncated = ( + encoded[:header] + _varint(4) + _varint(1) + encoded[values : values + 1] + ) + truncated += encoded[values + 4 :] + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(truncated) + + # a value outside [0, num_geographies) is not a usable tree index + out_of_range = bytearray(encoded) + out_of_range[values] = 100 + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(bytes(out_of_range)) + + # a huge shape count must be rejected before it is reserved + huge = encoded[:header] + _varint(4) + _varint(2**40) + encoded[values:] + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(huge) + + # ... as must a huge geography count + huge = encoded[:header] + _varint(2**63) + encoded[header + 1 :] + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(huge) + + +def _mixed_geographies(seed: int, num_points: int = 30) -> npt.NDArray[Any]: + # a mix of every geography kind, including polygons with holes, an + # antimeridian-crossing polygon, a multi-shape collection and an empty + # geography + rng = np.random.default_rng(seed) + lon = rng.uniform(-180, 180, num_points) + lat = rng.uniform(-85, 85, num_points) + geoms: list[Any] = list(np.asarray(spherely.points(lon, lat))) + for x, y in rng.uniform(-80, 80, (5, 2)): + geoms.append( + spherely.create_linestring([(x, y), (x + 2, y + 1), (x + 1, y + 3)]) + ) + for x, y in rng.uniform(-60, 60, (5, 2)): + geoms.append( + spherely.create_polygon( + [(x, y), (x + 8, y), (x + 8, y + 8), (x, y + 8)], + holes=[ + [(x + 3, y + 3), (x + 5, y + 3), (x + 5, y + 5), (x + 3, y + 5)] + ], + ) + ) + geoms.append( + # crosses the antimeridian + spherely.create_polygon([(178, -5), (-178, -5), (-178, 5), (178, 5)]) + ) + geoms.append( + spherely.create_collection( + [ + spherely.create_point(40, 40), + spherely.create_linestring([(41, 41), (42, 42)]), + ] + ) + ) + geoms.append(spherely.create_polygon(None)) # empty + return np.array(geoms) + + +@pytest.mark.parametrize("seed", [7, 21]) +def test_fat_encoded_randomized_parity(seed: int) -> None: + # a fat-decoded index must answer every query through the same code paths + # as the built index it came from: results must match exactly + geoms = _mixed_geographies(seed) + built = spherely.SpatialIndex(geoms) + fat = spherely.SpatialIndex.from_encoded(built.encode()) + + queries = [ + spherely.create_point(0, 0), + spherely.create_polygon([(-10, -10), (30, -10), (30, 30), (-10, 30)]), + # probes the antimeridian-crossing polygon from both sides + spherely.create_point(179.5, 0), + spherely.create_point(-179.5, 0), + spherely.create_polygon([(175, -8), (-175, -8), (-175, 8), (175, 8)]), + geoms[0], + ] + predicates = [ + "intersects", + "within", + "contains", + "covers", + "covered_by", + "touches", + "equals", + ] + for query in queries: + np.testing.assert_array_equal(built.query(query), fat.query(query)) + for predicate in predicates: + np.testing.assert_array_equal( + built.query(query, predicate=predicate), # type: ignore[call-overload] + fat.query(query, predicate=predicate), # type: ignore[call-overload] + ) + np.testing.assert_array_equal( + built.query_nearest(query), fat.query_nearest(query) + ) + for exclusive in [False, True]: + b_idx, b_dist = built.query_nearest( + query, exclusive=exclusive, return_distance=True + ) + f_idx, f_dist = fat.query_nearest( + query, exclusive=exclusive, return_distance=True + ) + np.testing.assert_array_equal(b_idx, f_idx) + np.testing.assert_array_equal(b_dist, f_dist) + + +def test_fat_encoded_geometries_round_trip() -> None: + # geometries reconstructs every entry in insertion order, preserving the + # geography type (e.g. a one-point MULTIPOINT stays a MULTIPOINT) and + # empty geographies + geoms = np.array( + [ + spherely.create_point(10, 10), + spherely.create_multipoint([(20, 20)]), # one-point multipoint + spherely.create_polygon(None), # empty + spherely.create_collection([spherely.create_point(30, 30)]), + spherely.create_polygon([(0, 0), (5, 0), (5, 5), (0, 5)]), + ] + ) + tree = spherely.SpatialIndex.from_encoded(spherely.SpatialIndex(geoms).encode()) + result = tree.geometries + assert result.shape == geoms.shape + np.testing.assert_array_equal( + spherely.get_type_id(result), spherely.get_type_id(geoms) + ) + np.testing.assert_array_equal(spherely.is_empty(result), spherely.is_empty(geoms)) + for a, b in zip(result, geoms): + if not spherely.is_empty(b): + assert spherely.equals(a, b) + # the array is cached: a second access returns the same objects + assert result[0] is tree.geometries[0] + + +def test_fat_encoded_lazy_large_tree() -> None: + # opening a large fat blob and running candidate-only queries must not + # reconstruct any geography at all + rng = np.random.default_rng(3) + geoms = np.asarray( + spherely.points(rng.uniform(-180, 180, 500), rng.uniform(-85, 85, 500)) + ) + tree = spherely.SpatialIndex.from_encoded(spherely.SpatialIndex(geoms).encode()) + assert tree._decoded_count == 0 + tree.query(spherely.create_polygon([(0, 0), (20, 0), (20, 20), (0, 20)])) + assert tree._decoded_count == 0 + + # a predicate query only reconstructs its candidates + candidates = tree.query( + spherely.create_polygon([(0, 0), (20, 0), (20, 20), (0, 20)]) + ) + tree.query( + spherely.create_polygon([(0, 0), (20, 0), (20, 20), (0, 20)]), + predicate="contains", + ) + assert tree._decoded_count == len(candidates) + + # a distance query reconstructs the entries whose shapes it examines -- + # a local neighborhood of the target, not the whole tree. The bound is a + # near-constant one on purpose: a regression to a reconstruction count + # that grows with the tree has to fail here (the count is 43 as written, + # against 500 entries) + tree.query_nearest(spherely.create_point(12, 34)) + assert tree._decoded_count <= 64 + + +def _fat_layout(num_geographies: int, num_shapes: int) -> tuple[int, int]: + # (table_start, blocks_start) of a fat blob whose counts and values all + # fit single-byte varints + table = 6 + 2 + num_shapes + return table, table + 16 * num_geographies + + +def test_from_encoded_unknown_flags() -> None: + encoded = spherely.SpatialIndex([spherely.create_point(0, 0)]).encode() + corrupt = encoded[:5] + bytes([encoded[5] | 0x02]) + encoded[6:] + with pytest.raises(ValueError, match="unsupported encoded SpatialIndex flags"): + spherely.SpatialIndex.from_encoded(corrupt) + + +def test_from_encoded_corrupt_block_table() -> None: + geoms = [spherely.create_point(0, 0), spherely.create_point(1, 1)] + encoded = spherely.SpatialIndex(geoms).encode() + table, blocks = _fat_layout(2, 2) + off0, len0, off1, len1 = struct.unpack_from(" bytes: + out = bytearray(encoded) + struct.pack_into(" None: + # a fat blob's shape -> entry map must be one contiguous ascending run + # per entry (shape addressing within a geography relies on it) + geoms = [spherely.create_point(0, 0), spherely.create_point(1, 1)] + encoded = spherely.SpatialIndex(geoms).encode() + # header (6 bytes) + the two single-byte counts, then one value per shape + values = 8 + assert encoded[values : values + 2] == b"\x00\x01" + corrupt = bytearray(encoded) + corrupt[values : values + 2] = b"\x01\x00" + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(bytes(corrupt)) + + +def test_from_encoded_corrupt_geography_block() -> None: + geoms = [spherely.create_point(0, 0), spherely.create_point(1, 1)] + encoded = spherely.SpatialIndex(geoms).encode() + _, blocks = _fat_layout(2, 2) + + # a block that no longer decodes as a geography fails cleanly on first + # touch (the open itself stays lazy and does not notice) + corrupt = bytearray(encoded) + corrupt[blocks + 2] = 0xFF # clobber the s2geography tag of block 0 + tree = spherely.SpatialIndex.from_encoded(bytes(corrupt)) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex geography"): + tree.geometries # noqa: B018 + + # ... as does an out-of-range geography type byte + corrupt = bytearray(encoded) + corrupt[blocks] = 100 + tree = spherely.SpatialIndex.from_encoded(bytes(corrupt)) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex geography"): + tree.geometries # noqa: B018 + + +def test_from_encoded_unwrappable_geography_kind() -> None: + # the tagged encoding's kind byte selects the s2geography subclass to + # build: the kinds a spherely Geography cannot wrap (here SHAPE_INDEX) + # must be rejected rather than cached in a slot and handed out by + # ``geometries``, where every accessor raises "Unsupported Geography + # subclass" (and an ENCODED_SHAPE_INDEX would outlive its own buffer) + geoms = [spherely.create_point(0, 0), spherely.create_point(1, 1)] + encoded = spherely.SpatialIndex(geoms).encode() + _, blocks = _fat_layout(2, 2) + + corrupt = bytearray(encoded) + corrupt[blocks + 2] = 5 # s2geography::GeographyKind::SHAPE_INDEX + tree = spherely.SpatialIndex.from_encoded(bytes(corrupt)) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex geography"): + tree.geometries # noqa: B018 + + +def test_from_encoded_block_type_byte_mismatch() -> None: + # nothing implies the block's type byte from its tagged encoding: without + # a cross-check a relabelled block yields a geography whose get_type_id + # and WKT disagree + geoms = [spherely.create_point(0, 0), spherely.create_point(1, 1)] + encoded = spherely.SpatialIndex(geoms).encode() + _, blocks = _fat_layout(2, 2) + assert encoded[blocks] == 0 # GeographyType::Point + + # LineString, Polygon, MultiLineString, MultiPolygon, GeometryCollection + # and None: none of them can describe a POINT-kind block + for geog_type in [1, 2, 4, 5, 6, 0xFF]: + corrupt = bytearray(encoded) + corrupt[blocks] = geog_type + tree = spherely.SpatialIndex.from_encoded(bytes(corrupt)) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex geography"): + tree.geometries # noqa: B018 + + # a singular type and its MULTI* form share a kind, so swapping the two + # labels is not detectable from the block alone -- the geometry itself is + # unaffected (see type_matches_kind) + corrupt = bytearray(encoded) + corrupt[blocks] = 3 # GeographyType::MultiPoint + tree = spherely.SpatialIndex.from_encoded(bytes(corrupt)) + assert spherely.to_wkt(tree.geometries[0]) == "POINT (0 0)" + + +def test_from_encoded_corrupt_block_touched_by_query_nearest( + geographies: npt.NDArray[Any], +) -> None: + # a corrupt block first touched from inside query_nearest -- i.e. from the + # ShapeFactory, deep inside s2's distance search rather than from spherely + # code -- surfaces as the same clean ValueError, and leaves the index + # usable: the entries decoded before it stay cached, candidate queries + # (which reconstruct nothing) keep working, and a repeated distance query + # raises identically rather than doing something new + built = spherely.SpatialIndex(geographies) + encoded = built.encode() + table, blocks = _fat_layout(4, 4) + offset, _ = struct.unpack_from(" None: + # a fully decoded fat index is read-only from here on: concurrent queries + # must agree with the single-threaded answers and reconstruct nothing more + geoms = _mixed_geographies(11) + built = spherely.SpatialIndex(geoms) + tree = spherely.SpatialIndex.from_encoded(built.encode()) + tree.geometries # noqa: B018 + decoded = tree._decoded_count + assert decoded == len(geoms) + + queries = [ + spherely.create_point(0, 0), + spherely.create_polygon([(-10, -10), (30, -10), (30, 30), (-10, 30)]), + spherely.create_point(179.5, 0), + spherely.create_linestring([(20, 20), (25, 25)]), + ] + expected = [ + ( + built.query(q, predicate="intersects"), + built.query_nearest(q), + built.query_nearest(q, exclusive=True), + ) + for q in queries + ] + + def run(i: int) -> list[tuple[Any, Any, Any]]: + return [ + ( + tree.query(q, predicate="intersects"), + tree.query_nearest(q), + tree.query_nearest(q, exclusive=True), + ) + for q in queries + ] + + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(run, range(4))) + + for got in results: + for (a, b, c), (x, y, z) in zip(got, expected): + np.testing.assert_array_equal(a, x) + np.testing.assert_array_equal(b, y) + np.testing.assert_array_equal(c, z) + assert tree._decoded_count == decoded + + +def _collection_geographies() -> npt.NDArray[Any]: + # collections mixing all three dimensions: their entries own several + # shapes each, which is what exercises the shape id -> (entry, sub-shape) + # addressing of a fat blob + geoms: list[Any] = [] + for x, y in [(0, 0), (20, 20), (-40, 10), (60, -30)]: + geoms.append( + spherely.create_collection( + [ + spherely.create_point(x, y), + spherely.create_linestring([(x + 1, y), (x + 2, y + 1)]), + spherely.create_polygon( + [(x + 3, y), (x + 5, y), (x + 5, y + 2), (x + 3, y + 2)] + ), + ] + ) + ) + geoms.append(spherely.create_collection([spherely.create_point(80, 80)])) + geoms.append(spherely.create_collection([])) # empty collection + geoms.append(spherely.create_point(0.5, 0.5)) # inside the first collection + return np.array(geoms) + + +def test_collections_full_parity() -> None: + geoms = _collection_geographies() + built = spherely.SpatialIndex(geoms) + fat = spherely.SpatialIndex.from_encoded(built.encode()) + thin = spherely.SpatialIndex.from_encoded(built.encode(include_geographies=False)) + + queries = [ + spherely.create_point(0, 0), + spherely.create_point(4, 1), + spherely.create_linestring([(1, 0), (2, 1)]), + spherely.create_polygon([(-1, -1), (6, -1), (6, 3), (-1, 3)]), + geoms[0], + geoms[-1], + ] + predicates = [ + "intersects", + "within", + "contains", + "covers", + "covered_by", + "touches", + "equals", + ] + for query in queries: + # candidate sets and distance queries agree on all three flavors + for other in [fat, thin]: + np.testing.assert_array_equal(built.query(query), other.query(query)) + b_idx, b_dist = built.query_nearest(query, return_distance=True) + o_idx, o_dist = other.query_nearest(query, return_distance=True) + np.testing.assert_array_equal(b_idx, o_idx) + np.testing.assert_array_equal(b_dist, o_dist) + # predicates and exclusive nearest need the geographies (thin raises) + for predicate in predicates: + np.testing.assert_array_equal( + built.query(query, predicate=predicate), # type: ignore[call-overload] + fat.query(query, predicate=predicate), # type: ignore[call-overload] + ) + b_idx, b_dist = built.query_nearest(query, exclusive=True, return_distance=True) + f_idx, f_dist = fat.query_nearest(query, exclusive=True, return_distance=True) + np.testing.assert_array_equal(b_idx, f_idx) + np.testing.assert_array_equal(b_dist, f_dist) + + +def test_from_encoded_block_shape_count_mismatch() -> None: + # entry 0 has one shape (a point), entry 1 has none (an empty polygon); + # swapping their blocks (and table lengths, keeping the section + # well-formed) must be caught by the shape count cross-check on touch + geoms = [spherely.create_point(0, 0), spherely.create_polygon(None)] + encoded = spherely.SpatialIndex(geoms).encode() + table, blocks = _fat_layout(2, 1) + off0, len0, off1, len1 = struct.unpack_from(" None: + geoms = [spherely.create_point(0, 0), spherely.create_point(1, 1)] + fat_bytes = spherely.SpatialIndex(geoms).encode() + thin_bytes = spherely.SpatialIndex(geoms).encode(include_geographies=False) + + fat = spherely.SpatialIndex.from_encoded(fat_bytes) + thin = spherely.SpatialIndex.from_encoded(thin_bytes) + # a decoded index re-encodes to the exact blob it was loaded from + assert fat.encode() == fat_bytes + assert thin.encode(include_geographies=False) == thin_bytes + # ... and cannot fabricate the other flavor + with pytest.raises(ValueError, match="include_geographies"): + fat.encode(include_geographies=False) + with pytest.raises(ValueError, match="include_geographies"): + thin.encode() + + +def _short_id_map(num_geographies: int) -> bytes: + # a fat blob of ``num_geographies`` single-shape points whose header claims + # one shape fewer than the index structure it still carries verbatim: the + # id map no longer covers the last shape id the index cells reference + geoms = [spherely.create_point(i, 0) for i in range(num_geographies)] + encoded = spherely.SpatialIndex(geoms).encode() + assert encoded[6] == num_geographies # num_geographies + assert encoded[7] == num_geographies # num_shapes (one per point) + values = 8 + assert encoded[values : values + num_geographies] == bytes(range(num_geographies)) + return ( + encoded[:6] + + bytes([num_geographies, num_geographies - 1]) + + encoded[values : values + num_geographies - 1] + + encoded[values + num_geographies :] + ) + + +def test_from_encoded_id_map_shorter_than_index() -> None: + # the shape count check in from_encoded compares the id map with + # num_shape_ids(), which for a fat blob is the id map's own size: without + # a cross-check against the index cells this blob is accepted, and the + # shape id its cells reference past the end of the map is then read out of + # bounds inside the distance queries (a segfault, not an exception) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(_short_id_map(4)) + + +def test_from_encoded_id_map_longer_than_index() -> None: + # the mirror image: an id map with more (still monotonic) entries than the + # index has shapes makes num_shape_ids() over-report, so it is rejected too + geoms = [spherely.create_point(i, 0) for i in range(3)] + encoded = spherely.SpatialIndex(geoms).encode() + values = 8 + assert encoded[values : values + 3] == b"\x00\x01\x02" + corrupt = ( + encoded[:6] + + bytes([3, 4]) + + b"\x00\x01\x02\x02" # one extra entry, still ascending + + encoded[values + 3 :] + ) + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(corrupt) + + +def test_from_encoded_shape_id_check_is_not_a_decode() -> None: + # the cell walk that checks the shape ids must not reconstruct anything: + # opening a fat blob stays lazy + geoms = np.asarray(spherely.points(np.arange(-80.0, 80.0), np.zeros(160))) + tree = spherely.SpatialIndex.from_encoded(spherely.SpatialIndex(geoms).encode()) + assert tree._decoded_count == 0 + + +def test_from_encoded_huge_geography_count_not_allocated() -> None: + # a 12 byte fat blob claiming 2**31-1 geographies must be rejected on the + # counts alone: sizing the per-geography vectors on it first would ask for + # tens of gigabytes before any bounds check ran + corrupt = b"SPIX" + bytes([1, 1]) + _varint(2**31 - 1) + _varint(0) + assert len(corrupt) < 16 + with pytest.raises(ValueError, match="invalid encoded SpatialIndex"): + spherely.SpatialIndex.from_encoded(corrupt) From e53a63eb50f3e281956a46ae56bda9a96df2f62f Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 19:04:47 -0700 Subject: [PATCH 11/12] build the shape id map here to avoid quadratic GeographyIndex::Add --- src/spatial_index.cpp | 41 +++++++++++++++++++++++++++++-------- tests/test_spatial_index.py | 33 +++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index e7a8fc0..8caa04d 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -86,8 +86,25 @@ class SpatialIndex { auto n = m_geographies.size(); auto* data = static_cast(m_geographies.request().ptr); + // this is what s2geography::GeographyIndex::Add does, except that it + // keeps the shape id -> value mapping in a vector it grows with + // ``reserve(size() + num_shapes)`` per geography. libc++ reserves + // exactly and the ``resize`` that follows leaves capacity == size, so + // every call reallocates and copies the whole vector: quadratic, and + // 15 s of pure vector growth for 555k polygons (0.04 s here). + // Reserving once and adding the shapes to the shape index directly + // builds exactly the same index. This is a workaround for the + // dependency -- it can go away once s2geography grows that vector + // geometrically. + auto& mutable_index = m_index->MutableShapeIndex(); + m_values.reserve(static_cast(n)); for (py::ssize_t i = 0; i < n; i++) { - m_index->Add(data[i].as_geog_ptr()->geog(), static_cast(i)); + const auto& geog = data[i].as_geog_ptr()->geog(); + for (int k = 0; k < geog.num_shapes(); k++) { + int shape_id = mutable_index.Add(geog.Shape(k)); + m_values.resize(static_cast(shape_id) + 1); + m_values[static_cast(shape_id)] = static_cast(i); + } } } @@ -240,7 +257,8 @@ class SpatialIndex { encoder.put_varint64(num_geographies); encoder.put_varint64(static_cast(num_shapes)); for (int i = 0; i < num_shapes; i++) { - encoder.put_varint64(static_cast(m_index->value(i))); + encoder.put_varint64( + static_cast(m_values.at(static_cast(i)))); } if (include_geographies) { @@ -502,16 +520,22 @@ class SpatialIndex { // destroyed in reverse declaration order, and an S2ShapeIndex touches its // shapes while being torn down. + // shape id -> tree index, filled by the constructor (built mode) or read + // back out of the blob by ``from_encoded`` (encoded mode) + std::vector m_values; + // built mode (constructed from an array of geographies) // m_geographies keeps the Geography objects alive (the index borrows // their S2Shapes) py::array_t m_geographies; + // wrapped for its S2ShapeIndex only: the constructor fills m_values + // itself, so this GeographyIndex's own shape id -> value table stays + // empty and neither its value() nor its Iterator may be used std::unique_ptr m_index; // encoded mode (loaded with from_encoded); the encoded bytes are kept // alive here as EncodedS2ShapeIndex works directly on them std::string m_encoded; - std::vector m_values; py::ssize_t m_num_geographies = 0; // whether the blob carries per-geography blocks, and where each block // lives inside ``m_encoded`` (absolute offset, length) @@ -836,12 +860,11 @@ class SpatialIndex { // tree index of the geography that owns the given shape int tree_value(int shape_id) const { - if (is_encoded()) { - // ``from_encoded`` checks the table covers every shape id; ``at`` - // keeps a mismatch a Python exception rather than a bad read - return m_values.at(static_cast(shape_id)); - } - return m_index->value(shape_id); + // the constructor covers every shape of a built index and + // ``from_encoded`` checks that the decoded table covers every shape + // id; ``at`` keeps a mismatch a Python exception rather than a bad + // read + return m_values.at(static_cast(shape_id)); } std::optional make_predicate(const std::optional& name) const { diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 4681f5f..7908733 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -802,6 +802,39 @@ def test_encode_concurrent() -> None: assert all(result == expected for result in results) +def test_shape_ids_map_to_tree_indices() -> None: + # the shape id -> tree index map is filled by the constructor rather than + # by s2geography::GeographyIndex::Add: geographies with several shapes + # (and ones with none) must still line up with their input positions + geoms = [ + spherely.create_point(0, 0), + spherely.create_collection( # two shapes + [ + spherely.create_point(30, 30), + spherely.create_linestring([(31, 31), (32, 32)]), + ] + ), + spherely.create_polygon(None), # empty: no shape + spherely.create_point(-60, -60), + ] + tree = spherely.SpatialIndex(geoms) + assert len(tree) == 4 + + for i, geog in enumerate(geoms): + if i == 2: # the empty geography is indexed but never returned + continue + np.testing.assert_array_equal(tree.query(geog, predicate="intersects"), [i]) + np.testing.assert_array_equal(tree.query_nearest(geog), [i]) + + # both shapes of the collection resolve to its tree index, and only once + np.testing.assert_array_equal( + tree.query(spherely.create_point(31.5, 31.5), predicate="intersects"), [] + ) + np.testing.assert_array_equal( + tree.query_nearest(spherely.create_point(31.5, 31.5)), [1] + ) + + def test_from_encoded_preserves_indices() -> None: # multi-shape and empty geographies make the shape id <-> tree index # mapping non-trivial: it must survive the encode / from_encoded round From e9758edbee6f1bd754d3ee1f4ad76380396c10d7 Mon Sep 17 00:00:00 2001 From: espg Date: Sat, 8 Aug 2026 21:01:27 -0700 Subject: [PATCH 12/12] portability fixes for s2 0.14 and LP64, plus a mypy list annotation --- src/spatial_index.cpp | 14 +++++++++++--- tests/test_spatial_index.py | 9 ++++++++- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp index 8caa04d..fb085e2 100644 --- a/src/spatial_index.cpp +++ b/src/spatial_index.cpp @@ -37,6 +37,14 @@ using namespace spherely; namespace { +// s2 0.11 declares Decoder::get_varint64(uint64*) -- `unsigned long long` -- while +// 0.14 declares get_varint64(uint64_t*), which on LP64 is `unsigned long`. The two +// are distinct types, so a pointer to the wrong one will not convert. Deduce +// whichever this build's s2 wants rather than naming either. +template +T varint64_value_type(bool (Decoder::*)(T*)); +using VarintU64 = decltype(varint64_value_type(&Decoder::get_varint64)); + /* ** Sets s2geometry's process-wide index-build temporary memory budget ** (``FLAGS_s2shape_index_tmp_memory_budget``) for the lifetime of the guard and @@ -304,8 +312,8 @@ class SpatialIndex { } self->m_has_geographies = (flags & kFlagHasGeographies) != 0; - std::uint64_t num_geographies; - std::uint64_t num_shapes; + VarintU64 num_geographies; + VarintU64 num_shapes; if (!decoder.get_varint64(&num_geographies) || !decoder.get_varint64(&num_shapes)) { fail(); } @@ -330,7 +338,7 @@ class SpatialIndex { self->m_num_geographies = static_cast(num_geographies); self->m_values.reserve(num_shapes); for (std::uint64_t i = 0; i < num_shapes; i++) { - std::uint64_t value; + VarintU64 value; if (!decoder.get_varint64(&value) || value >= num_geographies) { fail(); } diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py index 7908733..58e12ac 100644 --- a/tests/test_spatial_index.py +++ b/tests/test_spatial_index.py @@ -879,7 +879,14 @@ def test_from_encoded_bytes_like() -> None: ) encoded = tree.encode() query = spherely.create_point(0.9, 0.9) - for buffer in [encoded, bytearray(encoded), memoryview(encoded)]: + # annotated: an unannotated literal joins to Sequence[int], which + # from_encoded does not accept + buffers: list[bytes | bytearray | memoryview] = [ + encoded, + bytearray(encoded), + memoryview(encoded), + ] + for buffer in buffers: decoded = spherely.SpatialIndex.from_encoded(buffer) np.testing.assert_array_equal(decoded.query_nearest(query), [1])