enhance: [GoSDK] Refine search params and add some examples (#38523)

Related to #31293

This PR
- Add some example test code for some basic operations
- Refine search params and add some predefined one
- Split search & hybrid search option

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
pull/38542/head
congqixia 2024-12-17 20:52:44 +08:00 committed by GitHub
parent f0096ec292
commit 1ec858434f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 1352 additions and 67 deletions

45
client/index/ann_param.go Normal file
View File

@ -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),
},
}
}

View File

@ -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)
}

View File

@ -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
}

View File

@ -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
}

View File

@ -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
}

View File

@ -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,
}
}

View File

@ -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
}

View File

@ -0,0 +1,139 @@
// 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.
// nolint
package milvusclient_test
import (
"context"
"fmt"
"github.com/milvus-io/milvus/client/v2/milvusclient"
)
func ExampleClient_CreateAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
err = cli.CreateAlias(ctx, milvusclient.NewCreateAliasOption("customized_setup_2", "bob"))
if err != nil {
// handle error
}
err = cli.CreateAlias(ctx, milvusclient.NewCreateAliasOption("customized_setup_2", "alice"))
if err != nil {
// handle error
}
}
func ExampleClient_ListAliases() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
aliases, err := cli.ListAliases(ctx, milvusclient.NewListAliasesOption("customized_setup_2"))
if err != nil {
// handle error
}
fmt.Println(aliases)
}
func ExampleClient_DescribeAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
alias, err := cli.DescribeAlias(ctx, milvusclient.NewDescribeAliasOption("bob"))
if err != nil {
// handle error
}
fmt.Println(alias)
}
func ExampleClient_AlterAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
err = cli.AlterAlias(ctx, milvusclient.NewAlterAliasOption("alice", "customized_setup_1"))
if err != nil {
// handle error
}
aliases, err := cli.ListAliases(ctx, milvusclient.NewListAliasesOption("customized_setup_1"))
if err != nil {
// handle error
}
fmt.Println(aliases)
aliases, err = cli.ListAliases(ctx, milvusclient.NewListAliasesOption("customized_setup_2"))
if err != nil {
// handle error
}
fmt.Println(aliases)
}
func ExampleClient_DropAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
err = cli.DropAlias(ctx, milvusclient.NewDropAliasOption("alice"))
if err != nil {
// handle error
}
}

View File

@ -0,0 +1,336 @@
// 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.
// nolint
package milvusclient_test
import (
"context"
"fmt"
"log"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/client/v2/index"
"github.com/milvus-io/milvus/client/v2/milvusclient"
"github.com/milvus-io/milvus/pkg/common"
)
const (
milvusAddr = `127.0.0.1:19530`
)
func ExampleClient_CreateCollection_normal() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := `customized_setup_1`
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle err
}
indexOptions := []milvusclient.CreateIndexOption{
milvusclient.NewCreateIndexOption(collectionName, "my_vector", index.NewAutoIndex(entity.COSINE)).WithIndexName("my_vector"),
milvusclient.NewCreateIndexOption(collectionName, "my_id", index.NewSortedIndex()).WithIndexName("my_id"),
}
schema := entity.NewSchema().WithDynamicFieldEnabled(true).
WithField(entity.NewField().WithName("my_id").WithIsAutoID(true).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
WithField(entity.NewField().WithName("my_vector").WithDataType(entity.FieldTypeFloatVector).WithDim(5)).
WithField(entity.NewField().WithName("my_varchar").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512))
err = cli.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
WithIndexOptions(indexOptions...),
)
if err != nil {
// handle error
}
}
func ExampleClient_CreateCollection_quick() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := `quick_setup`
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle err
}
err = cli.CreateCollection(ctx, milvusclient.SimpleCreateCollectionOptions(collectionName, 5))
if err != nil {
// handle error
}
}
func ExampleClient_CreateCollection_shardNum() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := `customized_setup_3`
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle err
}
schema := entity.NewSchema().WithDynamicFieldEnabled(true).
WithField(entity.NewField().WithName("my_id").WithIsAutoID(true).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
WithField(entity.NewField().WithName("my_vector").WithDataType(entity.FieldTypeFloatVector).WithDim(5)).
WithField(entity.NewField().WithName("my_varchar").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512))
err = cli.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).WithShardNum(1))
if err != nil {
// handle error
}
}
func ExampleClient_CreateCollection_enableMmap() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := `customized_setup_4`
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle err
}
schema := entity.NewSchema().WithDynamicFieldEnabled(true).
WithField(entity.NewField().WithName("my_id").WithIsAutoID(true).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
WithField(entity.NewField().WithName("my_vector").WithDataType(entity.FieldTypeFloatVector).WithDim(5)).
WithField(entity.NewField().WithName("my_varchar").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512))
err = cli.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).WithProperty(common.MmapEnabledKey, true))
if err != nil {
// handle error
}
}
func ExampleClient_CreateCollection_ttl() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := `customized_setup_5`
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle err
}
schema := entity.NewSchema().WithDynamicFieldEnabled(true).
WithField(entity.NewField().WithName("my_id").WithIsAutoID(true).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
WithField(entity.NewField().WithName("my_vector").WithDataType(entity.FieldTypeFloatVector).WithDim(5)).
WithField(entity.NewField().WithName("my_varchar").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512))
err = cli.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).WithProperty(common.CollectionTTLConfigKey, 86400))
if err != nil {
// handle error
}
}
func ExampleClient_CreateCollection_consistencyLevel() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := `customized_setup_5`
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle err
}
schema := entity.NewSchema().WithDynamicFieldEnabled(true).
WithField(entity.NewField().WithName("my_id").WithIsAutoID(true).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
WithField(entity.NewField().WithName("my_vector").WithDataType(entity.FieldTypeFloatVector).WithDim(5)).
WithField(entity.NewField().WithName("my_varchar").WithDataType(entity.FieldTypeVarChar).WithMaxLength(512))
err = cli.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).WithConsistencyLevel(entity.ClBounded))
if err != nil {
// handle error
}
}
func ExampleClient_ListCollections() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
collectionNames, err := cli.ListCollections(ctx, milvusclient.NewListCollectionOption())
if err != nil {
// handle error
}
fmt.Println(collectionNames)
}
func ExampleClient_DescribeCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
collection, err := cli.DescribeCollection(ctx, milvusclient.NewDescribeCollectionOption("quick_setup"))
if err != nil {
// handle error
}
fmt.Println(collection)
}
func ExampleClient_RenameCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
err = cli.RenameCollection(ctx, milvusclient.NewRenameCollectionOption("my_collection", "my_new_collection"))
if err != nil {
// handle error
}
}
func ExampleClient_AlterCollection_setTTL() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
err = cli.AlterCollection(ctx, milvusclient.NewAlterCollectionOption("my_collection").WithProperty(common.CollectionTTLConfigKey, 60))
if err != nil {
// handle error
}
}
func ExampleClient_LoadCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
loadTask, err := cli.LoadCollection(ctx, milvusclient.NewLoadCollectionOption("customized_setup_1"))
if err != nil {
// handle error
}
// sync wait collection to be loaded
err = loadTask.Await(ctx)
if err != nil {
// handle error
}
}
func ExampleClient_ReleaseCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
err = cli.ReleaseCollection(ctx, milvusclient.NewReleaseCollectionOption("custom_quick_setup"))
if err != nil {
// handle error
}
}
func ExampleClient_DropCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
err = cli.DropCollection(ctx, milvusclient.NewDropCollectionOption("customized_setup_2"))
if err != nil {
// handle err
}
}

View File

@ -0,0 +1,164 @@
// 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.
// nolint
package milvusclient_test
import (
"context"
"fmt"
"github.com/milvus-io/milvus/client/v2/milvusclient"
)
func ExampleClient_ListPartitions() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
defer cli.Close(ctx)
partitionNames, err := cli.ListPartitions(ctx, milvusclient.NewListPartitionOption("quick_setup"))
if err != nil {
// handle error
}
fmt.Println(partitionNames)
}
func ExampleClient_CreatePartition() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
defer cli.Close(ctx)
err = cli.CreatePartition(ctx, milvusclient.NewCreatePartitionOption("quick_setup", "partitionA"))
if err != nil {
// handle error
}
partitionNames, err := cli.ListPartitions(ctx, milvusclient.NewListPartitionOption("quick_setup"))
if err != nil {
// handle error
}
fmt.Println(partitionNames)
}
func ExampleClient_HasPartition() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
defer cli.Close(ctx)
result, err := cli.HasPartition(ctx, milvusclient.NewHasPartitionOption("quick_setup", "partitionA"))
if err != nil {
// handle error
}
fmt.Println(result)
}
func ExampleClient_LoadPartitions() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
defer cli.Close(ctx)
task, err := cli.LoadPartitions(ctx, milvusclient.NewLoadPartitionsOption("quick_setup", "partitionA"))
// sync wait collection to be loaded
err = task.Await(ctx)
if err != nil {
// handle error
}
}
func ExampleClient_ReleasePartitions() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
defer cli.Close(ctx)
err = cli.ReleasePartitions(ctx, milvusclient.NewReleasePartitionsOptions("quick_setup", "partitionA"))
if err != nil {
// handle error
}
}
func ExampleClient_DropPartition() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
})
if err != nil {
// handle error
}
defer cli.Close(ctx)
err = cli.DropPartition(ctx, milvusclient.NewDropPartitionOption("quick_setup", "partitionA"))
if err != nil {
// handle error
}
}

View File

@ -49,7 +49,7 @@ func (c *Client) Search(ctx context.Context, option SearchOption, callOptions ..
if err != nil {
return err
}
resultSets, err = c.handleSearchResult(collection.Schema, req.GetOutputFields(), int(req.GetNq()), resp)
resultSets, err = c.handleSearchResult(collection.Schema, req.GetOutputFields(), int(resp.GetResults().GetNumQueries()), resp)
return err
})
@ -189,6 +189,33 @@ func (c *Client) Query(ctx context.Context, option QueryOption, callOptions ...g
return resultSet, err
}
func (c *Client) HybridSearch(ctx context.Context, option HybridSearchOption, callOptions ...grpc.CallOption) ([]ResultSet, error) {
req, err := option.HybridRequest()
if err != nil {
return nil, err
}
collection, err := c.getCollection(ctx, req.GetCollectionName())
if err != nil {
return nil, err
}
var resultSets []ResultSet
err = c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.HybridSearch(ctx, req, callOptions...)
err = merr.CheckRPCCall(resp, err)
if err != nil {
return err
}
resultSets, err = c.handleSearchResult(collection.Schema, req.GetOutputFields(), int(resp.GetResults().GetNumQueries()), resp)
return err
})
return resultSets, err
}
func expandWildcard(schema *entity.Schema, outputFields []string) ([]string, bool) {
wildcard := false
for _, outputField := range outputFields {

View File

@ -0,0 +1,204 @@
// 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.
// nolint
package milvusclient_test
import (
"context"
"log"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/client/v2/milvusclient"
)
func ExampleClient_Search_basic() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
token := "root:Milvus"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
APIKey: token,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
queryVector := []float32{0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592}
resultSets, err := cli.Search(ctx, milvusclient.NewSearchOption(
"quick_setup", // collectionName
3, // limit
[]entity.Vector{entity.FloatVector(queryVector)},
))
if err != nil {
log.Fatal("failed to perform basic ANN search collection: ", err.Error())
}
for _, resultSet := range resultSets {
log.Println("IDs: ", resultSet.IDs)
log.Println("Scores: ", resultSet.Scores)
}
}
func ExampleClient_Search_multivectors() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
token := "root:Milvus"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
APIKey: token,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
queryVectors := []entity.Vector{
entity.FloatVector([]float32{0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592}),
entity.FloatVector([]float32{0.19886812562848388, 0.06023560599112088, 0.6976963061752597, 0.2614474506242501, 0.838729485096104}),
}
resultSets, err := cli.Search(ctx, milvusclient.NewSearchOption(
"quick_setup", // collectionName
3, // limit
queryVectors,
))
if err != nil {
log.Fatal("failed to perform basic ANN search collection: ", err.Error())
}
for _, resultSet := range resultSets {
log.Println("IDs: ", resultSet.IDs)
log.Println("Scores: ", resultSet.Scores)
}
}
func ExampleClient_Search_partition() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
token := "root:Milvus"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
APIKey: token,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
queryVector := []float32{0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592}
resultSets, err := cli.Search(ctx, milvusclient.NewSearchOption(
"quick_setup", // collectionName
3, // limit
[]entity.Vector{entity.FloatVector(queryVector)},
).WithPartitions("partitionA"))
if err != nil {
log.Fatal("failed to perform basic ANN search collection: ", err.Error())
}
for _, resultSet := range resultSets {
log.Println("IDs: ", resultSet.IDs)
log.Println("Scores: ", resultSet.Scores)
}
}
func ExampleClient_Search_outputFields() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
token := "root:Milvus"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
APIKey: token,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
queryVector := []float32{0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592}
resultSets, err := cli.Search(ctx, milvusclient.NewSearchOption(
"quick_setup", // collectionName
3, // limit
[]entity.Vector{entity.FloatVector(queryVector)},
).WithOutputFields("color"))
if err != nil {
log.Fatal("failed to perform basic ANN search collection: ", err.Error())
}
for _, resultSet := range resultSets {
log.Println("IDs: ", resultSet.IDs)
log.Println("Scores: ", resultSet.Scores)
log.Println("Colors: ", resultSet.GetColumn("color"))
}
}
func ExampleClient_Search_offsetLimit() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
milvusAddr := "127.0.0.1:19530"
token := "root:Milvus"
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: milvusAddr,
APIKey: token,
})
if err != nil {
log.Fatal("failed to connect to milvus server: ", err.Error())
}
defer cli.Close(ctx)
queryVector := []float32{0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592}
resultSets, err := cli.Search(ctx, milvusclient.NewSearchOption(
"quick_setup", // collectionName
3, // limit
[]entity.Vector{entity.FloatVector(queryVector)},
).WithOffset(10))
if err != nil {
log.Fatal("failed to perform basic ANN search collection: ", err.Error())
}
for _, resultSet := range resultSets {
log.Println("IDs: ", resultSet.IDs)
log.Println("Scores: ", resultSet.Scores)
}
}
// func ExampleClient_Search_useLevel() {
// }

View File

@ -26,6 +26,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/client/v2/index"
)
const (
@ -47,82 +48,139 @@ type SearchOption interface {
var _ SearchOption = (*searchOption)(nil)
type searchOption struct {
annRequest *annRequest
collectionName string
partitionNames []string
topK int
offset int
outputFields []string
consistencyLevel entity.ConsistencyLevel
useDefaultConsistencyLevel bool
ignoreGrowing bool
expr string
// normal search request
request *annRequest
// TODO add sub request when support hybrid search
}
type annRequest struct {
vectors []entity.Vector
annField string
metricsType entity.MetricType
searchParam map[string]string
groupByField string
annField string
metricsType entity.MetricType
searchParam map[string]string
groupByField string
annParam index.AnnParam
ignoreGrowing bool
expr string
topK int
offset int
}
func (opt *searchOption) Request() (*milvuspb.SearchRequest, error) {
// TODO check whether search is hybrid after logic merged
return opt.prepareSearchRequest(opt.request)
func NewAnnRequest(annField string, limit int, vectors ...entity.Vector) *annRequest {
return &annRequest{
annField: annField,
vectors: vectors,
topK: limit,
}
}
func (opt *searchOption) prepareSearchRequest(annRequest *annRequest) (*milvuspb.SearchRequest, error) {
func (r *annRequest) searchRequest() (*milvuspb.SearchRequest, error) {
request := &milvuspb.SearchRequest{
CollectionName: opt.collectionName,
PartitionNames: opt.partitionNames,
Dsl: opt.expr,
DslType: commonpb.DslType_BoolExprV1,
ConsistencyLevel: commonpb.ConsistencyLevel(opt.consistencyLevel),
OutputFields: opt.outputFields,
Nq: int64(len(r.vectors)),
Dsl: r.expr,
DslType: commonpb.DslType_BoolExprV1,
}
if annRequest != nil {
// nq
request.Nq = int64(len(annRequest.vectors))
// search param
bs, _ := json.Marshal(annRequest.searchParam)
params := map[string]string{
spAnnsField: annRequest.annField,
spTopK: strconv.Itoa(opt.topK),
spOffset: strconv.Itoa(opt.offset),
spParams: string(bs),
spMetricsType: string(annRequest.metricsType),
spRoundDecimal: "-1",
spIgnoreGrowing: strconv.FormatBool(opt.ignoreGrowing),
}
if annRequest.groupByField != "" {
params[spGroupBy] = annRequest.groupByField
}
request.SearchParams = entity.MapKvPairs(params)
var err error
// placeholder group
request.PlaceholderGroup, err = vector2PlaceholderGroupBytes(annRequest.vectors)
if err != nil {
return nil, err
}
var err error
// placeholder group
request.PlaceholderGroup, err = vector2PlaceholderGroupBytes(r.vectors)
if err != nil {
return nil, err
}
params := map[string]string{
spAnnsField: r.annField,
spTopK: strconv.Itoa(r.topK),
spOffset: strconv.Itoa(r.offset),
spMetricsType: string(r.metricsType),
spRoundDecimal: "-1",
spIgnoreGrowing: strconv.FormatBool(r.ignoreGrowing),
}
if r.groupByField != "" {
params[spGroupBy] = r.groupByField
}
// ann param
if r.annParam != nil {
bs, _ := json.Marshal(r.annParam.Params())
params[spParams] = string(bs)
} else {
params[spParams] = "{}"
}
// use custom search param to overwrite
for k, v := range r.searchParam {
params[k] = v
}
request.SearchParams = entity.MapKvPairs(params)
return request, nil
}
func (r *annRequest) WithANNSField(annsField string) *annRequest {
r.annField = annsField
return r
}
func (r *annRequest) WithGroupByField(groupByField string) *annRequest {
r.groupByField = groupByField
return r
}
func (r *annRequest) WithSearchParam(key, value string) *annRequest {
r.searchParam[key] = value
return r
}
func (r *annRequest) WithAnnParam(ap index.AnnParam) *annRequest {
r.annParam = ap
return r
}
func (r *annRequest) WithFilter(expr string) *annRequest {
r.expr = expr
return r
}
func (r *annRequest) WithOffset(offset int) *annRequest {
r.offset = offset
return r
}
func (r *annRequest) WithIgnoreGrowing(ignoreGrowing bool) *annRequest {
r.ignoreGrowing = ignoreGrowing
return r
}
func (opt *searchOption) Request() (*milvuspb.SearchRequest, error) {
request, err := opt.annRequest.searchRequest()
if err != nil {
return nil, err
}
request.CollectionName = opt.collectionName
request.PartitionNames = opt.partitionNames
request.ConsistencyLevel = commonpb.ConsistencyLevel(opt.consistencyLevel)
request.UseDefaultConsistency = opt.useDefaultConsistencyLevel
request.OutputFields = opt.outputFields
return request, nil
}
func (opt *searchOption) WithPartitions(partitionNames ...string) *searchOption {
opt.partitionNames = partitionNames
return opt
}
func (opt *searchOption) WithFilter(expr string) *searchOption {
opt.expr = expr
opt.annRequest.WithFilter(expr)
return opt
}
func (opt *searchOption) WithOffset(offset int) *searchOption {
opt.offset = offset
opt.annRequest.WithOffset(offset)
return opt
}
@ -138,27 +196,38 @@ func (opt *searchOption) WithConsistencyLevel(consistencyLevel entity.Consistenc
}
func (opt *searchOption) WithANNSField(annsField string) *searchOption {
opt.request.annField = annsField
return opt
}
func (opt *searchOption) WithPartitions(partitionNames ...string) *searchOption {
opt.partitionNames = partitionNames
opt.annRequest.WithANNSField(annsField)
return opt
}
func (opt *searchOption) WithGroupByField(groupByField string) *searchOption {
opt.request.groupByField = groupByField
opt.annRequest.WithGroupByField(groupByField)
return opt
}
func (opt *searchOption) WithIgnoreGrowing(ignoreGrowing bool) *searchOption {
opt.annRequest.WithIgnoreGrowing(ignoreGrowing)
return opt
}
func (opt *searchOption) WithAnnParam(ap index.AnnParam) *searchOption {
opt.annRequest.WithAnnParam(ap)
return opt
}
func (opt *searchOption) WithSearchParam(key, value string) *searchOption {
opt.annRequest.WithSearchParam(key, value)
return opt
}
func NewSearchOption(collectionName string, limit int, vectors []entity.Vector) *searchOption {
return &searchOption{
collectionName: collectionName,
topK: limit,
request: &annRequest{
vectors: vectors,
annRequest: &annRequest{
vectors: vectors,
searchParam: make(map[string]string),
topK: limit,
},
collectionName: collectionName,
useDefaultConsistencyLevel: true,
consistencyLevel: entity.ClBounded,
}
@ -211,6 +280,66 @@ func vector2Placeholder(vectors []entity.Vector) (*commonpb.PlaceholderValue, er
return ph, nil
}
type HybridSearchOption interface {
HybridRequest() (*milvuspb.HybridSearchRequest, error)
}
type hybridSearchOption struct {
collectionName string
partitionNames []string
reqs []*annRequest
outputFields []string
useDefaultConsistency bool
consistencyLevel entity.ConsistencyLevel
}
func (opt *hybridSearchOption) WithConsistencyLevel(cl entity.ConsistencyLevel) *hybridSearchOption {
opt.consistencyLevel = cl
opt.useDefaultConsistency = false
return opt
}
func (opt *hybridSearchOption) WithPartitons(partitions ...string) *hybridSearchOption {
opt.partitionNames = partitions
return opt
}
func (opt *hybridSearchOption) WithOutputFields(outputFields ...string) *hybridSearchOption {
opt.outputFields = outputFields
return opt
}
func (opt *hybridSearchOption) HybridRequest() (*milvuspb.HybridSearchRequest, error) {
requests := make([]*milvuspb.SearchRequest, 0, len(opt.reqs))
for _, annRequest := range opt.reqs {
req, err := annRequest.searchRequest()
if err != nil {
return nil, err
}
requests = append(requests, req)
}
return &milvuspb.HybridSearchRequest{
CollectionName: opt.collectionName,
PartitionNames: opt.partitionNames,
Requests: requests,
UseDefaultConsistency: opt.useDefaultConsistency,
ConsistencyLevel: commonpb.ConsistencyLevel(opt.consistencyLevel),
OutputFields: opt.outputFields,
}, nil
}
func NewHybridSearchOption(collectionName string, annRequests ...*annRequest) *hybridSearchOption {
return &hybridSearchOption{
collectionName: collectionName,
reqs: annRequests,
useDefaultConsistency: true,
}
}
type QueryOption interface {
Request() *milvuspb.QueryRequest
}

View File

@ -13,6 +13,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/client/v2/index"
"github.com/milvus-io/milvus/pkg/util/merr"
)
@ -45,7 +46,6 @@ func (s *ReadSuite) TestSearch() {
s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
s.Equal(collectionName, sr.GetCollectionName())
s.ElementsMatch([]string{partitionName}, sr.GetPartitionNames())
// s.Equal(s)
return &milvuspb.SearchResults{
Status: merr.Success(),
@ -68,11 +68,17 @@ func (s *ReadSuite) TestSearch() {
}, nil
}).Once()
ap := index.NewCustomAnnParam()
ap.WithExtraParam("custom_level", 1)
_, err := s.client.Search(ctx, NewSearchOption(collectionName, 10, []entity.Vector{
entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
return rand.Float32()
})),
}).WithPartitions(partitionName).WithGroupByField("group_by"))
}).WithPartitions(partitionName).
WithGroupByField("group_by").
WithSearchParam("ignore_growing", "true").
WithAnnParam(ap),
)
s.NoError(err)
})
@ -134,6 +140,72 @@ func (s *ReadSuite) TestSearch() {
})
}
func (s *ReadSuite) TestHybridSearch() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
s.Run("success", func() {
collectionName := fmt.Sprintf("coll_%s", s.randString(6))
partitionName := fmt.Sprintf("part_%s", s.randString(6))
s.setupCache(collectionName, s.schema)
s.mock.EXPECT().HybridSearch(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, hsr *milvuspb.HybridSearchRequest) (*milvuspb.SearchResults, error) {
s.Equal(collectionName, hsr.GetCollectionName())
s.ElementsMatch([]string{partitionName}, hsr.GetPartitionNames())
s.ElementsMatch([]string{"*"}, hsr.GetOutputFields())
return &milvuspb.SearchResults{
Status: merr.Success(),
Results: &schemapb.SearchResultData{
NumQueries: 1,
TopK: 2,
FieldsData: []*schemapb.FieldData{
s.getInt64FieldData("ID", []int64{1, 2}),
s.getJSONBytesFieldData("$meta", [][]byte{
[]byte(`{"A": 123, "B": "456"}`),
[]byte(`{"B": "abc", "A": 456}`),
}, true),
},
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{1, 2},
},
},
},
Scores: make([]float32, 2),
Topks: []int64{2},
},
}, nil
}).Once()
_, err := s.client.HybridSearch(ctx, NewHybridSearchOption(collectionName, NewAnnRequest("vector", 10, entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
return rand.Float32()
}))).WithFilter("ID > 100"), NewAnnRequest("vector", 10, entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
return rand.Float32()
})))).WithConsistencyLevel(entity.ClStrong).WithPartitons(partitionName).WithOutputFields("*"))
s.NoError(err)
})
s.Run("failure", func() {
collectionName := fmt.Sprintf("coll_%s", s.randString(6))
s.setupCache(collectionName, s.schemaDyn)
_, err := s.client.HybridSearch(ctx, NewHybridSearchOption(collectionName, NewAnnRequest("vector", 10, nonSupportData{})))
s.Error(err)
s.mock.EXPECT().HybridSearch(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, hsr *milvuspb.HybridSearchRequest) (*milvuspb.SearchResults, error) {
return nil, merr.WrapErrServiceInternal("mocked")
}).Once()
_, err = s.client.HybridSearch(ctx, NewHybridSearchOption(collectionName, NewAnnRequest("vector", 10, entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
return rand.Float32()
}))).WithFilter("ID > 100"), NewAnnRequest("vector", 10, entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
return rand.Float32()
})))))
s.Error(err)
})
}
func (s *ReadSuite) TestQuery() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()