mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-06 17:18:35 +08:00
f23df95a77
https://github.com/milvus-io/milvus/issues/35856 #41312 Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
435 lines
14 KiB
Go
435 lines
14 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 rerank
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/suite"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/testutils"
|
|
)
|
|
|
|
func TestDecayFunction(t *testing.T) {
|
|
suite.Run(t, new(DecayFunctionSuite))
|
|
}
|
|
|
|
type DecayFunctionSuite struct {
|
|
suite.Suite
|
|
schema *schemapb.CollectionSchema
|
|
providers []string
|
|
}
|
|
|
|
func (s *DecayFunctionSuite) TestNewDecayErrors() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
|
|
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
|
|
{
|
|
FieldID: 102, Name: "vector", DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "4"},
|
|
},
|
|
},
|
|
{FieldID: 102, Name: "ts", DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
functionSchema := &schemapb.FunctionSchema{
|
|
Name: "test",
|
|
Type: schemapb.FunctionType_Rerank,
|
|
InputFieldNames: []string{"ts"},
|
|
OutputFieldNames: []string{"text"},
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: originKey, Value: "4"},
|
|
{Key: scaleKey, Value: "4"},
|
|
{Key: offsetKey, Value: "4"},
|
|
{Key: decayKey, Value: "0.5"},
|
|
{Key: functionKey, Value: "gauss"},
|
|
},
|
|
}
|
|
|
|
{
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Rerank function output field names should be empty")
|
|
}
|
|
|
|
{
|
|
functionSchema.OutputFieldNames = []string{}
|
|
functionSchema.InputFieldNames = []string{""}
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Rerank input field name cannot be empty string")
|
|
}
|
|
|
|
{
|
|
functionSchema.InputFieldNames = []string{"ts", "ts"}
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Each function input field should be used exactly once in the same function")
|
|
}
|
|
|
|
{
|
|
functionSchema.InputFieldNames = []string{"ts", "pk"}
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay function only supoorts single input, but gets")
|
|
}
|
|
|
|
{
|
|
functionSchema.InputFieldNames = []string{"notExists"}
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Function input field not found:")
|
|
}
|
|
|
|
{
|
|
functionSchema.InputFieldNames = []string{"vector"}
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay rerank: unsupported input field type")
|
|
}
|
|
|
|
{
|
|
functionSchema.InputFieldNames = []string{"ts"}
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.NoError(err)
|
|
}
|
|
|
|
{
|
|
for i := 0; i < 4; i++ {
|
|
functionSchema.Params[i].Value = "NotNum"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "is not a number")
|
|
functionSchema.Params[i].Value = "0.9"
|
|
}
|
|
}
|
|
|
|
{
|
|
fs := []string{gaussFunction, linerFunction, expFunction}
|
|
for i := 0; i < 3; i++ {
|
|
functionSchema.Params[4].Value = fs[i]
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.NoError(err)
|
|
}
|
|
functionSchema.Params[4].Value = "NotExist"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Invaild decay function:")
|
|
functionSchema.Params[4].Value = "exp"
|
|
}
|
|
|
|
{
|
|
newSchema := proto.Clone(schema).(*schemapb.CollectionSchema)
|
|
newSchema.Fields[0].IsPrimaryKey = false
|
|
_, err := newDecayFunction(newSchema, functionSchema)
|
|
s.ErrorContains(err, " can not found pk field")
|
|
}
|
|
}
|
|
|
|
func (s *DecayFunctionSuite) TestAllTypesInput() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
|
|
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
|
|
{
|
|
FieldID: 102, Name: "vector", DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "4"},
|
|
},
|
|
},
|
|
{FieldID: 102, Name: "ts", DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
|
|
functionSchema := &schemapb.FunctionSchema{
|
|
Name: "test",
|
|
Type: schemapb.FunctionType_Rerank,
|
|
InputFieldNames: []string{"ts"},
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: functionKey, Value: "gauss"},
|
|
{Key: originKey, Value: "4"},
|
|
{Key: scaleKey, Value: "4"},
|
|
{Key: offsetKey, Value: "4"},
|
|
{Key: decayKey, Value: "0.5"},
|
|
},
|
|
}
|
|
inputTypes := []schemapb.DataType{schemapb.DataType_Int64, schemapb.DataType_Int32, schemapb.DataType_Int16, schemapb.DataType_Int8, schemapb.DataType_Float, schemapb.DataType_Double, schemapb.DataType_Bool}
|
|
for i, inputType := range inputTypes {
|
|
schema.Fields[3].DataType = inputType
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
if i < len(inputTypes)-1 {
|
|
s.NoError(err)
|
|
} else {
|
|
s.ErrorContains(err, "Decay rerank: unsupported input field type")
|
|
}
|
|
}
|
|
|
|
schema.Fields[0].DataType = schemapb.DataType_String
|
|
for i, inputType := range inputTypes {
|
|
schema.Fields[3].DataType = inputType
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
if i < len(inputTypes)-1 {
|
|
s.NoError(err)
|
|
} else {
|
|
s.ErrorContains(err, "Decay rerank: unsupported input field type")
|
|
}
|
|
}
|
|
|
|
schema.Fields[3].DataType = schemapb.DataType_Double
|
|
|
|
{
|
|
functionSchema.Params[1].Key = "N"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay function lost param: origin")
|
|
functionSchema.Params[1].Key = originKey
|
|
}
|
|
{
|
|
functionSchema.Params[2].Key = "N"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay function lost param: scale")
|
|
functionSchema.Params[2].Key = scaleKey
|
|
}
|
|
|
|
{
|
|
functionSchema.Params[2].Value = "-1"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay function param: scale must > 0,")
|
|
functionSchema.Params[2].Value = "0.5"
|
|
}
|
|
|
|
{
|
|
functionSchema.Params[3].Value = "-1"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay function param: offset must => 0")
|
|
functionSchema.Params[3].Value = "0.5"
|
|
}
|
|
|
|
{
|
|
functionSchema.Params[4].Value = "10"
|
|
_, err := newDecayFunction(schema, functionSchema)
|
|
s.ErrorContains(err, "Decay function param: decay must 0 < decay < 1 0")
|
|
functionSchema.Params[2].Value = "0.5"
|
|
}
|
|
}
|
|
|
|
func (s *DecayFunctionSuite) TestRerankProcess() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
|
|
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
|
|
{
|
|
FieldID: 102, Name: "vector", DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "4"},
|
|
},
|
|
},
|
|
{FieldID: 102, Name: "ts", DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
functionSchema := &schemapb.FunctionSchema{
|
|
Name: "test",
|
|
Type: schemapb.FunctionType_Rerank,
|
|
InputFieldNames: []string{"ts"},
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: functionKey, Value: "gauss"},
|
|
{Key: originKey, Value: "0"},
|
|
{Key: scaleKey, Value: "4"},
|
|
{Key: offsetKey, Value: "2"},
|
|
},
|
|
}
|
|
|
|
// empty
|
|
{
|
|
nq := int64(1)
|
|
f, err := newDecayFunction(schema, functionSchema)
|
|
s.NoError(err)
|
|
ret, err := f.Process(context.Background(), &SearchParams{nq, 3, 2, -1, -1, 1, false}, []*schemapb.SearchResultData{})
|
|
s.NoError(err)
|
|
s.Equal(int64(3), ret.TopK)
|
|
s.Equal([]int64{}, ret.Topks)
|
|
}
|
|
|
|
// no input field exist
|
|
{
|
|
nq := int64(1)
|
|
f, err := newDecayFunction(schema, functionSchema)
|
|
data := genSearchResultData(nq, 10, schemapb.DataType_Int64, "noExist", 1000)
|
|
s.NoError(err)
|
|
_, err = f.Process(context.Background(), &SearchParams{nq, 3, 2, -1, -1, 1, false}, []*schemapb.SearchResultData{data})
|
|
s.ErrorContains(err, "Rerank decay function can not find input field, name")
|
|
}
|
|
|
|
// singleSearchResultData
|
|
// nq = 1
|
|
{
|
|
nq := int64(1)
|
|
f, err := newDecayFunction(schema, functionSchema)
|
|
s.NoError(err)
|
|
data := genSearchResultData(nq, 10, schemapb.DataType_Int64, "ts", 102)
|
|
ret, err := f.Process(context.Background(), &SearchParams{nq, 3, 2, -1, -1, 1, false}, []*schemapb.SearchResultData{data})
|
|
s.NoError(err)
|
|
s.Equal([]int64{3}, ret.Topks)
|
|
s.Equal(int64(3), ret.TopK)
|
|
s.Equal([]int64{2, 3, 4}, ret.Ids.GetIntId().Data)
|
|
}
|
|
// nq = 3
|
|
{
|
|
nq := int64(3)
|
|
f, err := newDecayFunction(schema, functionSchema)
|
|
s.NoError(err)
|
|
data := genSearchResultData(nq, 10, schemapb.DataType_Int64, "ts", 102)
|
|
ret, err := f.Process(context.Background(), &SearchParams{nq, 3, 2, -1, -1, 1, false}, []*schemapb.SearchResultData{data})
|
|
s.NoError(err)
|
|
s.Equal([]int64{3, 3, 3}, ret.Topks)
|
|
s.Equal(int64(3), ret.TopK)
|
|
s.Equal([]int64{2, 3, 4, 12, 13, 14, 22, 23, 24}, ret.Ids.GetIntId().Data)
|
|
}
|
|
|
|
// multipSearchResultData
|
|
// nq = 1
|
|
functionSchema2 := &schemapb.FunctionSchema{
|
|
Name: "test",
|
|
Type: schemapb.FunctionType_Rerank,
|
|
InputFieldNames: []string{"ts"},
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: functionKey, Value: "gauss"},
|
|
{Key: originKey, Value: "5"},
|
|
{Key: scaleKey, Value: "4"},
|
|
{Key: offsetKey, Value: "2"},
|
|
},
|
|
}
|
|
{
|
|
nq := int64(1)
|
|
f, err := newDecayFunction(schema, functionSchema2)
|
|
s.NoError(err)
|
|
// ts/id data: 0 - 9
|
|
data1 := genSearchResultData(nq, 10, schemapb.DataType_Int64, "ts", 102)
|
|
// ts/id data: 0 - 3
|
|
data2 := genSearchResultData(nq, 4, schemapb.DataType_Int64, "ts", 102)
|
|
ret, err := f.Process(context.Background(), &SearchParams{nq, 3, 2, -1, -1, 1, false}, []*schemapb.SearchResultData{data1, data2})
|
|
s.NoError(err)
|
|
s.Equal([]int64{3}, ret.Topks)
|
|
s.Equal(int64(3), ret.TopK)
|
|
s.Equal([]int64{5, 6, 7}, ret.Ids.GetIntId().Data)
|
|
}
|
|
// nq = 3
|
|
{
|
|
nq := int64(3)
|
|
f, err := newDecayFunction(schema, functionSchema2)
|
|
s.NoError(err)
|
|
// nq1 ts/id data: 0 - 9
|
|
// nq2 ts/id data: 10 - 19
|
|
// nq3 ts/id data: 20 - 29
|
|
data1 := genSearchResultData(nq, 10, schemapb.DataType_Int64, "ts", 102)
|
|
// nq1 ts/id data: 0 - 3
|
|
// nq2 ts/id data: 4 - 7
|
|
// nq3 ts/id data: 8 - 11
|
|
data2 := genSearchResultData(nq, 4, schemapb.DataType_Int64, "ts", 102)
|
|
ret, err := f.Process(context.Background(), &SearchParams{nq, 3, 2, 1, -1, 1, false}, []*schemapb.SearchResultData{data1, data2})
|
|
s.NoError(err)
|
|
s.Equal([]int64{3, 3, 3}, ret.Topks)
|
|
s.Equal(int64(3), ret.TopK)
|
|
s.Equal([]int64{5, 6, 7, 6, 7, 10, 10, 11, 20}, ret.Ids.GetIntId().Data)
|
|
}
|
|
}
|
|
|
|
func (s *DecayFunctionSuite) TestDecay() {
|
|
s.Equal(gaussianDecay(0, 1, 0.5, 5, 4), 1.0)
|
|
s.Equal(gaussianDecay(0, 1, 0.5, 5, 5), 1.0)
|
|
s.Less(gaussianDecay(0, 1, 0.5, 5, 6), 1.0)
|
|
|
|
s.Equal(expDecay(0, 1, 0.5, 5, 4), 1.0)
|
|
s.Equal(expDecay(0, 1, 0.5, 5, 5), 1.0)
|
|
s.Less(expDecay(0, 1, 0.5, 5, 6), 1.0)
|
|
|
|
s.Equal(linearDecay(0, 1, 0.5, 5, 4), 1.0)
|
|
s.Equal(linearDecay(0, 1, 0.5, 5, 5), 1.0)
|
|
s.Less(linearDecay(0, 1, 0.5, 5, 6), 1.0)
|
|
}
|
|
|
|
func (s *DecayFunctionSuite) TestUtil() {
|
|
inputTypes := []schemapb.DataType{schemapb.DataType_Int64, schemapb.DataType_Int32, schemapb.DataType_Int16, schemapb.DataType_Int8, schemapb.DataType_Float, schemapb.DataType_Double}
|
|
for _, tp := range inputTypes {
|
|
field := genSearchResultData(2, 10, tp, "test", 100)
|
|
num, err := getNumberic(field.FieldsData[0])
|
|
s.NoError(err)
|
|
switch tp {
|
|
case schemapb.DataType_Int32, schemapb.DataType_Int16, schemapb.DataType_Int8:
|
|
s.True(len(num.(*numberField[int32]).data) == 20)
|
|
case schemapb.DataType_Int64:
|
|
s.True(len(num.(*numberField[int64]).data) == 20)
|
|
case schemapb.DataType_Float:
|
|
s.True(len(num.(*numberField[float32]).data) == 20)
|
|
case schemapb.DataType_Double:
|
|
s.True(len(num.(*numberField[float64]).data) == 20)
|
|
}
|
|
}
|
|
|
|
field := genSearchResultData(2, 10, schemapb.DataType_Bool, "test", 100)
|
|
_, err := getNumberic(field.FieldsData[0])
|
|
s.ErrorContains(err, "only support numberic field")
|
|
|
|
{
|
|
ids := &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: testutils.GenerateInt64Array(10),
|
|
},
|
|
},
|
|
}
|
|
mid := newMilvusIDs(ids, schemapb.DataType_Int64).(milvusIDs[int64])
|
|
s.True(len(mid.data) == 10)
|
|
}
|
|
{
|
|
ids := &schemapb.IDs{
|
|
IdField: &schemapb.IDs_StrId{
|
|
StrId: &schemapb.StringArray{
|
|
Data: testutils.GenerateStringArray(10),
|
|
},
|
|
},
|
|
}
|
|
mid := newMilvusIDs(ids, schemapb.DataType_String).(milvusIDs[string])
|
|
s.True(len(mid.data) == 10)
|
|
}
|
|
}
|
|
|
|
func genSearchResultData(nq int64, topk int64, dType schemapb.DataType, fieldName string, fieldId int64) *schemapb.SearchResultData {
|
|
tops := make([]int64, nq)
|
|
for i := 0; i < int(nq); i++ {
|
|
tops[i] = topk
|
|
}
|
|
data := &schemapb.SearchResultData{
|
|
NumQueries: nq,
|
|
TopK: topk,
|
|
Scores: testutils.GenerateFloat32Array(int(nq * topk)),
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: testutils.GenerateInt64Array(int(nq * topk)),
|
|
},
|
|
},
|
|
},
|
|
Topks: tops,
|
|
FieldsData: []*schemapb.FieldData{testutils.GenerateScalarFieldData(dType, fieldName, int(nq*topk))},
|
|
}
|
|
data.FieldsData[0].FieldId = fieldId
|
|
return data
|
|
}
|