diff --git a/internal/datanode/importv2/executor_test.go b/internal/datanode/importv2/executor_test.go index ccd4605f77..0ef11f50c8 100644 --- a/internal/datanode/importv2/executor_test.go +++ b/internal/datanode/importv2/executor_test.go @@ -123,6 +123,9 @@ func createInsertData(t *testing.T, schema *schemapb.CollectionSchema, rowCount insertData, err := storage.NewInsertData(schema) assert.NoError(t, err) for _, field := range schema.GetFields() { + if field.GetAutoID() && field.GetIsPrimaryKey() { + continue + } switch field.GetDataType() { case schemapb.DataType_Bool: boolData := make([]bool, 0) diff --git a/internal/datanode/importv2/util.go b/internal/datanode/importv2/util.go index 4d3ea60e50..ad836c1378 100644 --- a/internal/datanode/importv2/util.go +++ b/internal/datanode/importv2/util.go @@ -19,6 +19,7 @@ package importv2 import ( "context" "fmt" + "strconv" "github.com/samber/lo" "go.uber.org/zap" @@ -168,7 +169,15 @@ func AppendSystemFieldsData(task *ImportTask, data *storage.InsertData) error { } idRange.Begin += int64(rowNum) if pkField.GetAutoID() { - data.Data[pkField.GetFieldID()] = &storage.Int64FieldData{Data: ids} + switch pkField.GetDataType() { + case schemapb.DataType_Int64: + data.Data[pkField.GetFieldID()] = &storage.Int64FieldData{Data: ids} + case schemapb.DataType_VarChar: + strIDs := lo.Map(ids, func(id int64, _ int) string { + return strconv.FormatInt(id, 10) + }) + data.Data[pkField.GetFieldID()] = &storage.StringFieldData{Data: strIDs} + } } data.Data[common.RowIDField] = &storage.Int64FieldData{Data: ids} tss := make([]int64, rowNum) diff --git a/internal/datanode/importv2/util_test.go b/internal/datanode/importv2/util_test.go new file mode 100644 index 0000000000..395e303ca0 --- /dev/null +++ b/internal/datanode/importv2/util_test.go @@ -0,0 +1,91 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package importv2 + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + "github.com/milvus-io/milvus/internal/proto/datapb" + "github.com/milvus-io/milvus/pkg/common" +) + +func Test_AppendSystemFieldsData(t *testing.T) { + const count = 100 + + pkField := &schemapb.FieldSchema{ + FieldID: 100, + Name: "pk", + IsPrimaryKey: true, + AutoID: true, + } + vecField := &schemapb.FieldSchema{ + FieldID: 101, + Name: "vec", + DataType: schemapb.DataType_FloatVector, + TypeParams: []*commonpb.KeyValuePair{ + { + Key: common.DimKey, + Value: "4", + }, + }, + } + int64Field := &schemapb.FieldSchema{ + FieldID: 102, + Name: "int64", + DataType: schemapb.DataType_Int64, + } + + schema := &schemapb.CollectionSchema{} + task := &ImportTask{ + req: &datapb.ImportRequest{ + Ts: 1000, + AutoIDRange: &datapb.AutoIDRange{ + Begin: 0, + End: count, + }, + }, + schema: schema, + } + + pkField.DataType = schemapb.DataType_Int64 + schema.Fields = []*schemapb.FieldSchema{pkField, vecField, int64Field} + insertData := createInsertData(t, schema, count) + assert.Equal(t, 0, insertData.Data[pkField.GetFieldID()].RowNum()) + assert.Nil(t, insertData.Data[common.RowIDField]) + assert.Nil(t, insertData.Data[common.TimeStampField]) + err := AppendSystemFieldsData(task, insertData) + assert.NoError(t, err) + assert.Equal(t, count, insertData.Data[pkField.GetFieldID()].RowNum()) + assert.Equal(t, count, insertData.Data[common.RowIDField].RowNum()) + assert.Equal(t, count, insertData.Data[common.TimeStampField].RowNum()) + + pkField.DataType = schemapb.DataType_VarChar + schema.Fields = []*schemapb.FieldSchema{pkField, vecField, int64Field} + insertData = createInsertData(t, schema, count) + assert.Equal(t, 0, insertData.Data[pkField.GetFieldID()].RowNum()) + assert.Nil(t, insertData.Data[common.RowIDField]) + assert.Nil(t, insertData.Data[common.TimeStampField]) + err = AppendSystemFieldsData(task, insertData) + assert.NoError(t, err) + assert.Equal(t, count, insertData.Data[pkField.GetFieldID()].RowNum()) + assert.Equal(t, count, insertData.Data[common.RowIDField].RowNum()) + assert.Equal(t, count, insertData.Data[common.TimeStampField].RowNum()) +}