Fix get_items

README.md

@@ -79,7 +79,7 @@ For other spaces use the nmslib library https://github.com/nmslib/nmslib.
 
 * `set_num_threads(num_threads)` set the default number of cpu threads used during data insertion/querying.
 
-* `get_items(ids)` - returns a numpy array (shape:`N*dim`) of vectors that have integer identifiers specified in `ids` numpy vector (shape:`N`). Note that for cosine similarity it currently returns **normalized** vectors.
+* `get_items(ids, return_type = 'numpy')` - returns a numpy array (shape:`N*dim`) of vectors that have integer identifiers specified in `ids` numpy vector (shape:`N`) if `return_type` is `list` return list of lists. Note that for cosine similarity it currently returns **normalized** vectors.
 
 * `get_ids_list()` - returns a list of all elements' ids.
 

python_bindings/bindings.cpp

@@ -304,7 +304,11 @@ class Index {
  }
 
 
- std::vector<std::vector<data_t>> getDataReturnList(py::object ids_ = py::none()) {
+ py::object getData(py::object ids_ = py::none(), std::string return_type = "numpy") {
+ std::vector<std::string> return_types{"numpy", "list"};
+ if (std::find(std::begin(return_types), std::end(return_types), return_type) == std::end(return_types)) {
+ throw std::invalid_argument("return_type should be \"numpy\" or \"list\"");
+ }
  std::vector<size_t> ids;
  if (!ids_.is_none()) {
  py::array_t < size_t, py::array::c_style | py::array::forcecast > items(ids_);
@@ -325,7 +329,12 @@ class Index {
  for (auto id : ids) {
  data.push_back(appr_alg->template getDataByLabel<data_t>(id));
  }
- return data;
+ if (return_type == "list") {
+ return py::cast(data);
+ }
+ if (return_type == "numpy") {
+ return py::array_t< data_t, py::array::c_style | py::array::forcecast >(py::cast(data));
+ }
  }
 
 
@@ -925,7 +934,7 @@ PYBIND11_PLUGIN(hnswlib) {
  py::arg("ids") = py::none(),
  py::arg("num_threads") = -1,
  py::arg("replace_deleted") = false)
- .def("get_items", &Index<float, float>::getDataReturnList, py::arg("ids") = py::none())
+ .def("get_items", &Index<float>::getData, py::arg("ids") = py::none(), py::arg("return_type") = "numpy")
  .def("get_ids_list", &Index<float>::getIdsList)
  .def("set_ef", &Index<float>::set_ef, py::arg("ef"))
  .def("set_num_threads", &Index<float>::set_num_threads, py::arg("num_threads"))

tests/python/bindings_test_getdata.py

@@ -45,5 +45,11 @@ def testGettingItems(self):
  self.assertRaises(ValueError, lambda: p.get_items(labels[0]))
 
  # After adding them, all labels should be retrievable
- returned_items = p.get_items(labels)
- self.assertSequenceEqual(data.tolist(), returned_items)
+ returned_items_np = p.get_items(labels)
+ self.assertTrue((data == returned_items_np).all())
+
+ # check returned type of get_items
+ self.assertTrue(isinstance(returned_items_np, np.ndarray))
+ returned_items_list = p.get_items(labels, return_type="list")
+ self.assertTrue(isinstance(returned_items_list, list))
+ self.assertTrue(isinstance(returned_items_list[0], list))

tests/python/bindings_test_replace.py

@@ -94,10 +94,10 @@ def testRandomSelf(self):
  remaining_data = comb_data[remaining_labels_list]
 
  returned_items = hnsw_index.get_items(remaining_labels_list)
- self.assertSequenceEqual(remaining_data.tolist(), returned_items)
+ self.assertTrue((remaining_data == returned_items).all())
 
  returned_items = hnsw_index.get_items(labels3_tr)
- self.assertSequenceEqual(data3_tr.tolist(), returned_items)
+ self.assertTrue((data3_tr == returned_items).all())
 
  # Check index serialization
  # Delete batch 3