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 <tangchenjie1210@gmail.com>

internal/core/src/exec/operator/search-groupby/SearchGroupByOperator.cpp

@@ -18,7 +18,6 @@
 #include <algorithm>
 #include <cstdint>
 #include <tuple>
-#include <unordered_set>
 #include <variant>
 
 #include "common/QueryInfo.h"
@@ -319,10 +318,6 @@ GroupIteratorResult(const std::shared_ptr<VectorIterator>& iterator,
     //note it may enumerate all data inside a segment and can block following
     //query and search possibly
     std::vector<std::tuple<int64_t, float, std::optional<T>>> 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<int64_t> seen_rows;
     while (iterator->HasNext() && !groupMap.IsGroupResEnough()) {
         auto offset_dis_pair = iterator->Next();
         AssertInfo(
@@ -332,17 +327,12 @@ GroupIteratorResult(const std::shared_ptr<VectorIterator>& 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<int32_t>(offset))
                     .first;
-            // Skip if we already saw this row (closest element wins)
-            if (!seen_rows.insert(row_offset).second) {
-                continue;
-            }
         }
 
         std::optional<T> row_data = data_getter->Get(row_offset);

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)
 			}

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()

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)
+	})
+}

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
 	}
 }
 

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
 	})

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
 }