Merge branch 'feat/add-search-radius' into 'main' (merge request !70)

feat: add search radius

Author: rogersqsliu

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## v1.5.4
+* feat: support search radius for searching in collection
+ 
 ## v1.5.3
 * fix: optimize warning message for modifing vector indexes interface
 

tcvectordb/api/document/api.go

@@ -178,6 +178,7 @@ type SearchCond struct {
 	Vectors        [][]float32   `json:"vectors,omitempty"`
 	Filter         string        `json:"filter,omitempty"`
 	EmbeddingItems []string      `json:"embeddingItems,omitempty"`
+	Radius         *float32      `json:"radius,omitempty"` // 距离阈值,范围搜索时有效
 }
 
 type SearchParams struct {

tcvectordb/base_document.go

@@ -181,12 +181,17 @@ func (i *implementerDocument) Query(ctx context.Context, documentIds []string, p
 //     If RetrieveVector is true, the vector values will be returned.
 //   - OutputFields: (Optional) Return columns specified by the list of column names.
 //   - Limit: (Required) Limit the number of documents returned (default to 1).
+//   - Radius: (Optional) Specifies the radius range for similarity retrieval.
+//     IP: return when score >= radius, value range (-∞, +∞).
+//     COSINE: return when score >= radius, value range [-1, 1].
+//     L2: return when score <= radius, value range [0, +∞).
 type SearchDocumentParams struct {
 	Filter         *Filter
 	Params         *SearchDocParams
 	RetrieveVector bool
 	OutputFields   []string
 	Limit          int64
+	Radius         *float32
 }
 
 // [SearchDocParams] holds the parameters for searching documents to a collection.
@@ -196,10 +201,10 @@ type SearchDocumentParams struct {
 //     of vectors to be accessed. Valid range is [1, nlist], and nlist is defined by creating collection.
 //   - Ef: (Optional) HNSW type index requires configuration parameter ef to specify the number
 //     of vectors to be accessed (default to 10). Valid range is [1, 32768]
-//   - Radius: (Optional) Specifies the radius range for similarity retrieval.
 type SearchDocParams struct {
-	Nprobe uint32  `json:"nprobe,omitempty"` // 搜索时查找的聚类数量，使用索引默认值即可
-	Ef     uint32  `json:"ef,omitempty"`     // HNSW
+	Nprobe uint32 `json:"nprobe,omitempty"` // 搜索时查找的聚类数量，使用索引默认值即可
+	Ef     uint32 `json:"ef,omitempty"`     // HNSW
+	// Deprecated:  Radius is deprecated and should not be used, which is supported in [SearchDocumentParams].
 	Radius float32 `json:"radius,omitempty"` // 距离阈值,范围搜索时有效
 }
 
@@ -694,11 +699,13 @@ func (i *implementerFlatDocument) search(ctx context.Context, databaseName, coll
 		req.Search.OutputFields = param.OutputFields
 		req.Search.Limit = param.Limit
 
+		req.Search.Params = new(document.SearchParams)
 		if param.Params != nil {
-			req.Search.Params = new(document.SearchParams)
 			req.Search.Params.Nprobe = param.Params.Nprobe
 			req.Search.Params.Ef = param.Params.Ef
-			req.Search.Params.Radius = param.Params.Radius
+		}
+		if param.Radius != nil {
+			req.Search.Radius = param.Radius
 		}
 	}
 

tcvectordb/rpc_base_document.go

@@ -689,8 +689,15 @@ func (r *rpcImplementerFlatDocument) search(ctx context.Context, databaseName, c
 			req.Search.Params = &olama.SearchParams{
 				Nprobe: param.Params.Nprobe,
 				Ef:     param.Params.Ef,
-				Radius: param.Params.Radius,
 			}
+
+		}
+		if param.Radius != nil {
+			if req.Search.Params == nil {
+				req.Search.Params = new(olama.SearchParams)
+			}
+			req.Search.Range = true
+			req.Search.Params.Radius = *param.Radius
 		}
 	}
 	res, err := r.rpcClient.Search(ctx, req)

tcvectordb/version.go

@@ -18,4 +18,4 @@
 
 package tcvectordb
 
-const SDKVersion = "v1.5.3"
+const SDKVersion = "v1.5.4"

test/normal_test.go

@@ -298,13 +298,17 @@ func TestQuery(t *testing.T) {
 func TestSearch(t *testing.T) {
 	col := cli.Database(database).Collection(collectionName)
 
+	radius := float32(0.76)
 	searchRes, err := col.Search(ctx, [][]float32{
-		{0.3123, 0.43, 0.213},
+		//{0.3123, 0.43, 0.213},
 		{0.233, 0.12, 0.97},
 	}, &tcvectordb.SearchDocumentParams{
-		Params:         &tcvectordb.SearchDocParams{Ef: 100},
+		Params: &tcvectordb.SearchDocParams{
+			Ef: 100,
+		},
+		Radius:         &radius,
 		RetrieveVector: false,
-		Limit:          10,
+		Limit:          2,
 	})
 	printErr(err)
 	log.Printf("search by vector-----------------")