enhance: [GoSDK] Support limit for search iterator (#43732)

Related to #37548

Add `WithIteratorLimit` option to limit search iterator overall return
limit.

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

Author: congqixia

client/milvusclient/iterator.go

@@ -38,6 +38,9 @@ const (
 	IteratorSearchLastBoundKey = "search_iter_last_bound"
 	IteratorSearchIDKey        = "search_iter_id"
 	CollectionIDKey            = `collection_id`
+
+	// Unlimited
+	Unlimited int64 = -1
 )
 
 var ErrServerVersionIncompatible = errors.New("server version incompatible")
@@ -53,10 +56,29 @@ type searchIteratorV2 struct {
 	client *Client
 	option SearchIteratorOption
 	schema *entity.Schema
+	limit  int64
 }
 
 func (it *searchIteratorV2) Next(ctx context.Context) (ResultSet, error) {
-	return it.next(ctx)
+	// limit reached, return EOF
+	if it.limit == 0 {
+		return ResultSet{}, io.EOF
+	}
+
+	rs, err := it.next(ctx)
+	if err != nil {
+		return rs, err
+	}
+
+	if it.limit == Unlimited {
+		return rs, err
+	}
+
+	if int64(rs.Len()) > it.limit {
+		rs = rs.Slice(0, int(it.limit))
+	}
+	it.limit -= int64(rs.Len())
+	return rs, nil
 }
 
 func (it *searchIteratorV2) next(ctx context.Context) (ResultSet, error) {
@@ -144,6 +166,7 @@ func newSearchIteratorV2(ctx context.Context, client *Client, option SearchItera
 	iter := &searchIteratorV2{
 		client: client,
 		option: option,
+		limit:  option.Limit(),
 	}
 	if err := iter.setupCollectionID(ctx); err != nil {
 		return nil, err

client/milvusclient/iterator_option.go

@@ -24,12 +24,16 @@ import (
 )
 
 type SearchIteratorOption interface {
+	// SearchOption returns the search option when iterate search
 	SearchOption() *searchOption
+	// Limit returns the overall limit of entries to iterate
+	Limit() int64
 }
 
 type searchIteratorOption struct {
 	*searchOption
-	batchSize int
+	batchSize     int
+	iteratorLimit int64
 }
 
 func (opt *searchIteratorOption) SearchOption() *searchOption {
@@ -38,6 +42,10 @@ func (opt *searchIteratorOption) SearchOption() *searchOption {
 	return opt.searchOption
 }
 
+func (opt *searchIteratorOption) Limit() int64 {
+	return opt.iteratorLimit
+}
+
 func (opt *searchIteratorOption) WithBatchSize(batchSize int) *searchIteratorOption {
 	opt.batchSize = batchSize
 	return opt
@@ -109,11 +117,17 @@ func (opt *searchIteratorOption) WithSearchParam(key, value string) *searchItera
 	return opt
 }
 
+func (opt *searchIteratorOption) WithIteratorLimit(limit int64) *searchIteratorOption {
+	opt.iteratorLimit = limit
+	return opt
+}
+
 func NewSearchIteratorOption(collectionName string, vector entity.Vector) *searchIteratorOption {
 	return &searchIteratorOption{
-		searchOption: NewSearchOption(collectionName, 100, []entity.Vector{vector}).
+		searchOption: NewSearchOption(collectionName, 1000, []entity.Vector{vector}).
 			WithSearchParam(IteratorKey, "true").
 			WithSearchParam(IteratorSearchV2Key, "true"),
-		batchSize: 1000,
+		batchSize:     1000,
+		iteratorLimit: Unlimited,
 	}
 }

client/milvusclient/iterator_test.go

@@ -298,6 +298,112 @@ func (s *SearchIteratorSuite) TestNext() {
 	s.ErrorIs(err, io.EOF)
 }
 
+func (s *SearchIteratorSuite) TestNextWithLimit() {
+	ctx := context.Background()
+	collectionName := fmt.Sprintf("coll_%s", s.randString(6))
+
+	token := fmt.Sprintf("iter_token_%s", s.randString(8))
+
+	s.mock.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
+		CollectionID: 1,
+		Schema:       s.schema.ProtoMessage(),
+	}, nil).Once()
+	s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
+		s.Equal(collectionName, sr.GetCollectionName())
+		checkSearchParam := func(kvs []*commonpb.KeyValuePair, key string, value string) bool {
+			for _, kv := range kvs {
+				if kv.GetKey() == key && kv.GetValue() == value {
+					return true
+				}
+			}
+			return false
+		}
+
+		s.True(checkSearchParam(sr.GetSearchParams(), IteratorKey, "true"))
+		s.True(checkSearchParam(sr.GetSearchParams(), IteratorSearchV2Key, "true"))
+		return &milvuspb.SearchResults{
+			Status: merr.Success(),
+			Results: &schemapb.SearchResultData{
+				NumQueries: 1,
+				TopK:       1,
+				FieldsData: []*schemapb.FieldData{
+					s.getInt64FieldData("ID", []int64{1}),
+				},
+				Ids: &schemapb.IDs{
+					IdField: &schemapb.IDs_IntId{
+						IntId: &schemapb.LongArray{
+							Data: []int64{1},
+						},
+					},
+				},
+				Scores:  make([]float32, 1),
+				Topks:   []int64{5},
+				Recalls: []float32{1},
+				SearchIteratorV2Results: &schemapb.SearchIteratorV2Results{
+					Token: token,
+				},
+			},
+		}, nil
+	}).Once()
+
+	iter, err := s.client.SearchIterator(ctx, NewSearchIteratorOption(collectionName, entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
+		return rand.Float32()
+	}))).WithIteratorLimit(6).WithBatchSize(5))
+	s.Require().NoError(err)
+	s.Require().NotNil(iter)
+
+	s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
+		s.Equal(collectionName, sr.GetCollectionName())
+		checkSearchParam := func(kvs []*commonpb.KeyValuePair, key string, value string) bool {
+			for _, kv := range kvs {
+				if kv.GetKey() == key && kv.GetValue() == value {
+					return true
+				}
+			}
+			return false
+		}
+
+		s.True(checkSearchParam(sr.GetSearchParams(), IteratorKey, "true"))
+		s.True(checkSearchParam(sr.GetSearchParams(), IteratorSearchV2Key, "true"))
+		return &milvuspb.SearchResults{
+			Status: merr.Success(),
+			Results: &schemapb.SearchResultData{
+				NumQueries: 1,
+				TopK:       1,
+				FieldsData: []*schemapb.FieldData{
+					s.getInt64FieldData("ID", []int64{1, 2, 3, 4, 5}),
+				},
+				Ids: &schemapb.IDs{
+					IdField: &schemapb.IDs_IntId{
+						IntId: &schemapb.LongArray{
+							Data: []int64{1, 2, 3, 4, 5},
+						},
+					},
+				},
+				Scores:  []float32{0.5, 0.4, 0.3, 0.2, 0.1},
+				Topks:   []int64{5},
+				Recalls: []float32{1},
+				SearchIteratorV2Results: &schemapb.SearchIteratorV2Results{
+					Token:     token,
+					LastBound: 0.5,
+				},
+			},
+		}, nil
+	}).Times(2)
+
+	rs, err := iter.Next(ctx)
+	s.NoError(err)
+	s.EqualValues(5, rs.IDs.Len(), "first batch, return all results")
+
+	rs, err = iter.Next(ctx)
+	s.NoError(err)
+	s.EqualValues(1, rs.IDs.Len(), "second batch, return sliced results")
+
+	_, err = iter.Next(ctx)
+	s.Error(err)
+	s.ErrorIs(err, io.EOF, "limit reached, return EOF")
+}
+
 func TestSearchIterator(t *testing.T) {
 	suite.Run(t, new(SearchIteratorSuite))
 }

client/milvusclient/results.go

@@ -21,6 +21,7 @@ import (
 	"runtime/debug"
 
 	"github.com/cockroachdb/errors"
+	"github.com/samber/lo"
 
 	"github.com/milvus-io/milvus/client/v2/column"
 	"github.com/milvus-io/milvus/client/v2/entity"
@@ -51,6 +52,31 @@ func (rs *ResultSet) GetColumn(fieldName string) column.Column {
 	return nil
 }
 
+func (rs ResultSet) Len() int {
+	return rs.ResultCount
+}
+
+func (rs ResultSet) Slice(start, end int) ResultSet {
+	result := ResultSet{
+		sch: rs.sch,
+		IDs: rs.IDs.Slice(start, end),
+		Fields: lo.Map(rs.Fields, func(column column.Column, _ int) column.Column {
+			return column.Slice(start, end)
+		}),
+		// Recall will not be sliced
+		Err: rs.Err,
+	}
+
+	if rs.GroupByValue != nil {
+		result.GroupByValue = rs.GroupByValue.Slice(start, end)
+	}
+
+	result.ResultCount = result.IDs.Len()
+	result.Scores = rs.Scores[start : start+result.ResultCount]
+
+	return result
+}
+
 // Unmarshal puts dataset into receiver in row based way.
 // `receiver` shall be a slice of pointer of model struct
 // eg, []*Records, in which type `Record` defines the row data.