From b1d7b7427a7f6b54733291df6d32a979190a2a52 Mon Sep 17 00:00:00 2001 From: Spade A <71589810+SpadeA-Tang@users.noreply.github.com> Date: Tue, 7 Apr 2026 14:55:38 +0800 Subject: [PATCH] feat: impl StructArray -- support group by primary key for element-level search (#48417) issue: https://github.com/milvus-io/milvus/issues/42148 design doc: https://github.com/milvus-io/milvus-design-docs/blob/main/design_docs/20260306-struct.md This PR: - Allow element-level search to use group by primary key for doc-level dedup (non pk is not allowed yet) - Block embedding list (multi-search-multi) with group by; hybrid search with ArrayOfVector group by remains blocked TODO: - support element-level query to use group by Signed-off-by: SpadeA --- .../search-groupby/SearchGroupByOperator.cpp | 12 +- internal/proxy/search_util.go | 5 +- internal/proxy/task_search.go | 42 ++++- internal/proxy/task_search_test.go | 150 +++++++++++++++++- internal/proxy/vector_type_convert.go | 24 +-- internal/proxy/vector_type_convert_test.go | 18 +-- pkg/util/typeutil/schema.go | 7 + 7 files changed, 216 insertions(+), 42 deletions(-) diff --git a/internal/core/src/exec/operator/search-groupby/SearchGroupByOperator.cpp b/internal/core/src/exec/operator/search-groupby/SearchGroupByOperator.cpp index ef34a67cb8..1c1cd91ba5 100644 --- a/internal/core/src/exec/operator/search-groupby/SearchGroupByOperator.cpp +++ b/internal/core/src/exec/operator/search-groupby/SearchGroupByOperator.cpp @@ -18,7 +18,6 @@ #include #include #include -#include #include #include "common/QueryInfo.h" @@ -319,10 +318,6 @@ GroupIteratorResult(const std::shared_ptr& iterator, //note it may enumerate all data inside a segment and can block following //query and search possibly std::vector>> res; - // For element-level search, multiple elements from the same row may be - // returned by the iterator. We must deduplicate by row_offset so that - // the same document does not consume multiple slots in a group. - std::unordered_set seen_rows; while (iterator->HasNext() && !groupMap.IsGroupResEnough()) { auto offset_dis_pair = iterator->Next(); AssertInfo( @@ -332,17 +327,12 @@ GroupIteratorResult(const std::shared_ptr& iterator, auto offset = offset_dis_pair.value().first; auto dis = offset_dis_pair.value().second; - // When array_offsets is present, iterator returns element_id, - // but DataGetter expects row_id. Convert element_id to row_id. + // For element-level search, the offset is the element_id, we need to convert it to the row_id. int64_t row_offset = offset; if (is_element_id) { row_offset = array_offsets->ElementIDToRowID(static_cast(offset)) .first; - // Skip if we already saw this row (closest element wins) - if (!seen_rows.insert(row_offset).second) { - continue; - } } std::optional row_data = data_getter->Get(row_offset); diff --git a/internal/proxy/search_util.go b/internal/proxy/search_util.go index c4a9952597..bf08d55c13 100644 --- a/internal/proxy/search_util.go +++ b/internal/proxy/search_util.go @@ -544,7 +544,10 @@ func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb "range search is not supported for vector array (embedding list) fields, fieldName:", annsFieldName) } - if groupByFieldId > 0 { + // For hybrid search (rankInfo != nil), group by on vector array is not supported. + // For simple search, group by validation is handled in initSearchRequest() + // where placeholder type is available to distinguish element-level vs embedding list. + if groupByFieldId > 0 && rankParams != nil { return nil, merr.WrapErrParameterInvalid("", "", "group by search is not supported for vector array (embedding list) fields, fieldName:", annsFieldName) } diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go index e631554295..d281afb7a0 100644 --- a/internal/proxy/task_search.go +++ b/internal/proxy/task_search.go @@ -551,7 +551,7 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error { metrics.ProxySearchSparseNumNonZeros.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), t.collectionName, metrics.HybridSearchLabel, strconv.FormatInt(internalSubReq.FieldId, 10)).Observe(float64(typeutil.EstimateSparseVectorNNZFromPlaceholderGroup(internalSubReq.PlaceholderGroup, int(internalSubReq.GetNq())))) } // Convert placeholder group vector type if needed (e.g., fp32 -> fp16/bf16) - internalSubReq.PlaceholderGroup, err = t.convertPlaceholderIfNeeded(subReq.GetPlaceholderGroup(), internalSubReq.FieldId) + internalSubReq.PlaceholderGroup, _, err = t.convertPlaceholderIfNeeded(subReq.GetPlaceholderGroup(), internalSubReq.FieldId) if err != nil { return err } @@ -763,10 +763,30 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error { metrics.ProxySearchSparseNumNonZeros.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), t.collectionName, metrics.SearchLabel, strconv.FormatInt(t.SearchRequest.FieldId, 10)).Observe(float64(typeutil.EstimateSparseVectorNNZFromPlaceholderGroup(t.request.GetPlaceholderGroup(), int(t.request.GetNq())))) } // Convert placeholder group vector type if needed (e.g., fp32 -> fp16/bf16) - t.SearchRequest.PlaceholderGroup, err = t.convertPlaceholderIfNeeded(t.request.GetPlaceholderGroup(), t.SearchRequest.FieldId) + var placeholderType commonpb.PlaceholderType + t.SearchRequest.PlaceholderGroup, placeholderType, err = t.convertPlaceholderIfNeeded(t.request.GetPlaceholderGroup(), t.SearchRequest.FieldId) if err != nil { return err } + + // For ArrayOfVector fields with group by: + // 1. Embedding list search does not support group by + // 2. Element-level search only supports group by PK (doc-level dedup) + if queryInfo.GetGroupByFieldId() > 0 { + annsField := typeutil.GetField(t.schema.CollectionSchema, t.SearchRequest.FieldId) + if annsField != nil && annsField.GetDataType() == schemapb.DataType_ArrayOfVector { + if isEmbeddingListPlaceholderType(placeholderType) { + return merr.WrapErrParameterInvalid("", "", + "group by is not supported for multi-search-multi on embedding list fields") + } + pkField, _ := typeutil.GetPrimaryFieldSchema(t.schema.CollectionSchema) + if pkField == nil || queryInfo.GetGroupByFieldId() != pkField.GetFieldID() { + return merr.WrapErrParameterInvalid("", "", + "only group by primary key is supported for element-level search on embedding list fields") + } + } + } + t.SearchRequest.Topk = queryInfo.GetTopk() t.SearchRequest.MetricType = queryInfo.GetMetricType() t.queryInfos = append(t.queryInfos, queryInfo) @@ -796,10 +816,11 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error { } // convertPlaceholderIfNeeded converts fp32 vectors to fp16/bf16 if the target field uses lower precision. -func (t *searchTask) convertPlaceholderIfNeeded(phgBytes []byte, fieldID int64) ([]byte, error) { +// Returns converted bytes and the original placeholder type (before any conversion). +func (t *searchTask) convertPlaceholderIfNeeded(phgBytes []byte, fieldID int64) ([]byte, commonpb.PlaceholderType, error) { field := typeutil.GetFieldByID(t.schema.CollectionSchema, fieldID) if field == nil { - return phgBytes, nil + return phgBytes, 0, nil } return ConvertPlaceholderGroup(phgBytes, field) } @@ -917,6 +938,19 @@ func getLastBound(result *milvuspb.SearchResults, incomingLastBound *float32, me return -math.MaxFloat32 } +func isEmbeddingListPlaceholderType(pt commonpb.PlaceholderType) bool { + switch pt { + case commonpb.PlaceholderType_EmbListFloatVector, + commonpb.PlaceholderType_EmbListFloat16Vector, + commonpb.PlaceholderType_EmbListBFloat16Vector, + commonpb.PlaceholderType_EmbListBinaryVector, + commonpb.PlaceholderType_EmbListInt8Vector: + return true + default: + return false + } +} + func (t *searchTask) PostExecute(ctx context.Context) error { ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search-PostExecute") defer sp.End() diff --git a/internal/proxy/task_search_test.go b/internal/proxy/task_search_test.go index 73c801fba1..1f259d4c69 100644 --- a/internal/proxy/task_search_test.go +++ b/internal/proxy/task_search_test.go @@ -3710,22 +3710,21 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { assert.Contains(t, err.Error(), "range search is not supported for vector array (embedding list) fields") }) - t.Run("vector array with group by", func(t *testing.T) { + t.Run("vector array with group by passes parseSearchInfo", func(t *testing.T) { + // Group by validation for ArrayOfVector is now handled in initSearchRequest(), + // so parseSearchInfo should pass regardless of group by field. schema := createSchemaWithVectorArray("embeddings_list") params := createSearchParams("embeddings_list") - // Add group by parameter + // Add group by parameter (non-PK field) params = append(params, &commonpb.KeyValuePair{ Key: GroupByFieldKey, Value: "group_field", }) searchInfo, err := parseSearchInfo(params, schema, nil, false) - assert.Error(t, err) - assert.Nil(t, searchInfo) - assert.ErrorIs(t, err, merr.ErrParameterInvalid) - assert.Contains(t, err.Error(), "group by search is not supported for vector array (embedding list) fields") - assert.Contains(t, err.Error(), "embeddings_list") + assert.NoError(t, err) + assert.NotNil(t, searchInfo) }) t.Run("vector array with iterator", func(t *testing.T) { @@ -5681,3 +5680,140 @@ func TestSearchTask_InitSearchRequestWithHighlighter(t *testing.T) { assert.Empty(t, plan.DynamicFields) }) } + +func TestIsEmbeddingListPlaceholderType(t *testing.T) { + embListTypes := []commonpb.PlaceholderType{ + commonpb.PlaceholderType_EmbListFloatVector, + commonpb.PlaceholderType_EmbListFloat16Vector, + commonpb.PlaceholderType_EmbListBFloat16Vector, + commonpb.PlaceholderType_EmbListBinaryVector, + commonpb.PlaceholderType_EmbListInt8Vector, + } + for _, pt := range embListTypes { + assert.True(t, isEmbeddingListPlaceholderType(pt), "expected true for %s", pt.String()) + } + + nonEmbListTypes := []commonpb.PlaceholderType{ + commonpb.PlaceholderType_FloatVector, + commonpb.PlaceholderType_BinaryVector, + commonpb.PlaceholderType_Float16Vector, + commonpb.PlaceholderType_BFloat16Vector, + commonpb.PlaceholderType_SparseFloatVector, + commonpb.PlaceholderType_Int8Vector, + commonpb.PlaceholderType_VarChar, + commonpb.PlaceholderType(0), + } + for _, pt := range nonEmbListTypes { + assert.False(t, isEmbeddingListPlaceholderType(pt), "expected false for %s", pt.String()) + } +} + +func TestSearchTask_ArrayOfVectorGroupBy(t *testing.T) { + paramtable.Init() + ctx := context.Background() + + schema := &schemapb.CollectionSchema{ + Name: "test_collection", + Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, Name: "regular_vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}}, + {FieldID: 102, Name: "scalar_field", DataType: schemapb.DataType_VarChar}, + }, + StructArrayFields: []*schemapb.StructArrayFieldSchema{ + { + FieldID: 103, + Name: "struct_array", + Fields: []*schemapb.FieldSchema{ + {FieldID: 104, Name: "emb_vec", DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}}, + }, + }, + }, + } + schemaInfo := newSchemaInfo(schema) + + makePlaceholderGroup := func(phType commonpb.PlaceholderType) []byte { + phg := &commonpb.PlaceholderGroup{ + Placeholders: []*commonpb.PlaceholderValue{{ + Tag: "$0", + Type: phType, + Values: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}}, // 4 floats + }}, + } + bs, _ := proto.Marshal(phg) + return bs + } + + makeTask := func(annsField string, groupByField string, phType commonpb.PlaceholderType) *searchTask { + params := []*commonpb.KeyValuePair{ + {Key: AnnsFieldKey, Value: annsField}, + {Key: TopKKey, Value: "10"}, + {Key: common.MetricTypeKey, Value: metric.L2}, + {Key: ParamsKey, Value: `{"nprobe": 10}`}, + } + if groupByField != "" { + params = append(params, &commonpb.KeyValuePair{Key: GroupByFieldKey, Value: groupByField}) + } + + phgBytes := makePlaceholderGroup(phType) + return &searchTask{ + ctx: ctx, + collectionName: "test_collection", + SearchRequest: &internalpb.SearchRequest{ + CollectionID: 1, + PartitionIDs: []int64{1}, + OutputFieldsId: []int64{100}, + PlaceholderGroup: nil, + DslType: commonpb.DslType_BoolExprV1, + }, + request: &milvuspb.SearchRequest{ + CollectionName: "test_collection", + OutputFields: []string{"pk"}, + SearchParams: params, + SearchInput: &milvuspb.SearchRequest_PlaceholderGroup{ + PlaceholderGroup: phgBytes, + }, + Nq: 1, + ConsistencyLevel: commonpb.ConsistencyLevel_Session, + }, + schema: schemaInfo, + translatedOutputFields: []string{"pk"}, + tr: timerecord.NewTimeRecorder("test"), + queryInfos: []*planpb.QueryInfo{{}}, + } + } + + t.Run("element-level search with group by PK should succeed", func(t *testing.T) { + task := makeTask("emb_vec", "pk", commonpb.PlaceholderType_FloatVector) + err := task.initSearchRequest(ctx) + assert.NoError(t, err) + }) + + t.Run("element-level search with group by non-PK should fail", func(t *testing.T) { + task := makeTask("emb_vec", "scalar_field", commonpb.PlaceholderType_FloatVector) + err := task.initSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "only group by primary key is supported") + }) + + t.Run("emblist search with group by should fail", func(t *testing.T) { + task := makeTask("emb_vec", "pk", commonpb.PlaceholderType_EmbListFloatVector) + err := task.initSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "not supported for multi-search-multi") + }) + + t.Run("element-level search without group by should succeed", func(t *testing.T) { + task := makeTask("emb_vec", "", commonpb.PlaceholderType_FloatVector) + err := task.initSearchRequest(ctx) + assert.NoError(t, err) + }) + + t.Run("regular vector with group by non-PK should succeed", func(t *testing.T) { + // Group by on non-ArrayOfVector field has no PK restriction + task := makeTask("regular_vec", "scalar_field", commonpb.PlaceholderType_FloatVector) + err := task.initSearchRequest(ctx) + assert.NoError(t, err) + }) +} diff --git a/internal/proxy/vector_type_convert.go b/internal/proxy/vector_type_convert.go index 2563cafb26..bd5d5407b6 100644 --- a/internal/proxy/vector_type_convert.go +++ b/internal/proxy/vector_type_convert.go @@ -105,51 +105,55 @@ func validateFloat32ForBFloat16(values []float32) error { // If the placeholder type matches the field type, returns the original bytes unchanged. // If the placeholder is fp32 and field is fp16/bf16, converts the vectors. // Otherwise returns an error for incompatible types. -func ConvertPlaceholderGroup(phgBytes []byte, fieldSchema *schemapb.FieldSchema) ([]byte, error) { +// Also returns the original placeholder type (before conversion) for callers that need it. +func ConvertPlaceholderGroup(phgBytes []byte, fieldSchema *schemapb.FieldSchema) ([]byte, commonpb.PlaceholderType, error) { var phg commonpb.PlaceholderGroup if err := proto.Unmarshal(phgBytes, &phg); err != nil { - return nil, fmt.Errorf("failed to unmarshal placeholder group: %w", err) + return nil, 0, fmt.Errorf("failed to unmarshal placeholder group: %w", err) } if len(phg.Placeholders) == 0 { - return phgBytes, nil + return phgBytes, 0, nil } placeholder := phg.Placeholders[0] + phType := placeholder.Type fieldType := fieldSchema.GetDataType() // Check if types already match if isVectorTypeMatch(placeholder.Type, fieldType) { - return phgBytes, nil + return phgBytes, phType, nil } // If placeholder is not a vector type (e.g., VarChar for text embedding), pass through. // Let downstream logic handle non-vector placeholders. if _, isVectorType := placeholderTypeToDataType[placeholder.Type]; !isVectorType { - return phgBytes, nil + return phgBytes, phType, nil } // Only handle fp32 -> fp16/bf16 conversion. // For other field types (e.g., SparseFloatVector, BinaryVector), pass through // and let downstream logic handle the type mismatch with appropriate error messages. if fieldType != schemapb.DataType_Float16Vector && fieldType != schemapb.DataType_BFloat16Vector { - return phgBytes, nil + return phgBytes, phType, nil } // Check if conversion is supported (fp32 -> fp16/bf16) if placeholder.Type != commonpb.PlaceholderType_FloatVector { - return nil, fmt.Errorf("vector type must be the same: field type %s, search type %s", + return nil, phType, fmt.Errorf("vector type must be the same: field type %s, search type %s", fieldType.String(), placeholder.Type.String()) } switch fieldType { case schemapb.DataType_Float16Vector: - return convertPlaceholder(&phg, validateFloat32ForFloat16, typeutil.Float32ArrayToFloat16Bytes, commonpb.PlaceholderType_Float16Vector) + converted, err := convertPlaceholder(&phg, validateFloat32ForFloat16, typeutil.Float32ArrayToFloat16Bytes, commonpb.PlaceholderType_Float16Vector) + return converted, phType, err case schemapb.DataType_BFloat16Vector: - return convertPlaceholder(&phg, validateFloat32ForBFloat16, typeutil.Float32ArrayToBFloat16Bytes, commonpb.PlaceholderType_BFloat16Vector) + converted, err := convertPlaceholder(&phg, validateFloat32ForBFloat16, typeutil.Float32ArrayToBFloat16Bytes, commonpb.PlaceholderType_BFloat16Vector) + return converted, phType, err default: // This should never be reached due to the check above, but keep for safety - return phgBytes, nil + return phgBytes, phType, nil } } diff --git a/internal/proxy/vector_type_convert_test.go b/internal/proxy/vector_type_convert_test.go index 2b691cd733..e1dadc57ad 100644 --- a/internal/proxy/vector_type_convert_test.go +++ b/internal/proxy/vector_type_convert_test.go @@ -193,7 +193,7 @@ func TestConvertPlaceholderGroupToFloat16(t *testing.T) { DataType: schemapb.DataType_Float16Vector, } - convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + convertedBytes, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) var resultPhg commonpb.PlaceholderGroup @@ -214,7 +214,7 @@ func TestConvertPlaceholderGroupToBFloat16(t *testing.T) { DataType: schemapb.DataType_BFloat16Vector, } - convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + convertedBytes, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) var resultPhg commonpb.PlaceholderGroup @@ -241,7 +241,7 @@ func TestConvertPlaceholderGroupTypeMismatch(t *testing.T) { DataType: schemapb.DataType_BFloat16Vector, } - _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + _, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.Error(t, err) assert.Contains(t, err.Error(), "vector type must be the same") } @@ -262,7 +262,7 @@ func TestConvertPlaceholderGroupPassThrough(t *testing.T) { DataType: schemapb.DataType_FloatVector, } - result, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + result, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) assert.Equal(t, phgBytes, result) } @@ -275,7 +275,7 @@ func TestConvertPlaceholderGroupNoConversionNeeded(t *testing.T) { DataType: schemapb.DataType_FloatVector, } - convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + convertedBytes, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) assert.Equal(t, phgBytes, convertedBytes) } @@ -288,7 +288,7 @@ func TestConvertPlaceholderGroupOutOfRange(t *testing.T) { DataType: schemapb.DataType_Float16Vector, } - _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + _, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.Error(t, err) assert.Contains(t, err.Error(), "exceeds float16 range") } @@ -327,7 +327,7 @@ func TestConvertPlaceholderGroupIntegration(t *testing.T) { DataType: schemapb.DataType_Float16Vector, } - convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + convertedBytes, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) var resultPhg commonpb.PlaceholderGroup @@ -344,7 +344,7 @@ func TestConvertPlaceholderGroupIntegration(t *testing.T) { DataType: schemapb.DataType_BFloat16Vector, } - convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + convertedBytes, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) var resultPhg commonpb.PlaceholderGroup @@ -361,7 +361,7 @@ func TestConvertPlaceholderGroupIntegration(t *testing.T) { DataType: schemapb.DataType_FloatVector, } - convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) + convertedBytes, _, err := ConvertPlaceholderGroup(phgBytes, fieldSchema) assert.NoError(t, err) assert.Equal(t, phgBytes, convertedBytes) // Should be unchanged }) diff --git a/pkg/util/typeutil/schema.go b/pkg/util/typeutil/schema.go index 2ce1a95801..d0240abc0c 100644 --- a/pkg/util/typeutil/schema.go +++ b/pkg/util/typeutil/schema.go @@ -2347,6 +2347,13 @@ func GetFieldByID(schema *schemapb.CollectionSchema, fieldID int64) *schemapb.Fi return field } } + for _, structField := range schema.GetStructArrayFields() { + for _, field := range structField.GetFields() { + if field.GetFieldID() == fieldID { + return field + } + } + } return nil }