Skip to content

Commit

Permalink
feat: add clustering key in create/describe collection (#29506)
Browse files Browse the repository at this point in the history 
#28410
/kind feature

Signed-off-by: wayblink <[email protected]>
  • Loading branch information
@wayblink
wayblink authored Jan 7, 2024
1 parent 156a0dd commit 635a7f7
Show file tree
Hide file tree
Showing 7 changed files with 299 additions and 65 deletions.
2 changes: 1 addition & 1 deletion go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ require (
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/klauspost/compress v1.16.7
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20231228051838-b5442d755fa4
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20231229025438-39bce6abb18f
github.com/minio/minio-go/v7 v7.0.61
github.com/prometheus/client_golang v1.14.0
github.com/prometheus/client_model v0.3.0
Expand Down
4 changes: 2 additions & 2 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -583,8 +583,8 @@ github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/le
github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20231228051838-b5442d755fa4 h1:nxIohfJOCMbixFAC3q4Lclmv0xg/8q6D8T7D8l258To=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20231228051838-b5442d755fa4/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20231229025438-39bce6abb18f h1:8lNcRqhQgUROtmtiIEdpQHGW82KMI5oASVKxkaZ/tBg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20231229025438-39bce6abb18f/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
github.com/milvus-io/milvus-storage/go v0.0.0-20231109072809-1cd7b0866092 h1:UYJ7JB+QlMOoFHNdd8mUa3/lV63t9dnBX7ILXmEEWPY=
github.com/milvus-io/milvus-storage/go v0.0.0-20231109072809-1cd7b0866092/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho=
github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A=
Expand Down
105 changes: 55 additions & 50 deletions internal/metastore/model/field.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,19 +7,20 @@ import (
)

type Field struct {
FieldID int64
Name string
IsPrimaryKey bool
Description string
DataType schemapb.DataType
TypeParams []*commonpb.KeyValuePair
IndexParams []*commonpb.KeyValuePair
AutoID bool
State schemapb.FieldState
IsDynamic bool
IsPartitionKey bool // partition key mode, multi logic partitions share a physical partition
DefaultValue *schemapb.ValueField
ElementType schemapb.DataType
FieldID int64
Name string
IsPrimaryKey bool
Description string
DataType schemapb.DataType
TypeParams []*commonpb.KeyValuePair
IndexParams []*commonpb.KeyValuePair
AutoID bool
State schemapb.FieldState
IsDynamic bool
IsPartitionKey bool // partition key mode, multi logic partitions share a physical partition
IsClusteringKey bool
DefaultValue *schemapb.ValueField
ElementType schemapb.DataType
}

func (f *Field) Available() bool {
Expand All @@ -28,19 +29,20 @@ func (f *Field) Available() bool {

func (f *Field) Clone() *Field {
return &Field{
FieldID: f.FieldID,
Name: f.Name,
IsPrimaryKey: f.IsPrimaryKey,
Description: f.Description,
DataType: f.DataType,
TypeParams: common.CloneKeyValuePairs(f.TypeParams),
IndexParams: common.CloneKeyValuePairs(f.IndexParams),
AutoID: f.AutoID,
State: f.State,
IsDynamic: f.IsDynamic,
IsPartitionKey: f.IsPartitionKey,
DefaultValue: f.DefaultValue,
ElementType: f.ElementType,
FieldID: f.FieldID,
Name: f.Name,
IsPrimaryKey: f.IsPrimaryKey,
Description: f.Description,
DataType: f.DataType,
TypeParams: common.CloneKeyValuePairs(f.TypeParams),
IndexParams: common.CloneKeyValuePairs(f.IndexParams),
AutoID: f.AutoID,
State: f.State,
IsDynamic: f.IsDynamic,
IsPartitionKey: f.IsPartitionKey,
IsClusteringKey: f.IsClusteringKey,
DefaultValue: f.DefaultValue,
ElementType: f.ElementType,
}
}

Expand Down Expand Up @@ -68,6 +70,7 @@ func (f *Field) Equal(other Field) bool {
f.AutoID == other.AutoID &&
f.IsPartitionKey == other.IsPartitionKey &&
f.IsDynamic == other.IsDynamic &&
f.IsClusteringKey == other.IsClusteringKey &&
f.DefaultValue == other.DefaultValue &&
f.ElementType == other.ElementType
}
Expand All @@ -91,18 +94,19 @@ func MarshalFieldModel(field *Field) *schemapb.FieldSchema {
}

return &schemapb.FieldSchema{
FieldID: field.FieldID,
Name: field.Name,
IsPrimaryKey: field.IsPrimaryKey,
Description: field.Description,
DataType: field.DataType,
TypeParams: field.TypeParams,
IndexParams: field.IndexParams,
AutoID: field.AutoID,
IsDynamic: field.IsDynamic,
IsPartitionKey: field.IsPartitionKey,
DefaultValue: field.DefaultValue,
ElementType: field.ElementType,
FieldID: field.FieldID,
Name: field.Name,
IsPrimaryKey: field.IsPrimaryKey,
Description: field.Description,
DataType: field.DataType,
TypeParams: field.TypeParams,
IndexParams: field.IndexParams,
AutoID: field.AutoID,
IsDynamic: field.IsDynamic,
IsPartitionKey: field.IsPartitionKey,
IsClusteringKey: field.IsClusteringKey,
DefaultValue: field.DefaultValue,
ElementType: field.ElementType,
}
}

Expand All @@ -124,18 +128,19 @@ func UnmarshalFieldModel(fieldSchema *schemapb.FieldSchema) *Field {
}

return &Field{
FieldID: fieldSchema.FieldID,
Name: fieldSchema.Name,
IsPrimaryKey: fieldSchema.IsPrimaryKey,
Description: fieldSchema.Description,
DataType: fieldSchema.DataType,
TypeParams: fieldSchema.TypeParams,
IndexParams: fieldSchema.IndexParams,
AutoID: fieldSchema.AutoID,
IsDynamic: fieldSchema.IsDynamic,
IsPartitionKey: fieldSchema.IsPartitionKey,
DefaultValue: fieldSchema.DefaultValue,
ElementType: fieldSchema.ElementType,
FieldID: fieldSchema.FieldID,
Name: fieldSchema.Name,
IsPrimaryKey: fieldSchema.IsPrimaryKey,
Description: fieldSchema.Description,
DataType: fieldSchema.DataType,
TypeParams: fieldSchema.TypeParams,
IndexParams: fieldSchema.IndexParams,
AutoID: fieldSchema.AutoID,
IsDynamic: fieldSchema.IsDynamic,
IsPartitionKey: fieldSchema.IsPartitionKey,
IsClusteringKey: fieldSchema.IsClusteringKey,
DefaultValue: fieldSchema.DefaultValue,
ElementType: fieldSchema.ElementType,
}
}

Expand Down
60 changes: 48 additions & 12 deletions internal/proxy/task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,36 @@ func (t *createCollectionTask) validatePartitionKey() error {
return nil
}

func (t *createCollectionTask) validateClusteringKey() error {
idx := -1
for i, field := range t.schema.Fields {
if field.GetIsClusteringKey() {
if idx != -1 {
return merr.WrapErrCollectionIllegalSchema(t.CollectionName,
fmt.Sprintf("there are more than one clustering key, field name = %s, %s", t.schema.Fields[idx].Name, field.Name))
}

if field.GetIsPrimaryKey() {
return merr.WrapErrCollectionIllegalSchema(t.CollectionName,
fmt.Sprintf("the clustering key field must not be primary key field, field name = %s", field.Name))
}

if field.GetIsPartitionKey() {
return merr.WrapErrCollectionIllegalSchema(t.CollectionName,
fmt.Sprintf("the clustering key field must not be partition key field, field name = %s", field.Name))
}
idx = i
}
}

if idx != -1 {
log.Info("create collection with clustering key",
zap.String("collectionName", t.CollectionName),
zap.String("clusteringKeyField", t.schema.Fields[idx].Name))
}
return nil
}

func (t *createCollectionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_CreateCollection
t.Base.SourceID = paramtable.GetNodeID()
Expand Down Expand Up @@ -266,6 +296,11 @@ func (t *createCollectionTask) PreExecute(ctx context.Context) error {
return err
}

// validate clustering key
if err := t.validateClusteringKey(); err != nil {
return err
}

for _, field := range t.schema.Fields {
// validate field name
if err := validateFieldName(field.Name); err != nil {
Expand Down Expand Up @@ -572,18 +607,19 @@ func (t *describeCollectionTask) Execute(ctx context.Context) error {
}
if field.FieldID >= common.StartOfUserFieldID {
t.result.Schema.Fields = append(t.result.Schema.Fields, &schemapb.FieldSchema{
FieldID: field.FieldID,
Name: field.Name,
IsPrimaryKey: field.IsPrimaryKey,
AutoID: field.AutoID,
Description: field.Description,
DataType: field.DataType,
TypeParams: field.TypeParams,
IndexParams: field.IndexParams,
IsDynamic: field.IsDynamic,
IsPartitionKey: field.IsPartitionKey,
DefaultValue: field.DefaultValue,
ElementType: field.ElementType,
FieldID: field.FieldID,
Name: field.Name,
IsPrimaryKey: field.IsPrimaryKey,
AutoID: field.AutoID,
Description: field.Description,
DataType: field.DataType,
TypeParams: field.TypeParams,
IndexParams: field.IndexParams,
IsDynamic: field.IsDynamic,
IsPartitionKey: field.IsPartitionKey,
IsClusteringKey: field.IsClusteringKey,
DefaultValue: field.DefaultValue,
ElementType: field.ElementType,
})
}
}
Expand Down
Loading

0 comments on commit 635a7f7

Please sign in to comment.