Merge pull request #494 from dyashuni/get_items_numpy

Fix get_items

Author: yurymalkov

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