diff --git a/client/collection.go b/client/collection.go index 1cd9e4cc7..0992457a8 100644 --- a/client/collection.go +++ b/client/collection.go @@ -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 { diff --git a/client/options.go b/client/options.go index b1fa0b535..f992490dc 100644 --- a/client/options.go +++ b/client/options.go @@ -37,6 +37,13 @@ func WithCollectionProperty(key string, value string) CreateCollectionOption { } } +// WithPartitionNum returns a create collection options to set physical partition number when logical partition feature. +func WithPartitionNum(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) diff --git a/client/options_test.go b/client/options_test.go index b876238fc..e7f8ea06a 100644 --- a/client/options_test.go +++ b/client/options_test.go @@ -12,6 +12,7 @@ package client import ( + "math/rand" "testing" common "github.com/milvus-io/milvus-proto/go-api/commonpb" @@ -44,6 +45,19 @@ func TestLoadCollectionWithReplicaNumber(t *testing.T) { assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber()) } +func TestCreateCollectionWithPartitionNum(t *testing.T) { + partitionNum := rand.Int63n(1000) + 1 + opt := WithPartitionNum(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",