milvus/client/milvusclient/collection_options.go
congqixia 11f4fe0177
enhance: [GoSDK] move client pkg go files to sub one (#37492)
Related to #31293

Client source files under client pkg cannot be evaluate correctly by
codecov. This PR moves them to `milvusclient` sub-package to fix this
issue and follow go major version best practice.

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2024-11-08 07:32:26 +08:00

313 lines
8.9 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 milvusclient
import (
"fmt"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/client/v2/index"
)
// CreateCollectionOption is the interface builds CreateCollectionRequest.
type CreateCollectionOption interface {
// Request is the method returns the composed request.
Request() *milvuspb.CreateCollectionRequest
// Indexes is the method returns IndexOption to create
Indexes() []CreateIndexOption
IsFast() bool
}
// createCollectionOption contains all the parameters to create collection.
type createCollectionOption struct {
name string
shardNum int32
// fast create collection params
varcharPK bool
varcharPKMaxLength int
pkFieldName string
vectorFieldName string
dim int64
autoID bool
enabledDynamicSchema bool
// advanced create collection params
schema *entity.Schema
consistencyLevel entity.ConsistencyLevel
properties map[string]string
// partition key
numPartitions int64
indexOptions []CreateIndexOption
// is fast create collection
isFast bool
// fast creation with index
metricType entity.MetricType
}
func (opt *createCollectionOption) WithAutoID(autoID bool) *createCollectionOption {
opt.autoID = autoID
return opt
}
func (opt *createCollectionOption) WithShardNum(shardNum int32) *createCollectionOption {
opt.shardNum = shardNum
return opt
}
func (opt *createCollectionOption) WithDynamicSchema(dynamicSchema bool) *createCollectionOption {
opt.enabledDynamicSchema = dynamicSchema
return opt
}
func (opt *createCollectionOption) WithVarcharPK(varcharPK bool, maxLen int) *createCollectionOption {
opt.varcharPK = varcharPK
opt.varcharPKMaxLength = maxLen
return opt
}
func (opt *createCollectionOption) WithIndexOptions(indexOpts ...CreateIndexOption) *createCollectionOption {
opt.indexOptions = append(opt.indexOptions, indexOpts...)
return opt
}
func (opt *createCollectionOption) WithProperty(key string, value any) *createCollectionOption {
opt.properties[key] = fmt.Sprintf("%v", value)
return opt
}
func (opt *createCollectionOption) WithConsistencyLevel(cl entity.ConsistencyLevel) *createCollectionOption {
opt.consistencyLevel = cl
return opt
}
func (opt *createCollectionOption) Request() *milvuspb.CreateCollectionRequest {
// fast create collection
if opt.isFast {
var pkField *entity.Field
if opt.varcharPK {
pkField = entity.NewField().WithDataType(entity.FieldTypeVarChar).WithMaxLength(int64(opt.varcharPKMaxLength))
} else {
pkField = entity.NewField().WithDataType(entity.FieldTypeInt64)
}
pkField = pkField.WithName(opt.pkFieldName).WithIsPrimaryKey(true).WithIsAutoID(opt.autoID)
opt.schema = entity.NewSchema().
WithName(opt.name).
WithAutoID(opt.autoID).
WithDynamicFieldEnabled(opt.enabledDynamicSchema).
WithField(pkField).
WithField(entity.NewField().WithName(opt.vectorFieldName).WithDataType(entity.FieldTypeFloatVector).WithDim(opt.dim))
}
var schemaBytes []byte
if opt.schema != nil {
opt.schema.WithName(opt.name)
schemaProto := opt.schema.ProtoMessage()
schemaBytes, _ = proto.Marshal(schemaProto)
}
return &milvuspb.CreateCollectionRequest{
DbName: "", // reserved fields, not used for now
CollectionName: opt.name,
Schema: schemaBytes,
ShardsNum: opt.shardNum,
ConsistencyLevel: commonpb.ConsistencyLevel(opt.consistencyLevel),
NumPartitions: opt.numPartitions,
Properties: entity.MapKvPairs(opt.properties),
}
}
func (opt *createCollectionOption) Indexes() []CreateIndexOption {
// fast create
if opt.isFast {
return []CreateIndexOption{
NewCreateIndexOption(opt.name, opt.vectorFieldName, index.NewGenericIndex("", map[string]string{})),
}
}
return nil
}
func (opt *createCollectionOption) IsFast() bool {
return opt.isFast
}
// SimpleCreateCollectionOptions returns a CreateCollectionOption with default fast collection options.
func SimpleCreateCollectionOptions(name string, dim int64) *createCollectionOption {
return &createCollectionOption{
name: name,
shardNum: 1,
pkFieldName: "id",
vectorFieldName: "vector",
autoID: true,
dim: dim,
enabledDynamicSchema: true,
consistencyLevel: entity.DefaultConsistencyLevel,
properties: make(map[string]string),
isFast: true,
metricType: entity.COSINE,
}
}
// NewCreateCollectionOption returns a CreateCollectionOption with customized collection schema
func NewCreateCollectionOption(name string, collectionSchema *entity.Schema) *createCollectionOption {
return &createCollectionOption{
name: name,
shardNum: 1,
schema: collectionSchema,
consistencyLevel: entity.DefaultConsistencyLevel,
properties: make(map[string]string),
metricType: entity.COSINE,
}
}
type ListCollectionOption interface {
Request() *milvuspb.ShowCollectionsRequest
}
type listCollectionOption struct{}
func (opt *listCollectionOption) Request() *milvuspb.ShowCollectionsRequest {
return &milvuspb.ShowCollectionsRequest{}
}
func NewListCollectionOption() *listCollectionOption {
return &listCollectionOption{}
}
// DescribeCollectionOption is the interface builds DescribeCollection request.
type DescribeCollectionOption interface {
// Request is the method returns the composed request.
Request() *milvuspb.DescribeCollectionRequest
}
type describeCollectionOption struct {
name string
}
func (opt *describeCollectionOption) Request() *milvuspb.DescribeCollectionRequest {
return &milvuspb.DescribeCollectionRequest{
CollectionName: opt.name,
}
}
// NewDescribeCollectionOption composes a describeCollectionOption with provided collection name.
func NewDescribeCollectionOption(name string) *describeCollectionOption {
return &describeCollectionOption{
name: name,
}
}
// HasCollectionOption is the interface to build DescribeCollectionRequest.
type HasCollectionOption interface {
Request() *milvuspb.DescribeCollectionRequest
}
type hasCollectionOpt struct {
name string
}
func (opt *hasCollectionOpt) Request() *milvuspb.DescribeCollectionRequest {
return &milvuspb.DescribeCollectionRequest{
CollectionName: opt.name,
}
}
func NewHasCollectionOption(name string) HasCollectionOption {
return &hasCollectionOpt{
name: name,
}
}
// The DropCollectionOption interface builds DropCollectionRequest.
type DropCollectionOption interface {
Request() *milvuspb.DropCollectionRequest
}
type dropCollectionOption struct {
name string
}
func (opt *dropCollectionOption) Request() *milvuspb.DropCollectionRequest {
return &milvuspb.DropCollectionRequest{
CollectionName: opt.name,
}
}
func NewDropCollectionOption(name string) *dropCollectionOption {
return &dropCollectionOption{
name: name,
}
}
type RenameCollectionOption interface {
Request() *milvuspb.RenameCollectionRequest
}
type renameCollectionOption struct {
oldCollectionName string
newCollectionName string
}
func (opt *renameCollectionOption) Request() *milvuspb.RenameCollectionRequest {
return &milvuspb.RenameCollectionRequest{
OldName: opt.oldCollectionName,
NewName: opt.newCollectionName,
}
}
func NewRenameCollectionOption(oldName, newName string) *renameCollectionOption {
return &renameCollectionOption{
oldCollectionName: oldName,
newCollectionName: newName,
}
}
type AlterCollectionOption interface {
Request() *milvuspb.AlterCollectionRequest
}
type alterCollectionOption struct {
collectionName string
properties map[string]string
}
func (opt *alterCollectionOption) WithProperty(key string, value any) *alterCollectionOption {
opt.properties[key] = fmt.Sprintf("%v", value)
return opt
}
func (opt *alterCollectionOption) Request() *milvuspb.AlterCollectionRequest {
return &milvuspb.AlterCollectionRequest{
CollectionName: opt.collectionName,
Properties: entity.MapKvPairs(opt.properties),
}
}
func NewAlterCollectionOption(collection string) *alterCollectionOption {
return &alterCollectionOption{collectionName: collection, properties: make(map[string]string)}
}