it should work

cubao · Nov 12, 2023 · 8beb9bd · 8beb9bd
1 parent b848bf7
commit 8beb9bd
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 11 deletions.
diff --git a/headers b/headers
diff --git a/src/geobuf/geobuf_plus.hpp b/src/geobuf/geobuf_plus.hpp
@@ -1,19 +1,16 @@
 #pragma once
 #include "geobuf.hpp"
 #include "planet.hpp"
+#include <spdlog/spdlog.h>
 
 namespace cubao
 {
 
 struct GeobufPlus
 {
-
-    static std::string encode(const Planet &planet) { return ""; }
-
     GeobufPlus() = default;
 
-    void init(const std::string &header_bytes) {
-    }
+    void init(const std::string &header_bytes) {}
     mapbox::geojson::feature_collection decode(const std::string &bytes) const
     {
         return {};
@@ -30,9 +27,69 @@ struct GeobufPlus
         return {};
     }
 
-    mapbox::feature::property_map decode_non_features(const char *tail_start, const char *tail_end) const {
+    mapbox::feature::property_map
+    decode_non_features(const char *tail_start, const char *tail_end) const
+    {
+        return {};
     }
 
-    static std::string encode(const Planet &planet) { return ""; }
+    static std::string encode(const Planet &planet) {
+        return "";
+    }
+
+    static bool encode(const std::string &input_geojson_path, const std::string &output_geobuf_plus_path) {
+        spdlog::info("Loading {} ...", input_geojson_path);
+        auto json = mapbox::geobuf::load_json(input_geojson_path);
+        if (json.IsNull()) {
+            spdlog::error("Invalid input. Abort");
+            return false;
+        }
+        auto geojson = mapbox::geojson::convert(json);
+        if (geojson.is<mapbox::geojson::geometry>()) {
+            geojson = mapbox::geojson::feature_collection{
+                mapbox::geojson::feature{geojson.get<mapbox::geojson::geometry>()}
+            };
+        } else if (geojson.is<mapbox::geojson::feature>()) {
+            geojson = mapbox::geojson::feature_collection{
+                    geojson.get<mapbox::geojson::feature>()
+            };
+        }
+        auto &fc = geojson.get<mapbox::geojson::feature_collection>();
+        spdlog::info("encoding {} features...", fc.size());
+
+        auto planet = Planet(fc);
+        auto bytes = encode(planet);
+        FILE *fp = fopen(output_geobuf_plus_path.c_str(), "wb");
+        if (!fp) {
+            spdlog::error("failed to open {} for writing", output_geobuf_plus_path);
+            return false;
+        }
+        fwrite("GeobufPlus", 10, 1, fp);
+        int num_features = fc.size();
+        fwrite(&num_features, sizeof(num_features), 1, fp);
+
+        auto &rtree = planet.packed_rtree();
+        auto extent = rtree.getExtent();
+        fwrite(&extent, sizeof(extent), 1, fp);
+        int num_items = rtree.getNumItems();
+        fwrite(&num_items, sizeof(num_items), 1, fp);
+        int num_nodes = rtree.getNumNodes();
+        fwrite(&num_nodes, sizeof(num_nodes), 1, fp);
+        int node_size = rtree.getNodeSize();
+        fwrite(&node_size, sizeof(node_size), 1, fp);
+        int tree_size = rtree.size();
+        fwrite(&tree_size, sizeof(tree_size), 1, fp);
+        rtree.streamWrite([&](const uint8_t *data, size_t size) {
+            fwrite(data, 1, size, fp);
+        });
+
+        fclose(fp);
+        return true;
+        // write magic, geobuf_plus
+        // write num_features, bbox
+        // write num_tree_bytes, tree
+        // write num_offsets, tree
+        // geobuf
+    }
 };
-} // namespace cubao
+} // namespace cubao
diff --git a/src/geobuf/planet.hpp b/src/geobuf/planet.hpp
@@ -132,6 +132,10 @@ struct Planet
         return fc;
     }
 
+    const FlatGeobuf::PackedRTree &packed_rtree() const {
+        return rtree();
+    }
+
   private:
     FeatureCollection features_;
 

diff --git a/src/main.cpp b/src/main.cpp
@@ -11,8 +11,8 @@
 #include <pybind11/stl_bind.h>
 
 #include "geobuf/geobuf.hpp"
-#include "geobuf/planet.hpp"
 #include "geobuf/geobuf_plus.hpp"
+#include "geobuf/planet.hpp"
 #include "geobuf/pybind11_helpers.hpp"
 
 #include <optional>
@@ -278,10 +278,12 @@ PYBIND11_MODULE(_pybind11_geobuf, m)
         ;
 
     using GeobufPlus = cubao::GeobufPlus;
-    py::class_<GeobufPlus>(m, "GeobufPlus", py::module_local())    //
+    py::class_<GeobufPlus>(m, "GeobufPlus", py::module_local()) //
         .def(py::init<>())
         .def("init", &GeobufPlus::init, "header_bytes"_a)
         //
+        .def_static("encode", py::overload_cast<const std::string &, const std::string &>(&GeobufPlus::encode), "input_geojson_path"_a, "output_geobuf_plus_path"_a)
+        //
         ;
 
     cubao::bind_rapidjson(m);

diff --git a/tests/test_geobuf.py b/tests/test_geobuf.py
@@ -1889,6 +1889,11 @@ def test_query():
 
 
 if __name__ == "__main__":
+
+    from pybind11_geobuf import GeobufPlus
+    ipath = f"{__pwd}/../data/suzhoubeizhan.json"
+    print(GeobufPlus.encode(ipath, f'{__pwd}/../build/export.pbf'))
+
     np.set_printoptions(suppress=True)
     pwd = os.path.abspath(os.path.dirname(__file__))
     pytest_main(pwd, test_file=os.path.basename(__file__))