diff --git a/python/knowhere/__init__.py b/python/knowhere/__init__.py index 3b194a0c8..9d3c77d3b 100644 --- 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 index ee47cf328..ddcb5be57 100644 --- 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;