github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/test/common/utils_client.go (about)

     1  package common
     2  
     3  import (
     4  	"github.com/milvus-io/milvus-sdk-go/v2/entity"
     5  )
     6  
     7  // -- create field --
     8  
     9  // CreateFieldOption is an option that is used to modify entity.Schema
    10  type CreateFieldOption func(field *entity.Field)
    11  
    12  func WithIsPrimaryKey(isPrimaryKey bool) CreateFieldOption {
    13  	return func(field *entity.Field) {
    14  		field.WithIsPrimaryKey(isPrimaryKey)
    15  	}
    16  }
    17  
    18  func WithAutoID(autoID bool) CreateFieldOption {
    19  	return func(field *entity.Field) {
    20  		field.WithIsAutoID(autoID)
    21  	}
    22  }
    23  
    24  func WithFieldDescription(desc string) CreateFieldOption {
    25  	return func(field *entity.Field) {
    26  		field.WithDescription(desc)
    27  	}
    28  }
    29  
    30  func WithIsPartitionKey(isPartitionKey bool) CreateFieldOption {
    31  	return func(field *entity.Field) {
    32  		field.WithIsPartitionKey(isPartitionKey)
    33  	}
    34  }
    35  
    36  func WithDim(dim int64) CreateFieldOption {
    37  	return func(field *entity.Field) {
    38  		field.WithDim(dim)
    39  	}
    40  }
    41  
    42  func WithMaxLength(maxLen int64) CreateFieldOption {
    43  	return func(field *entity.Field) {
    44  		field.WithMaxLength(maxLen)
    45  	}
    46  }
    47  
    48  func WithElementType(eleType entity.FieldType) CreateFieldOption {
    49  	return func(field *entity.Field) {
    50  		field.WithElementType(eleType)
    51  	}
    52  }
    53  
    54  func WithMaxCapacity(maxCap int64) CreateFieldOption {
    55  	return func(field *entity.Field) {
    56  		field.WithMaxCapacity(maxCap)
    57  	}
    58  }
    59  
    60  func WithTypeParams(key string, value string) CreateFieldOption {
    61  	return func(field *entity.Field) {
    62  		field.WithTypeParams(key, value)
    63  	}
    64  }
    65  
    66  func GenField(name string, fieldType entity.FieldType, opts ...CreateFieldOption) *entity.Field {
    67  	if name == "" {
    68  		name = fieldType.Name() + GenRandomString(2)
    69  	}
    70  	field := entity.NewField().
    71  		WithName(name).
    72  		WithDataType(fieldType)
    73  
    74  	for _, opt := range opts {
    75  		opt(field)
    76  	}
    77  	return field
    78  }
    79  
    80  // -- create field --
    81  
    82  // -- create schema --
    83  
    84  // CreateSchemaOption is an option that is used to modify entity.Schema
    85  type CreateSchemaOption func(schema *entity.Schema)
    86  
    87  func WithDescription(desc string) CreateSchemaOption {
    88  	return func(schema *entity.Schema) {
    89  		schema.Description = desc
    90  	}
    91  }
    92  
    93  func WithEnableDynamicField(enableDF bool) CreateSchemaOption {
    94  	return func(schema *entity.Schema) {
    95  		schema.EnableDynamicField = enableDF
    96  	}
    97  }
    98  
    99  // GenSchema gen schema
   100  func GenSchema(name string, autoID bool, fields []*entity.Field, opts ...CreateSchemaOption) *entity.Schema {
   101  	schema := &entity.Schema{
   102  		CollectionName: name,
   103  		AutoID:         autoID,
   104  		Fields:         fields,
   105  	}
   106  	for _, opt := range opts {
   107  		opt(schema)
   108  	}
   109  	return schema
   110  }
   111  
   112  // -- create schema --
   113  
   114  // GenColumnDataOption -- create column data --
   115  type GenColumnDataOption func(opt *genDataOpt)
   116  type genDataOpt struct {
   117  	dim          int64
   118  	ElementType  entity.FieldType
   119  	capacity     int64
   120  	maxLenSparse int
   121  }
   122  
   123  func WithVectorDim(dim int64) GenColumnDataOption {
   124  	return func(opt *genDataOpt) {
   125  		opt.dim = dim
   126  	}
   127  }
   128  
   129  func WithArrayElementType(eleType entity.FieldType) GenColumnDataOption {
   130  	return func(opt *genDataOpt) {
   131  		opt.ElementType = eleType
   132  	}
   133  }
   134  
   135  func WithArrayCapacity(capacity int64) GenColumnDataOption {
   136  	return func(opt *genDataOpt) {
   137  		opt.capacity = capacity
   138  	}
   139  }
   140  
   141  func WithSparseVectorLen(length int) GenColumnDataOption {
   142  	return func(opt *genDataOpt) {
   143  		opt.maxLenSparse = length
   144  	}
   145  }
   146  
   147  // -- create column data --