add nprobe option

apis/docs/v1/docs.md

@@ -894,6 +894,7 @@ Represent search configuration.
 | min_num               | [uint32](#uint32)                                                      |       | Minimum number of result to be returned.     |
 | aggregation_algorithm | [Search.AggregationAlgorithm](#payload-v1-Search-AggregationAlgorithm) |       | Aggregation Algorithm                        |
 | ratio                 | [google.protobuf.FloatValue](#google-protobuf-FloatValue)              |       | Search ratio for agent return result number. |
+| nprobe                | [uint32](#uint32)                                                      |       | Search nprobe.                               |
 
 <a name="payload-v1-Search-IDRequest"></a>
 

apis/grpc/v1/payload/payload_vtproto.pb.go

@@ -184,6 +184,7 @@ func (m *Search_Config) CloneVT() *Search_Config {
 	r.MinNum = m.MinNum
 	r.AggregationAlgorithm = m.AggregationAlgorithm
 	r.Ratio = (*wrapperspb.FloatValue)((*wrapperspb1.FloatValue)(m.Ratio).CloneVT())
+	r.Nprobe = m.Nprobe
 	if len(m.unknownFields) > 0 {
 		r.unknownFields = make([]byte, len(m.unknownFields))
 		copy(r.unknownFields, m.unknownFields)
@@ -2223,6 +2224,9 @@ func (this *Search_Config) EqualVT(that *Search_Config) bool {
 	if !(*wrapperspb1.FloatValue)(this.Ratio).EqualVT((*wrapperspb1.FloatValue)(that.Ratio)) {
 		return false
 	}
+	if this.Nprobe != that.Nprobe {
+		return false
+	}
 	return string(this.unknownFields) == string(that.unknownFields)
 }
 
@@ -5166,6 +5170,11 @@ func (m *Search_Config) MarshalToSizedBufferVT(dAtA []byte) (int, error) {
 		i -= len(m.unknownFields)
 		copy(dAtA[i:], m.unknownFields)
 	}
+	if m.Nprobe != 0 {
+		i = protohelpers.EncodeVarint(dAtA, i, uint64(m.Nprobe))
+		i--
+		dAtA[i] = 0x58
+	}
 	if m.Ratio != nil {
 		size, err := (*wrapperspb1.FloatValue)(m.Ratio).MarshalToSizedBufferVT(dAtA[:i])
 		if err != nil {
@@ -9783,6 +9792,9 @@ func (m *Search_Config) SizeVT() (n int) {
 		l = (*wrapperspb1.FloatValue)(m.Ratio).SizeVT()
 		n += 1 + l + protohelpers.SizeOfVarint(uint64(l))
 	}
+	if m.Nprobe != 0 {
+		n += 1 + protohelpers.SizeOfVarint(uint64(m.Nprobe))
+	}
 	n += len(m.unknownFields)
 	return n
 }
@@ -12382,6 +12394,25 @@ func (m *Search_Config) UnmarshalVT(dAtA []byte) error {
 				return err
 			}
 			iNdEx = postIndex
+		case 11:
+			if wireType != 0 {
+				return fmt.Errorf("proto: wrong wireType = %d for field Nprobe", wireType)
+			}
+			m.Nprobe = 0
+			for shift := uint(0); ; shift += 7 {
+				if shift >= 64 {
+					return protohelpers.ErrIntOverflow
+				}
+				if iNdEx >= l {
+					return io.ErrUnexpectedEOF
+				}
+				b := dAtA[iNdEx]
+				iNdEx++
+				m.Nprobe |= uint32(b&0x7F) << shift
+				if b < 0x80 {
+					break
+				}
+			}
 		default:
 			iNdEx = preIndex
 			skippy, err := protohelpers.Skip(dAtA[iNdEx:])

apis/proto/v1/payload/payload.proto

@@ -96,6 +96,8 @@ message Search {
     AggregationAlgorithm aggregation_algorithm = 9;
     // Search ratio for agent return result number.
     google.protobuf.FloatValue ratio = 10;
+    // Search nprobe.
+    uint32 nprobe = 11;
   }
 
   // AggregationAlgorithm is enum of each aggregation algorithms

apis/swagger/v1/vald/filter.swagger.json

@@ -550,6 +550,11 @@
           "type": "number",
           "format": "float",
           "description": "Search ratio for agent return result number."
+        },
+        "nprobe": {
+          "type": "integer",
+          "format": "int64",
+          "description": "Search nprobe."
         }
       },
       "description": "Represent search configuration."

apis/swagger/v1/vald/search.swagger.json

@@ -464,6 +464,11 @@
           "type": "number",
           "format": "float",
           "description": "Search ratio for agent return result number."
+        },
+        "nprobe": {
+          "type": "integer",
+          "format": "int64",
+          "description": "Search nprobe."
         }
       },
       "description": "Represent search configuration."

internal/core/algorithm/faiss/Capi.cpp

@@ -47,10 +47,8 @@ FaissStruct* faiss_create_index(
   switch (method_type) {
     case IVFPQ:
       return faiss_create_index_ivfpq(d, nlist, m, nbits_per_idx, metric_type);
-      break;
     case BINARYIVF:
       return faiss_create_index_binaryivf(d*8, nlist);
-      break;
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";
@@ -75,10 +73,8 @@ FaissStruct* faiss_create_index_ivfpq(
     switch (metric_type) {
       case faiss::METRIC_INNER_PRODUCT:
         quantizer = new faiss::IndexFlat(d, faiss::METRIC_INNER_PRODUCT);
-        break;
       case faiss::METRIC_L2:
         quantizer = new faiss::IndexFlat(d, faiss::METRIC_L2);
-        break;
       default:
         std::stringstream ss;
         ss << "Capi : " << __FUNCTION__ << "() : Error: no metric type.";
@@ -134,10 +130,8 @@ FaissStruct* faiss_read_index(const char* fname, const int method_type) {
   switch (method_type) {
     case IVFPQ:
       return faiss_read_index_ivfpq(fname);
-      break;
     case BINARYIVF:
       return faiss_read_index_binaryindex(fname);
-      break;
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";
@@ -197,10 +191,8 @@ bool faiss_write_index(
   switch (method_type) {
     case IVFPQ:
       return faiss_write_index_ivfpq(st, fname);
-      break;
     case BINARYIVF:
       return faiss_write_index_binaryivf(st, fname);
-      break;
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";
@@ -261,10 +253,8 @@ bool faiss_train(
   switch (method_type) {
     case IVFPQ:
       return faiss_train_ivfpq(st, nb, xb);
-      break;
     case BINARYIVF:
       return faiss_train_binaryivf(st, nb, reinterpret_cast<const uint8_t*>(xb));
-      break;
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";
@@ -332,10 +322,8 @@ int faiss_add(
   switch (method_type) {
     case IVFPQ:
       return faiss_add_ivfpq(st, nb, xb, xids);
-      break;
     case BINARYIVF:
       return faiss_add_binaryivf(st, nb, reinterpret_cast<const uint8_t*>(xb), xids);
-      break;
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";
@@ -395,6 +383,7 @@ int faiss_add_binaryivf(
 bool faiss_search(
     const FaissStruct* st,
     const int k,
+    const int nprobe,
     const int nq,
     const float* xq,
     long* I,
@@ -406,11 +395,9 @@ bool faiss_search(
 
   switch (method_type) {
     case IVFPQ:
-      return faiss_search_ivfpq(st, k, nq, xq, I, D);
-      break;
+      return faiss_search_ivfpq(st, k, nprobe, nq, xq, I, D);
     case BINARYIVF:
-      return faiss_search_binaryivf(st, k, nq, reinterpret_cast<const uint8_t*>(xq), I, D);
-      break;
+      return faiss_search_binaryivf(st, k, nprobe, nq, reinterpret_cast<const uint8_t*>(xq), I, D);
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";
@@ -422,6 +409,7 @@ bool faiss_search(
 bool faiss_search_ivfpq(
     const FaissStruct* st,
     const int k,
+    const int nprobe,
     const int nq,
     const float* xq,
     long* I,
@@ -433,6 +421,7 @@ bool faiss_search_ivfpq(
   try {
     //printf("is_trained: %d\n", (static_cast<faiss::IndexIVFPQ*>(st->faiss_index))->is_trained);
     //printf("ntotal: %ld\n", (static_cast<faiss::IndexIVFPQ*>(st->faiss_index))->ntotal);
+    (static_cast<faiss::IndexIVFPQ*>(st->faiss_index))->nprobe = nprobe;
     (static_cast<faiss::IndexIVFPQ*>(st->faiss_index))->search(nq, xq, k, D, I);
     //printf("I=\n");
     //for(int i = 0; i < nq; i++) {
@@ -461,6 +450,7 @@ bool faiss_search_ivfpq(
 bool faiss_search_binaryivf(
     const FaissStruct* st,
     const int k,
+    const int nprobe,
     const int nq,
     const uint8_t* xq,
     long* I,
@@ -471,6 +461,7 @@ bool faiss_search_binaryivf(
 
   int32_t* tmpD = new int32_t[nq*k];
   try {
+    (static_cast<faiss::IndexBinaryIVF*>(st->faiss_index))->nprobe = nprobe;
     (static_cast<faiss::IndexBinaryIVF*>(st->faiss_index))->search(nq, xq, k, tmpD, I);
   } catch(std::exception &err) {
     std::stringstream ss;
@@ -497,10 +488,8 @@ int faiss_remove(
   switch (method_type) {
     case IVFPQ:
       return faiss_remove_ivfpq(st, size, ids);
-      break;
     case BINARYIVF:
       return faiss_remove_binaryivf(st, size, ids);
-      break;
     default:
       std::stringstream ss;
       ss << "Capi : " << __FUNCTION__ << "() : Error: no method type.";

internal/core/algorithm/faiss/Capi.h

@@ -90,6 +90,7 @@ extern "C" {
   bool faiss_search(
       const FaissStruct* st,
       const int k,
+      const int nprobe,
       const int nq,
       const float* xq,
       long* I,
@@ -98,13 +99,15 @@ extern "C" {
   bool faiss_search_ivfpq(
       const FaissStruct* st,
       const int k,
+      const int nprobe,
       const int nq,
       const float* xq,
       long* I,
       float* D);
   bool faiss_search_binaryivf(
       const FaissStruct* st,
       const int k,
+      const int nprobe,
       const int nq,
       const uint8_t* xq,
       long* I,

internal/core/algorithm/faiss/faiss.go

@@ -47,7 +47,7 @@ type (
 		Add(nb int, xb []float32, xids []int64) (int, error)
 
 		// Search returns search result as []algorithm.SearchResult.
-		Search(k, nq int, xq []float32) ([]algorithm.SearchResult, error)
+		Search(k, nprobe, nq int, xq []float32) ([]algorithm.SearchResult, error)
 
 		// Remove removes from faiss index.
 		Remove(size int, ids []int64) (int, error)
@@ -215,15 +215,15 @@ func (f *faiss) Add(nb int, xb []float32, xids []int64) (int, error) {
 }
 
 // Search returns search result as []algorithm.SearchResult.
-func (f *faiss) Search(k, nq int, xq []float32) ([]algorithm.SearchResult, error) {
+func (f *faiss) Search(k, nprobe, nq int, xq []float32) ([]algorithm.SearchResult, error) {
 	if len(xq) != nq*int(f.dimension) {
 		return nil, errors.ErrIncompatibleDimensionSize(len(xq), int(f.dimension))
 	}
 
 	I := make([]int64, k*nq)
 	D := make([]float32, k*nq)
 	f.mu.RLock()
-	ret := C.faiss_search(f.st, (C.int)(k), (C.int)(nq), (*C.float)(&xq[0]), (*C.long)(&I[0]), (*C.float)(&D[0]), C.int(f.methodType))
+	ret := C.faiss_search(f.st, (C.int)(k), (C.int)(nprobe), (C.int)(nq), (*C.float)(&xq[0]), (*C.long)(&I[0]), (*C.float)(&D[0]), C.int(f.methodType))
 	f.mu.RUnlock()
 	if ret == ErrorCode {
 		return nil, errors.NewFaissError("failed to faiss_search")

pkg/agent/core/faiss/handler/grpc/search.go

@@ -69,6 +69,7 @@ func (s *server) Search(
 	}
 	res, err = s.faiss.Search(
 		req.GetConfig().GetNum(),
+		req.GetConfig().GetNprobe(),
 		1,
 		req.GetVector())
 	if err == nil && res == nil {

pkg/agent/core/faiss/service/faiss.go

@@ -57,7 +57,7 @@ type (
 		CreateIndex(ctx context.Context) error
 		SaveIndex(ctx context.Context) error
 		CreateAndSaveIndex(ctx context.Context) error
-		Search(k, nq uint32, xq []float32) (*payload.Search_Response, error)
+		Search(k, nprobe, nq uint32, xq []float32) (*payload.Search_Response, error)
 		Delete(uuid string) error
 		DeleteWithTime(uuid string, t int64) error
 		Exists(uuid string) (uint32, bool)
@@ -1123,12 +1123,17 @@ func (f *faiss) CreateAndSaveIndex(ctx context.Context) error {
 	return f.SaveIndex(ctx)
 }
 
-func (f *faiss) Search(k, nq uint32, xq []float32) (res *payload.Search_Response, err error) {
+func (f *faiss) Search(
+	k, nprobe, nq uint32, xq []float32,
+) (res *payload.Search_Response, err error) {
 	if f.IsIndexing() {
 		return nil, errors.ErrCreateIndexingIsInProgress
 	}
+	if nprobe == 0 {
+		nprobe = 1
+	}
 
-	sr, err := f.core.Search(int(k), int(nq), xq)
+	sr, err := f.core.Search(int(k), int(nprobe), int(nq), xq)
 	if err != nil {
 		if f.IsIndexing() {
 			return nil, errors.ErrCreateIndexingIsInProgress

rust/libs/proto/src/payload.v1.rs

@@ -113,6 +113,9 @@ pub mod search {
         /// Search ratio for agent return result number.
         #[prost(message, optional, tag="10")]
         pub ratio: ::core::option::Option<f32>,
+        /// Search nprobe.
+        #[prost(uint32, tag="11")]
+        pub nprobe: u32,
     }
     /// Represent a search response.
     #[allow(clippy::derive_partial_eq_without_eq)]