add batch_query_array which flattens nested vec to two numpy arrays before returning to python (#35)

* add batch_query_array for numpy output instead of nested lists

* add batch_query_array for numpy output instead of nested lists

Author: LunarLanding

cpp/main.cc

@@ -33,6 +33,9 @@ PYBIND11_MODULE(PRTree, m)
       .def("batch_query", &PRTree<T, B, 2>::find_all, R"pbdoc(
           parallel query with multi-thread
         )pbdoc")
+      .def("batch_query_array", &PRTree<T, B, 2>::find_all_array, R"pbdoc(
+          parallel query with multi-thread with array output
+        )pbdoc")
       .def("erase", &PRTree<T, B, 2>::erase, R"pbdoc(
           Delete from prtree
         )pbdoc")
@@ -74,6 +77,9 @@ PYBIND11_MODULE(PRTree, m)
       .def("batch_query", &PRTree<T, B, 3>::find_all, R"pbdoc(
           parallel query with multi-thread
         )pbdoc")
+      .def("batch_query_array", &PRTree<T, B, 3>::find_all_array, R"pbdoc(
+          parallel query with multi-thread with array output
+        )pbdoc")
       .def("erase", &PRTree<T, B, 3>::erase, R"pbdoc(
           Delete from prtree
         )pbdoc")
@@ -115,6 +121,9 @@ PYBIND11_MODULE(PRTree, m)
       .def("batch_query", &PRTree<T, B, 4>::find_all, R"pbdoc(
           parallel query with multi-thread
         )pbdoc")
+      .def("batch_query_array", &PRTree<T, B, 4>::find_all_array, R"pbdoc(
+          parallel query with multi-thread with array output
+        )pbdoc")
       .def("erase", &PRTree<T, B, 4>::erase, R"pbdoc(
           Delete from prtree
         )pbdoc")

cpp/prtree.h

@@ -51,6 +51,38 @@ namespace py = pybind11;
 template <class T>
 using vec = std::vector<T>;
 
+template <typename Sequence >
+inline py::array_t<typename Sequence::value_type> as_pyarray(Sequence& seq) {
+
+  auto size = seq.size();
+  auto data = seq.data();
+  std::unique_ptr<Sequence> seq_ptr = std::make_unique<Sequence>(std::move(seq));
+  auto capsule = py::capsule(seq_ptr.get(), [](void *p) { std::unique_ptr<Sequence>(reinterpret_cast<Sequence*>(p)); });
+  seq_ptr.release();
+  return py::array(size, data, capsule);
+}
+
+template <typename T>
+auto list_list_to_arrays(vec<vec<T>> out_ll){
+  vec<T> out_s;
+  out_s.reserve(out_ll.size());
+  std::size_t sum = 0;
+  for (auto &&i : out_ll) {
+    out_s.push_back(i.size());
+    sum += i.size();
+  }
+  vec<T> out;
+  out.reserve(sum);
+  for(const auto &v: out_ll)
+    out.insert(out.end(), v.begin(), v.end());
+
+  return make_tuple(
+      std::move(as_pyarray(out_s))
+      ,
+      std::move(as_pyarray(out))
+  );
+}
+
 template <class T, size_t StaticCapacity>
 using svec = itlib::small_vector<T, StaticCapacity>;
 
@@ -1259,6 +1291,10 @@ class PRTree
     return out;
   }
 
+  auto find_all_array(const py::array_t<float> &x){
+    return list_list_to_arrays(std::move(find_all(x)));
+  }
+
   auto find_one(const vec<float> &x)
   {
     bool is_point = false;