Skip to content

Commit

Permalink
Add Logical partition num options for create collection
Browse files Browse the repository at this point in the history
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
  • Loading branch information
congqixia committed May 25, 2023
1 parent 096d105 commit b06cfbe
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 0 deletions.
1 change: 1 addition & 0 deletions client/collection.go
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ func (c *GrpcClient) CreateCollection(ctx context.Context, collSchema *entity.Sc
// default consistency level is strong
// to be consistent with previous version
ConsistencyLevel: common.ConsistencyLevel_Bounded,
NumPartitions: 0,
}
// apply options on request
for _, opt := range opts {
Expand Down
6 changes: 6 additions & 0 deletions client/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,12 @@ func WithCollectionProperty(key string, value string) CreateCollectionOption {
}
}

func WithLogicalPartitionNum(partitionNums int64) CreateCollectionOption {
return func(req *server.CreateCollectionRequest) {
req.NumPartitions = partitionNums
}
}

// LoadCollectionOption is an option that is used to modify LoadCollectionRequest
type LoadCollectionOption func(*server.LoadCollectionRequest)

Expand Down
14 changes: 14 additions & 0 deletions client/options_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
package client

import (
"math/rand"
"testing"

common "github.com/milvus-io/milvus-proto/go-api/commonpb"
Expand Down Expand Up @@ -44,6 +45,19 @@ func TestLoadCollectionWithReplicaNumber(t *testing.T) {
assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber())
}

func TestCreateCollectionWithLogicalPartitionNum(t *testing.T) {
partitionNum := rand.Int63n(1000) + 1
opt := WithLogicalPartitionNum(partitionNum)
assert.NotNil(t, opt)
req := &server.CreateCollectionRequest{}

assert.NotPanics(t, func() {
opt(req)
})

assert.Equal(t, partitionNum, req.GetNumPartitions())
}

func TestMakeSearchQueryOption(t *testing.T) {
c := &entity.Collection{
Name: "999",
Expand Down

0 comments on commit b06cfbe

Please sign in to comment.