Repository: milvus-io/milvus-sdk-go
Pull request #120 (2021-07-09 09:44:09)

Add partition load/release
Add partition load & release api
Add related unit test cases
Fix some lint issue

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/120#" 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/120#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/120#issuecomment-877076240" 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.go b/client/client.go
index 0628006..9e19a5b 100644
--- a/client/client.go
+++ b/client/client.go
@@ -55,6 +55,10 @@ type Client interface {
 	ShowPartitions(ctx context.Context, collName string) ([]*entity.Partition, error)
 	// HasPartition check whether partition exists in collection
 	HasPartition(ctx context.Context, collName string, partitionName string) (bool, error)
+	// LoadPartitions load partitions into memory
+	LoadPartitions(ctx context.Context, collName string, partitionNames []string, async bool) error
+	// ReleasePartitions release partitions
+	ReleasePartitions(ctx context.Context, collName string, partitionNames []string) error
 
 	// -- index --
 
diff --git a/client/client_grpc_dd.go b/client/client_grpc_dd.go
index 059b0fb..5ae17cc 100644
--- a/client/client_grpc_dd.go
+++ b/client/client_grpc_dd.go
@@ -15,6 +15,7 @@ package client
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/golang/protobuf/proto"
@@ -187,7 +188,7 @@ func validateSchema(sch entity.Schema) error {
 		}
 	}
 	if vectors <= 0 {
-		return errors.New("vector field not set!")
+		return errors.New("vector field not set")
 	}
 	return nil
 }
@@ -411,6 +412,96 @@ func (c *grpcClient) ShowPartitions(ctx context.Context, collName string) ([]*en
 	return partitions, nil
 }
 
+// LoadPartitions load collection paritions into memory
+func (c *grpcClient) LoadPartitions(ctx context.Context, collName string, partitionNames []string, async bool) error {
+	if c.service == nil {
+		return ErrClientNotReady
+	}
+
+	partitions, err := c.ShowPartitions(ctx, collName)
+	if err != nil {
+		return err
+	}
+	m := make(map[string]int64)
+	for _, partition := range partitions {
+		m[partition.Name] = partition.ID
+	}
+	// load partitions ids
+	ids := make(map[int64]struct{})
+	for _, partitionName := range partitionNames {
+		id, has := m[partitionName]
+		if !has {
+			return fmt.Errorf("Collection %s does not has partitions %s", collName, partitionName)
+		}
+		ids[id] = struct{}{}
+	}
+
+	req := &server.LoadPartitionsRequest{
+		DbName:         "", // reserved
+		CollectionName: collName,
+		PartitionNames: partitionNames,
+	}
+	resp, err := c.service.LoadPartitions(ctx, req)
+	if err != nil {
+		return err
+	}
+	if err := handleRespStatus(resp); err != nil {
+		return err
+	}
+
+	if !async {
+		segments, _ := c.GetPersistentSegmentInfo(ctx, collName)
+		target := make(map[int64]*entity.Segment)
+		for _, segment := range segments {
+			if segment.NumRows == 0 {
+				continue
+			}
+			if _, has := ids[segment.ParititionID]; !has {
+				// segment not belongs to partition
+				continue
+			}
+			target[segment.ID] = segment
+		}
+		for len(target) > 0 {
+			current, err := c.GetQuerySegmentInfo(ctx, collName)
+			if err == nil {
+				for _, segment := range current {
+					ts, has := target[segment.ID]
+					if has {
+						if segment.NumRows >= ts.NumRows {
+							delete(target, segment.ID)
+						}
+					}
+				}
+			}
+			time.Sleep(time.Millisecond * 100)
+		}
+	}
+
+	return nil
+}
+
+// ReleasePartitions release partitions
+func (c *grpcClient) ReleasePartitions(ctx context.Context, collName string, partitionNames []string) error {
+	if c.service == nil {
+		return ErrClientNotReady
+	}
+	req := &server.ReleasePartitionsRequest{
+		DbName:         "", // reserved
+		CollectionName: collName,
+		PartitionNames: partitionNames,
+	}
+	resp, err := c.service.ReleasePartitions(ctx, req)
+	if err != nil {
+		return err
+	}
+	if err := handleRespStatus(resp); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // CreateIndex create index for collection
 func (c *grpcClient) CreateIndex(ctx context.Context, collName string, fieldName string, idx entity.Index) error {
 	if c.service == nil {


Test patch:
diff --git a/client/client_grpc_dd_test.go b/client/client_grpc_dd_test.go
index 5ac825b..b31db41 100644
--- a/client/client_grpc_dd_test.go
+++ b/client/client_grpc_dd_test.go
@@ -23,7 +23,7 @@ func TestGrpcClientNil(t *testing.T) {
 	tp := reflect.TypeOf(c)
 	v := reflect.ValueOf(c)
 	ctx := context.Background()
-	c2 := testClient(t, ctx)
+	c2 := testClient(ctx, t)
 	v2 := reflect.ValueOf(c2)
 
 	ctxDone, cancel := context.WithCancel(context.Background())
@@ -38,7 +38,7 @@ func TestGrpcClientNil(t *testing.T) {
 			continue
 		}
 		ins := make([]reflect.Value, 0, mt.NumIn())
-		for j := 1; j < mt.NumIn(); j++ { // idx == 0, is the reciever v
+		for j := 1; j < mt.NumIn(); j++ { // idx == 0, is the receiver v
 			if j == 1 {
 				//non-general solution, hard code context!
 				ins = append(ins, reflect.ValueOf(ctx))
@@ -97,12 +97,12 @@ func TestGrpcClientClose(t *testing.T) {
 	ctx := context.Background()
 
 	t.Run("normal close", func(t *testing.T) {
-		c := testClient(t, ctx)
+		c := testClient(ctx, t)
 		assert.Nil(t, c.Close())
 	})
 
 	t.Run("double close", func(t *testing.T) {
-		c := testClient(t, ctx)
+		c := testClient(ctx, t)
 		assert.Nil(t, c.Close())
 		assert.Nil(t, c.Close())
 	})
@@ -110,7 +110,7 @@ func TestGrpcClientClose(t *testing.T) {
 
 func TestGrpcClientListCollections(t *testing.T) {
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	type testCase struct {
 		ids     []int64
@@ -171,7 +171,7 @@ func TestGrpcClientListCollections(t *testing.T) {
 
 func TestGrpcClientCreateCollection(t *testing.T) {
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 	t.Run("Test normal creation", func(t *testing.T) {
 		ds := defaultSchema()
 		shardsNum := int32(1)
@@ -260,7 +260,7 @@ func TestGrpcClientCreateCollection(t *testing.T) {
 
 func TestGrpcClientDropCollection(t *testing.T) {
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	mock.setInjection(mDropCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := (raw).(*server.DropCollectionRequest)
@@ -284,7 +284,7 @@ func TestGrpcClientDropCollection(t *testing.T) {
 
 func TestGrpcClientLoadCollection(t *testing.T) {
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 	// injection check collection name equals
 	mock.setInjection(mLoadCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := raw.(*server.LoadCollectionRequest)
@@ -355,7 +355,7 @@ func TestGrpcClientLoadCollection(t *testing.T) {
 func TestReleaseCollection(t *testing.T) {
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	mock.setInjection(mReleaseCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := raw.(*server.ReleaseCollectionRequest)
@@ -372,7 +372,7 @@ func TestReleaseCollection(t *testing.T) {
 func TestGrpcClientHasCollection(t *testing.T) {
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	mock.setInjection(mHasCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := raw.(*server.HasCollectionRequest)
@@ -398,7 +398,7 @@ func TestGrpcClientHasCollection(t *testing.T) {
 func TestGrpcClientDescribeCollection(t *testing.T) {
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	collectionID := rand.Int63()
 
@@ -433,7 +433,7 @@ func TestGrpcClientDescribeCollection(t *testing.T) {
 func TestGrpcClientGetCollectionStatistics(t *testing.T) {
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	stat := make(map[string]string)
 	stat["row_count"] = "0"
@@ -466,48 +466,229 @@ func TestGrpcClientGetCollectionStatistics(t *testing.T) {
 func TestGrpcClientCreatePartition(t *testing.T) {
 	//TODO add detail asserts
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 	assert.Nil(t, c.CreatePartition(ctx, testCollectionName, "_default_part"))
 }
 
 func TestGrpcClientDropPartition(t *testing.T) {
 	//TODO add detail asserts
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 	assert.Nil(t, c.DropPartition(ctx, testCollectionName, "_default_part"))
 }
 
 func TestGrpcClientHasPartition(t *testing.T) {
 	//TODO add detail asserts
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 	r, err := c.HasPartition(ctx, testCollectionName, "_default_part")
 	assert.Nil(t, err)
 	assert.False(t, r)
 }
 
+// default partition intercetion for ShowPartitions, generates testCollection related paritition data
+func getPartitionsInterception(t *testing.T, collName string, partitions ...*entity.Partition) func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+	return func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.ShowPartitionsRequest)
+		resp := &server.ShowPartitionsResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		assert.Equal(t, collName, req.GetCollectionName())
+		resp.PartitionIDs = make([]int64, 0, len(partitions))
+		resp.PartitionNames = make([]string, 0, len(partitions))
+		for _, part := range partitions {
+			resp.PartitionIDs = append(resp.PartitionIDs, part.ID)
+			resp.PartitionNames = append(resp.PartitionNames, part.Name)
+		}
+		s, err := successStatus()
+		resp.Status = s
+		return resp, err
+	}
+}
+
 func TestGrpcClientShowPartitions(t *testing.T) {
+
 	ctx := context.Background()
-	c := testClient(t, ctx)
-	r, err := c.ShowPartitions(ctx, testCollectionName)
-	assert.Nil(t, err)
-	assert.NotNil(t, r)
+	c := testClient(ctx, t)
+
+	type testCase struct {
+		collName      string
+		partitions    []*entity.Partition
+		shouldSuccess bool
+	}
+	cases := []testCase{
+		{
+			collName: testCollectionName,
+			partitions: []*entity.Partition{
+				{
+					ID:   1,
+					Name: "_part1",
+				},
+				{
+					ID:   2,
+					Name: "_part2",
+				},
+				{
+					ID:   3,
+					Name: "_part3",
+				},
+			},
+			shouldSuccess: true,
+		},
+	}
+	for _, tc := range cases {
+		mock.setInjection(mShowPartitions, getPartitionsInterception(t, tc.collName, tc.partitions...))
+		r, err := c.ShowPartitions(ctx, tc.collName)
+		if tc.shouldSuccess {
+			assert.Nil(t, err)
+			assert.NotNil(t, r)
+			if assert.Equal(t, len(tc.partitions), len(r)) {
+				for idx, part := range tc.partitions {
+					assert.Equal(t, part.ID, r[idx].ID)
+					assert.Equal(t, part.Name, r[idx].Name)
+				}
+			}
+		} else {
+			assert.NotNil(t, err)
+		}
+	}
 }
 
 func TestGrpcClientLoadPartitions(t *testing.T) {
-	//TODO add implementation
-	//ctx := context.Background()
-	//c := testClient(t, ctx)
-	//assert.Nil(t, c.LoadParitions(ctx, testCollectionName, "_default_part"))
+	ctx := context.Background()
+	c := testClient(ctx, t)
+
+	type testCase struct {
+		collName      string
+		partitions    []*entity.Partition
+		shouldSuccess bool
+		loadNames     []string
+	}
+	cases := []testCase{
+		{
+			collName: testCollectionName,
+			partitions: []*entity.Partition{
+				{
+					ID:   1,
+					Name: "_part1",
+				},
+				{
+					ID:   2,
+					Name: "_part2",
+				},
+				{
+					ID:   3,
+					Name: "_part3",
+				},
+			},
+			loadNames:     []string{"_part1", "_part2"},
+			shouldSuccess: true,
+		},
+		{
+			collName: testCollectionName,
+			partitions: []*entity.Partition{
+				{
+					ID:   1,
+					Name: "_part1",
+				},
+				{
+					ID:   2,
+					Name: "_part2",
+				},
+				{
+					ID:   3,
+					Name: "_part3",
+				},
+			},
+			loadNames:     []string{"_part4", "_part1"},
+			shouldSuccess: false,
+		},
+	}
+	for _, tc := range cases {
+		// one segment per paritions
+		segmentCount := len(tc.partitions)
+		start := time.Now()
+		loadTime := rand.Intn(1500) + 100
+		rowCounts := rand.Intn(1000) + 100
+		ok := false
+		mock.setInjection(mGetPersistentSegmentInfo, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			s, err := successStatus()
+			r := &server.GetPersistentSegmentInfoResponse{
+				Status: s,
+				Infos:  make([]*server.PersistentSegmentInfo, 0, segmentCount),
+			}
+			for i := 0; i < segmentCount; i++ {
+				r.Infos = append(r.Infos, &server.PersistentSegmentInfo{
+					SegmentID:   int64(i),
+					NumRows:     int64(rowCounts),
+					PartitionID: tc.partitions[i].ID,
+				})
+			}
+			r.Infos = append(r.Infos, &server.PersistentSegmentInfo{
+				SegmentID: int64(segmentCount),
+				NumRows:   0, // handcrafted empty segment
+			})
+			return r, err
+		})
+		mock.setInjection(mGetQuerySegmentInfo, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			s, err := successStatus()
+			r := &server.GetQuerySegmentInfoResponse{
+				Status: s,
+				Infos:  make([]*server.QuerySegmentInfo, 0, segmentCount),
+			}
+			rc := 0
+			if time.Since(start) > time.Duration(loadTime)*time.Millisecond {
+				rc = rowCounts // after load time, row counts set to full amount
+				ok = true
+			}
+			for i := 0; i < segmentCount; i++ {
+				r.Infos = append(r.Infos, &server.QuerySegmentInfo{
+					SegmentID:   int64(i),
+					NumRows:     int64(rc),
+					PartitionID: tc.partitions[i].ID,
+				})
+			}
+			return r, err
+		})
+
+		mock.setInjection(mShowPartitions, getPartitionsInterception(t, tc.collName, tc.partitions...))
+		err := c.LoadPartitions(ctx, tc.collName, tc.loadNames, false)
+		if tc.shouldSuccess {
+			assert.Nil(t, err)
+			assert.True(t, ok)
+		} else {
+			assert.NotNil(t, err)
+		}
+	}
 }
 
 func TestGrpcClientReleasePartitions(t *testing.T) {
-	//TODO add implementation
+	ctx := context.Background()
+
+	c := testClient(ctx, t)
+
+	parts := []string{"_part1", "_part2"}
+
+	mock.setInjection(mReleasePartitions, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.ReleasePartitionsRequest)
+		if !ok {
+			return badRequestStatus()
+		}
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		assert.ElementsMatch(t, parts, req.GetPartitionNames())
+
+		return successStatus()
+	})
+
+	assert.Nil(t, c.ReleasePartitions(ctx, testCollectionName, parts))
 }
 
 func TestGrpcClientCreateIndex(t *testing.T) {
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 	idx := entity.NewGenericIndex("", entity.Flat, map[string]string{
 		"nlist":       "1024",
 		"metric_type": "IP",
@@ -518,7 +699,7 @@ func TestGrpcClientCreateIndex(t *testing.T) {
 
 func TestGrpcClientDropIndex(t *testing.T) {
 	ctx := context.Background()
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	assert.Nil(t, c.DropIndex(ctx, testCollectionName, "vector"))
 }
@@ -526,7 +707,7 @@ func TestGrpcClientDropIndex(t *testing.T) {
 func TestGrpcClientDescribeIndex(t *testing.T) {
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	fieldName := "vector"
 
diff --git a/client/client_grpc_dm_test.go b/client/client_grpc_dm_test.go
index 0ff0a16..75d01ea 100644
--- a/client/client_grpc_dm_test.go
+++ b/client/client_grpc_dm_test.go
@@ -15,7 +15,7 @@ import (
 func TestGrpcClientInsert(t *testing.T) {
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	vector := generateFloatVector(4096, testVectorDim)
 	mock.setInjection(mInsert, func(_ context.Context, raw proto.Message) (proto.Message, error) {
@@ -50,7 +50,7 @@ func TestGrpcClientFlush(t *testing.T) {
 
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	assert.Nil(t, c.Flush(ctx, testCollectionName, false))
 	assert.Nil(t, c.Flush(ctx, testCollectionName, true))
@@ -60,7 +60,7 @@ func TestGrpcSearch(t *testing.T) {
 
 	ctx := context.Background()
 
-	c := testClient(t, ctx)
+	c := testClient(ctx, t)
 
 	vectors := generateFloatVector(4096, testVectorDim)
 
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
index 21f2f79..2e48be7 100644
--- a/client/client_mock_test.go
+++ b/client/client_mock_test.go
@@ -54,7 +54,7 @@ func bufDialer(context.Context, string) (net.Conn, error) {
 	return lis.Dial()
 }
 
-func testClient(t *testing.T, ctx context.Context) Client {
+func testClient(ctx context.Context, t *testing.T) Client {
 	c, err := NewGrpcClient(ctx, "bufnet", grpc.WithBlock(),
 		grpc.WithInsecure(), grpc.WithContextDialer(bufDialer))
 
diff --git a/tests/client_collection_test.go b/tests/client_collection_test.go
index fcecd81..782858c 100644
--- a/tests/client_collection_test.go
+++ b/tests/client_collection_test.go
@@ -2,7 +2,10 @@ package tests
 
 import (
 	"context"
+	"fmt"
+	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/milvus-io/milvus-sdk-go/client"
 	"github.com/milvus-io/milvus-sdk-go/entity"
@@ -31,6 +34,7 @@ func TestListCollections(t *testing.T) {
 }
 
 func TestCreateCollection(t *testing.T) {
+	rand.Seed(time.Now().UnixNano())
 	c, err := client.NewGrpcClient(context.Background(), "localhost:19530")
 	assert.Nil(t, err)
 	assert.NotNil(t, c)
@@ -43,7 +47,7 @@ func TestCreateCollection(t *testing.T) {
 		AutoID:         false,
 		Fields: []*entity.Field{
 			{
-				Name:       "int64",
+				Name:       fmt.Sprintf("int64_%v", rand.Intn(10)),
 				DataType:   entity.FieldTypeInt64,
 				PrimaryKey: true,
 			},
@@ -56,7 +60,8 @@ func TestCreateCollection(t *testing.T) {
 			},
 		},
 	}
-	err = c.CreateCollection(context.Background(), schema, 1)
+	shards := rand.Int31n(3)
+	err = c.CreateCollection(context.Background(), schema, shards)
 	assert.Nil(t, err)
 }
 
diff --git a/tests/client_index_test.go b/tests/client_index_test.go
index ca5af12..2ee608a 100644
--- a/tests/client_index_test.go
+++ b/tests/client_index_test.go
@@ -23,7 +23,7 @@ func TestCreateIndex(t *testing.T) {
 	})
 	err = c.CreateIndex(context.Background(), "test_go_sdk", "vector", idx)
 	assert.Nil(t, err)
-	indexes, err := c.DescribeIndex(context.Background(), "test_go_sdk")
+	indexes, err := c.DescribeIndex(context.Background(), "test_go_sdk", "vector")
 	if assert.Nil(t, err) {
 		for _, idx := range indexes {
 			t.Log(idx.IndexType())