Support Array datatype

Signed-off-by: wayblink <[email protected]>
zilliztech · Nov 1, 2023 · d70e3ed · d70e3ed
1 parent 8f3af76
commit d70e3ed
Show file tree

Hide file tree

Showing 6 changed files with 266 additions and 169 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
+
+	"github.com/zilliztech/milvus-backup/internal/log"
 )
 
 var (
@@ -28,4 +30,6 @@ func Execute() {
 
 func SetVersionInfo(version, commit, date string) {
 	rootCmd.Version = fmt.Sprintf("%s (Built on %s from Git SHA %s)", version, date, commit)
+	println(rootCmd.Version)
+	log.Info(fmt.Sprintf("Milvus backup version: %s", rootCmd.Version))
 }
diff --git a/core/backup_impl_create_backup.go b/core/backup_impl_create_backup.go
@@ -242,6 +242,7 @@ func (b *BackupContext) backupCollection(ctx context.Context, backupInfo *backup
 			IndexParams:    utils.MapToKVPair(field.IndexParams),
 			IsDynamic:      field.IsDynamic,
 			IsPartitionKey: field.IsPartitionKey,
+			ElementType:    backuppb.DataType(field.ElementType),
 		})
 	}
 	schema := &backuppb.CollectionSchema{

diff --git a/core/backup_impl_restore_backup.go b/core/backup_impl_restore_backup.go
@@ -391,6 +391,7 @@ func (b *BackupContext) executeRestoreCollectionTask(ctx context.Context, backup
 			IndexParams:    utils.KvPairsMap(field.GetIndexParams()),
 			IsDynamic:      field.GetIsDynamic(),
 			IsPartitionKey: field.GetIsPartitionKey(),
+			ElementType:    entity.FieldType(field.GetElementType()),
 		})
 		if field.GetIsPartitionKey() {
 			hasPartitionKey = true

diff --git a/core/milvus_sdk_test.go b/core/milvus_sdk_test.go
@@ -42,6 +42,87 @@ func TestProxyClient(t *testing.T) {
 	//assert.Empty(t, collections)
 }
 
+func TestCreateCollection(t *testing.T) {
+	ctx := context.Background()
+	milvusAddr := "local:19530"
+	//c, err := proxy.NewClient(context, milvusAddr)
+	//assert.NoError(t, err)
+
+	client, err := gomilvus.NewGrpcClient(ctx, milvusAddr)
+	assert.NoError(t, err)
+
+	//MILVUS_DATA_PATH := "/tmp/milvus/data/"
+
+	_COLLECTION_NAME := "create_demo"
+	_ID_FIELD_NAME := "id_field"
+	_VECTOR_FIELD_NAME := "float_vector_field"
+	_STR_FIELD_NAME := "str_field"
+	_ARRAY_FIELD_NAME := "arr_field"
+
+	// String field parameter
+	_MAX_LENGTH := "65535"
+	_ARRAY_MAX_CAPACITY := "100"
+	// Vector field parameter
+	_DIM := "8"
+
+	//field1 = FieldSchema(name=_ID_FIELD_NAME, dtype=DataType.INT64, description="int64", is_primary=True, auto_id=True)
+	//field2 = FieldSchema(name=_VECTOR_FIELD_NAME, dtype=DataType.FLOAT_VECTOR, description="float vector", dim=_DIM,
+	//	is_primary=False)
+	//field3 = FieldSchema(name=_STR_FIELD_NAME, dtype=DataType.VARCHAR, description="string",
+	//	max_length=_MAX_LENGTH, is_primary=False)
+	//schema = CollectionSchema(fields=[field1, field2, field3], description="collection description")
+	//collection = Collection(name=_COLLECTION_NAME, data=None, schema=schema)
+
+	field1 := &entity.Field{
+		Name:        _ID_FIELD_NAME,
+		DataType:    entity.FieldTypeInt64,
+		Description: "int64",
+		PrimaryKey:  true,
+		AutoID:      true,
+	}
+	field2 := &entity.Field{
+		Name:        _VECTOR_FIELD_NAME,
+		DataType:    entity.FieldTypeFloatVector,
+		Description: "float vector",
+		TypeParams: map[string]string{
+			entity.TypeParamDim: _DIM,
+		},
+		PrimaryKey: false,
+	}
+	field3 := &entity.Field{
+		Name:        _STR_FIELD_NAME,
+		DataType:    entity.FieldTypeVarChar,
+		Description: "string",
+		PrimaryKey:  false,
+		TypeParams: map[string]string{
+			entity.TypeParamMaxLength: _MAX_LENGTH,
+		},
+	}
+	field4 := &entity.Field{
+		Name:        _ARRAY_FIELD_NAME,
+		DataType:    entity.FieldTypeArray,
+		Description: "arr",
+		ElementType: entity.FieldTypeInt64,
+		TypeParams: map[string]string{
+			entity.TypeParamMaxCapacity: _ARRAY_MAX_CAPACITY,
+		},
+	}
+	schema := &entity.Schema{
+		CollectionName: _COLLECTION_NAME,
+		Description:    "demo bulkinsert",
+		AutoID:         true,
+		Fields:         []*entity.Field{field1, field2, field3, field4},
+	}
+	//client.DropCollection(ctx, _COLLECTION_NAME)
+
+	err = client.CreateCollection(ctx, schema, 2)
+	println(err)
+	_PART_1 := "part_1"
+	client.CreatePartition(ctx, _COLLECTION_NAME, _PART_1)
+
+	//client.DropCollection(ctx, _COLLECTION_NAME)
+}
+
 func TestBulkInsert(t *testing.T) {
 	ctx := context.Background()
 	milvusAddr := "localhost:19530"

diff --git a/core/proto/backup.proto b/core/proto/backup.proto
@@ -364,6 +364,9 @@ enum DataType {
   String = 20;
   VarChar = 21; // variable-length strings with a specified maximum length
 
+  Array = 22;
+  Json = 23;
+
   BinaryVector = 100;
   FloatVector = 101;
 }