mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-07 01:28:27 +08:00
7d13bdcf4c
Cherry-pick from master pr: #45444 Related to #45338 When using bulk vector search in hybrid search with rerank functions, the output field values for different queries were all equal to the values returned by the first query, instead of the correct values belonging to each document ID. The document IDs were correct, but the entity field values were wrong. In rerank functions (RRF, weighted, decay, model), when processing multiple queries in a batch, the `idLocations` stored only the relative offset within each result set (`idx`), not accounting for the absolute position within the entire batch. This caused `FillFieldData` to retrieve field data from the wrong positions, always using offsets relative to the first query. This fix ensures that when processing bulk searches with rerank functions, each result correctly retrieves its corresponding field data based on the absolute offset within the entire batch, resolving the issue where all queries returned the first query's field values. Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
104 lines
3.2 KiB
Go
104 lines
3.2 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"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
)
|
|
|
|
const (
|
|
RRFParamsKey string = "k"
|
|
|
|
defaultRRFParamsValue float64 = 60
|
|
)
|
|
|
|
type RRFFunction[T PKType] struct {
|
|
RerankBase
|
|
|
|
k float32
|
|
}
|
|
|
|
func newRRFFunction(collSchema *schemapb.CollectionSchema, funcSchema *schemapb.FunctionSchema) (Reranker, error) {
|
|
base, err := newRerankBase(collSchema, funcSchema, RRFName, true)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if len(base.GetInputFieldNames()) != 0 {
|
|
return nil, fmt.Errorf("The rrf function does not support input parameters, but got %s", base.GetInputFieldNames())
|
|
}
|
|
|
|
k := float64(defaultRRFParamsValue)
|
|
for _, param := range funcSchema.Params {
|
|
if strings.ToLower(param.Key) == RRFParamsKey {
|
|
if k, err = strconv.ParseFloat(param.Value, 64); err != nil {
|
|
return nil, fmt.Errorf("Param k:%s is not a number", param.Value)
|
|
}
|
|
}
|
|
}
|
|
if k <= 0 || k >= 16384 {
|
|
return nil, fmt.Errorf("The rank params k should be in range (0, %d)", 16384)
|
|
}
|
|
if base.pkType == schemapb.DataType_Int64 {
|
|
return &RRFFunction[int64]{RerankBase: *base, k: float32(k)}, nil
|
|
} else {
|
|
return &RRFFunction[string]{RerankBase: *base, k: float32(k)}, nil
|
|
}
|
|
}
|
|
|
|
func (rrf *RRFFunction[T]) processOneSearchData(ctx context.Context, searchParams *SearchParams, cols []*columns, idGroup map[any]any) (*IDScores[T], error) {
|
|
rrfScores := map[T]float32{}
|
|
idLocations := make(map[T]IDLoc)
|
|
for i, col := range cols {
|
|
if col.size == 0 {
|
|
continue
|
|
}
|
|
ids := col.ids.([]T)
|
|
for idx, id := range ids {
|
|
if score, ok := rrfScores[id]; !ok {
|
|
idLocations[id] = IDLoc{batchIdx: i, offset: idx + int(col.nqOffset)}
|
|
rrfScores[id] = 1 / (rrf.k + float32(idx+1))
|
|
} else {
|
|
rrfScores[id] = score + 1/(rrf.k+float32(idx+1))
|
|
}
|
|
}
|
|
}
|
|
if searchParams.isGrouping() {
|
|
return newGroupingIDScores(rrfScores, idLocations, searchParams, idGroup)
|
|
}
|
|
return newIDScores(rrfScores, idLocations, searchParams, true), nil
|
|
}
|
|
|
|
func (rrf *RRFFunction[T]) Process(ctx context.Context, searchParams *SearchParams, inputs *rerankInputs) (*rerankOutputs, error) {
|
|
outputs := newRerankOutputs(inputs, searchParams)
|
|
for _, cols := range inputs.data {
|
|
idScore, err := rrf.processOneSearchData(ctx, searchParams, cols, inputs.idGroupValue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
appendResult(inputs, outputs, idScore)
|
|
}
|
|
return outputs, nil
|
|
}
|