admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-06 17:18:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: [2.5] Handle empty FieldsData in reduce/rerank for requery scenario (#44919)

Browse Source 
issue: #44909
pr: #44917
When requery optimization is enabled, search results contain IDs but
empty FieldsData. During reduce/rerank operations, if the first shard
has empty FieldsData while others have data, PrepareResultFieldData
initializes an empty array, causing AppendFieldData to panic when
accessing array indices.

Changes:
- Find first non-empty FieldsData as template in 5 functions:
  reduceAdvanceGroupBY, reduceSearchResultDataWithGroupBy,
  reduceSearchResultDataNoGroupBy, rankSearchResultDataByGroup,
  rankSearchResultDataByPk
- Add length check before 4 AppendFieldData calls to prevent panic
- Add unit tests for empty and partial empty FieldsData scenarios

This fix handles both pure requery (all empty) and mixed scenarios
(some empty, some with data) without breaking normal search flow.

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
This commit is contained in:
wei liu 2025-10-21 19:48:04 +08:00 committed by GitHub
parent d43b030b4d
commit c633556fee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 305 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
internal/proxy/search_reduce_util.go
View File

@ -99,7 +99,13 @@ func reduceAdvanceGroupBY(ctx context.Context, subSearchResultData []*schemapb.S
} else {
ret.GetResults().AllSearchCount = allSearchCount
limit = int64(hitNum)
ret.GetResults().FieldsData = typeutil.PrepareResultFieldData(subSearchResultData[0].GetFieldsData(), limit)
// Find the first non-empty FieldsData as template
for _, result := range subSearchResultData {
if len(result.GetFieldsData()) > 0 {
ret.GetResults().FieldsData = typeutil.PrepareResultFieldData(result.GetFieldsData(), limit)
break
}
}
}
if err := setupIdListForSearchResult(ret, pkType, limit); err != nil {
@ -186,7 +192,7 @@ func reduceSearchResultDataWithGroupBy(ctx context.Context, subSearchResultData
Results: &schemapb.SearchResultData{
NumQueries: nq,
TopK: topk,
FieldsData: typeutil.PrepareResultFieldData(subSearchResultData[0].GetFieldsData(), limit),
FieldsData: []*schemapb.FieldData{},
Scores: []float32{},
Ids: &schemapb.IDs{},
Topks: []int64{},
@ -204,6 +210,14 @@ func reduceSearchResultDataWithGroupBy(ctx context.Context, subSearchResultData
ret.GetResults().AllSearchCount = allSearchCount
}
// Find the first non-empty FieldsData as template
for _, result := range subSearchResultData {
if len(result.GetFieldsData()) > 0 {
ret.GetResults().FieldsData = typeutil.PrepareResultFieldData(result.GetFieldsData(), limit)
break
}
}
var (
subSearchNum = len(subSearchResultData)
// for results of each subSearchResultData, storing the start offset of each query of nq queries
@ -280,7 +294,9 @@ func reduceSearchResultDataWithGroupBy(ctx context.Context, subSearchResultData
groupEntities := groupByValMap[groupVal]
for _, groupEntity := range groupEntities {
subResData := subSearchResultData[groupEntity.subSearchIdx]
if len(ret.Results.FieldsData) > 0 {
retSize += typeutil.AppendFieldData(ret.Results.FieldsData, subResData.FieldsData, groupEntity.resultIdx)
}
typeutil.AppendPKs(ret.Results.Ids, groupEntity.id)
ret.Results.Scores = append(ret.Results.Scores, groupEntity.score)
if err := typeutil.AppendGroupByValue(ret.Results, groupVal, subResData.GetGroupByFieldValue().GetType()); err != nil {
@ -330,7 +346,7 @@ func reduceSearchResultDataNoGroupBy(ctx context.Context, subSearchResultData []
Results: &schemapb.SearchResultData{
NumQueries: nq,
TopK: topk,
FieldsData: typeutil.PrepareResultFieldData(subSearchResultData[0].GetFieldsData(), limit),
FieldsData: []*schemapb.FieldData{},
Scores: []float32{},
Ids: &schemapb.IDs{},
Topks: []int64{},
@ -348,6 +364,14 @@ func reduceSearchResultDataNoGroupBy(ctx context.Context, subSearchResultData []
ret.GetResults().AllSearchCount = allSearchCount
}
// Find the first non-empty FieldsData as template
for _, result := range subSearchResultData {
if len(result.GetFieldsData()) > 0 {
ret.GetResults().FieldsData = typeutil.PrepareResultFieldData(result.GetFieldsData(), limit)
break
}
}
subSearchNum := len(subSearchResultData)
if subSearchNum == 1 && offset == 0 {
// sorting is not needed if there is only one shard and no offset, assigning the result directly.
@ -401,7 +425,9 @@ func reduceSearchResultDataNoGroupBy(ctx context.Context, subSearchResultData []
}
score := subSearchResultData[subSearchIdx].Scores[resultDataIdx]
if len(ret.Results.FieldsData) > 0 {
retSize += typeutil.AppendFieldData(ret.Results.FieldsData, subSearchResultData[subSearchIdx].FieldsData, resultDataIdx)
}
typeutil.CopyPk(ret.Results.Ids, subSearchResultData[subSearchIdx].GetIds(), int(resultDataIdx))
ret.Results.Scores = append(ret.Results.Scores, score)
cursors[subSearchIdx]++
@ -515,8 +541,13 @@ func rankSearchResultDataByGroup(ctx context.Context,
return ret, nil
}
// init FieldsData
ret.Results.FieldsData = typeutil.PrepareResultFieldData(searchResults[0].GetResults().GetFieldsData(), limit)
// Find the first non-empty FieldsData as template
for _, result := range searchResults {
if len(result.GetResults().GetFieldsData()) > 0 {
ret.Results.FieldsData = typeutil.PrepareResultFieldData(result.GetResults().GetFieldsData(), limit)
break
}
}
totalCount := limit * groupSize
if err := setupIdListForSearchResult(ret, pkType, totalCount); err != nil {
@ -643,7 +674,9 @@ func rankSearchResultDataByGroup(ctx context.Context,
}
ret.Results.Scores = append(ret.Results.Scores, score)
loc := pk2DataOffset[i][group.idList[idx]]
if len(ret.Results.FieldsData) > 0 {
typeutil.AppendFieldData(ret.Results.FieldsData, searchResults[loc.resultIdx].GetResults().GetFieldsData(), int64(loc.offset))
}
typeutil.AppendGroupByValue(ret.Results, group.groupVal, groupByDataType)
}
returnedRowNum += len(group.idList)
@ -712,8 +745,13 @@ func rankSearchResultDataByPk(ctx context.Context,
return ret, nil
}
// init FieldsData
ret.Results.FieldsData = typeutil.PrepareResultFieldData(searchResults[0].GetResults().GetFieldsData(), limit)
// Find the first non-empty FieldsData as template
for _, result := range searchResults {
if len(result.GetResults().GetFieldsData()) > 0 {
ret.Results.FieldsData = typeutil.PrepareResultFieldData(result.GetResults().GetFieldsData(), limit)
break
}
}
if err := setupIdListForSearchResult(ret, pkType, limit); err != nil {
return ret, nil
@ -783,9 +821,11 @@ func rankSearchResultDataByPk(ctx context.Context,
}
ret.Results.Scores = append(ret.Results.Scores, score)
loc := pk2DataOffset[i][keys[index]]
if len(ret.Results.FieldsData) > 0 {
typeutil.AppendFieldData(ret.Results.FieldsData, searchResults[loc.resultIdx].GetResults().GetFieldsData(), loc.offset)
}
}
}
return ret, nil
}

254
internal/proxy/search_reduce_util_test.go
View File

@ -142,6 +142,260 @@ func (struts *SearchReduceUtilTestSuite) TestReduceSearchResult() {
}
}
// TestReduceWithEmptyFieldsData tests reduce functions when FieldsData is empty (requery scenario)
func (struts *SearchReduceUtilTestSuite) TestReduceWithEmptyFieldsData() {
ctx := context.Background()
nq := int64(1)
topK := int64(5)
offset := int64(0)
// Create search results with empty FieldsData (simulating requery scenario)
searchResultData1 := &schemapb.SearchResultData{
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{1, 2, 3, 4, 5},
},
},
},
Scores: []float32{0.9, 0.8, 0.7, 0.6, 0.5},
Topks: []int64{5},
NumQueries: nq,
TopK: topK,
FieldsData: []*schemapb.FieldData{}, // Empty FieldsData for requery
}
searchResultData2 := &schemapb.SearchResultData{
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{6, 7, 8, 9, 10},
},
},
},
Scores: []float32{0.85, 0.75, 0.65, 0.55, 0.45},
Topks: []int64{5},
NumQueries: nq,
TopK: topK,
FieldsData: []*schemapb.FieldData{}, // Empty FieldsData for requery
}
// Test reduceSearchResultDataNoGroupBy with empty FieldsData
{
results, err := reduceSearchResultDataNoGroupBy(ctx, []*schemapb.SearchResultData{searchResultData1, searchResultData2}, nq, topK, "L2", schemapb.DataType_Int64, offset)
struts.NoError(err)
struts.NotNil(results)
// Should have merged results without panic
struts.Equal(int64(5), results.Results.Topks[0])
// FieldsData should be empty since all inputs were empty
struts.Equal(0, len(results.Results.FieldsData))
}
// Test reduceSearchResultDataWithGroupBy with empty FieldsData
{
// Add GroupByFieldValue to support group by
searchResultData1.GroupByFieldValue = &schemapb.FieldData{
Type: schemapb.DataType_VarChar,
FieldName: "group",
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"a", "b", "c", "a", "b"},
},
},
},
},
}
searchResultData2.GroupByFieldValue = &schemapb.FieldData{
Type: schemapb.DataType_VarChar,
FieldName: "group",
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"c", "a", "b", "c", "a"},
},
},
},
},
}
results, err := reduceSearchResultDataWithGroupBy(ctx, []*schemapb.SearchResultData{searchResultData1, searchResultData2}, nq, topK, "L2", schemapb.DataType_Int64, offset, int64(2))
struts.NoError(err)
struts.NotNil(results)
// FieldsData should be empty since all inputs were empty
struts.Equal(0, len(results.Results.FieldsData))
}
// Test reduceAdvanceGroupBY with empty FieldsData
{
results, err := reduceAdvanceGroupBY(ctx, []*schemapb.SearchResultData{searchResultData1, searchResultData2}, nq, topK, schemapb.DataType_Int64, "L2")
struts.NoError(err)
struts.NotNil(results)
// FieldsData should be empty since all inputs were empty
struts.Equal(0, len(results.Results.FieldsData))
}
}
// TestReduceWithPartialEmptyFieldsData tests when first result has empty FieldsData but second has data
func (struts *SearchReduceUtilTestSuite) TestReduceWithPartialEmptyFieldsData() {
ctx := context.Background()
nq := int64(1)
topK := int64(3)
offset := int64(0)
// First result with empty FieldsData
searchResultData1 := &schemapb.SearchResultData{
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{1, 2, 3},
},
},
},
Scores: []float32{0.9, 0.8, 0.7},
Topks: []int64{3},
NumQueries: nq,
TopK: topK,
FieldsData: []*schemapb.FieldData{}, // Empty
}
// Second result with non-empty FieldsData
searchResultData2 := &schemapb.SearchResultData{
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{4, 5, 6},
},
},
},
Scores: []float32{0.85, 0.75, 0.65},
Topks: []int64{3},
NumQueries: nq,
TopK: topK,
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64,
FieldName: "field1",
FieldId: 100,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{40, 50, 60},
},
},
},
},
},
},
}
// Test: Should use the non-empty FieldsData from second result
results, err := reduceSearchResultDataNoGroupBy(ctx, []*schemapb.SearchResultData{searchResultData1, searchResultData2}, nq, topK, "L2", schemapb.DataType_Int64, offset)
struts.NoError(err)
struts.NotNil(results)
// Should have initialized FieldsData from second result
struts.Greater(len(results.Results.FieldsData), 0)
}
// TestRankWithEmptyFieldsData tests rank functions when FieldsData is empty
func (struts *SearchReduceUtilTestSuite) TestRankWithEmptyFieldsData() {
ctx := context.Background()
nq := int64(1)
limit := int64(3)
offset := int64(0)
roundDecimal := int64(-1)
rankParams := &rankParams{limit: limit, offset: offset, roundDecimal: roundDecimal}
// Create search results with empty FieldsData
searchResult1 := &milvuspb.SearchResults{
Results: &schemapb.SearchResultData{
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{1, 2, 3},
},
},
},
Scores: []float32{0.9, 0.8, 0.7},
Topks: []int64{3},
NumQueries: nq,
FieldsData: []*schemapb.FieldData{}, // Empty
},
}
searchResult2 := &milvuspb.SearchResults{
Results: &schemapb.SearchResultData{
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{2, 4, 5},
},
},
},
Scores: []float32{0.85, 0.75, 0.65},
Topks: []int64{3},
NumQueries: nq,
FieldsData: []*schemapb.FieldData{}, // Empty
},
}
searchResults := []*milvuspb.SearchResults{searchResult1, searchResult2}
// Test rankSearchResultDataByPk with empty FieldsData
{
results, err := rankSearchResultDataByPk(ctx, nq, rankParams, schemapb.DataType_Int64, searchResults)
struts.NoError(err)
struts.NotNil(results)
// FieldsData should be empty since all inputs were empty
struts.Equal(0, len(results.Results.FieldsData))
}
// Test rankSearchResultDataByGroup with empty FieldsData
{
// Add group by values
searchResult1.Results.GroupByFieldValue = &schemapb.FieldData{
Type: schemapb.DataType_VarChar,
FieldName: "group",
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"a", "b", "c"},
},
},
},
},
}
searchResult2.Results.GroupByFieldValue = &schemapb.FieldData{
Type: schemapb.DataType_VarChar,
FieldName: "group",
FieldId: 101,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"b", "a", "c"},
},
},
},
},
}
groupScorer, _ := GetGroupScorer("max")
results, err := rankSearchResultDataByGroup(ctx, nq, rankParams, schemapb.DataType_Int64, searchResults, groupScorer, int64(2))
struts.NoError(err)
struts.NotNil(results)
// FieldsData should be empty since all inputs were empty
struts.Equal(0, len(results.Results.FieldsData))
}
}
func TestSearchReduceUtilTestSuite(t *testing.T) {
suite.Run(t, new(SearchReduceUtilTestSuite))
}