Repository: milvus-io/milvus-sdk-go
Pull request #550 (2023-08-08 03:40:47)

Support GPU index type and search param
Add `GPUIvfFlat` and `GPUIvfPQ` index support
/kind improvement

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/550#" 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":[]} -->
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/550#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/550#pullrequestreview-1566512461" title="LGTM">yah01</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,yah01]

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/entity/genidx/genidx.go b/entity/genidx/genidx.go
index 9c6552f..379979e 100644
--- a/entity/genidx/genidx.go
+++ b/entity/genidx/genidx.go
@@ -638,6 +638,76 @@ func main() {
 					"level = -1",
 				},
 			},
+			{
+				IdxName: "GPUIvfFlat",
+				IdxType: entity.GPUIvfFlat,
+				ConstructParams: []idxParam{
+					{
+						Name:           "nlist",
+						ValidationRule: "[1, 65536]",
+					},
+				},
+				SearchParams: []idxParam{
+					{
+						Name:           "nprobe",
+						ValidationRule: "[1, 65536]", // [1, nlist], refer to index construct param, not supported yet
+					},
+				},
+				ValidExamples: []string{
+					"nlist = 10",
+				},
+				InvalidExamples: []string{
+					"nlist = 0",
+					"nlist = 65537",
+				},
+				ValidSearchParams: []string{
+					"nprobe = 10",
+				},
+				InvalidSearchParams: []string{
+					"nprobe = 0",
+					"nprobe = 65537",
+				},
+			},
+			{
+				IdxName: "GPUIvfPQ",
+				IdxType: entity.GPUIvfPQ,
+				ConstructParams: []idxParam{
+					{
+						Name:           "nlist",
+						ValidationRule: "[1, 65536]",
+					},
+					{
+						Name:           "m",
+						ValidationRule: "dim===0 (mod self)",
+					},
+					{
+						Name:           "nbits",
+						ValidationRule: "[1, 64]",
+					},
+				},
+				SearchParams: []idxParam{
+					{
+						Name:           "nprobe",
+						ValidationRule: "[1, 65536]", // [1, nlist], refer to index construct param, not supported yet
+					},
+				},
+				ValidExamples: []string{
+					"nlist, m, nbits = 10, 8, 8",
+				},
+				InvalidExamples: []string{
+					"nlist, m, nbits = 0, 8, 8",
+					"nlist, m, nbits = 65537, 8, 8",
+					"nlist, m, nbits = 10, 8, 0",
+					"nlist, m, nbits = 10, 8, 65",
+				},
+				ValidSearchParams: []string{
+					"nprobe = 10",
+				},
+				InvalidSearchParams: []string{
+					"nprobe = 0",
+					"nprobe = 65537",
+				},
+			},
 		},
 	}
 
diff --git a/entity/index.go b/entity/index.go
index b656072..bc92bff 100644
--- a/entity/index.go
+++ b/entity/index.go
@@ -36,6 +36,9 @@ const (
 	IvfHNSW    IndexType = "IVF_HNSW"
 	AUTOINDEX  IndexType = "AUTOINDEX"
 	DISKANN    IndexType = "DISKANN"
+
+	GPUIvfFlat IndexType = "GPU_IVF_FLAT"
+	GPUIvfPQ   IndexType = "GPU_IVF_PQ"
 )
 
 // Metric Constants
diff --git a/entity/indexes_gen.go b/entity/indexes_gen.go
index 32ae649..acbd696 100755
--- a/entity/indexes_gen.go
+++ b/entity/indexes_gen.go
@@ -591,3 +591,136 @@ func NewIndexAUTOINDEX(metricType MetricType, ) (*IndexAUTOINDEX, error) {
 	}, nil
 }
 
+
+var _ Index = &IndexGPUIvfFlat{}
+
+// IndexGPUIvfFlat idx type for GPU_IVF_FLAT
+type IndexGPUIvfFlat struct { //auto generated fields
+	nlist int
+	metricType MetricType
+}
+
+// Name returns index type name, implementing Index interface
+func(i *IndexGPUIvfFlat) Name() string {
+	return "GPUIvfFlat"
+}
+
+// IndexType returns IndexType, implementing Index interface
+func(i *IndexGPUIvfFlat) IndexType() IndexType {
+	return IndexType("GPU_IVF_FLAT")
+}
+
+// SupportBinary returns whether index type support binary vector
+func(i *IndexGPUIvfFlat) SupportBinary() bool {
+	return 0 & 2 > 0
+}
+
+// Params returns index construction params, implementing Index interface
+func(i *IndexGPUIvfFlat) Params() map[string]string {
+	params := map[string]string {//auto generated mapping 
+		"nlist": fmt.Sprintf("%v",i.nlist),
+	}
+	bs, _ := json.Marshal(params)
+	return map[string]string {
+		"params": string(bs),
+		"index_type": string(i.IndexType()),
+		"metric_type": string(i.metricType),
+	}
+}
+
+// NewIndexGPUIvfFlat create index with construction parameters
+func NewIndexGPUIvfFlat(metricType MetricType, 
+	nlist int,
+) (*IndexGPUIvfFlat, error) {
+	// auto generate parameters validation code, if any
+	if nlist < 1 {
+		return nil, errors.New("nlist has to be in range [1, 65536]")
+	}
+	if nlist > 65536 {
+		return nil, errors.New("nlist has to be in range [1, 65536]")
+	}
+	
+	return &IndexGPUIvfFlat{ 
+	//auto generated setting
+	nlist: nlist,
+	metricType: metricType,
+	}, nil
+}
+
+
+var _ Index = &IndexGPUIvfPQ{}
+
+// IndexGPUIvfPQ idx type for GPU_IVF_PQ
+type IndexGPUIvfPQ struct { //auto generated fields
+	nlist int
+	m int
+	nbits int
+	metricType MetricType
+}
+
+// Name returns index type name, implementing Index interface
+func(i *IndexGPUIvfPQ) Name() string {
+	return "GPUIvfPQ"
+}
+
+// IndexType returns IndexType, implementing Index interface
+func(i *IndexGPUIvfPQ) IndexType() IndexType {
+	return IndexType("GPU_IVF_PQ")
+}
+
+// SupportBinary returns whether index type support binary vector
+func(i *IndexGPUIvfPQ) SupportBinary() bool {
+	return 0 & 2 > 0
+}
+
+// Params returns index construction params, implementing Index interface
+func(i *IndexGPUIvfPQ) Params() map[string]string {
+	params := map[string]string {//auto generated mapping 
+		"nlist": fmt.Sprintf("%v",i.nlist),
+		"m": fmt.Sprintf("%v",i.m),
+		"nbits": fmt.Sprintf("%v",i.nbits),
+	}
+	bs, _ := json.Marshal(params)
+	return map[string]string {
+		"params": string(bs),
+		"index_type": string(i.IndexType()),
+		"metric_type": string(i.metricType),
+	}
+}
+
+// NewIndexGPUIvfPQ create index with construction parameters
+func NewIndexGPUIvfPQ(metricType MetricType, 
+	nlist int,
+
+	m int,
+
+	nbits int,
+) (*IndexGPUIvfPQ, error) {
+	// auto generate parameters validation code, if any
+	if nlist < 1 {
+		return nil, errors.New("nlist has to be in range [1, 65536]")
+	}
+	if nlist > 65536 {
+		return nil, errors.New("nlist has to be in range [1, 65536]")
+	}
+	
+	
+	
+	if nbits < 1 {
+		return nil, errors.New("nbits has to be in range [1, 64]")
+	}
+	if nbits > 64 {
+		return nil, errors.New("nbits has to be in range [1, 64]")
+	}
+	
+	return &IndexGPUIvfPQ{ 
+	//auto generated setting
+	nlist: nlist,
+	//auto generated setting
+	m: m,
+	//auto generated setting
+	nbits: nbits,
+	metricType: metricType,
+	}, nil
+}
+
diff --git a/entity/indexes_search_param_gen.go b/entity/indexes_search_param_gen.go
index 2ee4458..bcff101 100755
--- a/entity/indexes_search_param_gen.go
+++ b/entity/indexes_search_param_gen.go
@@ -318,3 +318,65 @@ func NewIndexAUTOINDEXSearchParam(
 	return sp, nil
 }
 
+var _ SearchParam = &IndexGPUIvfFlatSearchParam{}
+
+// IndexGPUIvfFlatSearchParam search param struct for index type GPU_IVF_FLAT
+type IndexGPUIvfFlatSearchParam struct { //auto generated fields
+	baseSearchParams
+	
+	nprobe int
+}
+
+// NewIndexGPUIvfFlatSearchParam create index search param
+func NewIndexGPUIvfFlatSearchParam(
+	nprobe int,
+) (*IndexGPUIvfFlatSearchParam, error) {
+	// auto generate parameters validation code, if any
+	if nprobe < 1 {
+		return nil, errors.New("nprobe has to be in range [1, 65536]")
+	}
+	if nprobe > 65536 {
+		return nil, errors.New("nprobe has to be in range [1, 65536]")
+	}
+	
+	sp := &IndexGPUIvfFlatSearchParam{
+		baseSearchParams: newBaseSearchParams(),
+	}
+	
+	//auto generated setting
+	sp.params["nprobe"] = nprobe
+
+	return sp, nil
+}
+
+var _ SearchParam = &IndexGPUIvfPQSearchParam{}
+
+// IndexGPUIvfPQSearchParam search param struct for index type GPU_IVF_PQ
+type IndexGPUIvfPQSearchParam struct { //auto generated fields
+	baseSearchParams
+	
+	nprobe int
+}
+
+// NewIndexGPUIvfPQSearchParam create index search param
+func NewIndexGPUIvfPQSearchParam(
+	nprobe int,
+) (*IndexGPUIvfPQSearchParam, error) {
+	// auto generate parameters validation code, if any
+	if nprobe < 1 {
+		return nil, errors.New("nprobe has to be in range [1, 65536]")
+	}
+	if nprobe > 65536 {
+		return nil, errors.New("nprobe has to be in range [1, 65536]")
+	}
+	
+	sp := &IndexGPUIvfPQSearchParam{
+		baseSearchParams: newBaseSearchParams(),
+	}
+	
+	//auto generated setting
+	sp.params["nprobe"] = nprobe
+
+	return sp, nil
+}
+


Test patch:
diff --git a/entity/indexes_gen_test.go b/entity/indexes_gen_test.go
index d635ec9..4de1bbc 100755
--- a/entity/indexes_gen_test.go
+++ b/entity/indexes_gen_test.go
@@ -462,3 +462,111 @@ func TestIndexAUTOINDEX(t *testing.T){
 	})
 }
 
+func TestIndexGPUIvfFlat(t *testing.T){
+	
+	var nlist int
+
+	mt := L2
+	
+
+	t.Run("valid usage case", func(t *testing.T){
+		
+		nlist = 10
+		idx0, err := NewIndexGPUIvfFlat(mt, 
+			nlist,
+		)
+		assert.Nil(t, err)
+		assert.NotNil(t, idx0)
+		assert.Equal(t, "GPUIvfFlat", idx0.Name())
+		assert.EqualValues(t, "GPU_IVF_FLAT", idx0.IndexType())
+		assert.NotNil(t, idx0.Params())
+		assert.False(t, idx0.SupportBinary())
+		
+	})
+
+	t.Run("invalid usage case", func(t *testing.T){
+		
+		nlist = 0
+		idx0, err := NewIndexGPUIvfFlat(mt, 
+			nlist,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx0)
+		
+		nlist = 65537
+		idx1, err := NewIndexGPUIvfFlat(mt, 
+			nlist,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx1)
+		
+	})
+}
+
+func TestIndexGPUIvfPQ(t *testing.T){
+	
+	var nlist int
+	var m int
+	var nbits int
+
+	mt := L2
+	
+
+	t.Run("valid usage case", func(t *testing.T){
+		
+		nlist, m, nbits = 10, 8, 8
+		idx0, err := NewIndexGPUIvfPQ(mt, 
+			nlist,
+			m,
+			nbits,
+		)
+		assert.Nil(t, err)
+		assert.NotNil(t, idx0)
+		assert.Equal(t, "GPUIvfPQ", idx0.Name())
+		assert.EqualValues(t, "GPU_IVF_PQ", idx0.IndexType())
+		assert.NotNil(t, idx0.Params())
+		assert.False(t, idx0.SupportBinary())
+		
+	})
+
+	t.Run("invalid usage case", func(t *testing.T){
+		
+		nlist, m, nbits = 0, 8, 8
+		idx0, err := NewIndexGPUIvfPQ(mt, 
+			nlist,
+			m,
+			nbits,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx0)
+		
+		nlist, m, nbits = 65537, 8, 8
+		idx1, err := NewIndexGPUIvfPQ(mt, 
+			nlist,
+			m,
+			nbits,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx1)
+		
+		nlist, m, nbits = 10, 8, 0
+		idx2, err := NewIndexGPUIvfPQ(mt, 
+			nlist,
+			m,
+			nbits,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx2)
+		
+		nlist, m, nbits = 10, 8, 65
+		idx3, err := NewIndexGPUIvfPQ(mt, 
+			nlist,
+			m,
+			nbits,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx3)
+		
+	})
+}
+
diff --git a/entity/indexes_search_param_gen_test.go b/entity/indexes_search_param_gen_test.go
index 762b7ff..f196d33 100755
--- a/entity/indexes_search_param_gen_test.go
+++ b/entity/indexes_search_param_gen_test.go
@@ -377,3 +377,75 @@ func TestIndexAUTOINDEXSearchParam(t *testing.T) {
 	
 }
 
+func TestIndexGPUIvfFlatSearchParam(t *testing.T) {
+	
+	var nprobe int
+
+	t.Run("valid usage case", func(t *testing.T){
+		
+		nprobe = 10
+		idx0, err := NewIndexGPUIvfFlatSearchParam(
+			nprobe,
+		)
+		assert.Nil(t, err)
+		assert.NotNil(t, idx0)
+		assert.NotNil(t, idx0.Params())
+		
+	})
+	
+	t.Run("invalid usage case", func(t *testing.T){
+		
+		nprobe = 0
+		idx0, err := NewIndexGPUIvfFlatSearchParam(
+			nprobe,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx0)
+		
+		nprobe = 65537
+		idx1, err := NewIndexGPUIvfFlatSearchParam(
+			nprobe,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx1)
+		
+	})
+	
+}
+
+func TestIndexGPUIvfPQSearchParam(t *testing.T) {
+	
+	var nprobe int
+
+	t.Run("valid usage case", func(t *testing.T){
+		
+		nprobe = 10
+		idx0, err := NewIndexGPUIvfPQSearchParam(
+			nprobe,
+		)
+		assert.Nil(t, err)
+		assert.NotNil(t, idx0)
+		assert.NotNil(t, idx0.Params())
+		
+	})
+	
+	t.Run("invalid usage case", func(t *testing.T){
+		
+		nprobe = 0
+		idx0, err := NewIndexGPUIvfPQSearchParam(
+			nprobe,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx0)
+		
+		nprobe = 65537
+		idx1, err := NewIndexGPUIvfPQSearchParam(
+			nprobe,
+		)
+		assert.NotNil(t, err)
+		assert.Nil(t, idx1)
+		
+	})
+	
+}
+