Fix col-major array bug in Python binding

Fixes nmslib#230.
ViperCraft · Aug 29, 2017 · 4788374 · 4788374
1 parent d327e79
commit 4788374
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/python_bindings/nmslib.cc b/python_bindings/nmslib.cc
@@ -208,7 +208,7 @@ struct IndexWrapper {
 
     } else if (data_type == DATATYPE_DENSE_VECTOR) {
       // allow numpy arrays to be returned here too
-      py::array_t<dist_t> items(input);
+      py::array_t<dist_t, py::array::c_style | py::array::forcecast> items(input);
       auto buffer = items.request();
       if (buffer.ndim != 2) throw std::runtime_error("data must be a 2d array");
 

diff --git a/python_bindings/tests/bindings_test.py b/python_bindings/tests/bindings_test.py
@@ -47,6 +47,13 @@ def testKnnQueryBatch(self):
         for query, (ids, distances) in zip(queries, results):
             self.assertTrue(get_hitrate(get_exact_cosine(query, data), ids) >= 5)
 
+        # test col-major arrays
+        queries = np.asfortranarray(queries)
+        results = index.knnQueryBatch(queries, k=10)
+        for query, (ids, distances) in zip(queries, results):
+            self.assertTrue(get_hitrate(get_exact_cosine(query, data), ids) >= 5)
+
+
     def testReloadIndex(self):
         np.random.seed(23)
         data = np.random.randn(1000, 10).astype(np.float32)