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

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_<symbol> 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

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)

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

src/geography.hpp

@@ -79,7 +79,7 @@ class Geography {
         return reinterpret_cast<const T*>(&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<s2geog::ShapeIndexGeography>(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;
 

src/predicates.hpp

@@ -0,0 +1,30 @@
+#ifndef SPHERELY_PREDICATES_H_
+#define SPHERELY_PREDICATES_H_
+
+#include <s2geography.h>
+
+#include <functional>
+#include <string>
+
+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<bool(const s2geog::ShapeIndexGeography&, const s2geog::ShapeIndexGeography&)>;
+
+// 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_

src/predicates_common.cpp

@@ -0,0 +1,101 @@
+#include <pybind11/pybind11.h>
+#include <s2/s2boolean_operation.h>
+#include <s2geography.h>
+
+#include <unordered_map>
+
+#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<std::string, PredicateId> 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

src/s2_tmp_memory_budget.cpp

@@ -0,0 +1,52 @@
+#include "s2_tmp_memory_budget.hpp"
+
+#include <absl/flags/flag.h>
+
+/*
+** ``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<std::int64_t> FLAGS_s2shape_index_tmp_memory_budget;
+#else
+#include <s2/mutable_s2shape_index.h>
+#endif
+
+namespace spherely {
+
+std::int64_t get_s2_tmp_memory_budget() {
+    return static_cast<std::int64_t>(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

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 <cstdint>
+
+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_