diff --git a/internal/proxy/task_upsert.go b/internal/proxy/task_upsert.go index 57501393ef..116c247f61 100644 --- a/internal/proxy/task_upsert.go +++ b/internal/proxy/task_upsert.go @@ -294,11 +294,41 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { if len(it.upsertMsg.InsertMsg.GetFieldsData()) == 0 { return merr.WrapErrParameterInvalidMsg("upsert field data is empty") } + fieldsDataToCheckAligned := make([]*schemapb.FieldData, 0, len(it.upsertMsg.InsertMsg.GetFieldsData())) for _, fieldData := range it.upsertMsg.InsertMsg.GetFieldsData() { fieldName := fieldData.GetFieldName() if fieldData.GetIsDynamic() { fieldName = "$meta" } + if typeutil.IsStructSubField(fieldName) { + return merr.WrapErrParameterInvalidMsg("partial struct update is not supported for struct sub-field '%s'; use the whole struct field instead", fieldName) + } + if structSchema := it.schema.schemaHelper.GetStructArrayFieldFromName(fieldName); structSchema != nil { + fieldData.FieldId = structSchema.GetFieldID() + fieldData.FieldName = fieldName + if err := validateWholeStructFieldDataForPartialUpdate(it.schema.schemaHelper, structSchema, fieldData, upsertIDSize); err != nil { + return err + } + for _, subField := range fieldData.GetStructArrays().GetFields() { + if len(subField.GetValidData()) != 0 && subFieldHasData(subField) { + subFieldSchema, err := it.schema.schemaHelper.GetFieldFromName(storedStructSubFieldName(structSchema.GetName(), subField.GetFieldName())) + if err != nil { + log.Info("get struct sub-field schema failed", zap.Error(err)) + return err + } + if subFieldSchema.GetDefaultValue() != nil { + err = FillWithDefaultValue(subField, subFieldSchema, int(it.upsertMsg.InsertMsg.NRows())) + } else { + err = FillWithNullValue(subField, subFieldSchema, int(it.upsertMsg.InsertMsg.NRows())) + } + if err != nil { + log.Info("unify struct field data format failed", zap.Error(err)) + return err + } + } + } + continue + } fieldSchema, err := it.schema.schemaHelper.GetFieldFromName(fieldName) if err != nil { log.Info("get field schema failed", zap.Error(err)) @@ -335,10 +365,11 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { return err } } + fieldsDataToCheckAligned = append(fieldsDataToCheckAligned, fieldData) } // Validate field data alignment before processing to prevent index out of range panic - if err := newValidateUtil().checkAligned(it.upsertMsg.InsertMsg.GetFieldsData(), it.schema.schemaHelper, uint64(upsertIDSize)); err != nil { + if err := newValidateUtil().checkAligned(fieldsDataToCheckAligned, it.schema.schemaHelper, uint64(upsertIDSize)); err != nil { log.Warn("check field data aligned failed", zap.Error(err)) return err } @@ -396,8 +427,34 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { } existIndices[i] = int64(idx) } + upsertRows := make([]int64, len(updateIdxInUpsert)) + for i, upsertIdx := range updateIdxInUpsert { + upsertRows[i] = int64(upsertIdx) + } for fieldIdx, existField := range existFieldData { + dstField := it.insertFieldData[fieldIdx] + upsertField := upsertFieldMap[existField.FieldId] + + op := schemapb.FieldPartialUpdateOp_REPLACE + if fieldOpMap != nil { + if resolved, ok := fieldOpMap[existField.GetFieldName()]; ok { + op = resolved + } + } + + if it.schema.schemaHelper.GetStructArrayFieldFromName(existField.GetFieldName()) != nil { + if upsertField != nil { + if op != schemapb.FieldPartialUpdateOp_REPLACE { + return merr.WrapErrParameterInvalidMsg("op %s is not supported for struct field %q", op.String(), existField.GetFieldName()) + } + typeutil.AppendFieldDataByColumn(dstField, upsertField, upsertRows) + } else { + typeutil.AppendFieldDataByColumn(dstField, existField, existIndices) + } + continue + } + fieldSchema, err := it.schema.schemaHelper.GetFieldFromName(existField.GetFieldName()) if err != nil { log.Info("get field schema failed", zap.Error(err)) @@ -406,7 +463,7 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { // Note: For fields containing default values, default values need to be set according to valid data during insertion, // but query results fields do not set valid data when returning default value fields, - // therefore valid data needs to be manually set to true + // therefore valid data needs to be manually set to true. if fieldSchema.GetDefaultValue() != nil && len(existField.GetValidData()) == 0 { existField.ValidData = make([]bool, existIDsLen) for i := range existField.ValidData { @@ -414,15 +471,12 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { } } - dstField := it.insertFieldData[fieldIdx] - upsertField := upsertFieldMap[existField.FieldId] isNullableVector := typeutil.IsCompactNullableVectorFieldData(existField) existComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{existField}) existSrcIndices := make([]int64, len(updateIdxInUpsert)) for i, existIdx := range existIndices { existSrcIndices[i] = existComputer.Compute(existIdx)[0] } - if upsertField != nil { upsertComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{upsertField}) upsertSrcIndices := make([]int64, len(updateIdxInUpsert)) @@ -446,12 +500,6 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { for i := range dstIndices { dstIndices[i] = int64(i) } - op := schemapb.FieldPartialUpdateOp_REPLACE - if fieldOpMap != nil { - if resolved, ok := fieldOpMap[existField.GetFieldName()]; ok { - op = resolved - } - } if op == schemapb.FieldPartialUpdateOp_REPLACE { if err := typeutil.UpdateFieldDataByColumn(dstField, upsertField, dstIndices, upsertSrcIndices); err != nil { return err @@ -489,9 +537,10 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { return lackOfFieldErr } - // if the nullable field has not passed in upsert request, which means the len(upsertFieldData) < len(it.insertFieldData) - // we need to generate the nullable field data before append as insert - insertWithNullField := make([]*schemapb.FieldData, 0) + // Build field data for rows that take insert semantics. Request-provided + // fields are kept as-is; missing nullable/default fields are synthesized + // so the insert rows still have a complete field set. + insertFieldsToAppend := make([]*schemapb.FieldData, 0) upsertFieldMap := lo.SliceToMap(it.upsertMsg.InsertMsg.GetFieldsData(), func(field *schemapb.FieldData) (string, *schemapb.FieldData) { return field.GetFieldName(), field }) @@ -503,10 +552,15 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { log.Info("generate nullable field data failed", zap.Error(err)) return err } - insertWithNullField = append(insertWithNullField, fieldData) + insertFieldsToAppend = append(insertFieldsToAppend, fieldData) } } else { - insertWithNullField = append(insertWithNullField, fieldData) + insertFieldsToAppend = append(insertFieldsToAppend, fieldData) + } + } + for _, structSchema := range it.schema.GetStructArrayFields() { + if fieldData, ok := upsertFieldMap[structSchema.GetName()]; ok { + insertFieldsToAppend = append(insertFieldsToAppend, fieldData) } } @@ -517,7 +571,7 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { } // Process insert data by column (field), similar to update path - for _, srcField := range insertWithNullField { + for _, srcField := range insertFieldsToAppend { isNullableVector := typeutil.IsCompactNullableVectorFieldData(srcField) srcComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{srcField}) @@ -554,6 +608,13 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { } for _, fieldData := range it.insertFieldData { + if fieldData.GetType() == schemapb.DataType_ArrayOfStruct { + if err := ToCompressedFormatNullableStructField(fieldData); err != nil { + log.Info("convert struct field data to compressed format nullable failed", zap.Error(err)) + return err + } + continue + } if len(fieldData.GetValidData()) > 0 { err := ToCompressedFormatNullable(fieldData) if err != nil { @@ -565,7 +626,7 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error { return nil } -// ToCompressedFormatNullable converts the field data from full format nullable to compressed format nullable +// ToCompressedFormatNullable converts nullable field data from full format to compressed format. func ToCompressedFormatNullable(field *schemapb.FieldData) error { if getValidNumber(field.GetValidData()) == len(field.GetValidData()) { return nil @@ -882,7 +943,8 @@ func GenNullableFieldData(field *schemapb.FieldSchema, upsertIDSize int) (*schem Scalars: &schemapb.ScalarField{ Data: &schemapb.ScalarField_ArrayData{ ArrayData: &schemapb.ArrayArray{ - Data: make([]*schemapb.ScalarField, upsertIDSize), + Data: make([]*schemapb.ScalarField, upsertIDSize), + ElementType: field.GetElementType(), }, }, }, @@ -1037,6 +1099,213 @@ func GenNullableFieldData(field *schemapb.FieldSchema, upsertIDSize int) (*schem } } +func GenNullableStructArrayFieldData(structField *schemapb.StructArrayFieldSchema, rowCount int) *schemapb.FieldData { + fields := make([]*schemapb.FieldData, 0, len(structField.GetFields())) + for _, subField := range structField.GetFields() { + fieldName := subField.GetName() + if typeutil.IsStructSubField(fieldName) { + if originalName, err := typeutil.ExtractStructFieldName(fieldName); err == nil { + fieldName = originalName + } + } + fields = append(fields, &schemapb.FieldData{ + FieldName: fieldName, + FieldId: subField.GetFieldID(), + Type: subField.GetDataType(), + // The source field is indexed by original upsert row IDs before + // appending the insert subset, so ValidData uses the same row domain. + ValidData: make([]bool, rowCount), + }) + } + return &schemapb.FieldData{ + FieldId: structField.GetFieldID(), + FieldName: structField.GetName(), + Type: schemapb.DataType_ArrayOfStruct, + Field: &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{Fields: fields}, + }, + } +} + +func storedStructSubFieldName(structName string, fieldName string) string { + if typeutil.IsStructSubField(fieldName) { + return fieldName + } + return typeutil.ConcatStructFieldName(structName, fieldName) +} + +func subFieldHasData(fieldData *schemapb.FieldData) bool { + switch fieldData.GetType() { + case schemapb.DataType_Array: + scalars := fieldData.GetScalars() + return scalars != nil && scalars.GetArrayData() != nil && len(scalars.GetArrayData().GetData()) > 0 + case schemapb.DataType_ArrayOfVector: + vectors := fieldData.GetVectors() + return vectors != nil && vectors.GetVectorArray() != nil && len(vectors.GetVectorArray().GetData()) > 0 + default: + return false + } +} + +// ToCompressedFormatNullableStructField converts nullable struct sub-fields from full format to compressed format. +func ToCompressedFormatNullableStructField(fieldData *schemapb.FieldData) error { + structArrays := fieldData.GetStructArrays() + if structArrays == nil { + return nil + } + for _, subField := range structArrays.GetFields() { + if len(subField.GetValidData()) == 0 { + continue + } + + validData := subField.GetValidData() + validRows := getValidNumber(validData) + if validRows == 0 && !subFieldHasData(subField) { + continue + } + payloadRows, ok := structSubFieldPayloadRowsForPartialUpdate(subField) + if !ok { + return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' has invalid field data", subField.GetFieldName()) + } + if payloadRows == validRows { + continue + } + if payloadRows != len(validData) { + return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' payload row count %d does not match valid rows %d or logical rows %d", + subField.GetFieldName(), payloadRows, validRows, len(validData)) + } + if subField.GetType() == schemapb.DataType_ArrayOfVector { + vectorArray := subField.GetVectors().GetVectorArray() + compactRows := make([]*schemapb.VectorField, 0, validRows) + for row, valid := range validData { + if valid { + compactRows = append(compactRows, vectorArray.GetData()[row]) + } + } + vectorArray.Data = compactRows + continue + } + if err := ToCompressedFormatNullable(subField); err != nil { + return err + } + } + return nil +} + +// validateWholeStructFieldDataForPartialUpdate validates only the logical +// top-level struct payload used by partial-update whole struct REPLACE. +// It checks: +// 1. ArrayOfStruct wrapper and sub-field count/name/type. +// 2. All-or-none sub-field payload presence and nullable struct null-mask consistency. +// 3. Payload row counts needed for safe row append. +// It does not flatten struct data or merge sub-fields; final insert +// pre-execute still runs checkAndFlattenStructFieldData before storage writes. +func validateWholeStructFieldDataForPartialUpdate(schemaHelper *typeutil.SchemaHelper, structSchema *schemapb.StructArrayFieldSchema, fieldData *schemapb.FieldData, rowCount int) error { + if fieldData.GetType() != schemapb.DataType_ArrayOfStruct { + return merr.WrapErrParameterInvalidMsg("field %q expects ArrayOfStruct payload, got %s", fieldData.GetFieldName(), fieldData.GetType().String()) + } + structArrays := fieldData.GetStructArrays() + if structArrays == nil { + return merr.WrapErrParameterInvalidMsg("field %q expects non-nil StructArrays payload", fieldData.GetFieldName()) + } + if len(structArrays.GetFields()) != len(structSchema.GetFields()) { + return merr.WrapErrParameterInvalidMsg("length of fields of struct field mismatch length of the fields in schema, fieldName: %s, fieldData fields length:%d, schema fields length:%d", + fieldData.GetFieldName(), len(structArrays.GetFields()), len(structSchema.GetFields())) + } + + seenSubFieldIDs := make(map[int64]struct{}, len(structArrays.GetFields())) + for _, subField := range structArrays.GetFields() { + subFieldSchema, err := schemaHelper.GetFieldFromName(storedStructSubFieldName(structSchema.GetName(), subField.GetFieldName())) + if err != nil { + return err + } + subField.FieldId = subFieldSchema.GetFieldID() + if _, ok := seenSubFieldIDs[subFieldSchema.GetFieldID()]; ok { + return merr.WrapErrParameterInvalidMsg("duplicated sub-field '%s' in struct '%s'", subField.GetFieldName(), fieldData.GetFieldName()) + } + seenSubFieldIDs[subFieldSchema.GetFieldID()] = struct{}{} + if subField.GetType() != subFieldSchema.GetDataType() { + return merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s' expects type %s, got %s", + subField.GetFieldName(), fieldData.GetFieldName(), subFieldSchema.GetDataType().String(), subField.GetType().String()) + } + if len(subField.GetValidData()) > 0 && len(subField.GetValidData()) != rowCount { + return merr.WrapErrParameterInvalidMsg("sub-field ValidData length mismatch in struct '%s': '%s' has %d, expected %d", + fieldData.GetFieldName(), subField.GetFieldName(), len(subField.GetValidData()), rowCount) + } + } + + hasDataCount := 0 + for _, subField := range structArrays.GetFields() { + if subFieldHasData(subField) { + hasDataCount++ + } + } + // All sub-fields must be either present together or absent together. + // A partial payload would mean partial struct replace, which is not + // supported by this path. + if hasDataCount == 0 { + for _, subField := range structArrays.GetFields() { + if len(subField.GetValidData()) != 0 && len(subField.GetValidData()) != rowCount { + return merr.WrapErrParameterInvalidMsg("sub-field ValidData length mismatch in struct '%s': '%s' has %d, expected %d", + fieldData.GetFieldName(), subField.GetFieldName(), len(subField.GetValidData()), rowCount) + } + for j, valid := range subField.GetValidData() { + if valid { + return merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s' claims row %d is valid but no payload is provided", + subField.GetFieldName(), fieldData.GetFieldName(), j) + } + } + } + return nil + } + if hasDataCount != len(structArrays.GetFields()) { + return merr.WrapErrParameterInvalidMsg("inconsistent sub-field data in struct '%s': %d of %d sub-fields have data, all must be present or all absent", + fieldData.GetFieldName(), hasDataCount, len(structArrays.GetFields())) + } + + for _, subField := range structArrays.GetFields() { + // Struct sub-fields without ValidData use dense logical rows. With + // ValidData, the payload must be compact so FillWithNullValue can expand + // it consistently before merge. + payloadRows, ok := structSubFieldPayloadRowsForPartialUpdate(subField) + if !ok { + return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' has invalid field data", subField.GetFieldName()) + } + if len(subField.GetValidData()) == 0 { + if payloadRows != rowCount { + return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' payload row count %d does not match logical rows %d", + subField.GetFieldName(), payloadRows, rowCount) + } + continue + } + validRows := getValidNumber(subField.GetValidData()) + if payloadRows != validRows { + return merr.WrapErrParameterInvalidMsg("nullable struct sub-field '%s' payload must be compact: payload row count %d does not match valid rows %d", + subField.GetFieldName(), payloadRows, validRows) + } + } + return nil +} + +func structSubFieldPayloadRowsForPartialUpdate(fieldData *schemapb.FieldData) (int, bool) { + switch fieldData.GetType() { + case schemapb.DataType_Array: + scalars := fieldData.GetScalars() + if scalars == nil || scalars.GetArrayData() == nil { + return 0, false + } + return len(scalars.GetArrayData().GetData()), true + case schemapb.DataType_ArrayOfVector: + vectors := fieldData.GetVectors() + if vectors == nil || vectors.GetVectorArray() == nil { + return 0, false + } + return len(vectors.GetVectorArray().GetData()), true + default: + return 0, false + } +} + func (it *upsertTask) insertPreExecute(ctx context.Context) error { ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert-insertPreExecute") defer sp.End() diff --git a/internal/proxy/task_upsert_partial_op.go b/internal/proxy/task_upsert_partial_op.go index 1292e51ec8..fb05503134 100644 --- a/internal/proxy/task_upsert_partial_op.go +++ b/internal/proxy/task_upsert_partial_op.go @@ -102,6 +102,10 @@ func validateFieldPartialUpdateOps(req *milvuspb.UpsertRequest, schema *schemapb if len(fieldOps) == 0 { return false, nil } + schemaHelper, err := typeutil.CreateSchemaHelper(schema) + if err != nil { + return false, err + } // Precompute PK names and a lookup table of FieldData by name so we // can validate payload alignment in O(1) per op. @@ -131,6 +135,10 @@ func validateFieldPartialUpdateOps(req *milvuspb.UpsertRequest, schema *schemapb op := opMsg.GetOp() if op == schemapb.FieldPartialUpdateOp_REPLACE { + if typeutil.IsStructSubField(name) { + return false, merr.WrapErrParameterInvalidMsg( + "partial struct update is not supported for struct sub-field %q; use the whole struct field instead", name) + } // An explicit REPLACE is legal but indistinguishable from no // op at all. Accept silently — no further validation needed. continue @@ -141,6 +149,14 @@ func validateFieldPartialUpdateOps(req *milvuspb.UpsertRequest, schema *schemapb return false, merr.WrapErrParameterInvalidMsg( fmt.Sprintf("field %q is the primary key and cannot carry a partial-update op", name)) } + if schemaHelper.GetStructArrayFieldFromName(name) != nil { + return false, merr.WrapErrParameterInvalidMsg( + "op %s is not supported for struct field %q", op.String(), name) + } + if typeutil.IsStructSubField(name) { + return false, merr.WrapErrParameterInvalidMsg( + "op %s is not supported for struct field %q", op.String(), name) + } fieldSchema, err := findFieldSchemaByName(schema, name) if err != nil { diff --git a/internal/proxy/task_upsert_partial_op_test.go b/internal/proxy/task_upsert_partial_op_test.go index 4ecf3e21e2..7a0378a810 100644 --- a/internal/proxy/task_upsert_partial_op_test.go +++ b/internal/proxy/task_upsert_partial_op_test.go @@ -197,6 +197,34 @@ func TestValidateFieldPartialUpdateOps_RejectsNonArrayField(t *testing.T) { assert.Contains(t, err.Error(), "Array field") } +func TestValidateFieldPartialUpdateOps_RejectsStructOps(t *testing.T) { + schema := &schemapb.CollectionSchema{ + StructArrayFields: []*schemapb.StructArrayFieldSchema{ + { + Name: "profile", + Fields: []*schemapb.FieldSchema{ + {Name: "profile[age]", DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_Int64}, + }, + }, + }, + } + + req := &milvuspb.UpsertRequest{ + FieldsData: []*schemapb.FieldData{{FieldName: "profile", Type: schemapb.DataType_ArrayOfStruct}}, + FieldOps: []*schemapb.FieldPartialUpdateOp{op("profile", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)}, + } + _, err := validateFieldPartialUpdateOps(req, schema) + require.Error(t, err) + assert.Contains(t, err.Error(), "not supported for struct field") + + req = &milvuspb.UpsertRequest{ + FieldOps: []*schemapb.FieldPartialUpdateOp{op("profile[age]", schemapb.FieldPartialUpdateOp_REPLACE)}, + } + _, err = validateFieldPartialUpdateOps(req, schema) + require.Error(t, err) + assert.Contains(t, err.Error(), "partial struct update is not supported") +} + func TestValidateFieldPartialUpdateOps_RejectsElementTypeMismatch(t *testing.T) { schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{ {Name: "tags", DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_VarChar}, diff --git a/internal/proxy/task_upsert_test.go b/internal/proxy/task_upsert_test.go index 41f2f22dc2..5af727b2ac 100644 --- a/internal/proxy/task_upsert_test.go +++ b/internal/proxy/task_upsert_test.go @@ -1486,6 +1486,197 @@ func TestUpsertTask_queryPreExecute_ArrayPartialOpSkipsGenericReplace(t *testing assert.Equal(t, []string{"new1"}, noteField.GetScalars().GetStringData().GetData()) } +func TestUpsertTask_queryPreExecute_StructWholeReplace(t *testing.T) { + schema := newSchemaInfo(&schemapb.CollectionSchema{ + Name: "test_struct_partial_update", + Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64}, + {FieldID: 101, Name: "value", DataType: schemapb.DataType_Int32}, + }, + StructArrayFields: []*schemapb.StructArrayFieldSchema{ + { + FieldID: 200, + Name: "profile", + Fields: []*schemapb.FieldSchema{ + {FieldID: 201, Name: "profile[age]", DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_Int32}, + }, + }, + }, + }) + + idField := func(ids ...int64) *schemapb.FieldData { + return &schemapb.FieldData{ + FieldName: "id", FieldId: 100, Type: schemapb.DataType_Int64, + Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: ids}}}}, + } + } + valueField := func(values ...int32) *schemapb.FieldData { + return &schemapb.FieldData{ + FieldName: "value", FieldId: 101, Type: schemapb.DataType_Int32, + Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: values}}}}, + } + } + profileField := func(values ...int32) *schemapb.FieldData { + rows := make([]*schemapb.ScalarField, 0, len(values)) + for _, value := range values { + rows = append(rows, &schemapb.ScalarField{ + Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{value}}}, + }) + } + return &schemapb.FieldData{ + FieldName: "profile", + FieldId: 200, + Type: schemapb.DataType_ArrayOfStruct, + Field: &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{ + Fields: []*schemapb.FieldData{ + { + FieldName: "age", + FieldId: 201, + Type: schemapb.DataType_Array, + Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{ + ElementType: schemapb.DataType_Int32, + Data: rows, + }}}}, + }, + }, + }, + }, + } + } + queryResult := func() *milvuspb.QueryResults { + return &milvuspb.QueryResults{ + Status: merr.Success(), + FieldsData: []*schemapb.FieldData{ + idField(1, 2), + valueField(10, 20), + profileField(11, 22), + }, + } + } + run := func(fields []*schemapb.FieldData) (*upsertTask, error) { + task := &upsertTask{ + ctx: context.Background(), + schema: schema, + req: &milvuspb.UpsertRequest{ + FieldsData: fields, + NumRows: 2, + PartialUpdate: true, + CollectionName: "test_struct_partial_update", + }, + upsertMsg: &msgstream.UpsertMsg{InsertMsg: &msgstream.InsertMsg{ + InsertRequest: &msgpb.InsertRequest{ + FieldsData: fields, + NumRows: 2, + Version: msgpb.InsertDataVersion_ColumnBased, + }, + }}, + node: &Proxy{}, + } + mockRetrieve := mockey.Mock(retrieveByPKs).Return(queryResult(), segcore.StorageCost{}, nil).Build() + defer mockRetrieve.UnPatch() + return task, task.queryPreExecute(context.Background()) + } + structValues := func(field *schemapb.FieldData) []int32 { + rows := field.GetStructArrays().GetFields()[0].GetScalars().GetArrayData().GetData() + values := make([]int32, 0, len(rows)) + for _, row := range rows { + values = append(values, row.GetIntData().GetData()[0]) + } + return values + } + findProfile := func(task *upsertTask) *schemapb.FieldData { + for _, field := range task.insertFieldData { + if field.GetFieldName() == "profile" { + return field + } + } + return nil + } + + t.Run("omitted struct preserves old value", func(t *testing.T) { + task, err := run([]*schemapb.FieldData{ + idField(1, 2), + valueField(100, 200), + }) + assert.NoError(t, err) + profile := findProfile(task) + if assert.NotNil(t, profile) { + assert.Equal(t, []int32{11, 22}, structValues(profile)) + } + }) + + t.Run("provided top-level struct replaces whole struct", func(t *testing.T) { + task, err := run([]*schemapb.FieldData{ + idField(1, 2), + valueField(100, 200), + profileField(111, 222), + }) + assert.NoError(t, err) + profile := findProfile(task) + if assert.NotNil(t, profile) { + assert.Equal(t, []int32{111, 222}, structValues(profile)) + } + }) + + t.Run("mixed insert update keeps request struct rows", func(t *testing.T) { + task := &upsertTask{ + ctx: context.Background(), + schema: schema, + req: &milvuspb.UpsertRequest{ + FieldsData: []*schemapb.FieldData{ + idField(1, 3), + valueField(100, 300), + profileField(111, 333), + }, + NumRows: 2, + PartialUpdate: true, + CollectionName: "test_struct_partial_update", + }, + upsertMsg: &msgstream.UpsertMsg{InsertMsg: &msgstream.InsertMsg{ + InsertRequest: &msgpb.InsertRequest{ + FieldsData: []*schemapb.FieldData{ + idField(1, 3), + valueField(100, 300), + profileField(111, 333), + }, + NumRows: 2, + Version: msgpb.InsertDataVersion_ColumnBased, + }, + }}, + node: &Proxy{}, + } + mockRetrieve := mockey.Mock(retrieveByPKs).Return(&milvuspb.QueryResults{ + Status: merr.Success(), + FieldsData: []*schemapb.FieldData{ + idField(1), + valueField(10), + profileField(11), + }, + }, segcore.StorageCost{}, nil).Build() + defer mockRetrieve.UnPatch() + + err := task.queryPreExecute(context.Background()) + assert.NoError(t, err) + profile := findProfile(task) + if assert.NotNil(t, profile) { + assert.Equal(t, []int32{111, 333}, structValues(profile)) + } + }) + + t.Run("direct struct sub-field update is rejected", func(t *testing.T) { + subField := profileField(111, 222).GetStructArrays().GetFields()[0] + subField.FieldName = "profile[age]" + _, err := run([]*schemapb.FieldData{ + idField(1, 2), + valueField(100, 200), + subField, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "partial struct update is not supported") + }) +} + func TestCheckDynamicFieldDataForPartialUpdate(t *testing.T) { t.Run("preserves $meta keys matching static field names after schema evolution", func(t *testing.T) { schema := &schemapb.CollectionSchema{ diff --git a/internal/proxy/util.go b/internal/proxy/util.go index db041cb3a6..d7f89a2aa9 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -2070,6 +2070,12 @@ func LackOfFieldsDataBySchema(schema *schemapb.CollectionSchema, fieldsData []*s return merr.WrapErrParameterInvalidMsg("fieldSchema(%s) has no corresponding fieldData pass in", fieldSchema.GetName()) } } + for _, structSchema := range schema.GetStructArrayFields() { + if _, ok := dataNameMap[structSchema.GetName()]; !ok { + log.Info("no corresponding struct fieldData pass in", zap.String("structFieldSchema", structSchema.GetName())) + return merr.WrapErrParameterInvalidMsg("structFieldSchema(%s) has no corresponding fieldData pass in", structSchema.GetName()) + } + } return nil } @@ -2408,7 +2414,6 @@ func verifyDynamicFieldData(schema *schemapb.CollectionSchema, insertMsg *msgstr } } } - return nil } diff --git a/pkg/util/typeutil/schema.go b/pkg/util/typeutil/schema.go index 2f3559f749..9dd6aaa1c8 100644 --- a/pkg/util/typeutil/schema.go +++ b/pkg/util/typeutil/schema.go @@ -851,6 +851,12 @@ func PrepareResultFieldData(sample []*schemapb.FieldData, topK int64) []*schemap } } fd.Field = vectors + case *schemapb.FieldData_StructArrays: + fd.Field = &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{ + Fields: PrepareResultFieldData(fieldData.GetStructArrays().GetFields(), topK), + }, + } } result = append(result, fd) } @@ -1452,7 +1458,53 @@ func AppendFieldDataByColumn(dst, src *schemapb.FieldData, dataIndices []int64, dstVector.Data.(*schemapb.VectorField_Int8Vector).Int8Vector, srcVector.Int8Vector[start:end]...) } + case *schemapb.VectorField_VectorArray: + if srcVector.VectorArray == nil { + return + } + if dstVector.GetVectorArray() == nil { + dstVector.Data = &schemapb.VectorField_VectorArray{ + VectorArray: &schemapb.VectorArray{ + Data: make([]*schemapb.VectorField, 0, len(dataIndices)), + Dim: srcVector.VectorArray.Dim, + ElementType: srcVector.VectorArray.ElementType, + }, + } + } + for _, idx := range dataIndices { + dstVector.GetVectorArray().Data = append(dstVector.GetVectorArray().Data, srcVector.VectorArray.Data[idx]) + } } + case *schemapb.FieldData_StructArrays: + appendStructFieldDataByColumn(dst, srcField.StructArrays, dataIndices, rowIndices...) + } +} + +func appendStructFieldDataByColumn(dst *schemapb.FieldData, src *schemapb.StructArrayField, dataIndices []int64, rowIndices ...[]int64) { + if src == nil { + return + } + if dst.GetStructArrays() == nil { + dst.Field = &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{ + Fields: PrepareResultFieldData(src.GetFields(), int64(len(dataIndices))), + }, + } + } + dstStruct := dst.GetStructArrays() + dstSubFields := make(map[int64]*schemapb.FieldData, len(dstStruct.GetFields())) + for _, dstSubField := range dstStruct.GetFields() { + dstSubFields[dstSubField.GetFieldId()] = dstSubField + } + for _, srcSubField := range src.GetFields() { + dstSubField := dstSubFields[srcSubField.GetFieldId()] + if dstSubField == nil { + newFields := PrepareResultFieldData([]*schemapb.FieldData{srcSubField}, int64(len(dataIndices))) + dstSubField = newFields[0] + dstStruct.Fields = append(dstStruct.Fields, dstSubField) + dstSubFields[dstSubField.GetFieldId()] = dstSubField + } + AppendFieldDataByColumn(dstSubField, srcSubField, dataIndices, rowIndices...) } } @@ -3319,6 +3371,12 @@ func ConcatStructFieldName(structName string, fieldName string) string { return fmt.Sprintf("%s[%s]", structName, fieldName) } +// IsStructSubField checks if a field name follows the "structName[fieldName]" convention, +// indicating it is a sub-field within a StructArrayField. +func IsStructSubField(fieldName string) bool { + return strings.Contains(fieldName, "[") +} + func ExtractStructFieldName(fieldName string) (string, error) { parts := strings.Split(fieldName, "[") if len(parts) == 1 { diff --git a/pkg/util/typeutil/schema_test.go b/pkg/util/typeutil/schema_test.go index 006c73cf1b..9c3a30cb33 100644 --- a/pkg/util/typeutil/schema_test.go +++ b/pkg/util/typeutil/schema_test.go @@ -5332,6 +5332,127 @@ func TestAppendFieldDataByColumn(t *testing.T) { AppendFieldDataByColumn(dst, src, []int64{0, 2}) assert.Equal(t, [][]byte{{1, 2}, {5, 6}}, dst.GetVectors().GetSparseFloatVector().Contents) }) + + t.Run("nullable array of vector copies dense row data", func(t *testing.T) { + validData := []bool{true, false, true, false} + src := newArrayOfVectorFieldData(200, "arrvec", 1, schemapb.DataType_FloatVector, validData) + dst := &schemapb.FieldData{Type: schemapb.DataType_ArrayOfVector} + + AppendFieldDataByColumn(dst, src, []int64{0, 1, 2, 3}, []int64{0, 1, 2, 3}) + + got := dst.GetVectors().GetVectorArray() + require.NotNil(t, got) + require.Len(t, got.GetData(), 4) + assert.Equal(t, validData, dst.GetValidData()) + assert.Equal(t, []float32{1}, got.GetData()[0].GetFloatVector().GetData()) + assert.Empty(t, got.GetData()[1].GetFloatVector().GetData()) + assert.Equal(t, []float32{3}, got.GetData()[2].GetFloatVector().GetData()) + assert.Empty(t, got.GetData()[3].GetFloatVector().GetData()) + assert.EqualValues(t, 1, got.GetDim()) + assert.Equal(t, schemapb.DataType_FloatVector, got.GetElementType()) + }) + + t.Run("struct array appends dense nullable sub-field rows", func(t *testing.T) { + src := &schemapb.FieldData{ + FieldName: "profile", + FieldId: 200, + Type: schemapb.DataType_ArrayOfStruct, + Field: &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{ + Fields: []*schemapb.FieldData{ + { + FieldName: "age", + FieldId: 201, + Type: schemapb.DataType_Array, + ValidData: []bool{true, false, true}, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_ArrayData{ + ArrayData: &schemapb.ArrayArray{ + ElementType: schemapb.DataType_Int32, + Data: []*schemapb.ScalarField{ + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{10}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{30}}}}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + } + dst := PrepareResultFieldData([]*schemapb.FieldData{src}, 3)[0] + + AppendFieldDataByColumn(dst, src, []int64{0, 1, 2}) + + require.NotNil(t, dst.GetStructArrays()) + require.Len(t, dst.GetStructArrays().GetFields(), 1) + subField := dst.GetStructArrays().GetFields()[0] + assert.Equal(t, []bool{true, false, true}, subField.GetValidData()) + got := subField.GetScalars().GetArrayData().GetData() + require.Len(t, got, 3) + assert.Equal(t, []int32{10}, got[0].GetIntData().GetData()) + assert.Empty(t, got[1].GetIntData().GetData()) + assert.Equal(t, []int32{30}, got[2].GetIntData().GetData()) + }) + + t.Run("struct array matches sub-fields by id", func(t *testing.T) { + dst := &schemapb.FieldData{ + FieldName: "profile", + FieldId: 200, + Type: schemapb.DataType_ArrayOfStruct, + Field: &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{ + Fields: []*schemapb.FieldData{ + { + FieldName: "profile[age]", + FieldId: 201, + Type: schemapb.DataType_Array, + Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{ + ElementType: schemapb.DataType_Int32, + Data: []*schemapb.ScalarField{}, + }}}}, + }, + }, + }, + }, + } + src := &schemapb.FieldData{ + FieldName: "profile", + FieldId: 200, + Type: schemapb.DataType_ArrayOfStruct, + Field: &schemapb.FieldData_StructArrays{ + StructArrays: &schemapb.StructArrayField{ + Fields: []*schemapb.FieldData{ + { + FieldName: "age", + FieldId: 201, + Type: schemapb.DataType_Array, + Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{ + ElementType: schemapb.DataType_Int32, + Data: []*schemapb.ScalarField{ + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{42}}}}, + }, + }}}}, + }, + }, + }, + }, + } + + AppendFieldDataByColumn(dst, src, []int64{0}) + + require.NotNil(t, dst.GetStructArrays()) + require.Len(t, dst.GetStructArrays().GetFields(), 1) + subField := dst.GetStructArrays().GetFields()[0] + assert.Equal(t, "profile[age]", subField.GetFieldName()) + got := subField.GetScalars().GetArrayData().GetData() + require.Len(t, got, 1) + assert.Equal(t, []int32{42}, got[0].GetIntData().GetData()) + }) } func TestUpdateFieldDataByColumn(t *testing.T) { @@ -5639,9 +5760,30 @@ func TestUpdateFieldDataByColumn(t *testing.T) { }) } -func newArrayOfVectorFieldData(fieldID int64, fieldName string, dim int64, elementType schemapb.DataType, rowCount int) *schemapb.FieldData { +func newArrayOfVectorFieldData(fieldID int64, fieldName string, dim int64, elementType schemapb.DataType, rows any) *schemapb.FieldData { + var rowCount int + var validData []bool + switch v := rows.(type) { + case int: + rowCount = v + case []bool: + validData = v + rowCount = len(v) + default: + panic("unexpected ArrayOfVector row spec") + } + data := make([]*schemapb.VectorField, rowCount) for i := 0; i < rowCount; i++ { + if len(validData) > 0 && !validData[i] { + data[i] = &schemapb.VectorField{ + Dim: dim, + Data: &schemapb.VectorField_FloatVector{ + FloatVector: &schemapb.FloatArray{}, + }, + } + continue + } data[i] = &schemapb.VectorField{ Dim: dim, Data: &schemapb.VectorField_FloatVector{ @@ -5649,7 +5791,7 @@ func newArrayOfVectorFieldData(fieldID int64, fieldName string, dim int64, eleme }, } } - return &schemapb.FieldData{ + fd := &schemapb.FieldData{ Type: schemapb.DataType_ArrayOfVector, FieldName: fieldName, FieldId: fieldID, @@ -5666,6 +5808,10 @@ func newArrayOfVectorFieldData(fieldID int64, fieldName string, dim int64, eleme }, }, } + if len(validData) > 0 { + fd.ValidData = validData + } + return fd } func TestFieldDataIdxComputer_ArrayOfVectorIsNonVector(t *testing.T) {