Skip to content

Commit

Permalink
enhance: Support specifying field schema to be cluster key (#794)
Browse files Browse the repository at this point in the history
See also milvus-io/milvus#32939

Signed-off-by: Congqi Xia <[email protected]>
  • Loading branch information
congqixia authored Jul 25, 2024
1 parent d45cab5 commit d26d3d7
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 22 deletions.
52 changes: 30 additions & 22 deletions entity/schema.go
Original file line number Diff line number Diff line change
Expand Up @@ -141,33 +141,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() *schema.FieldSchema {
return &schema.FieldSchema{
FieldID: f.ID,
Name: f.Name,
Description: f.Description,
IsPrimaryKey: f.PrimaryKey,
AutoID: f.AutoID,
DataType: schema.DataType(f.DataType),
TypeParams: MapKvPairs(f.TypeParams),
IndexParams: MapKvPairs(f.IndexParams),
IsDynamic: f.IsDynamic,
IsPartitionKey: f.IsPartitionKey,
ElementType: schema.DataType(f.ElementType),
FieldID: f.ID,
Name: f.Name,
Description: f.Description,
IsPrimaryKey: f.PrimaryKey,
AutoID: f.AutoID,
DataType: schema.DataType(f.DataType),
TypeParams: MapKvPairs(f.TypeParams),
IndexParams: MapKvPairs(f.IndexParams),
IsDynamic: f.IsDynamic,
IsPartitionKey: f.IsPartitionKey,
IsClusteringKey: f.IsClusteringKey,
ElementType: schema.DataType(f.ElementType),
}
}

Expand Down Expand Up @@ -214,6 +216,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 = &schema.ValueField{
Expand Down Expand Up @@ -318,6 +325,7 @@ func (f *Field) ReadProto(p *schema.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
Expand Down
3 changes: 3 additions & 0 deletions entity/schema_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ func TestFieldSchema(t *testing.T) {
NewField().WithName("array_field").WithDataType(FieldTypeArray).WithElementType(FieldTypeBool).WithMaxCapacity(128),
NewField().WithName("fp16_field").WithDataType(FieldTypeFloat16Vector).WithDim(128),
NewField().WithName("bf16_field").WithDataType(FieldTypeBFloat16Vector).WithDim(128),
NewField().WithName("cluster_key").WithDataType(FieldTypeInt32).WithIsClusteringKey(true),
/*
NewField().WithName("default_value_bool").WithDataType(FieldTypeBool).WithDefaultValueBool(true),
NewField().WithName("default_value_int").WithDataType(FieldTypeInt32).WithDefaultValueInt(1),
Expand All @@ -45,6 +46,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()))
Expand All @@ -60,6 +62,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)
}
Expand Down

0 comments on commit d26d3d7

Please sign in to comment.