mirror of https://github.com/milvus-io/milvus.git
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
GitHub
parent
f0096ec292
commit
1ec858434f
|
|
@ -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),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ const (
|
|||
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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
||||
// }
|
||||
|
|
@ -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,19 +48,12 @@ 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 {
|
||||
|
|
@ -69,60 +63,124 @@ type annRequest struct {
|
|||
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,
|
||||
Nq: int64(len(r.vectors)),
|
||||
Dsl: r.expr,
|
||||
DslType: commonpb.DslType_BoolExprV1,
|
||||
ConsistencyLevel: commonpb.ConsistencyLevel(opt.consistencyLevel),
|
||||
OutputFields: opt.outputFields,
|
||||
}
|
||||
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)
|
||||
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{
|
||||
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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue