export PackedRTree (#28)

* update headers

* export packed rtree

* export

* test packed_rtree

* ready to release

* lint

* lint

* fix

---------

Co-authored-by: TANG ZHIXIONG <zhixiong.tang@momenta.ai>

Author: district10

.github/workflows/wheels.yml

@@ -49,7 +49,7 @@ jobs:
       with:
         platforms: all
 
-    - uses: pypa/cibuildwheel@v2.12.0
+    - uses: pypa/cibuildwheel@v2.17.0
       env:
         # CIBW_ARCHS: auto64
         CIBW_ARCHS_LINUX: x86_64 aarch64

docs/about/release-notes.md

@@ -10,6 +10,10 @@ To upgrade `pybind11-geobuf` to the latest version, use pip:
 pip install -U pybind11-geobuf
 ```
 
+## Version 0.2.1 (2024-03-14)
+
+*   Export PackedRTree
+
 ## Version 0.2.0 (2023-11-18)
 
 *   Indexing geobuf in protobuf format; Spec: geobuf_index.proto

headers

@@ -127,7 +127,7 @@ def build_extension(self, ext):
 # logic and declaration, and simpler if you include description/version in a file.
 setup(
     name="pybind11_geobuf",
-    version="0.2.0",
+    version="0.2.1",
     author="tzx",
     author_email="dvorak4tzx@gmail.com",
     url="https://geobuf-cpp.readthedocs.io",

setup.py

@@ -269,6 +269,69 @@ PYBIND11_MODULE(_pybind11_geobuf, m)
     auto geojson = m.def_submodule("geojson");
     cubao::bind_geojson(geojson);
 
+    using namespace FlatGeobuf;
+    py::class_<NodeItem>(m, "NodeItem", py::module_local())
+        .def_property_readonly("min_x",
+                               [](const NodeItem &self) { return self.minX; })
+        .def_property_readonly("min_y",
+                               [](const NodeItem &self) { return self.minY; })
+        .def_property_readonly("max_x",
+                               [](const NodeItem &self) { return self.maxX; })
+        .def_property_readonly("max_y",
+                               [](const NodeItem &self) { return self.maxY; })
+        .def_property_readonly("offset",
+                               [](const NodeItem &self) { return self.offset; })
+        .def_property_readonly(
+            "width", [](const NodeItem &self) { return self.width(); })
+        .def_property_readonly(
+            "height", [](const NodeItem &self) { return self.height(); })
+        //
+        .def("expand", &NodeItem::expand, "other"_a)
+        .def("intersects", &NodeItem::intersects, "other"_a)
+        .def(py::self == py::self)
+        .def(py::self != py::self)
+        .def("to_numpy",
+             [](const NodeItem &self) -> Eigen::Vector4d {
+                 return {self.minX, self.minY, self.maxX, self.maxY};
+             })
+        //
+        ;
+
+    using PackedRTree = FlatGeobuf::PackedRTree;
+    py::class_<PackedRTree>(m, "PackedRTree", py::module_local())
+        .def(
+            "search",
+            [](const PackedRTree &self, double minX, double minY, double maxX,
+               double maxY) {
+                auto hits = self.search(minX, minY, maxX, maxY);
+                std::vector<size_t> ret;
+                ret.reserve(hits.size());
+                for (auto &h : hits) {
+                    ret.push_back(h.offset);
+                }
+                return ret;
+            },
+            "min_x"_a, "min_y"_a, "max_x"_a, "max_y"_a)
+        .def_property_readonly(
+            "size", [](const PackedRTree &self) { return self.size(); })
+        .def_property_readonly("extent",
+                               [](const PackedRTree &self) {
+                                   auto bbox = self.getExtent();
+                                   return Eigen::Vector4d(bbox.minX, bbox.minY,
+                                                          bbox.maxX, bbox.maxY);
+                               })
+        .def_property_readonly(
+            "num_items",
+            [](const PackedRTree &self) { return self.getNumItems(); })
+        .def_property_readonly(
+            "num_nodes",
+            [](const PackedRTree &self) { return self.getNumNodes(); })
+        .def_property_readonly(
+            "node_size",
+            [](const PackedRTree &self) { return self.getNodeSize(); })
+        //
+        ;
+
     using Planet = cubao::Planet;
     py::class_<Planet>(m, "Planet", py::module_local())
         .def(py::init<>())
@@ -281,6 +344,7 @@ PYBIND11_MODULE(_pybind11_geobuf, m)
         .def("build", &Planet::build, py::kw_only(),
              "per_line_segment"_a = false, "force"_a = false)
         .def("query", &Planet::query, "min"_a, "max"_a)
+        .def("packed_rtree", &Planet::packed_rtree, rvp::reference_internal)
         .def("copy", &Planet::copy)
         .def("crop", &Planet::crop, "polygon"_a, py::kw_only(),
              "clipping_mode"_a = "longest", //
@@ -303,6 +367,15 @@ PYBIND11_MODULE(_pybind11_geobuf, m)
             "offsets", [](const GeobufIndex &self) { return self.offsets; })
         .def_property_readonly("ids",
                                [](const GeobufIndex &self) { return self.ids; })
+        .def_property_readonly(
+            "packed_rtree",
+            [](const GeobufIndex &self) -> const FlatGeobuf::PackedRTree * {
+                if (!self.packed_rtree) {
+                    return nullptr;
+                }
+                return &*self.packed_rtree;
+            },
+            rvp::reference_internal)
         //
         .def("init", py::overload_cast<const std::string &>(&GeobufIndex::init),
              "index_bytes"_a)

src/main.cpp

@@ -1896,6 +1896,19 @@ def test_query():
     path = f"{__pwd}/../data/suzhoubeizhan.pbf"
     fc = geojson.FeatureCollection().load(path)
     planet = Planet(fc)
+
+    tree = planet.packed_rtree()
+    assert id(tree) == id(planet.packed_rtree())
+    assert tree.extent.tolist() == [
+        120.623512,
+        31.4001743,
+        120.67509,
+        31.4458499,
+    ]
+    assert tree.num_items == 1016
+    assert tree.num_nodes == 1085
+    assert tree.node_size == 16
+
     hits = planet.query([120.64094, 31.41515], [120.64137, 31.41534])
     assert len(hits) == 4
     assert hits.tolist() == [529, 530, 536, 659]
@@ -1935,6 +1948,17 @@ def test_geobuf_index():
     assert len(offsets) == 1018
     assert offsets[:5] == [56, 318, 817, 999, 1174]
     assert indexer.ids["24"] == 0
+    tree = indexer.packed_rtree
+    assert id(tree) == id(indexer.packed_rtree)
+    assert tree.extent.tolist() == [
+        120.623512,
+        31.4001743,
+        120.67509,
+        31.4458499,
+    ]
+    assert tree.num_items == 1016
+    assert tree.num_nodes == 1085
+    assert tree.node_size == 16
 
     expected = {
         "type": "Feature",
@@ -2025,6 +2049,15 @@ def test_geobuf_index():
 
     hits = indexer.query([120.64094, 31.41515], [120.64137, 31.41534])
     assert len(hits) == 4
+    tree = indexer.packed_rtree
+    assert sorted(hits) == sorted(
+        tree.search(
+            120.64094,
+            31.41515,
+            120.64137,
+            31.41534,
+        )
+    )
     hits = [indexer.decode_feature(h) for h in hits]
     hits = [h.properties()["id"]() for h in hits]
     assert hits == ["943", "936", "937", "1174"]