diff --git a/CMakeLists.txt b/CMakeLists.txt index 3cd1d11..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? @@ -80,7 +105,12 @@ set(CPP_SOURCES 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/s2_tmp_memory_budget.cpp + src/spatial_index.cpp) endif() add_library(spherely MODULE ${CPP_SOURCES}) @@ -91,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/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..ac37458 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -6,6 +6,14 @@ :toctree: _api_generated/ Geography + SpatialIndex + SpatialIndex.build + SpatialIndex.encode + SpatialIndex.from_encoded + SpatialIndex.geometries + SpatialIndex.is_built + SpatialIndex.query + SpatialIndex.query_nearest Projection Projection.lnglat Projection.pseudo_mercator 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.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..2d6701b --- /dev/null +++ b/src/predicates_common.cpp @@ -0,0 +1,101 @@ +#include +#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; +} + +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; + + 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/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 new file mode 100644 index 0000000..fb085e2 --- /dev/null +++ b/src/spatial_index.cpp @@ -0,0 +1,1430 @@ +#include +#include +#include +#include +#include +#include +#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" +#include "s2_tmp_memory_budget.hpp" + +namespace py = pybind11; +namespace s2geog = s2geography; +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 +** 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 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_previous(get_s2_tmp_memory_budget()) { + set_s2_tmp_memory_budget(bytes); + } + + ~TmpMemoryBudgetGuard() { + set_s2_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 +** lookup similar to shapely's STRtree. +*/ +class SpatialIndex { +public: + SpatialIndex(const py::array_t& geographies) { + if (geographies.ndim() != 1) { + throw py::type_error("geographies must be a 1-dimensional array"); + } + + // 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 = 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++) { + 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); + } + } + } + + // 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(); + } + + // 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)"); + } + + // 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& index = m_index->ShapeIndex(); + if (tmp_memory_budget.has_value()) { + TmpMemoryBudgetGuard guard(*tmp_memory_budget); + index.ForceBuild(); + } else { + 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 { + 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 + // 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. 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; + 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_values.at(static_cast(i)))); + } + + 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); + } + index.Encode(&encoder); + } + 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; + + VarintU64 num_geographies; + VarintU64 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++) { + VarintU64 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 { + auto pred = make_predicate(predicate); + auto results = query_one(geography, pred ? &*pred : nullptr); + return to_index_array(results); + } + + // 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 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(), pred ? &*pred : nullptr); + for (int t : results) { + input_idx.push_back(i); + tree_idx.push_back(static_cast(t)); + } + } + + 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"); + } + + // 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); + 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++) { + dist_data(static_cast(j)) = dists[j]; + } + return py::make_tuple(pairs, distances); + } + +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); + // 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. + + // 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; + 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 + // 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()); + } + } + + // 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 { + // 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 { + 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 { + std::unordered_set hits; + + auto region = query_geog.geog().Region(); + S2RegionCoverer coverer; + std::vector covering; + coverer.GetCovering(*region, &covering); + + 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(); + for (int candidate : hits) { + auto* cand_geog = candidate_geog(candidate); + if ((*pred)(query_index, cand_geog->geog_index())) { + results.push_back(candidate); + } + } + } + + std::sort(results.begin(), results.end()); + 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(tree_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 (exclusive) { + check_has_geographies("query_nearest with exclusive=True"); + } + if (query_geog.is_empty()) { + return {}; + } + + 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); + 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(); + std::vector non_equal; + for (int t : zero_hits) { + if (!equals_pred(query_index, candidate_geog(t)->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)); + 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; + } + + 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) { + 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``. + + 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 instead. + + 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"), + "__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 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`) -- including on an + index loaded with :py:meth:`SpatialIndex.from_encoded`, which has no + build to force. + + 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 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 + 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 + ----- + 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. + + 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 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, + R"pbdoc(encode(include_geographies=True) + + Serialize the index to bytes. + + By default the Geography objects themselves are serialized along with + 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 + 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_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, + py::const_), + 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") + .def("query", + py::overload_cast&, std::optional>( + &SpatialIndex::query, py::const_), + py::arg("geography"), + 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. + 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.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..1a78c6f 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -232,6 +232,106 @@ 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 + +_IndexPredicate = Literal[ + "intersects", "within", "contains", "covers", "covered_by", "touches", "equals" +] + +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: ... + # 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 + ) -> npt.NDArray[np.intp]: ... + @overload + 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: ... + # 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..58e12ac --- /dev/null +++ b/tests/test_spatial_index.py @@ -0,0 +1,1566 @@ +import struct +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable + +import numpy as np +import numpy.typing as npt +import pytest + +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 + 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 + return np.array( + [ + spherely.create_point(0, 0), + spherely.create_point(1, 1), + spherely.create_point(2, 2), + spherely.create_point(50, 50), + ] + ) + + +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], make_full_tree: MakeTree +) -> None: + tree = make_full_tree(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_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_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_matches( + 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 + + # 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: + # 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 + + 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)]) +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, 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: + 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], 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) + 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(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 + geographies = np.array( + [ + spherely.create_point(1, 1), + spherely.create_point(4, 4), + spherely.create_point(50, 50), + ] + ) + tree = make_full_tree(geographies) + triangle = spherely.create_polygon([(0, 0), (5, 0), (0, 5), (0, 0)]) + + coarse = tree.query(triangle) + refined = tree.query(triangle, predicate="contains") + + # 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], 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], make_full_tree: MakeTree) -> None: + tree = make_full_tree(geographies) + queries = np.array( + [ + spherely.create_point(1, 1), + spherely.create_point(50, 50), + ] + ) + result = tree.query(queries, predicate="intersects") + # (input_index, tree_index) pairs + np.testing.assert_array_equal(result, [[0, 1], [1, 3]]) + + +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 = 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], 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], 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], 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(make_tree: MakeTree) -> None: + # two points exactly symmetric about the query point on the equator + tree = make_tree( + [ + 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(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 = make_tree( + [ + 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(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 = make_full_tree( + [ + 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 = make_full_tree( + [ + 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(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 = 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 = 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(make_tree: MakeTree) -> None: + tree = make_tree( + [ + 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 = 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], 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) + 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], make_tree: MakeTree +) -> None: + tree = make_tree(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(make_tree: MakeTree) -> None: + tree = make_tree( + [ + 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], make_tree: MakeTree +) -> None: + tree = make_tree(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(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 = make_tree(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], make_tree: MakeTree, max_distance: float +) -> None: + # NaN in particular used to be silently accepted as "no bound" + 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): + 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], 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 = 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): + 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(make_full_tree: MakeTree) -> None: + tree = make_full_tree( + [ + 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 = 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(make_full_tree: MakeTree) -> None: + # all geographies at distance zero are equal to the query + 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 = 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) + 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( + 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 = 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 = make_full_tree(equal) + assert tree.query_nearest(query, exclusive=True).shape == (0,) + + +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 = 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) + 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) + + +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) + + +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 + # 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) + # 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]) + + 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)