Add test cases for fp16, bf16, hybrid search

Signed-off-by: ThreadDao <[email protected]>
milvus-io · Mar 29, 2024 · f21ccdd · f21ccdd
1 parent d43cd42
commit f21ccdd
Show file tree

Hide file tree

Showing 19 changed files with 1,086 additions and 173 deletions.
diff --git a/client/ann_request.go b/client/ann_request.go
@@ -20,11 +20,12 @@ type ANNSearchRequest struct {
 	limit       int
 }
 
-func NewANNSearchRequest(fieldName string, metricsType entity.MetricType, vectors []entity.Vector, searchParam entity.SearchParam, limit int, options ...SearchQueryOptionFunc) *ANNSearchRequest {
+func NewANNSearchRequest(fieldName string, metricsType entity.MetricType, expr string, vectors []entity.Vector, searchParam entity.SearchParam, limit int, options ...SearchQueryOptionFunc) *ANNSearchRequest {
 	return &ANNSearchRequest{
 		fieldName:   fieldName,
 		vectors:     vectors,
 		metricType:  metricsType,
+		expr:        expr,
 		searchParam: searchParam,
 		limit:       limit,
 		options:     options,

diff --git a/examples/multivectors/main.go b/examples/multivectors/main.go
@@ -139,8 +139,8 @@ func main() {
 
 	result, err = c.HybridSearch(ctx, collectionName, nil, topK, []string{keyCol, embeddingCol1, embeddingCol2},
 		client.NewRRFReranker(), []*client.ANNSearchRequest{
-			client.NewANNSearchRequest(embeddingCol1, entity.L2, vec2search1, sp, topK),
-			client.NewANNSearchRequest(embeddingCol2, entity.L2, vec2search2, sp, topK),
+			client.NewANNSearchRequest(embeddingCol1, entity.L2, "", vec2search1, sp, topK),
+			client.NewANNSearchRequest(embeddingCol2, entity.L2, "", vec2search2, sp, topK),
 		})
 	if err != nil {
 		log.Fatalf("failed to search collection, err: %v", err)

diff --git a/test/base/milvus_client.go b/test/base/milvus_client.go
@@ -410,6 +410,17 @@ func (mc *MilvusClient) Search(ctx context.Context, collName string, partitions
 	return searchResult, err
 }
 
+func (mc *MilvusClient) HybridSearch(ctx context.Context, collName string, partitions []string, limit int, outputFields []string,
+	reranker client.Reranker, subRequests []*client.ANNSearchRequest) ([]client.SearchResult, error) {
+	funcName := "HybridSearch"
+	preRequest(funcName, ctx, collName, partitions, limit, outputFields, reranker, subRequests)
+
+	searchResult, err := mc.mClient.HybridSearch(ctx, collName, partitions, limit, outputFields, reranker, subRequests)
+	postResponse(funcName, err, searchResult)
+
+	return searchResult, err
+}
+
 // QueryByPks query from collection
 func (mc *MilvusClient) QueryByPks(ctx context.Context, collName string, partitions []string, ids entity.Column,
 	outputFields []string, opts ...client.SearchQueryOptionFunc) (client.ResultSet, error) {

diff --git a/test/common/response_check.go b/test/common/response_check.go
@@ -3,6 +3,7 @@ package common
 import (
 	"fmt"
 	"log"
+	"strings"
 	"testing"
 
 	"github.com/milvus-io/milvus-sdk-go/v2/client"
@@ -15,11 +16,22 @@ func CheckErr(t *testing.T, actualErr error, expErrNil bool, expErrorMsg ...stri
 	if expErrNil {
 		require.NoError(t, actualErr)
 	} else {
-		if len(expErrorMsg) == 0 {
+		switch len(expErrorMsg) {
+		case 0:
 			log.Fatal("expect error message should not be empty")
+		case 1:
+			require.ErrorContains(t, actualErr, expErrorMsg[0])
+		default:
+			contains := false
+			for i := 0; i < len(expErrorMsg); i++ {
+				if strings.Contains(actualErr.Error(), expErrorMsg[i]) {
+					contains = true
+				}
+			}
+			if !contains {
+				t.FailNow()
+			}
 		}
-		require.NotEmpty(t, expErrorMsg[0])
-		require.ErrorContains(t, actualErr, expErrorMsg[0])
 	}
 }