diff --git a/src/geobuf/geobuf_index.hpp b/src/geobuf/geobuf_index.hpp index 832fecf..5670253 100644 --- a/src/geobuf/geobuf_index.hpp +++ b/src/geobuf/geobuf_index.hpp @@ -309,6 +309,19 @@ struct GeobufIndex return {}; } + std::set query(const Eigen::Vector2d &min, + const Eigen::Vector2d &max) const + { + if (!packed_rtree) { + return {}; + } + std::set 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 feature_id = "@", diff --git a/src/main.cpp b/src/main.cpp index c37f03a..f9f080a 100644 --- 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(&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( + &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 index 2dd5621..29eaf17 100644 --- 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()