Binary vector from/to np.uint8 array in pyknowhere api (#550)

Signed-off-by: chasingegg <[email protected]>
zilliztech · May 8, 2024 · 7db6314 · 7db6314
1 parent 8c91025
commit 7db6314
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/python/knowhere/__init__.py b/python/knowhere/__init__.py
@@ -73,7 +73,7 @@ def ArrayToDataSet(arr):
     if arr.ndim == 1:
         return swigknowhere.Array2DataSetIds(arr)
     if arr.ndim == 2:
-        if arr.dtype == np.int32:
+        if arr.dtype == np.uint8:
             return swigknowhere.Array2DataSetI(arr)
         if arr.dtype == np.float32:
             return swigknowhere.Array2DataSetF(arr)
@@ -85,7 +85,7 @@ def ArrayToDataSet(arr):
             return swigknowhere.Array2DataSetBF16(arr)
     raise ValueError(
         """
-        ArrayToDataSet only support numpy array dtype float32,int32,float16 and bfloat16.
+        ArrayToDataSet only support numpy array dtype float32,uint8,float16 and bfloat16.
         """
     )
 
@@ -156,9 +156,9 @@ def GetBFloat16VectorDataSetToArray(ans):
     return data
 
 def GetBinaryVectorDataSetToArray(ans):
-    dim = int(swigknowhere.DataSet_Dim(ans) / 32)
+    dim = int(swigknowhere.DataSet_Dim(ans) / 8)
     rows = swigknowhere.DataSet_Rows(ans)
-    data = np.zeros([rows, dim]).astype(np.int32)
+    data = np.zeros([rows, dim]).astype(np.uint8)
     swigknowhere.BinaryDataSetTensor2Array(ans, data)
     return data
 

diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i
@@ -73,6 +73,8 @@ import_array();
 %apply (float* IN_ARRAY2, int DIM1, int DIM2) {(float* xb, int nb, int dim)}
 %apply (int* IN_ARRAY2, int DIM1, int DIM2) {(int* xb, int nb, int dim)}
 %apply (uint8_t *IN_ARRAY1, int DIM1) {(uint8_t *block, int size)}
+%apply (uint8_t* IN_ARRAY2, int DIM1, int DIM2) {(uint8_t* xb, int nb, int dim)}
+%apply (uint8_t* INPLACE_ARRAY2, int DIM1, int DIM2) {(uint8_t *data,int rows,int dim)}
 %apply (int *IN_ARRAY1, int DIM1) {(int *lims, int len)}
 %apply (int *IN_ARRAY1, int DIM1) {(int *ids, int len)}
 %apply (float *IN_ARRAY1, int DIM1) {(float *dis, int len)}
@@ -354,11 +356,11 @@ CurrentVersion() {
 }
 
 knowhere::DataSetPtr
-Array2DataSetI(int *xb, int nb, int dim){
+Array2DataSetI(uint8_t* xb, int nb, int dim) {
     auto ds = std::make_shared<DataSet>();
     ds->SetIsOwner(false);
     ds->SetRows(nb);
-    ds->SetDim(dim*32);
+    ds->SetDim(dim*8);
     ds->SetTensor(xb);
     return ds;
 };
@@ -469,12 +471,12 @@ BFloat16DataSetTensor2Array(knowhere::DataSetPtr result, float* data, int rows,
 }
 
 void
-BinaryDataSetTensor2Array(knowhere::DataSetPtr result, int32_t* data, int rows, int dim) {
+BinaryDataSetTensor2Array(knowhere::DataSetPtr result, uint8_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);
+            *(data + i * dim + j) = *((uint8_t*)(data_) + i * dim + j);
         }
     }
 }