Add pyknowhere API to support binary GetVectorByIds (#158)

Signed-off-by: Yudong Cai <[email protected]>
zilliztech · Oct 20, 2023 · f4c1757 · f4c1757
1 parent ce65e5e
commit f4c1757
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/python/knowhere/__init__.py b/python/knowhere/__init__.py
@@ -87,3 +87,11 @@ def GetVectorDataSetToArray(ans):
     data = np.zeros([rows, dim]).astype(np.float32)
     swigknowhere.DataSetTensor2Array(ans, data)
     return data
+
+
+def GetBinaryVectorDataSetToArray(ans):
+    dim = int(swigknowhere.DataSet_Dim(ans) / 32)
+    rows = swigknowhere.DataSet_Rows(ans)
+    data = np.zeros([rows, dim]).astype(np.int32)
+    swigknowhere.BinaryDataSetTensor2Array(ans, data)
+    return data
diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i
@@ -69,6 +69,7 @@ import_array();
 %apply (float* INPLACE_ARRAY2, int DIM1, int DIM2){(float *dis,int nq_1,int k_1)}
 %apply (int *INPLACE_ARRAY2, int DIM1, int DIM2){(int *ids,int nq_2,int k_2)}
 %apply (float* INPLACE_ARRAY2, int DIM1, int DIM2){(float *data,int rows,int dim)}
+%apply (int32_t *INPLACE_ARRAY2, int DIM1, int DIM2){(int32_t *data,int rows,int dim)}
 
 %typemap(in, numinputs=0) knowhere::Status& status(knowhere::Status tmp) %{
     $1 = &tmp;
@@ -329,6 +330,17 @@ DataSetTensor2Array(knowhere::DataSetPtr result, float* data, int rows, int dim)
     }
 }
 
+void
+BinaryDataSetTensor2Array(knowhere::DataSetPtr result, int32_t* data, int rows, int dim) {
+    GILReleaser rel;
+    auto data_ = result->GetTensor();
+    for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < dim; ++j) {
+            *(data + i * dim + j) = *((int32_t*)(data_) + i * dim + j);
+        }
+    }
+}
+
 void
 DumpRangeResultIds(knowhere::DataSetPtr result, int* ids, int len) {
     GILReleaser rel;