diff --git a/python_bindings/nmslib.cc b/python_bindings/nmslib.cc index 4ed85f7..8d51afd 100644 --- a/python_bindings/nmslib.cc +++ b/python_bindings/nmslib.cc @@ -150,8 +150,7 @@ struct IndexWrapper { py::gil_scoped_release l; index->Search(&rq, -1); } - std::unique_ptr> res(rq.Result());//->Clone() - return convertResult(res.get()); + return convertResult_rq(rq.ResultSet()); } py::object knnQueryBatch(py::object input, size_t k, int num_threads) { @@ -197,6 +196,22 @@ struct IndexWrapper { } return py::make_tuple(ids, distances); } + + py::object convertResult_rq(std::set res) { + // Create numpy arrays for the output + size_t size = res->Size(); + py::array_t ids(size); + py::array_t distances(size); + + std::set::iterator it = setOfStr.begin(); + while (it != setOfStr.end()) + { + ids.mutable_at(i) = it->id(); + //distances.mutable_at(size) = res->at(i)->id(); + it++; + } + return py::make_tuple(ids, distances); + } const Object * readObject(py::object input, int id = 0) { switch (data_type) {