Repository: milvus-io/milvus-sdk-go
Pull request #168 (2021-07-30 04:04:49)

Split Search request according to row size
Split Search request according to row size

Fix issue: #167 

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/168#" title="Author self-approved">congqixia</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
sunby: /lgtm
/approve
sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/168#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/168#issuecomment-889656971" title="Approved">sunby</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia,sunby]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/client/client_grpc_data.go b/client/client_grpc_data.go
index a99a222..4426fb2 100644
--- a/client/client_grpc_data.go
+++ b/client/client_grpc_data.go
@@ -16,6 +16,9 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
+	"strconv"
+	"sync"
 	"time"
 
 	"github.com/golang/protobuf/proto"
@@ -190,6 +193,9 @@ func (c *grpcClient) Search(ctx context.Context, collName string, partitions []s
 	if err != nil {
 		return nil, err
 	}
+	if coll.Schema.CollectionName == "" {
+		coll.Schema.CollectionName = collName
+	}
 	mNameField := make(map[string]*entity.Field)
 	for _, field := range coll.Schema.Fields {
 		mNameField[field.Name] = field
@@ -223,6 +229,69 @@ func (c *grpcClient) Search(ctx context.Context, collName string, partitions []s
 	}
 
 	// 2. Request milvus service
+	reqs := splitSearchRequest(coll.Schema, partitions, expr, outputFields, vectors, vectorField, metricType, topK, sp)
+	if len(reqs) == 0 {
+		return nil, errors.New("empty request generated")
+	}
+	wg := &sync.WaitGroup{}
+	wg.Add(len(reqs))
+	var batchErr error
+	sr := make([]SearchResult, 0, len(vectors))
+	mut := sync.Mutex{}
+	for _, req := range reqs {
+		go func(req *server.SearchRequest) {
+			defer wg.Done()
+			resp, err := c.service.Search(ctx, req)
+			if err != nil {
+				batchErr = err
+				return
+			}
+			if err := handleRespStatus(resp.GetStatus()); err != nil {
+				batchErr = err
+				return
+			}
+			// 3. parse result into result
+			results := resp.GetResults()
+			offset := 0
+			fieldDataList := results.GetFieldsData()
+			for i := 0; i < int(results.GetNumQueries()); i++ {
+
+				rc := int(results.GetTopks()[i]) // result entry count for current query
+				entry := SearchResult{
+					ResultCount: rc,
+					Scores:      results.GetScores()[offset:rc],
+				}
+				entry.IDs, entry.Err = entity.IDColumns(results.GetIds(), offset, offset+rc)
+				if entry.Err != nil {
+					offset += rc
+					continue
+				}
+				entry.Fields = make([]entity.Column, 0, len(fieldDataList))
+				for _, fieldData := range fieldDataList {
+					column, err := entity.FieldDataColumn(fieldData, offset, offset+rc)
+					if err != nil {
+						entry.Err = err
+						continue
+					}
+					entry.Fields = append(entry.Fields, column)
+				}
+				mut.Lock()
+				sr = append(sr, entry)
+				mut.Unlock()
+				offset += rc
+			}
+		}(req)
+	}
+	wg.Wait()
+	if batchErr != nil {
+		return []SearchResult{}, batchErr
+	}
+	return sr, nil
+}
+
+func splitSearchRequest(sch *entity.Schema, partitions []string,
+	expr string, outputFields []string, vectors []entity.Vector, vectorField string,
+	metricType entity.MetricType, topK int, sp entity.SearchParam) []*server.SearchRequest {
 	params := sp.Params()
 	bs, _ := json.Marshal(params)
 	searchParams := entity.MapKvPairs(map[string]string{
@@ -232,53 +301,29 @@ func (c *grpcClient) Search(ctx context.Context, collName string, partitions []s
 		"metric_type": string(metricType),
 	})
 
-	req := &server.SearchRequest{
-		DbName:           "",
-		CollectionName:   collName,
-		PartitionNames:   partitions,
-		SearchParams:     searchParams,
-		Dsl:              expr,
-		DslType:          common.DslType_BoolExprV1,
-		OutputFields:     outputFields,
-		PlaceholderGroup: vector2PlaceholderGroupBytes(vectors),
-	}
-	resp, err := c.service.Search(ctx, req)
-	if err != nil {
-		return []SearchResult{}, err
-	}
-	if err := handleRespStatus(resp.GetStatus()); err != nil {
-		return []SearchResult{}, err
-	}
-	// 3. parse result into result
-	results := resp.GetResults()
-	offset := 0
-	sr := make([]SearchResult, 0, results.GetNumQueries())
-	fieldDataList := results.GetFieldsData()
-	for i := 0; i < int(results.GetNumQueries()); i++ {
-
-		rc := int(results.GetTopks()[i]) // result entry count for current query
-		entry := SearchResult{
-			ResultCount: rc,
-			Scores:      results.GetScores()[offset:rc],
+	ers := estRowSize(sch, outputFields)
+	maxBatch := int(math.Ceil(float64(5*1024*1024) / float64(ers*int64(topK))))
+	result := []*server.SearchRequest{}
+	for i := 0; i*maxBatch < len(vectors); i++ {
+		start := i * maxBatch
+		end := (i + 1) * maxBatch
+		if end > len(vectors) {
+			end = len(vectors)
 		}
-		entry.IDs, entry.Err = entity.IDColumns(results.GetIds(), offset, offset+rc)
-		if entry.Err != nil {
-			offset += rc
-			continue
+		batchVectors := vectors[start:end]
+		req := &server.SearchRequest{
+			DbName:           "",
+			CollectionName:   sch.CollectionName,
+			PartitionNames:   partitions,
+			SearchParams:     searchParams,
+			Dsl:              expr,
+			DslType:          common.DslType_BoolExprV1,
+			OutputFields:     outputFields,
+			PlaceholderGroup: vector2PlaceholderGroupBytes(batchVectors),
 		}
-		entry.Fields = make([]entity.Column, 0, len(fieldDataList))
-		for _, fieldData := range fieldDataList {
-			column, err := entity.FieldDataColumn(fieldData, offset, offset+rc)
-			if err != nil {
-				entry.Err = err
-				continue
-			}
-			entry.Fields = append(entry.Fields, column)
-		}
-		sr = append(sr, entry)
-		offset += rc
+		result = append(result, req)
 	}
-	return sr, nil
+	return result
 }
 
 // GetPersistentSegmentInfo get persistent segment info
@@ -514,3 +559,49 @@ func isCollectionPrimaryKey(coll *entity.Collection, column entity.Column) bool
 	}
 	return false
 }
+
+// estRowSize estimate size per row for the specified schema
+func estRowSize(sch *entity.Schema, selected []string) int64 {
+	var total int64
+	for _, field := range sch.Fields {
+		if len(selected) > 0 {
+			found := false
+			for _, sel := range selected {
+				if field.Name == sel {
+					found = true
+					break
+				}
+			}
+			if !found {
+				continue
+			}
+		}
+		switch field.DataType {
+		case entity.FieldTypeBool:
+			total += 1
+		case entity.FieldTypeInt8:
+			total += 1
+		case entity.FieldTypeInt16:
+			total += 2
+		case entity.FieldTypeInt32:
+			total += 4
+		case entity.FieldTypeInt64:
+			total += 8
+		case entity.FieldTypeFloat:
+			total += 4
+		case entity.FieldTypeDouble:
+			total += 8
+		case entity.FieldTypeString:
+			//TODO string need varchar[max] syntax like limitation
+		case entity.FieldTypeFloatVector:
+			dimStr := field.TypeParams["dim"]
+			dim, _ := strconv.ParseInt(dimStr, 10, 64)
+			total += 4 * dim
+		case entity.FieldTypeBinaryVector:
+			dimStr := field.TypeParams["dim"]
+			dim, _ := strconv.ParseInt(dimStr, 10, 64)
+			total += 4 * dim / 8
+		}
+	}
+	return total
+}


Test patch:
diff --git a/client/client_grpc_data_test.go b/client/client_grpc_data_test.go
index d2d3453..db495ff 100644
--- a/client/client_grpc_data_test.go
+++ b/client/client_grpc_data_test.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"math/rand"
 	"testing"
 	"time"
@@ -547,6 +548,125 @@ func TestIsCollectionPrimaryKey(t *testing.T) {
 	})
 }
 
+func TestEstRowSize(t *testing.T) {
+	// a schema contains all supported vector
+	sch := &entity.Schema{
+		CollectionName: testCollectionName,
+		AutoID:         false,
+		Fields: []*entity.Field{
+			{
+				Name:       testPrimaryField,
+				DataType:   entity.FieldTypeInt64,
+				PrimaryKey: true,
+				AutoID:     true,
+			},
+			{
+				Name:     "attr1",
+				DataType: entity.FieldTypeInt8,
+			},
+			{
+				Name:     "attr2",
+				DataType: entity.FieldTypeInt16,
+			},
+			{
+				Name:     "attr3",
+				DataType: entity.FieldTypeInt32,
+			},
+			{
+				Name:     "attr4",
+				DataType: entity.FieldTypeFloat,
+			},
+			{
+				Name:     "attr5",
+				DataType: entity.FieldTypeDouble,
+			},
+			{
+				Name:     "attr6",
+				DataType: entity.FieldTypeBool,
+			},
+			{
+				Name:     testVectorField,
+				DataType: entity.FieldTypeFloatVector,
+				TypeParams: map[string]string{
+					"dim": fmt.Sprintf("%d", testVectorDim),
+				},
+			},
+			{
+				Name:     "binary_vector",
+				DataType: entity.FieldTypeBinaryVector,
+				TypeParams: map[string]string{
+					"dim": fmt.Sprintf("%d", testVectorDim),
+				},
+			},
+		},
+	}
+
+	// one row
+	columnID := entity.NewColumnInt64(testPrimaryField, []int64{0})
+	columnAttr1 := entity.NewColumnInt8("attr1", []int8{0})
+	columnAttr2 := entity.NewColumnInt16("attr2", []int16{0})
+	columnAttr3 := entity.NewColumnInt32("attr3", []int32{0})
+	columnAttr4 := entity.NewColumnFloat("attr4", []float32{0})
+	columnAttr5 := entity.NewColumnDouble("attr5", []float64{0})
+	columnAttr6 := entity.NewColumnBool("attr6", []bool{true})
+	columnFv := entity.NewColumnFloatVector(testVectorField, testVectorDim, generateFloatVector(1, testVectorDim))
+	columnBv := entity.NewColumnBinaryVector("binary_vector", testVectorDim, generateBinaryVector(1, testVectorDim))
+
+	sr := &server.SearchResults{
+		Results: &schema.SearchResultData{
+			FieldsData: []*schema.FieldData{
+				columnID.FieldData(),
+				columnAttr1.FieldData(),
+				columnAttr2.FieldData(),
+				columnAttr3.FieldData(),
+				columnAttr4.FieldData(),
+				columnAttr5.FieldData(),
+				columnAttr6.FieldData(),
+				columnFv.FieldData(),
+				columnBv.FieldData(),
+			},
+		},
+	}
+	bs, err := proto.Marshal(sr)
+	assert.Nil(t, err)
+	sr1l := len(bs)
+	// 2Row
+	columnID = entity.NewColumnInt64(testPrimaryField, []int64{0, 1})
+	columnAttr1 = entity.NewColumnInt8("attr1", []int8{0, 1})
+	columnAttr2 = entity.NewColumnInt16("attr2", []int16{0, 1})
+	columnAttr3 = entity.NewColumnInt32("attr3", []int32{0, 1})
+	columnAttr4 = entity.NewColumnFloat("attr4", []float32{0, 1})
+	columnAttr5 = entity.NewColumnDouble("attr5", []float64{0, 1})
+	columnAttr6 = entity.NewColumnBool("attr6", []bool{true, true})
+	columnFv = entity.NewColumnFloatVector(testVectorField, testVectorDim, generateFloatVector(2, testVectorDim))
+	columnBv = entity.NewColumnBinaryVector("binary_vector", testVectorDim, generateBinaryVector(2, testVectorDim))
+
+	sr = &server.SearchResults{
+		Results: &schema.SearchResultData{
+			FieldsData: []*schema.FieldData{
+				columnID.FieldData(),
+				columnAttr1.FieldData(),
+				columnAttr2.FieldData(),
+				columnAttr3.FieldData(),
+				columnAttr4.FieldData(),
+				columnAttr5.FieldData(),
+				columnAttr6.FieldData(),
+				columnFv.FieldData(),
+				columnBv.FieldData(),
+			},
+		},
+	}
+	bs, err = proto.Marshal(sr)
+	assert.Nil(t, err)
+	sr2l := len(bs)
+
+	t.Log(sr1l, sr2l, sr2l-sr1l)
+	est := estRowSize(sch, []string{})
+	t.Log(est)
+
+	assert.Greater(t, est, int64(sr2l-sr1l))
+}
+
 func generateFloatVector(num, dim int) [][]float32 {
 	r := make([][]float32, 0, num)
 	for i := 0; i < num; i++ {