mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-11-29 18:38:44 +08:00
11f4fe0177
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>
313 lines
8.9 KiB
Go
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)}
|
|
}
|