enhance: [GoSDK] Refine search params and add some examples

Related to milvus-io#31293

Signed-off-by: Congqi Xia <[email protected]>

client/index/ann_param.go

@@ -0,0 +1,45 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package index
+
+type AnnParam interface {
+	Params() map[string]any
+}
+
+type baseAnnParam struct {
+	params map[string]any
+}
+
+func (b baseAnnParam) WithExtraParam(key string, value any) {
+	b.params[key] = value
+}
+
+func (b baseAnnParam) Params() map[string]any {
+	return b.params
+}
+
+type CustomAnnParam struct {
+	baseAnnParam
+}
+
+func NewCustomAnnParam() CustomAnnParam {
+	return CustomAnnParam{
+		baseAnnParam: baseAnnParam{
+			params: make(map[string]any),
+		},
+	}
+}

client/index/ann_param_test.go

@@ -0,0 +1,52 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package index
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAnnParams(t *testing.T) {
+	ivfAP := NewIvfAnnParam(16)
+	result := ivfAP.Params()
+	v, ok := result[ivfNprobeKey]
+	assert.True(t, ok)
+	assert.Equal(t, 16, v)
+
+	hnswAP := NewHNSWAnnParam(16)
+	result = hnswAP.Params()
+	v, ok = result[hnswEfKey]
+	assert.True(t, ok)
+	assert.Equal(t, 16, v)
+
+	diskAP := NewDiskAnnParam(10)
+	result = diskAP.Params()
+	v, ok = result[diskANNSearchListKey]
+	assert.True(t, ok)
+	assert.Equal(t, 10, v)
+
+	scannAP := NewSCANNAnnParam(32, 50)
+	result = scannAP.Params()
+	v, ok = result[scannNProbeKey]
+	assert.True(t, ok)
+	assert.Equal(t, 32, v)
+	v, ok = result[scannReorderKKey]
+	assert.True(t, ok)
+	assert.Equal(t, 50, v)
+}

client/index/auto.go

@@ -16,6 +16,10 @@
 
 package index
 
+const (
+	autoLevelKey = `level`
+)
+
 var _ Index = autoIndex{}
 
 type autoIndex struct {
@@ -37,3 +41,23 @@ func NewAutoIndex(metricType MetricType) Index {
 		},
 	}
 }
+
+type autoAnnParam struct {
+	baseAnnParam
+	level int
+}
+
+func NewAutoAnnParam(level int) autoAnnParam {
+	return autoAnnParam{
+		baseAnnParam: baseAnnParam{
+			params: make(map[string]any),
+		},
+		level: level,
+	}
+}
+
+func (ap autoAnnParam) Params() map[string]any {
+	result := ap.baseAnnParam.params
+	result[autoLevelKey] = ap.level
+	return result
+}

client/index/disk_ann.go

@@ -16,6 +16,10 @@
 
 package index
 
+const (
+	diskANNSearchListKey = `search_list`
+)
+
 var _ Index = diskANNIndex{}
 
 type diskANNIndex struct {
@@ -37,3 +41,23 @@ func NewDiskANNIndex(metricType MetricType) Index {
 		},
 	}
 }
+
+type diskANNParam struct {
+	baseAnnParam
+	searchList int
+}
+
+func NewDiskAnnParam(searchList int) diskANNParam {
+	return diskANNParam{
+		baseAnnParam: baseAnnParam{
+			params: make(map[string]any),
+		},
+		searchList: searchList,
+	}
+}
+
+func (ap diskANNParam) Params() map[string]any {
+	result := ap.baseAnnParam.params
+	result[diskANNSearchListKey] = ap.searchList
+	return result
+}

client/index/hnsw.go

@@ -16,11 +16,14 @@
 
 package index
 
-import "strconv"
+import (
+	"strconv"
+)
 
 const (
 	hnswMKey           = `M`
 	hsnwEfConstruction = `efConstruction`
+	hnswEfKey          = `ef`
 )
 
 var _ Index = hnswIndex{}
@@ -51,3 +54,23 @@ func NewHNSWIndex(metricType MetricType, m int, efConstruction int) Index {
 		efConstruction: efConstruction,
 	}
 }
+
+type hsnwAnnParam struct {
+	baseAnnParam
+	ef int
+}
+
+func NewHNSWAnnParam(ef int) hsnwAnnParam {
+	return hsnwAnnParam{
+		baseAnnParam: baseAnnParam{
+			params: make(map[string]any),
+		},
+		ef: ef,
+	}
+}
+
+func (ap hsnwAnnParam) Params() map[string]any {
+	result := ap.baseAnnParam.params
+	result[hnswEfKey] = ap.ef
+	return result
+}

client/index/ivf.go

@@ -19,9 +19,10 @@ package index
 import "strconv"
 
 const (
-	ivfNlistKey = `nlist`
-	ivfPQMKey   = `m`
-	ivfPQNbits  = `nbits`
+	ivfNlistKey  = `nlist`
+	ivfPQMKey    = `m`
+	ivfPQNbits   = `nbits`
+	ivfNprobeKey = `nprobe`
 )
 
 var _ Index = ivfFlatIndex{}
@@ -137,3 +138,23 @@ func NewBinIvfFlatIndex(metricType MetricType, nlist int) Index {
 		nlist: nlist,
 	}
 }
+
+type ivfAnnParam struct {
+	baseAnnParam
+	nprobe int
+}
+
+func (ap ivfAnnParam) Params() map[string]any {
+	result := ap.baseAnnParam.Params()
+	result[ivfNprobeKey] = ap.nprobe
+	return result
+}
+
+func NewIvfAnnParam(nprobe int) ivfAnnParam {
+	return ivfAnnParam{
+		baseAnnParam: baseAnnParam{
+			params: make(map[string]any),
+		},
+		nprobe: nprobe,
+	}
+}

client/index/scann.go

@@ -21,6 +21,8 @@ import "strconv"
 const (
 	scannNlistKey       = `nlist`
 	scannWithRawDataKey = `with_raw_data`
+	scannNProbeKey      = `nprobe`
+	scannReorderKKey    = `reorder_k`
 )
 
 type scannIndex struct {
@@ -49,3 +51,26 @@ func NewSCANNIndex(metricType MetricType, nlist int, withRawData bool) Index {
 		withRawData: withRawData,
 	}
 }
+
+type scannAnnParam struct {
+	baseAnnParam
+	nprobe   int
+	reorderK int
+}
+
+func NewSCANNAnnParam(nprobe int, reorderK int) scannAnnParam {
+	return scannAnnParam{
+		baseAnnParam: baseAnnParam{
+			params: make(map[string]any),
+		},
+		nprobe:   nprobe,
+		reorderK: reorderK,
+	}
+}
+
+func (ap scannAnnParam) Params() map[string]any {
+	result := ap.baseAnnParam.params
+	result[scannNProbeKey] = ap.nprobe
+	result[scannReorderKKey] = ap.reorderK
+	return result
+}