enhance: [GoSDK]Support specifying field schema to be cluster key

See also milvus-io#32939 milvus-io/milvus-sdk-go#794 Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
congqixia · Jul 26, 2024 · 8878370 · 8878370
1 parent 4ee6c69
commit 8878370
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 22 deletions.
diff --git a/client/entity/schema.go b/client/entity/schema.go
@@ -151,33 +151,35 @@ func (s *Schema) PKField() *Field {
 
 // Field represent field schema in milvus
 type Field struct {
-	ID             int64  // field id, generated when collection is created, input value is ignored
-	Name           string // field name
-	PrimaryKey     bool   // is primary key
-	AutoID         bool   // is auto id
-	Description    string
-	DataType       FieldType
-	TypeParams     map[string]string
-	IndexParams    map[string]string
-	IsDynamic      bool
-	IsPartitionKey bool
-	ElementType    FieldType
+	ID              int64  // field id, generated when collection is created, input value is ignored
+	Name            string // field name
+	PrimaryKey      bool   // is primary key
+	AutoID          bool   // is auto id
+	Description     string
+	DataType        FieldType
+	TypeParams      map[string]string
+	IndexParams     map[string]string
+	IsDynamic       bool
+	IsPartitionKey  bool
+	IsClusteringKey bool
+	ElementType     FieldType
 }
 
 // ProtoMessage generates corresponding FieldSchema
 func (f *Field) ProtoMessage() *schemapb.FieldSchema {
 	return &schemapb.FieldSchema{
-		FieldID:        f.ID,
-		Name:           f.Name,
-		Description:    f.Description,
-		IsPrimaryKey:   f.PrimaryKey,
-		AutoID:         f.AutoID,
-		DataType:       schemapb.DataType(f.DataType),
-		TypeParams:     MapKvPairs(f.TypeParams),
-		IndexParams:    MapKvPairs(f.IndexParams),
-		IsDynamic:      f.IsDynamic,
-		IsPartitionKey: f.IsPartitionKey,
-		ElementType:    schemapb.DataType(f.ElementType),
+		FieldID:         f.ID,
+		Name:            f.Name,
+		Description:     f.Description,
+		IsPrimaryKey:    f.PrimaryKey,
+		AutoID:          f.AutoID,
+		DataType:        schemapb.DataType(f.DataType),
+		TypeParams:      MapKvPairs(f.TypeParams),
+		IndexParams:     MapKvPairs(f.IndexParams),
+		IsDynamic:       f.IsDynamic,
+		IsPartitionKey:  f.IsPartitionKey,
+		IsClusteringKey: f.IsClusteringKey,
+		ElementType:     schemapb.DataType(f.ElementType),
 	}
 }
 
@@ -224,6 +226,11 @@ func (f *Field) WithIsPartitionKey(isPartitionKey bool) *Field {
 	return f
 }
 
+func (f *Field) WithIsClusteringKey(isClusteringKey bool) *Field {
+	f.IsClusteringKey = isClusteringKey
+	return f
+}
+
 /*
 func (f *Field) WithDefaultValueBool(defaultValue bool) *Field {
 	f.DefaultValue = &schemapb.ValueField{
@@ -340,6 +347,7 @@ func (f *Field) ReadProto(p *schemapb.FieldSchema) *Field {
 	f.IndexParams = KvPairsMap(p.GetIndexParams())
 	f.IsDynamic = p.GetIsDynamic()
 	f.IsPartitionKey = p.GetIsPartitionKey()
+	f.IsClusteringKey = p.GetIsClusteringKey()
 	f.ElementType = FieldType(p.GetElementType())
 
 	return f

diff --git a/client/entity/schema_test.go b/client/entity/schema_test.go
@@ -43,6 +43,7 @@ func TestFieldSchema(t *testing.T) {
 		NewField().WithName("string_field").WithDataType(FieldTypeString).WithIsAutoID(false).WithIsPrimaryKey(true).WithIsDynamic(false).WithTypeParams("max_len", "32").WithDescription("string_field desc"),
 		NewField().WithName("partition_key").WithDataType(FieldTypeInt32).WithIsPartitionKey(true),
 		NewField().WithName("array_field").WithDataType(FieldTypeArray).WithElementType(FieldTypeBool).WithMaxCapacity(128),
+		NewField().WithName("clustering_key").WithDataType(FieldTypeInt32).WithIsClusteringKey(true),
 		/*
 			NewField().WithName("default_value_bool").WithDataType(FieldTypeBool).WithDefaultValueBool(true),
 			NewField().WithName("default_value_int").WithDataType(FieldTypeInt32).WithDefaultValueInt(1),
@@ -60,6 +61,7 @@ func TestFieldSchema(t *testing.T) {
 		assert.Equal(t, field.AutoID, fieldSchema.GetAutoID())
 		assert.Equal(t, field.PrimaryKey, fieldSchema.GetIsPrimaryKey())
 		assert.Equal(t, field.IsPartitionKey, fieldSchema.GetIsPartitionKey())
+		assert.Equal(t, field.IsClusteringKey, fieldSchema.GetIsClusteringKey())
 		assert.Equal(t, field.IsDynamic, fieldSchema.GetIsDynamic())
 		assert.Equal(t, field.Description, fieldSchema.GetDescription())
 		assert.Equal(t, field.TypeParams, KvPairsMap(fieldSchema.GetTypeParams()))
@@ -75,6 +77,7 @@ func TestFieldSchema(t *testing.T) {
 		assert.Equal(t, field.Description, nf.Description)
 		assert.Equal(t, field.IsDynamic, nf.IsDynamic)
 		assert.Equal(t, field.IsPartitionKey, nf.IsPartitionKey)
+		assert.Equal(t, field.IsClusteringKey, nf.IsClusteringKey)
 		assert.EqualValues(t, field.TypeParams, nf.TypeParams)
 		assert.EqualValues(t, field.ElementType, nf.ElementType)
 	}