mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
Bingyi SunandGitHub
70f144a91c
Shard data by namespace. issue: #50154 --------- Signed-off-by: sunby <sunbingyi1992@gmail.com>
624 lines
18 KiB
Go
624 lines
18 KiB
Go
// 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 milvusclient
|
|
|
|
import (
|
|
"fmt"
|
|
"math/rand"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus/client/v2/column"
|
|
"github.com/milvus-io/milvus/client/v2/entity"
|
|
)
|
|
|
|
type SearchOptionSuite struct {
|
|
suite.Suite
|
|
}
|
|
|
|
type nonSupportData struct{}
|
|
|
|
func (d nonSupportData) Serialize() []byte {
|
|
return []byte{}
|
|
}
|
|
|
|
func (d nonSupportData) Dim() int {
|
|
return 0
|
|
}
|
|
|
|
func (d nonSupportData) FieldType() entity.FieldType {
|
|
return entity.FieldType(0)
|
|
}
|
|
|
|
func (s *SearchOptionSuite) TestBasic() {
|
|
collName := "search_opt_basic"
|
|
|
|
topK := rand.Intn(100) + 1
|
|
opt := NewSearchOption(collName, topK, []entity.Vector{entity.FloatVector([]float32{0.1, 0.2})})
|
|
|
|
rerankerFunction := entity.NewFunction().WithName("time_decay").WithInputFields("timestamp").WithType(entity.FunctionTypeRerank).
|
|
WithParam("reranker", "decay").
|
|
WithParam("function", "gauss").
|
|
WithParam("origin", 1754995249).
|
|
WithParam("scale", 7*24*60*60).
|
|
WithParam("offset", 24*60*60).
|
|
WithParam("decay", 0.5)
|
|
|
|
opt = opt.WithANNSField("test_field").WithOutputFields("ID", "Value").WithConsistencyLevel(entity.ClStrong).WithFilter("ID > 1000").WithGroupByField("group_field").WithGroupSize(10).WithStrictGroupSize(true).WithFunctionReranker(rerankerFunction)
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
|
|
s.Equal(collName, req.GetCollectionName())
|
|
s.Equal("ID > 1000", req.GetDsl())
|
|
s.ElementsMatch([]string{"ID", "Value"}, req.GetOutputFields())
|
|
searchParams := entity.KvPairsMap(req.GetSearchParams())
|
|
annField, ok := searchParams[spAnnsField]
|
|
s.Require().True(ok)
|
|
s.Equal("test_field", annField)
|
|
groupField, ok := searchParams[spGroupBy]
|
|
s.Require().True(ok)
|
|
s.Equal("group_field", groupField)
|
|
groupSize, ok := searchParams[spGroupSize]
|
|
s.Require().True(ok)
|
|
s.Equal("10", groupSize)
|
|
spStrictGroupSize, ok := searchParams[spStrictGroupSize]
|
|
s.Require().True(ok)
|
|
s.Equal("true", spStrictGroupSize)
|
|
|
|
functionScore := req.GetFunctionScore()
|
|
s.Len(functionScore.GetFunctions(), 1)
|
|
|
|
opt = NewSearchOption(collName, topK, []entity.Vector{nonSupportData{}})
|
|
_, err = opt.Request()
|
|
s.Error(err)
|
|
}
|
|
|
|
func (s *SearchOptionSuite) TestWithNamespace() {
|
|
collName := "namespace_read_option"
|
|
namespace := "tenant_a"
|
|
|
|
searchReq, err := NewSearchOption(collName, 10, []entity.Vector{entity.FloatVector([]float32{0.1, 0.2})}).
|
|
WithNamespace(namespace).
|
|
Request()
|
|
s.Require().NoError(err)
|
|
s.Equal(namespace, searchReq.GetNamespace())
|
|
|
|
queryReq, err := NewQueryOption(collName).
|
|
WithNamespace(namespace).
|
|
Request()
|
|
s.Require().NoError(err)
|
|
s.Equal(namespace, queryReq.GetNamespace())
|
|
|
|
hybridReq, err := NewHybridSearchOption(collName, 10, NewAnnRequest("vector", 10, entity.FloatVector([]float32{0.1, 0.2}))).
|
|
WithNamespace(namespace).
|
|
HybridRequest()
|
|
s.Require().NoError(err)
|
|
s.Equal(namespace, hybridReq.GetNamespace())
|
|
}
|
|
|
|
func (s *SearchOptionSuite) TestPlaceHolder() {
|
|
type testCase struct {
|
|
tag string
|
|
input []entity.Vector
|
|
expectError bool
|
|
expectType commonpb.PlaceholderType
|
|
}
|
|
|
|
sparse, err := entity.NewSliceSparseEmbedding([]uint32{0, 10, 12}, []float32{0.1, 0.2, 0.3})
|
|
s.Require().NoError(err)
|
|
|
|
cases := []*testCase{
|
|
{
|
|
tag: "empty_input",
|
|
input: nil,
|
|
expectType: commonpb.PlaceholderType_None,
|
|
},
|
|
{
|
|
tag: "float_vector",
|
|
input: []entity.Vector{entity.FloatVector([]float32{0.1, 0.2, 0.3})},
|
|
expectType: commonpb.PlaceholderType_FloatVector,
|
|
},
|
|
{
|
|
tag: "sparse_vector",
|
|
input: []entity.Vector{sparse},
|
|
expectType: commonpb.PlaceholderType_SparseFloatVector,
|
|
},
|
|
{
|
|
tag: "fp16_vector",
|
|
input: []entity.Vector{entity.Float16Vector([]byte{})},
|
|
expectType: commonpb.PlaceholderType_Float16Vector,
|
|
},
|
|
{
|
|
tag: "bf16_vector",
|
|
input: []entity.Vector{entity.BFloat16Vector([]byte{})},
|
|
expectType: commonpb.PlaceholderType_BFloat16Vector,
|
|
},
|
|
{
|
|
tag: "binary_vector",
|
|
input: []entity.Vector{entity.BinaryVector([]byte{})},
|
|
expectType: commonpb.PlaceholderType_BinaryVector,
|
|
},
|
|
{
|
|
tag: "int8_vector",
|
|
input: []entity.Vector{entity.Int8Vector([]int8{})},
|
|
expectType: commonpb.PlaceholderType_Int8Vector,
|
|
},
|
|
{
|
|
tag: "text",
|
|
input: []entity.Vector{entity.Text("abc")},
|
|
expectType: commonpb.PlaceholderType_VarChar,
|
|
},
|
|
{
|
|
tag: "emb_list_float",
|
|
input: []entity.Vector{entity.FloatVectorArray{entity.FloatVector([]float32{0.1, 0.2})}},
|
|
expectType: commonpb.PlaceholderType_EmbListFloatVector,
|
|
},
|
|
{
|
|
tag: "emb_list_fp16",
|
|
input: []entity.Vector{entity.Float16VectorArray{entity.Float16Vector([]byte{})}},
|
|
expectType: commonpb.PlaceholderType_EmbListFloat16Vector,
|
|
},
|
|
{
|
|
tag: "emb_list_bf16",
|
|
input: []entity.Vector{entity.BFloat16VectorArray{entity.BFloat16Vector([]byte{})}},
|
|
expectType: commonpb.PlaceholderType_EmbListBFloat16Vector,
|
|
},
|
|
{
|
|
tag: "emb_list_binary",
|
|
input: []entity.Vector{entity.BinaryVectorArray{entity.BinaryVector([]byte{})}},
|
|
expectType: commonpb.PlaceholderType_EmbListBinaryVector,
|
|
},
|
|
{
|
|
tag: "emb_list_int8",
|
|
input: []entity.Vector{entity.Int8VectorArray{entity.Int8Vector([]int8{})}},
|
|
expectType: commonpb.PlaceholderType_EmbListInt8Vector,
|
|
},
|
|
{
|
|
tag: "non_supported",
|
|
input: []entity.Vector{nonSupportData{}},
|
|
expectError: true,
|
|
},
|
|
}
|
|
for _, tc := range cases {
|
|
s.Run(tc.tag, func() {
|
|
phv, err := vector2Placeholder(tc.input)
|
|
if tc.expectError {
|
|
s.Error(err)
|
|
} else {
|
|
s.NoError(err)
|
|
s.Equal(tc.expectType, phv.GetType())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *SearchOptionSuite) TestSearchAggregationOption() {
|
|
opt := NewSearchOption("coll", 100, []entity.Vector{entity.FloatVector([]float32{0.1, 0.2})}).
|
|
WithSearchAggregation(
|
|
NewSearchAggregation([]string{"brand"}, 3).
|
|
WithMetric("count_all", "count", "*").
|
|
WithTopHits(NewTopHits(2).WithSort("_score", "asc")),
|
|
)
|
|
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
s.NotNil(req.GetSearchAggregation())
|
|
s.Equal([]string{"brand"}, req.GetSearchAggregation().GetFields())
|
|
s.EqualValues(3, req.GetSearchAggregation().GetSize())
|
|
s.Equal("count", req.GetSearchAggregation().GetMetrics()["count_all"].GetOp())
|
|
}
|
|
|
|
func (s *SearchOptionSuite) TestSearchAggregationRejectsConflictingOptions() {
|
|
base := func() *searchOption {
|
|
return NewSearchOption("coll", 10, []entity.Vector{entity.FloatVector([]float32{0.1, 0.2})}).
|
|
WithSearchAggregation(NewSearchAggregation([]string{"brand"}, 3))
|
|
}
|
|
|
|
cases := []struct {
|
|
name string
|
|
opt *searchOption
|
|
msg string
|
|
}{
|
|
{
|
|
name: "legacy group by field",
|
|
opt: base().WithGroupByField("brand"),
|
|
msg: "search_aggregation and group_by_field/group_size are mutually exclusive",
|
|
},
|
|
{
|
|
name: "legacy group size",
|
|
opt: base().WithGroupSize(2),
|
|
msg: "search_aggregation and group_by_field/group_size are mutually exclusive",
|
|
},
|
|
{
|
|
name: "legacy strict group size",
|
|
opt: base().WithStrictGroupSize(true),
|
|
msg: "search_aggregation and group_by_field/group_size are mutually exclusive",
|
|
},
|
|
{
|
|
name: "offset",
|
|
opt: base().WithOffset(1),
|
|
msg: "offset is not supported with search_aggregation",
|
|
},
|
|
{
|
|
name: "raw offset param",
|
|
opt: base().WithSearchParam("offset", "1"),
|
|
msg: "offset is not supported with search_aggregation",
|
|
},
|
|
{
|
|
name: "raw group_by_field param",
|
|
opt: base().WithSearchParam("group_by_field", "brand"),
|
|
msg: "group_by_field and search_aggregation cannot be used simultaneously",
|
|
},
|
|
{
|
|
name: "raw group_by_fields param",
|
|
opt: base().WithSearchParam("group_by_fields", "brand,color"),
|
|
msg: "group_by_fields and search_aggregation cannot be used simultaneously",
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
s.Run(tc.name, func() {
|
|
_, err := tc.opt.Request()
|
|
s.Require().Error(err)
|
|
s.Contains(err.Error(), tc.msg)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSearchOption(t *testing.T) {
|
|
suite.Run(t, new(SearchOptionSuite))
|
|
}
|
|
|
|
func TestAny2TmplValue(t *testing.T) {
|
|
t.Run("primitives", func(t *testing.T) {
|
|
t.Run("int", func(t *testing.T) {
|
|
v := rand.Int()
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, v, val.GetInt64Val())
|
|
})
|
|
|
|
t.Run("int32", func(t *testing.T) {
|
|
v := rand.Int31()
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, v, val.GetInt64Val())
|
|
})
|
|
|
|
t.Run("int64", func(t *testing.T) {
|
|
v := rand.Int63()
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, v, val.GetInt64Val())
|
|
})
|
|
|
|
t.Run("float32", func(t *testing.T) {
|
|
v := rand.Float32()
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, v, val.GetFloatVal())
|
|
})
|
|
|
|
t.Run("float64", func(t *testing.T) {
|
|
v := rand.Float64()
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, v, val.GetFloatVal())
|
|
})
|
|
|
|
t.Run("bool", func(t *testing.T) {
|
|
val, err := any2TmplValue(true)
|
|
assert.NoError(t, err)
|
|
assert.True(t, val.GetBoolVal())
|
|
})
|
|
|
|
t.Run("string", func(t *testing.T) {
|
|
v := fmt.Sprintf("%v", rand.Int())
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, v, val.GetStringVal())
|
|
})
|
|
})
|
|
|
|
t.Run("slice", func(t *testing.T) {
|
|
t.Run("int", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]int, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, rand.Int())
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetLongData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
|
|
t.Run("int32", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]int32, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, rand.Int31())
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetLongData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
|
|
t.Run("int64", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]int64, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, rand.Int63())
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetLongData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
|
|
t.Run("float32", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]float32, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, rand.Float32())
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetDoubleData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
|
|
t.Run("float64", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]float64, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, rand.Float64())
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetDoubleData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
|
|
t.Run("bool", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]bool, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, rand.Int()%2 == 0)
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetBoolData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
|
|
t.Run("string", func(t *testing.T) {
|
|
l := rand.Intn(10) + 1
|
|
v := make([]string, 0, l)
|
|
for i := 0; i < l; i++ {
|
|
v = append(v, fmt.Sprintf("%v", rand.Int()))
|
|
}
|
|
val, err := any2TmplValue(v)
|
|
assert.NoError(t, err)
|
|
data := val.GetArrayVal().GetStringData().GetData()
|
|
assert.Equal(t, l, len(data))
|
|
for i, val := range data {
|
|
assert.EqualValues(t, v[i], val)
|
|
}
|
|
})
|
|
})
|
|
|
|
t.Run("unsupported", func(*testing.T) {
|
|
_, err := any2TmplValue(struct{}{})
|
|
assert.Error(t, err)
|
|
|
|
_, err = any2TmplValue([]struct{}{})
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func (s *SearchOptionSuite) TestSearchByIDs() {
|
|
collName := "search_by_ids_test"
|
|
limit := 10
|
|
|
|
s.Run("int64_ids", func() {
|
|
ids := column.NewColumnInt64("id", []int64{1, 2, 3, 4, 5})
|
|
opt := NewSearchByIDsOption(collName, limit, ids)
|
|
opt = opt.WithANNSField("vector_field").WithConsistencyLevel(entity.ClStrong)
|
|
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
|
|
s.Equal(collName, req.GetCollectionName())
|
|
s.Equal(int64(5), req.GetNq()) // 5 IDs
|
|
s.NotNil(req.GetIds())
|
|
s.Equal([]int64{1, 2, 3, 4, 5}, req.GetIds().GetIntId().GetData())
|
|
})
|
|
|
|
s.Run("varchar_ids", func() {
|
|
ids := column.NewColumnVarChar("id", []string{"a", "b", "c"})
|
|
opt := NewSearchByIDsOption(collName, limit, ids)
|
|
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
|
|
s.Equal(int64(3), req.GetNq()) // 3 IDs
|
|
s.NotNil(req.GetIds())
|
|
s.Equal([]string{"a", "b", "c"}, req.GetIds().GetStrId().GetData())
|
|
})
|
|
|
|
s.Run("empty_ids_error", func() {
|
|
ids := column.NewColumnInt64("id", []int64{})
|
|
opt := NewSearchByIDsOption(collName, limit, ids)
|
|
|
|
_, err := opt.Request()
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "cannot be empty")
|
|
})
|
|
|
|
s.Run("with_filter", func() {
|
|
ids := column.NewColumnInt64("id", []int64{1, 2, 3})
|
|
opt := NewSearchByIDsOption(collName, limit, ids).
|
|
WithFilter("status == 'active'").
|
|
WithANNSField("vector")
|
|
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
|
|
s.Equal("status == 'active'", req.GetDsl())
|
|
s.Equal(int64(3), req.GetNq())
|
|
})
|
|
|
|
s.Run("with_output_fields", func() {
|
|
ids := column.NewColumnInt64("id", []int64{1, 2})
|
|
opt := NewSearchByIDsOption(collName, limit, ids).
|
|
WithOutputFields("id", "name", "vector")
|
|
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
|
|
s.ElementsMatch([]string{"id", "name", "vector"}, req.GetOutputFields())
|
|
})
|
|
|
|
s.Run("single_id", func() {
|
|
ids := column.NewColumnInt64("id", []int64{42})
|
|
opt := NewSearchByIDsOption(collName, limit, ids)
|
|
|
|
req, err := opt.Request()
|
|
s.Require().NoError(err)
|
|
|
|
s.Equal(int64(1), req.GetNq())
|
|
s.Equal([]int64{42}, req.GetIds().GetIntId().GetData())
|
|
})
|
|
|
|
s.Run("unsupported_id_type_error", func() {
|
|
// Float column is not a valid primary key type
|
|
ids := column.NewColumnFloat("id", []float32{1.0, 2.0})
|
|
opt := NewSearchByIDsOption(collName, limit, ids)
|
|
|
|
_, err := opt.Request()
|
|
s.Error(err)
|
|
s.Contains(err.Error(), "failed to convert IDs column")
|
|
s.Contains(err.Error(), "unsupported primary key type")
|
|
})
|
|
}
|
|
|
|
func TestColumn2IDs(t *testing.T) {
|
|
t.Run("int64_column", func(t *testing.T) {
|
|
col := column.NewColumnInt64("pk", []int64{100, 200, 300})
|
|
ids, err := column2IDs(col)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, ids.GetIntId())
|
|
assert.Equal(t, []int64{100, 200, 300}, ids.GetIntId().GetData())
|
|
})
|
|
|
|
t.Run("varchar_column", func(t *testing.T) {
|
|
col := column.NewColumnVarChar("pk", []string{"id1", "id2", "id3"})
|
|
ids, err := column2IDs(col)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, ids.GetStrId())
|
|
assert.Equal(t, []string{"id1", "id2", "id3"}, ids.GetStrId().GetData())
|
|
})
|
|
|
|
t.Run("string_column", func(t *testing.T) {
|
|
col := column.NewColumnString("pk", []string{"str1", "str2"})
|
|
ids, err := column2IDs(col)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, ids.GetStrId())
|
|
assert.Equal(t, []string{"str1", "str2"}, ids.GetStrId().GetData())
|
|
})
|
|
|
|
t.Run("nil_column_error", func(t *testing.T) {
|
|
_, err := column2IDs(nil)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "cannot be nil")
|
|
})
|
|
|
|
t.Run("unsupported_type_error", func(t *testing.T) {
|
|
col := column.NewColumnFloat("pk", []float32{1.0, 2.0})
|
|
_, err := column2IDs(col)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "unsupported primary key type")
|
|
})
|
|
}
|
|
|
|
func TestAnnRequestWithIDs(t *testing.T) {
|
|
t.Run("basic_with_ids", func(t *testing.T) {
|
|
ids := column.NewColumnInt64("pk", []int64{1, 2, 3})
|
|
req := NewAnnRequest("vector_field", 10)
|
|
req.WithIDs(ids)
|
|
|
|
searchReq, err := req.searchRequest()
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, int64(3), searchReq.GetNq())
|
|
assert.NotNil(t, searchReq.GetIds())
|
|
})
|
|
|
|
t.Run("ids_override_vectors", func(t *testing.T) {
|
|
// When IDs are set, vectors should be ignored
|
|
vectors := []entity.Vector{entity.FloatVector([]float32{0.1, 0.2, 0.3})}
|
|
ids := column.NewColumnInt64("pk", []int64{1, 2})
|
|
|
|
req := NewAnnRequest("vector_field", 10, vectors...)
|
|
req.WithIDs(ids)
|
|
|
|
searchReq, err := req.searchRequest()
|
|
assert.NoError(t, err)
|
|
// Nq should be based on IDs count, not vectors count
|
|
assert.Equal(t, int64(2), searchReq.GetNq())
|
|
assert.NotNil(t, searchReq.GetIds())
|
|
assert.Nil(t, searchReq.GetPlaceholderGroup())
|
|
})
|
|
|
|
t.Run("with_search_params", func(t *testing.T) {
|
|
ids := column.NewColumnInt64("pk", []int64{1, 2, 3})
|
|
req := NewAnnRequest("vector_field", 10).
|
|
WithIDs(ids).
|
|
WithSearchParam("nprobe", "16").
|
|
WithFilter("category == 'A'")
|
|
|
|
searchReq, err := req.searchRequest()
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "category == 'A'", searchReq.GetDsl())
|
|
|
|
params := entity.KvPairsMap(searchReq.GetSearchParams())
|
|
assert.Equal(t, "16", params["nprobe"])
|
|
})
|
|
}
|