mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
342 lines
9.1 KiB
Go
342 lines
9.1 KiB
Go
|
// 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 entity
|
||
|
|
||
|
import (
|
||
|
"strconv"
|
||
|
|
||
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
||
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||
|
)
|
||
|
|
||
|
const (
|
||
|
// TypeParamDim is the const for field type param dimension
|
||
|
TypeParamDim = "dim"
|
||
|
|
||
|
// TypeParamMaxLength is the const for varchar type maximal length
|
||
|
TypeParamMaxLength = "max_length"
|
||
|
|
||
|
// TypeParamMaxCapacity is the const for array type max capacity
|
||
|
TypeParamMaxCapacity = `max_capacity`
|
||
|
|
||
|
// ClStrong strong consistency level
|
||
|
ClStrong ConsistencyLevel = ConsistencyLevel(commonpb.ConsistencyLevel_Strong)
|
||
|
// ClBounded bounded consistency level with default tolerance of 5 seconds
|
||
|
ClBounded ConsistencyLevel = ConsistencyLevel(commonpb.ConsistencyLevel_Bounded)
|
||
|
// ClSession session consistency level
|
||
|
ClSession ConsistencyLevel = ConsistencyLevel(commonpb.ConsistencyLevel_Session)
|
||
|
// ClEvenually eventually consistency level
|
||
|
ClEventually ConsistencyLevel = ConsistencyLevel(commonpb.ConsistencyLevel_Eventually)
|
||
|
// ClCustomized customized consistency level and users pass their own `guarantee_timestamp`.
|
||
|
ClCustomized ConsistencyLevel = ConsistencyLevel(commonpb.ConsistencyLevel_Customized)
|
||
|
)
|
||
|
|
||
|
// ConsistencyLevel enum type for collection Consistency Level
|
||
|
type ConsistencyLevel commonpb.ConsistencyLevel
|
||
|
|
||
|
// CommonConsistencyLevel returns corresponding commonpb.ConsistencyLevel
|
||
|
func (cl ConsistencyLevel) CommonConsistencyLevel() commonpb.ConsistencyLevel {
|
||
|
return commonpb.ConsistencyLevel(cl)
|
||
|
}
|
||
|
|
||
|
// Schema represents schema info of collection in milvus
|
||
|
type Schema struct {
|
||
|
CollectionName string
|
||
|
Description string
|
||
|
AutoID bool
|
||
|
Fields []*Field
|
||
|
EnableDynamicField bool
|
||
|
}
|
||
|
|
||
|
// NewSchema creates an empty schema object.
|
||
|
func NewSchema() *Schema {
|
||
|
return &Schema{}
|
||
|
}
|
||
|
|
||
|
// WithName sets the name value of schema, returns schema itself.
|
||
|
func (s *Schema) WithName(name string) *Schema {
|
||
|
s.CollectionName = name
|
||
|
return s
|
||
|
}
|
||
|
|
||
|
// WithDescription sets the description value of schema, returns schema itself.
|
||
|
func (s *Schema) WithDescription(desc string) *Schema {
|
||
|
s.Description = desc
|
||
|
return s
|
||
|
}
|
||
|
|
||
|
func (s *Schema) WithAutoID(autoID bool) *Schema {
|
||
|
s.AutoID = autoID
|
||
|
return s
|
||
|
}
|
||
|
|
||
|
func (s *Schema) WithDynamicFieldEnabled(dynamicEnabled bool) *Schema {
|
||
|
s.EnableDynamicField = dynamicEnabled
|
||
|
return s
|
||
|
}
|
||
|
|
||
|
// WithField adds a field into schema and returns schema itself.
|
||
|
func (s *Schema) WithField(f *Field) *Schema {
|
||
|
s.Fields = append(s.Fields, f)
|
||
|
return s
|
||
|
}
|
||
|
|
||
|
// ProtoMessage returns corresponding server.CollectionSchema
|
||
|
func (s *Schema) ProtoMessage() *schemapb.CollectionSchema {
|
||
|
r := &schemapb.CollectionSchema{
|
||
|
Name: s.CollectionName,
|
||
|
Description: s.Description,
|
||
|
AutoID: s.AutoID,
|
||
|
EnableDynamicField: s.EnableDynamicField,
|
||
|
}
|
||
|
r.Fields = make([]*schemapb.FieldSchema, 0, len(s.Fields))
|
||
|
for _, field := range s.Fields {
|
||
|
r.Fields = append(r.Fields, field.ProtoMessage())
|
||
|
}
|
||
|
return r
|
||
|
}
|
||
|
|
||
|
// ReadProto parses proto Collection Schema
|
||
|
func (s *Schema) ReadProto(p *schemapb.CollectionSchema) *Schema {
|
||
|
s.Description = p.GetDescription()
|
||
|
s.CollectionName = p.GetName()
|
||
|
s.Fields = make([]*Field, 0, len(p.GetFields()))
|
||
|
for _, fp := range p.GetFields() {
|
||
|
if fp.GetAutoID() {
|
||
|
s.AutoID = true
|
||
|
}
|
||
|
s.Fields = append(s.Fields, NewField().ReadProto(fp))
|
||
|
}
|
||
|
s.EnableDynamicField = p.GetEnableDynamicField()
|
||
|
return s
|
||
|
}
|
||
|
|
||
|
// PKFieldName returns pk field name for this schemapb.
|
||
|
func (s *Schema) PKFieldName() string {
|
||
|
for _, field := range s.Fields {
|
||
|
if field.PrimaryKey {
|
||
|
return field.Name
|
||
|
}
|
||
|
}
|
||
|
return ""
|
||
|
}
|
||
|
|
||
|
// 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
|
||
|
}
|
||
|
|
||
|
// 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),
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// NewField creates a new Field with map initialized.
|
||
|
func NewField() *Field {
|
||
|
return &Field{
|
||
|
TypeParams: make(map[string]string),
|
||
|
IndexParams: make(map[string]string),
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithName(name string) *Field {
|
||
|
f.Name = name
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDescription(desc string) *Field {
|
||
|
f.Description = desc
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDataType(dataType FieldType) *Field {
|
||
|
f.DataType = dataType
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithIsPrimaryKey(isPrimaryKey bool) *Field {
|
||
|
f.PrimaryKey = isPrimaryKey
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithIsAutoID(isAutoID bool) *Field {
|
||
|
f.AutoID = isAutoID
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithIsDynamic(isDynamic bool) *Field {
|
||
|
f.IsDynamic = isDynamic
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithIsPartitionKey(isPartitionKey bool) *Field {
|
||
|
f.IsPartitionKey = isPartitionKey
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
/*
|
||
|
func (f *Field) WithDefaultValueBool(defaultValue bool) *Field {
|
||
|
f.DefaultValue = &schemapb.ValueField{
|
||
|
Data: &schemapb.ValueField_BoolData{
|
||
|
BoolData: defaultValue,
|
||
|
},
|
||
|
}
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDefaultValueInt(defaultValue int32) *Field {
|
||
|
f.DefaultValue = &schemapb.ValueField{
|
||
|
Data: &schemapb.ValueField_IntData{
|
||
|
IntData: defaultValue,
|
||
|
},
|
||
|
}
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDefaultValueLong(defaultValue int64) *Field {
|
||
|
f.DefaultValue = &schemapb.ValueField{
|
||
|
Data: &schemapb.ValueField_LongData{
|
||
|
LongData: defaultValue,
|
||
|
},
|
||
|
}
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDefaultValueFloat(defaultValue float32) *Field {
|
||
|
f.DefaultValue = &schemapb.ValueField{
|
||
|
Data: &schemapb.ValueField_FloatData{
|
||
|
FloatData: defaultValue,
|
||
|
},
|
||
|
}
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDefaultValueDouble(defaultValue float64) *Field {
|
||
|
f.DefaultValue = &schemapb.ValueField{
|
||
|
Data: &schemapb.ValueField_DoubleData{
|
||
|
DoubleData: defaultValue,
|
||
|
},
|
||
|
}
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDefaultValueString(defaultValue string) *Field {
|
||
|
f.DefaultValue = &schemapb.ValueField{
|
||
|
Data: &schemapb.ValueField_StringData{
|
||
|
StringData: defaultValue,
|
||
|
},
|
||
|
}
|
||
|
return f
|
||
|
}*/
|
||
|
|
||
|
func (f *Field) WithTypeParams(key string, value string) *Field {
|
||
|
if f.TypeParams == nil {
|
||
|
f.TypeParams = make(map[string]string)
|
||
|
}
|
||
|
f.TypeParams[key] = value
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithDim(dim int64) *Field {
|
||
|
if f.TypeParams == nil {
|
||
|
f.TypeParams = make(map[string]string)
|
||
|
}
|
||
|
f.TypeParams[TypeParamDim] = strconv.FormatInt(dim, 10)
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithMaxLength(maxLen int64) *Field {
|
||
|
if f.TypeParams == nil {
|
||
|
f.TypeParams = make(map[string]string)
|
||
|
}
|
||
|
f.TypeParams[TypeParamMaxLength] = strconv.FormatInt(maxLen, 10)
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithElementType(eleType FieldType) *Field {
|
||
|
f.ElementType = eleType
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
func (f *Field) WithMaxCapacity(maxCap int64) *Field {
|
||
|
if f.TypeParams == nil {
|
||
|
f.TypeParams = make(map[string]string)
|
||
|
}
|
||
|
f.TypeParams[TypeParamMaxCapacity] = strconv.FormatInt(maxCap, 10)
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
// ReadProto parses FieldSchema
|
||
|
func (f *Field) ReadProto(p *schemapb.FieldSchema) *Field {
|
||
|
f.ID = p.GetFieldID()
|
||
|
f.Name = p.GetName()
|
||
|
f.PrimaryKey = p.GetIsPrimaryKey()
|
||
|
f.AutoID = p.GetAutoID()
|
||
|
f.Description = p.GetDescription()
|
||
|
f.DataType = FieldType(p.GetDataType())
|
||
|
f.TypeParams = KvPairsMap(p.GetTypeParams())
|
||
|
f.IndexParams = KvPairsMap(p.GetIndexParams())
|
||
|
f.IsDynamic = p.GetIsDynamic()
|
||
|
f.IsPartitionKey = p.GetIsPartitionKey()
|
||
|
f.ElementType = FieldType(p.GetElementType())
|
||
|
|
||
|
return f
|
||
|
}
|
||
|
|
||
|
// MapKvPairs converts map into commonpb.KeyValuePair slice
|
||
|
func MapKvPairs(m map[string]string) []*commonpb.KeyValuePair {
|
||
|
pairs := make([]*commonpb.KeyValuePair, 0, len(m))
|
||
|
for k, v := range m {
|
||
|
pairs = append(pairs, &commonpb.KeyValuePair{
|
||
|
Key: k,
|
||
|
Value: v,
|
||
|
})
|
||
|
}
|
||
|
return pairs
|
||
|
}
|
||
|
|
||
|
// KvPairsMap converts commonpb.KeyValuePair slices into map
|
||
|
func KvPairsMap(kvps []*commonpb.KeyValuePair) map[string]string {
|
||
|
m := make(map[string]string)
|
||
|
for _, kvp := range kvps {
|
||
|
m[kvp.Key] = kvp.Value
|
||
|
}
|
||
|
return m
|
||
|
}
|