From 8f870d9b410e8d741659d67d34eef6b0b134bcf6 Mon Sep 17 00:00:00 2001 From: Cai Zhang Date: Tue, 25 Jun 2024 14:55:52 +0800 Subject: [PATCH] Export indexed rows for descrieb_index Signed-off-by: Cai Zhang --- client/index.go | 6 ++++ client/index_test.go | 8 +++++ entity/genidx/genidx.go | 5 +++ entity/index.go | 28 +++++++++++++--- entity/index_cagra.go | 8 +++++ entity/index_scalar.go | 4 +++ entity/index_sparse.go | 8 +++++ entity/indexes_gen.go | 65 ++++++++++++++++++++++++++++++++++++ test/testcases/index_test.go | 4 +-- 9 files changed, 130 insertions(+), 6 deletions(-) diff --git a/client/index.go b/client/index.go index 97d3ef77..4e35b2ee 100644 --- a/client/index.go +++ b/client/index.go @@ -199,10 +199,16 @@ func (c *GrpcClient) DescribeIndex(ctx context.Context, collName string, fieldNa } params := entity.KvPairsMap(info.Params) it := params["index_type"] // TODO change to const + progress := make(map[string]string) + progress["total_rows"] = strconv.FormatInt(info.GetTotalRows(), 10) + progress["indexed_rows"] = strconv.FormatInt(info.GetIndexedRows(), 10) + progress["pending_index_rows"] = strconv.FormatInt(info.GetPendingIndexRows(), 10) + progress["state"] = info.GetState().String() idx := entity.NewGenericIndex( info.IndexName, entity.IndexType(it), params, + entity.WithProgress(progress), ) indexes = append(indexes, idx) } diff --git a/client/index_test.go b/client/index_test.go index 323ef7c6..86bd71a0 100644 --- a/client/index_test.go +++ b/client/index_test.go @@ -117,6 +117,10 @@ func TestGrpcClientDescribeIndex(t *testing.T) { "nlist": "1024", "metric_type": "IP", }), + TotalRows: 10000, + IndexedRows: 4000, + PendingIndexRows: 10000, + State: commonpb.IndexState_InProgress, }, } s, err := SuccessStatus() @@ -127,6 +131,10 @@ func TestGrpcClientDescribeIndex(t *testing.T) { idxes, err := c.DescribeIndex(ctx, testCollectionName, fieldName) assert.Nil(t, err) assert.NotNil(t, idxes) + assert.Equal(t, idxes[0].Progress()["total_rows"], "10000") + assert.Equal(t, idxes[0].Progress()["indexed_rows"], "4000") + assert.Equal(t, idxes[0].Progress()["pending_index_rows"], "10000") + assert.Equal(t, idxes[0].Progress()["state"], "InProgress") }) t.Run("Service return errors", func(t *testing.T) { diff --git a/entity/genidx/genidx.go b/entity/genidx/genidx.go index 4a4a63f3..bf1cb18a 100644 --- a/entity/genidx/genidx.go +++ b/entity/genidx/genidx.go @@ -73,6 +73,11 @@ func(i *Index{{.IdxName}}) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *Index{{.IdxName}}) Progress() map[string]string { + return map[string]string {} +} + // NewIndex{{.IdxName}} create index with construction parameters func NewIndex{{.IdxName}}(metricType MetricType, {{range .ConstructParams}}{{with .}} {{.Name}} {{.ParamType}}, diff --git a/entity/index.go b/entity/index.go index eafdadeb..796f9433 100644 --- a/entity/index.go +++ b/entity/index.go @@ -11,7 +11,9 @@ package entity -import common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" +import ( + common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" +) //go:generate go run genidx/genidx.go @@ -81,6 +83,7 @@ type Index interface { Name() string IndexType() IndexType Params() map[string]string + Progress() map[string]string } // SearchParam interface for index related search param @@ -136,7 +139,16 @@ func (b baseIndex) IndexType() IndexType { // no constraint for index is applied type GenericIndex struct { baseIndex - params map[string]string + params map[string]string + progress map[string]string +} + +type OptionFunc func(*GenericIndex) + +func WithProgress(progress map[string]string) OptionFunc { + return func(index *GenericIndex) { + index.progress = progress + } } // Params implements Index @@ -151,13 +163,21 @@ func (gi GenericIndex) Params() map[string]string { return m } +func (gi GenericIndex) Progress() map[string]string { + return gi.progress +} + // NewGenericIndex create generic index instance -func NewGenericIndex(name string, it IndexType, params map[string]string) Index { - return GenericIndex{ +func NewGenericIndex(name string, it IndexType, params map[string]string, opts ...OptionFunc) Index { + index := GenericIndex{ baseIndex: baseIndex{ it: it, name: name, }, params: params, } + for _, opt := range opts { + opt(&index) + } + return index } diff --git a/entity/index_cagra.go b/entity/index_cagra.go index afc711d8..7fa208d9 100644 --- a/entity/index_cagra.go +++ b/entity/index_cagra.go @@ -62,6 +62,10 @@ func (i *IndexGPUCagra) Params() map[string]string { } } +func (i *IndexGPUCagra) Progress() map[string]string { + return map[string]string{} +} + type IndexGPUCagraSearchParam struct { baseSearchParams } @@ -131,6 +135,10 @@ func (i *IndexGPUBruteForce) Params() map[string]string { } } +func (i *IndexGPUBruteForce) Progress() map[string]string { + return map[string]string{} +} + type IndexGPUBruteForceSearchParam struct { baseSearchParams } diff --git a/entity/index_scalar.go b/entity/index_scalar.go index ffaa5de9..fdc0d380 100644 --- a/entity/index_scalar.go +++ b/entity/index_scalar.go @@ -14,6 +14,10 @@ func (i *indexScalar) Params() map[string]string { return result } +func (i *indexScalar) Progress() map[string]string { + return map[string]string{} +} + func NewScalarIndex() Index { return &indexScalar{ baseIndex: baseIndex{}, diff --git a/entity/index_sparse.go b/entity/index_sparse.go index e1f47087..2cc74b16 100644 --- a/entity/index_sparse.go +++ b/entity/index_sparse.go @@ -35,6 +35,10 @@ func (i *IndexSparseInverted) Params() map[string]string { } } +func (i *IndexSparseInverted) Progress() map[string]string { + return map[string]string{} +} + type IndexSparseInvertedSearchParam struct { baseSearchParams } @@ -87,6 +91,10 @@ func (i *IndexSparseWAND) Params() map[string]string { } } +func (i *IndexSparseWAND) Progress() map[string]string { + return map[string]string{} +} + // IndexSparseWAND index type for SPARSE_WAND, weak-and func NewIndexSparseWAND(metricType MetricType, dropRatio float64) (*IndexSparseWAND, error) { if dropRatio < 0 || dropRatio >= 1.0 { diff --git a/entity/indexes_gen.go b/entity/indexes_gen.go index e4b84fe5..fdf7b634 100755 --- a/entity/indexes_gen.go +++ b/entity/indexes_gen.go @@ -45,6 +45,11 @@ func(i *IndexFlat) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexFlat) Progress() map[string]string { + return map[string]string {} +} + // NewIndexFlat create index with construction parameters func NewIndexFlat(metricType MetricType, ) (*IndexFlat, error) { // auto generate parameters validation code, if any @@ -90,6 +95,11 @@ func(i *IndexBinFlat) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexBinFlat) Progress() map[string]string { + return map[string]string {} +} + // NewIndexBinFlat create index with construction parameters func NewIndexBinFlat(metricType MetricType, nlist int, @@ -146,6 +156,11 @@ func(i *IndexIvfFlat) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexIvfFlat) Progress() map[string]string { + return map[string]string {} +} + // NewIndexIvfFlat create index with construction parameters func NewIndexIvfFlat(metricType MetricType, nlist int, @@ -202,6 +217,11 @@ func(i *IndexBinIvfFlat) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexBinIvfFlat) Progress() map[string]string { + return map[string]string {} +} + // NewIndexBinIvfFlat create index with construction parameters func NewIndexBinIvfFlat(metricType MetricType, nlist int, @@ -258,6 +278,11 @@ func(i *IndexIvfSQ8) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexIvfSQ8) Progress() map[string]string { + return map[string]string {} +} + // NewIndexIvfSQ8 create index with construction parameters func NewIndexIvfSQ8(metricType MetricType, nlist int, @@ -318,6 +343,11 @@ func(i *IndexIvfPQ) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexIvfPQ) Progress() map[string]string { + return map[string]string {} +} + // NewIndexIvfPQ create index with construction parameters func NewIndexIvfPQ(metricType MetricType, nlist int, @@ -393,6 +423,11 @@ func(i *IndexHNSW) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexHNSW) Progress() map[string]string { + return map[string]string {} +} + // NewIndexHNSW create index with construction parameters func NewIndexHNSW(metricType MetricType, M int, @@ -464,6 +499,11 @@ func(i *IndexIvfHNSW) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexIvfHNSW) Progress() map[string]string { + return map[string]string {} +} + // NewIndexIvfHNSW create index with construction parameters func NewIndexIvfHNSW(metricType MetricType, nlist int, @@ -540,6 +580,11 @@ func(i *IndexDISKANN) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexDISKANN) Progress() map[string]string { + return map[string]string {} +} + // NewIndexDISKANN create index with construction parameters func NewIndexDISKANN(metricType MetricType, ) (*IndexDISKANN, error) { // auto generate parameters validation code, if any @@ -583,6 +628,11 @@ func(i *IndexAUTOINDEX) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexAUTOINDEX) Progress() map[string]string { + return map[string]string {} +} + // NewIndexAUTOINDEX create index with construction parameters func NewIndexAUTOINDEX(metricType MetricType, ) (*IndexAUTOINDEX, error) { // auto generate parameters validation code, if any @@ -628,6 +678,11 @@ func(i *IndexGPUIvfFlat) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexGPUIvfFlat) Progress() map[string]string { + return map[string]string {} +} + // NewIndexGPUIvfFlat create index with construction parameters func NewIndexGPUIvfFlat(metricType MetricType, nlist int, @@ -688,6 +743,11 @@ func(i *IndexGPUIvfPQ) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexGPUIvfPQ) Progress() map[string]string { + return map[string]string {} +} + // NewIndexGPUIvfPQ create index with construction parameters func NewIndexGPUIvfPQ(metricType MetricType, nlist int, @@ -763,6 +823,11 @@ func(i *IndexSCANN) Params() map[string]string { } } +// Progress returns index progress information, only fill it after DescribeIndex for GenericIndex +func(i *IndexSCANN) Progress() map[string]string { + return map[string]string {} +} + // NewIndexSCANN create index with construction parameters func NewIndexSCANN(metricType MetricType, nlist int, diff --git a/test/testcases/index_test.go b/test/testcases/index_test.go index e90b5802..918fd769 100644 --- a/test/testcases/index_test.go +++ b/test/testcases/index_test.go @@ -995,7 +995,7 @@ func TestCreateIndexSparseVector2(t *testing.T) { // describe index idx2, err := mc.DescribeIndex(ctx, collName, common.DefaultSparseVecFieldName) expIndex := entity.NewGenericIndex(common.DefaultSparseVecFieldName, idx.IndexType(), idx.Params()) - require.EqualValues(t, expIndex, idx2[0]) + require.EqualValues(t, expIndex.Params(), idx2[0].Params()) common.CheckErr(t, err, true) common.CheckIndexResult(t, idx2, expIndex) } @@ -1448,7 +1448,7 @@ func TestDropIndexCreateIndex(t *testing.T) { // describe index ipIndexes, _ := mc.DescribeIndex(ctx, collName, common.DefaultFloatVecFieldName) - require.EqualValues(t, entity.NewGenericIndex(common.DefaultFloatVecFieldName, ipIdx.IndexType(), ipIdx.Params()), ipIndexes[0]) + require.EqualValues(t, entity.NewGenericIndex(common.DefaultFloatVecFieldName, ipIdx.IndexType(), ipIdx.Params()).Params(), ipIndexes[0].Params()) // describe collection t.Log("Issue: https://github.com/milvus-io/milvus-sdk-go/issues/385")