Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history 
Signed-off-by: ThreadDao <yufen.zong@zilliz.com>
  • Loading branch information
@ThreadDao
ThreadDao committed Nov 6, 2023
1 parent fcb8d4e commit 19f97bd
Show file tree
Hide file tree
Showing 3 changed files with 184 additions and 30 deletions.
81 changes: 51 additions & 30 deletions test/testcases/main_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package testcases
import (
"context"
"flag"
"fmt"
"log"
"math/rand"
"os"
Expand Down Expand Up @@ -242,8 +243,6 @@ func createVarcharCollectionWithDataIndex(ctx context.Context, t *testing.T, mc
return collName, ids
}

type CollectionFieldsType string

const (
Int64FloatVec CollectionFieldsType = "PkInt64FloatVec" // int64 + float + floatVec
Int64BinaryVec CollectionFieldsType = "Int64BinaryVec" // int64 + float + binaryVec
Expand All @@ -253,16 +252,6 @@ const (
CustomerFields CollectionFieldsType = "CustomerFields" // customer fields
)

type CollectionParams struct {
CollectionFieldsType CollectionFieldsType // collection fields type
AutoID bool // autoId
EnableDynamicField bool // enable dynamic field
ShardsNum int32
Fields []*entity.Field
Dim int64
MaxLength int64
}

func createCollection(ctx context.Context, t *testing.T, mc *base.MilvusClient, cp CollectionParams, opts ...client.CreateCollectionOption) string {
collName := common.GenRandomString(4)
var fields []*entity.Field
Expand Down Expand Up @@ -300,18 +289,7 @@ func createCollection(ctx context.Context, t *testing.T, mc *base.MilvusClient,
return collName
}

type DataParams struct {
CollectionName string // insert data into which collection
PartitionName string
CollectionFieldsType CollectionFieldsType // collection fields type
start int // start
nb int // insert how many data
dim int64
EnableDynamicField bool // whether insert dynamic field data
WithRows bool
Data []entity.Column
Rows []interface{}
}


func insertData(ctx context.Context, t *testing.T, mc *base.MilvusClient, dp DataParams) (entity.Column, error) {
// todo autoid
Expand Down Expand Up @@ -446,12 +424,6 @@ func createCollectionAllFields(ctx context.Context, t *testing.T, mc *base.Milvu
return collName, ids
}

type HelpPartitionColumns struct {
PartitionName string
IdsColumn entity.Column
VectorColumn entity.Column
}

func createInsertTwoPartitions(ctx context.Context, t *testing.T, mc *base.MilvusClient, collName string, nb int) (partitionName string, defaultPartition HelpPartitionColumns, newPartition HelpPartitionColumns) {
// create new partition
partitionName = "new"
Expand Down Expand Up @@ -486,6 +458,55 @@ func createInsertTwoPartitions(ctx context.Context, t *testing.T, mc *base.Milvu
return partitionName, defaultPartition, newPartition
}

func prepareCollection(ctx context.Context, t *testing.T, mc *base.MilvusClient, collParam CollectionParams, opts ...PrepareCollectionOption) {
// default insert nb entities
defaultDp := DataParams{CollectionName: "", PartitionName: "", CollectionFieldsType: Int64FloatVecJSON,
start: 0, nb: common.DefaultNb, dim: common.DefaultDim, EnableDynamicField: true, WithRows: false, DoInsert: true}
// default do flush
defaultFp := FlushParams{PartitionNames: []string{}, async: false, DoFlush: true}

opt := &ClientParamsOption{
DataParams: defaultDp,
FlushParams: defaultFp,
}
for _, o := range opts{
o(opt)
}
// create collection
collName := createCollection(ctx, t, mc, collParam)
fmt.Println(collName)

// insert
if opt.DataParams.EnableDynamicField != collParam.EnableDynamicField {
t.Fatalf("The EnableDynamicField of CollectionParams and DataParams should be equal.")
}
if opt.DataParams.DoInsert {
opt.DataParams.CollectionName = collName
insertData(ctx, t, mc, opt.DataParams)
}

// flush
if opt.FlushParams.DoFlush {
mc.Flush(ctx, collName, opt.FlushParams.async)
}

// index
if opt.IndexParams.DoIndex {
err := mc.CreateIndex(ctx, collName, opt.IndexParams.FieldName, opt.IndexParams.Index, opt.IndexParams.async, opt.IndexParams.IndexOption)
common.CheckErr(t, err, true)
}

// load
if opt.LoadParams.DoLoad {
if len (opt.LoadParams.PartitionNames) > 0 {
err := mc.LoadPartitions(ctx, collName, opt.LoadParams.PartitionNames, opt.LoadParams.async)
common.CheckErr(t, err, true)
}else {
err := mc.LoadCollection(ctx, collName, opt.LoadParams.async, opt.LoadParams.LoadOption)
common.CheckErr(t, err, true)
}
}
}
func TestMain(m *testing.M) {
flag.Parse()
log.Printf("parse addr=%s", *addr)
Expand Down
93 changes: 93 additions & 0 deletions test/testcases/option.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
package testcases

import (
"github.com/milvus-io/milvus-sdk-go/v2/client"
"github.com/milvus-io/milvus-sdk-go/v2/entity"
)

type HelpPartitionColumns struct {
PartitionName string
IdsColumn entity.Column
VectorColumn entity.Column
}

type CollectionFieldsType string

type CollectionParams struct {
CollectionFieldsType CollectionFieldsType // collection fields type
AutoID bool // autoId
EnableDynamicField bool // enable dynamic field
ShardsNum int32
Fields []*entity.Field
Dim int64
MaxLength int64
}

type DataParams struct {
CollectionName string // insert data into which collection
PartitionName string
CollectionFieldsType CollectionFieldsType // collection fields type
start int // start
nb int // insert how many data
dim int64
EnableDynamicField bool // whether insert dynamic field data
WithRows bool
Data []entity.Column
Rows []interface{}
DoInsert bool
}

type FlushParams struct {
DoFlush bool
PartitionNames []string
async bool
}

type IndexParams struct {
DoIndex bool
Index entity.Index
FieldName string
async bool
IndexOption client.IndexOption
}

type LoadParams struct {
DoLoad bool
PartitionNames []string
async bool
LoadOption client.LoadCollectionOption
}

type ClientParamsOption struct {
DataParams DataParams
FlushParams FlushParams
IndexParams IndexParams
LoadParams LoadParams
}

type PrepareCollectionOption func(opt *ClientParamsOption)


func WithDataParams(dp DataParams) PrepareCollectionOption {
return func(opt *ClientParamsOption) {
opt.DataParams = dp
}
}

func WithFlushParams(fp FlushParams) PrepareCollectionOption {
return func(opt *ClientParamsOption) {
opt.FlushParams = fp
}
}

func WithIndexParams(ip IndexParams) PrepareCollectionOption {
return func(opt *ClientParamsOption) {
opt.IndexParams = ip
}
}

func WithLoadParams(lp LoadParams) PrepareCollectionOption {
return func(opt *ClientParamsOption) {
opt.LoadParams = lp
}
}
40 changes: 40 additions & 0 deletions test/testcases/upsert_test.go
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package testcases

import (
"github.com/milvus-io/milvus-sdk-go/v2/entity"
"github.com/milvus-io/milvus-sdk-go/v2/test/common"
"testing"
"time"
Expand All @@ -13,6 +14,37 @@ func TestUpsert(t *testing.T) {
// delete some pks
// upsert part deleted(not exist) pk and part existed pk
// upsert all not exist pk
for _, enableDynamic := range []bool{true, false} {
ctx := createContext(t, time.Second*common.DefaultTimeout)
// connect
mc := createMilvusClient(ctx, t)

// create collection
cp := CollectionParams{CollectionFieldsType: Int64FloatVecJSON, AutoID: false, EnableDynamicField: enableDynamic,
ShardsNum: common.DefaultShards, Dim: common.DefaultDim}
collName := createCollection(ctx, t, mc, cp)

// insert
dp := DataParams{CollectionName: collName, PartitionName: "", CollectionFieldsType: Int64FloatVecJSON,
start: 0, nb: common.DefaultNb, dim: common.DefaultDim, EnableDynamicField: true, WithRows: false}
_, _ = insertData(ctx, t, mc, dp)
mc.Flush(ctx, collName, false)

// create scann index
indexScann, _ := entity.NewIndexSCANN(entity.L2, 16, false)
_ = mc.CreateIndex(ctx, collName, common.DefaultFloatVecFieldName, indexScann, false)

// load collection
errLoad := mc.LoadCollection(ctx, collName, false)
common.CheckErr(t, errLoad, true)

// query

}
}

func TestDebug(t *testing.T) {

}

func TestUpsertAutoID(t *testing.T) {
Expand Down Expand Up @@ -73,4 +105,12 @@ func TestUpsertDynamicField(t *testing.T) {
// enable dynamic field and insert dynamic column
// upsert exist pk without dynamic column
// upsert not exist pk with dynamic column
}

func TestUpsertPartitionKeyCollection(t *testing.T) {
// upsert data into collection that has partition key field
}

func TestUpsertWithoutLoading(t *testing.T) {
// test upsert without loading (because delete need loading)
}

0 comments on commit 19f97bd

Please sign in to comment.