From b06cfbe6c319bf6a395dcad7bf5bd08d4452582e Mon Sep 17 00:00:00 2001
From: Congqi Xia <congqi.xia@zilliz.com>
Date: Thu, 25 May 2023 15:43:44 +0800
Subject: [PATCH] Add Logical partition num options for create collection

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 client/collection.go   |  1 +
 client/options.go      |  6 ++++++
 client/options_test.go | 14 ++++++++++++++
 3 files changed, 21 insertions(+)

diff --git a/client/collection.go b/client/collection.go
index 1cd9e4cc..0992457a 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 b1fa0b53..8ea9505c 100644
--- a/client/options.go
+++ b/client/options.go
@@ -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)
 
diff --git a/client/options_test.go b/client/options_test.go
index b876238f..ed66c4f1 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 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",