not ready

cubao · Nov 18, 2023 · 685b2f4 · 685b2f4
1 parent fbd7df1
commit 685b2f4
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 16 deletions.
diff --git a/src/geobuf/geobuf_index.hpp b/src/geobuf/geobuf_index.hpp
@@ -309,6 +309,19 @@ struct GeobufIndex
         return {};
     }
 
+    std::set<uint32_t> query(const Eigen::Vector2d &min,
+                             const Eigen::Vector2d &max) const
+    {
+        if (!packed_rtree) {
+            return {};
+        }
+        std::set<uint32_t> hits;
+        for (auto h : packed_rtree->search(min[0], min[1], max[0], max[1])) {
+            hits.insert(h.offset);
+        }
+        return hits;
+    }
+
     static bool indexing(const std::string &input_geobuf_path,
                          const std::string &output_index_path,
                          const std::optional<std::string> feature_id = "@",

diff --git a/src/main.cpp b/src/main.cpp
@@ -303,21 +303,6 @@ 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",
-            [](GeobufIndex &self) -> FlatGeobuf::PackedRTree * {
-                if (!self.packed_rtree) {
-                    return nullptr;
-                }
-                return &(*self.packed_rtree);
-            },
-            rvp::reference_internal)
-        .def_property_readonly(
-            "decoder",
-            [](GeobufIndex &self) -> GeobufIndex::Decoder & {
-                return self.decoder;
-            },
-            rvp::reference_internal)
         //
         .def("init", py::overload_cast<const std::string &>(&GeobufIndex::init),
              "index_bytes"_a)
@@ -369,6 +354,10 @@ PYBIND11_MODULE(_pybind11_geobuf, m)
              "bytes"_a)
         .def("decode_non_features",
              py::overload_cast<>(&GeobufIndex::decode_non_features))
+        .def(
+            "query",
+            py::overload_cast<const Eigen::Vector2d &, const Eigen::Vector2d &>(
+                &GeobufIndex::query, py::const_))
         //
         .def_static("indexing", &GeobufIndex::indexing, //
                     "input_geobuf_path"_a,              //

diff --git a/tests/test_geobuf.py b/tests/test_geobuf.py
@@ -1920,9 +1920,20 @@ def test_geobuf_index():
     build_dir = os.path.abspath(f"{__pwd}/../build")
     opath = f"{build_dir}/suzhoubeizhan.idx"
 
-    assert GeobufIndex.indexing(ipath, opath)
     indexer = GeobufIndex()
+    assert indexer.header_size == indexer.num_features == 0
+    assert not indexer.offsets
+    assert indexer.ids is None
+
+    assert GeobufIndex.indexing(ipath, opath)
     assert indexer.mmap_init(opath, ipath)
+    assert indexer.header_size == 80
+    assert indexer.num_features == 1016
+    offsets = indexer.offsets
+    assert len(offsets) == 1018
+    assert offsets[:5] == [84, 346, 845, 1027, 1202]
+    assert indexer.ids["24"] == 0
+
     expected = {
         "type": "Feature",
         "geometry": {
@@ -2010,6 +2021,10 @@ def test_geobuf_index():
     assert indexer.decode_non_features()
     assert indexer.decode_non_features()() == expected_custom_props
 
+    hits = planet.query([120.64094, 31.41515], [120.64137, 31.41534])
+    assert len(hits) == 4
+    assert hits.tolist() == [529, 530, 536, 659]
+
 
 if __name__ == "__main__":
     test_geobuf_index()