*: use add field to avoid field mismatch (#887)

Signed-off-by: huanghaoyuanhhy <[email protected]>

Author: huanghaoyuanhhy

core/restore/collection_ddl_task.go

@@ -3,6 +3,7 @@ package restore
 import (
 	"context"
 	"fmt"
+	"sort"
 	"strings"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
@@ -89,36 +90,96 @@ func (ddlt *collectionDDLTask) dropExistedColl(ctx context.Context) error {
 	return nil
 }
 
+type sortableFields []*schemapb.FieldSchema
+
+func (s sortableFields) Len() int           { return len(s) }
+func (s sortableFields) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+func (s sortableFields) Less(i, j int) bool { return s[i].GetFieldID() < s[j].GetFieldID() }
+
+// fields returns two types of fields:
+// 1. fields that can be created in CreateCollection API
+// 2. fields that need to be created by addField API
+//
+// If the original collection had addField operations, the $meta field might not be at the end of the fields list.
+// However, the CreateCollection API always appends the $meta field to the end.
+// This discrepancy in field order between the new collection and the backup file can cause restore errors.
+// Therefore, we stop processing fields if the field IDs are not continuous.
+// The remaining fields will be imported via addField.
+func (ddlt *collectionDDLTask) fields() ([]*schemapb.FieldSchema, []*schemapb.FieldSchema, error) {
+	fields, err := ddlt.convFields(ddlt.collBackup.GetSchema().GetFields())
+	if err != nil {
+		return nil, nil, fmt.Errorf("collection: get fields: %w", err)
+	}
+
+	if !ddlt.collBackup.GetSchema().GetEnableDynamicField() {
+		return fields, nil, nil
+	}
+
+	sort.Sort(sortableFields(fields))
+
+	var createFields, addFields []*schemapb.FieldSchema
+	if len(fields) == 0 {
+		return createFields, addFields, nil
+	}
+
+	createFields = append(createFields, fields[0])
+	isContinuous := true
+	for i := 1; i < len(fields); i++ {
+		prevField := fields[i-1]
+		currField := fields[i]
+
+		if isContinuous && currField.GetFieldID() != prevField.GetFieldID()+1 {
+			isContinuous = false
+		}
+
+		if isContinuous {
+			createFields = append(createFields, currField)
+		} else {
+			addFields = append(addFields, currField)
+		}
+	}
+
+	ddlt.logger.Info("fields", zap.Any("create_fields", createFields), zap.Any("add_fields", addFields))
+	return createFields, addFields, nil
+}
+
+func (ddlt *collectionDDLTask) addFields(ctx context.Context, fields []*schemapb.FieldSchema) error {
+	// add fields
+	for _, field := range fields {
+		if err := ddlt.grpcCli.AddField(ctx, ddlt.targetNS.DBName(), ddlt.targetNS.CollName(), field); err != nil {
+			return fmt.Errorf("restore: add field %s id %d: %w", field.GetName(), field.GetFieldID(), err)
+		}
+	}
+
+	return nil
+}
+
 func (ddlt *collectionDDLTask) createColl(ctx context.Context) error {
 	if ddlt.option.SkipCreateCollection {
 		ddlt.logger.Info("skip create collection")
 		return nil
 	}
 
-	fields, err := ddlt.convFields(ddlt.collBackup.GetSchema().GetFields())
+	createFields, addFields, err := ddlt.fields()
 	if err != nil {
 		return fmt.Errorf("collection: get fields: %w", err)
 	}
-	ddlt.logger.Info("restore collection fields", zap.Any("fields", fields))
-
 	functions := conv.Functions(ddlt.collBackup.GetSchema().GetFunctions())
-	ddlt.logger.Info("restore collection functions", zap.Any("functions", functions))
-
 	structArrayFields, err := ddlt.structArrayFields()
 	if err != nil {
 		return fmt.Errorf("restore: conv struct array fields: %w", err)
 	}
-
 	schema := &schemapb.CollectionSchema{
 		Name:               ddlt.targetNS.CollName(),
 		Description:        ddlt.collBackup.GetSchema().GetDescription(),
 		AutoID:             ddlt.collBackup.GetSchema().GetAutoID(),
 		Functions:          functions,
-		Fields:             fields,
+		Fields:             createFields,
 		EnableDynamicField: ddlt.collBackup.GetSchema().GetEnableDynamicField(),
 		Properties:         pbconv.BakKVToMilvusKV(ddlt.collBackup.GetSchema().GetProperties()),
 		StructArrayFields:  structArrayFields,
 	}
+	ddlt.logger.Info("create collection", zap.Any("schema", schema))
 
 	opt := milvus.CreateCollectionInput{
 		DB:           ddlt.targetNS.DBName(),
@@ -132,6 +193,10 @@ func (ddlt *collectionDDLTask) createColl(ctx context.Context) error {
 		return fmt.Errorf("restore: call create collection api after retry: %w", err)
 	}
 
+	if err := ddlt.addFields(ctx, addFields); err != nil {
+		return fmt.Errorf("restore: add fields: %w", err)
+	}
+
 	return nil
 }
 
@@ -141,7 +206,7 @@ func (ddlt *collectionDDLTask) convFields(bakFields []*backuppb.FieldSchema) ([]
 	for _, bakField := range bakFields {
 		defaultValue, err := conv.DefaultValue(bakField)
 		if err != nil {
-			return nil, fmt.Errorf("restore: failed to get default value: %w", err)
+			return nil, fmt.Errorf("restore: get default value: %w", err)
 		}
 
 		fieldRestore := &schemapb.FieldSchema{
@@ -173,7 +238,7 @@ func (ddlt *collectionDDLTask) structArrayFields() ([]*schemapb.StructArrayField
 	for _, bakField := range bakFields {
 		fields, err := ddlt.convFields(bakField.GetFields())
 		if err != nil {
-			return nil, fmt.Errorf("restore: failed to convert struct array fields: %w", err)
+			return nil, fmt.Errorf("restore: convert struct array fields: %w", err)
 		}
 
 		structArrayField := &schemapb.StructArrayFieldSchema{

core/restore/collection_ddl_task_test.go

@@ -34,6 +34,56 @@ func TestCollectionDDLTask_shardNum(t *testing.T) {
 	})
 }
 
+func TestCollectionDDLTask_fields(t *testing.T) {
+	t.Run("WithDynamicField", func(t *testing.T) {
+		ddlt := newTestCollectionDDLTask()
+
+		ddlt.collBackup = &backuppb.CollectionBackupInfo{
+			Schema: &backuppb.CollectionSchema{
+				Fields: []*backuppb.FieldSchema{
+					{Name: "field1", FieldID: 1},
+					{Name: "field2", FieldID: 2},
+					{Name: "field3", FieldID: 3},
+					// $meta field's field id is 4
+					{Name: "field5", FieldID: 5},
+					{Name: "field6", FieldID: 6},
+				},
+				EnableDynamicField: true,
+			},
+		}
+
+		createFields, addFields, err := ddlt.fields()
+		assert.NoError(t, err)
+		assert.Len(t, createFields, 3)
+		assert.Len(t, addFields, 2)
+
+		assert.Equal(t, "field1", createFields[0].GetName())
+		assert.Equal(t, "field2", createFields[1].GetName())
+		assert.Equal(t, "field3", createFields[2].GetName())
+		assert.Equal(t, "field5", addFields[0].GetName())
+		assert.Equal(t, "field6", addFields[1].GetName())
+	})
+
+	t.Run("WithoutDynamicField", func(t *testing.T) {
+		ddlt := newTestCollectionDDLTask()
+
+		ddlt.collBackup = &backuppb.CollectionBackupInfo{
+			Schema: &backuppb.CollectionSchema{
+				Fields:             []*backuppb.FieldSchema{{Name: "field1", FieldID: 1}, {Name: "field2", FieldID: 2}},
+				EnableDynamicField: false,
+			},
+		}
+
+		createFields, addFields, err := ddlt.fields()
+		assert.NoError(t, err)
+		assert.Len(t, createFields, 2)
+		assert.Empty(t, addFields)
+
+		assert.Equal(t, "field1", createFields[0].GetName())
+		assert.Equal(t, "field2", createFields[1].GetName())
+	})
+}
+
 func TestCollectionTask_createColl(t *testing.T) {
 	t.Run("Skip", func(t *testing.T) {
 		ct := newTestCollectionDDLTask()

internal/client/milvus/grpc.go

@@ -119,6 +119,7 @@ type Grpc interface {
 	CreateCollection(ctx context.Context, input CreateCollectionInput) error
 	CreatePartition(ctx context.Context, db, collName, partitionName string) error
 	HasPartition(ctx context.Context, db, collName, partitionName string) (bool, error)
+	AddField(ctx context.Context, db, collName string, field *schemapb.FieldSchema) error
 	CreateIndex(ctx context.Context, input CreateIndexInput) error
 	DropIndex(ctx context.Context, db, collName, indexName string) error
 	BackupRBAC(ctx context.Context) (*milvuspb.BackupRBACMetaResponse, error)
@@ -886,6 +887,23 @@ func (g *GrpcClient) HasPartition(ctx context.Context, db, collName string, part
 	return resp.GetValue(), nil
 }
 
+func (g *GrpcClient) AddField(ctx context.Context, db, collName string, field *schemapb.FieldSchema) error {
+	ctx = g.newCtxWithDB(ctx, db)
+
+	bytes, err := proto.Marshal(field)
+	if err != nil {
+		return fmt.Errorf("client: add field marshal proto: %w", err)
+	}
+
+	in := &milvuspb.AddCollectionFieldRequest{CollectionName: collName, Schema: bytes}
+	resp, err := g.srv.AddCollectionField(ctx, in)
+	if err := checkResponse(resp, err); err != nil {
+		return fmt.Errorf("client: add field failed: %w", err)
+	}
+
+	return nil
+}
+
 func mapKvPairs(m map[string]string) []*commonpb.KeyValuePair {
 	pairs := make([]*commonpb.KeyValuePair, 0, len(m))
 	for k, v := range m {