pybind: replaced input type to bytearray - MOD-4597 (#304)

* replaced input type to bytearray

* added a module to convert numpy to bytearay

* moved mybytearray to python bindings lib

Author: meiravgri

setup.py

@@ -82,7 +82,6 @@ def build_extension(self, ext):
             ["cmake", "--build", "."] + build_args, cwd=self.build_temp
         )
 
-
 # The information here can also be placed in setup.cfg - better separation of
 # logic and declaration, and simpler if you include description/version in a file.
 setup(
@@ -93,5 +92,6 @@ def build_extension(self, ext):
     description="Python library around collection of vector similarity algorithm",
     long_description="",
     ext_modules=[CMakeExtension("VecSim", "src/python_bindings")],
+    py_modules=['src/python_bindings/Mybytearray'],
     cmdclass={"build_ext": CMakeBuild}
 )

src/python_bindings/Mybytearray.py

@@ -0,0 +1,3 @@
+
+def create_bytearray(np_arr):
+        return bytearray(np_arr)

src/python_bindings/bindings.cpp

@@ -69,24 +69,32 @@ class PyBatchIterator {
     void reset() { VecSimBatchIterator_Reset(batchIterator.get()); }
     virtual ~PyBatchIterator() {}
 };
+// @input or @query arguments are a py::object object. (numpy arrays are acceptable)
+
+// To convert input or query to a pointer use input_to_blob(input)
+// For example:
+// VecSimIndex_AddVector(index, input_to_blob(input), id);
 
 class PyVecSimIndex {
 public:
-    PyVecSimIndex() {}
-
-    PyVecSimIndex(const VecSimParams &params) { index = VecSimIndex_New(&params); }
+    PyVecSimIndex()
+        : create_bytearray(
+              py::module::import("src.python_bindings.Mybytearray").attr("create_bytearray")) {}
+
+    PyVecSimIndex(const VecSimParams &params)
+        : create_bytearray(
+              py::module::import("src.python_bindings.Mybytearray").attr("create_bytearray")) {
+        index = VecSimIndex_New(&params);
+    }
 
     void addVector(py::object input, size_t id) {
-        py::array_t<float, py::array::c_style | py::array::forcecast> items(input);
-        VecSimIndex_AddVector(index, (void *)items.data(0), id);
+        VecSimIndex_AddVector(index, input_to_blob(input), id);
     }
-
     void deleteVector(size_t id) { VecSimIndex_DeleteVector(index, id); }
 
     py::object knn(py::object input, size_t k, VecSimQueryParams *query_params) {
-        py::array_t<float, py::array::c_style | py::array::forcecast> items(input);
         VecSimQueryResult_List res =
-            VecSimIndex_TopKQuery(index, (void *)items.data(0), k, query_params, BY_SCORE);
+            VecSimIndex_TopKQuery(index, input_to_blob(input), k, query_params, BY_SCORE);
         if (VecSimQueryResult_Len(res) != k) {
             throw std::runtime_error("Cannot return the results in a contiguous 2D array. Probably "
                                      "ef or M is too small");
@@ -95,27 +103,32 @@ class PyVecSimIndex {
     }
 
     py::object range(py::object input, double radius, VecSimQueryParams *query_params) {
-        py::array_t<float, py::array::c_style | py::array::forcecast> items(input);
         VecSimQueryResult_List res =
-            VecSimIndex_RangeQuery(index, (void *)items.data(0), radius, query_params, BY_SCORE);
+            VecSimIndex_RangeQuery(index, input_to_blob(input), radius, query_params, BY_SCORE);
         return wrap_results(res, VecSimQueryResult_Len(res));
     }
 
     size_t indexSize() { return VecSimIndex_IndexSize(index); }
 
-    PyBatchIterator createBatchIterator(py::object &query_blob, VecSimQueryParams *query_params) {
-        py::array_t<float, py::array::c_style | py::array::forcecast> items(query_blob);
-        float *vector_data = (float *)items.data(0);
-        return PyBatchIterator(VecSimBatchIterator_New(index, vector_data, query_params));
+    PyBatchIterator createBatchIterator(py::object input, VecSimQueryParams *query_params) {
+        return PyBatchIterator(VecSimBatchIterator_New(index, input_to_blob(input), query_params));
     }
 
     virtual ~PyVecSimIndex() { VecSimIndex_Free(index); }
 
 protected:
     VecSimIndex *index;
+
+private:
+    // save the bytearray to keep its pointer valid
+    py::bytearray tmp_bytearray;
+    const py::function create_bytearray;
+    const char *input_to_blob(py::object input) {
+        tmp_bytearray = create_bytearray(input);
+        return PyByteArray_AS_STRING(tmp_bytearray.ptr());
+    }
 };
 
-// Currently supports only floats. TODO change after serializer refactoring
 class PyHNSWLibIndex : public PyVecSimIndex {
 public:
     PyHNSWLibIndex(const HNSWParams &hnsw_params) {

tests/flow/test_bruteforce.py

@@ -4,6 +4,60 @@
 
 from common import *
 
+def test_sanity_bf():
+    class TestData:
+        def __init__(self, data_type, metric, dist_func, np_fuc):
+            dim = 16
+            num_elements = 10
+            params = VecSimParams()
+            bfparams = BFParams()
+
+            params.algo = VecSimAlgo_BF
+            bfparams.initialCapacity = num_elements
+            bfparams.blockSize = num_elements
+            bfparams.dim = dim
+            bfparams.type = data_type
+            bfparams.metric = metric
+
+            params.bfParams = bfparams
+
+            self.index = VecSimIndex(params)
+
+            self.metric = metric
+            self.type = data_type
+            self.dist_func = dist_func
+
+            np.random.seed(47)
+            self.data = np_fuc(np.random.random((num_elements, dim)))
+            self.query = np_fuc(np.random.random((1, dim)))
+            self.vectors = []
+            for i, vector in enumerate(self.data):
+                self.vectors.append((i, vector))
+                self.index.add_vector(vector, i)
+
+        def measure_dists(self, k):
+            dists = [(self.dist_func(self.query.flat, vec), key) for key, vec in self.vectors]
+            dists = sorted(dists)[:k]
+            keys = [key for _, key in dists]
+            dists = [dist for dist, _ in dists]
+            return (keys, dists)       
+    
+    test_datas = []
+
+    dist_funcs = [(VecSimMetric_Cosine, spatial.distance.cosine), (VecSimMetric_L2, spatial.distance.sqeuclidean)]
+    types = [(VecSimType_FLOAT32, np.float32), (VecSimType_FLOAT64, np.float64)]
+    for type_name, np_type in types:
+        for dist_name, dist_func in dist_funcs:
+            test_datas.append(TestData(type_name, dist_name, dist_func, np_type))
+
+    k = 10
+    for test_data in test_datas:
+
+        keys, dists = test_data.measure_dists(k)
+        bf_labels, bf_distances = test_data.index.knn_query(test_data.query, k=k)
+        assert_allclose(bf_labels, [keys],  rtol=1e-5, atol=0)
+        assert_allclose(bf_distances, [dists],  rtol=1e-5, atol=0)
+        print(f"\nsanity test for {test_data.metric} and {test_data.type} pass")
 
 def test_bf_cosine():
     dim = 128