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

     1  //go:build L0
     2  
     3  package testcases
     4  
     5  import (
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/milvus-io/milvus-sdk-go/v2/client"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/milvus-io/milvus-sdk-go/v2/entity"
    14  	"github.com/milvus-io/milvus-sdk-go/v2/test/common"
    15  )
    16  
    17  const (
    18  	DefaultPkFieldName     = "id"
    19  	DefaultVectorFieldName = "vector"
    20  	nb                     = 10000
    21  )
    22  
    23  // test highlevel api new collection
    24  func TestNewCollection(t *testing.T) {
    25  	ctx := createContext(t, time.Second*common.DefaultTimeout)
    26  	mc := createMilvusClient(ctx, t)
    27  
    28  	// new collection
    29  	collName := common.GenRandomString(5)
    30  	err := mc.NewCollection(ctx, collName, common.DefaultDim, client.WithConsistencyLevel(entity.ClStrong))
    31  	common.CheckErr(t, err, true)
    32  
    33  	// describe collection and check
    34  	collection, _ := mc.DescribeCollection(ctx, collName)
    35  
    36  	pkField := common.GenField(DefaultPkFieldName, entity.FieldTypeInt64, common.WithIsPrimaryKey(true))
    37  	vecField := common.GenField(DefaultVectorFieldName, entity.FieldTypeFloatVector, common.WithDim(common.DefaultDim))
    38  	expSchema := common.GenSchema(collName, false, []*entity.Field{pkField, vecField}, common.WithEnableDynamicField(true))
    39  	common.CheckCollection(t, collection, collName, 1, expSchema, entity.ClStrong)
    40  
    41  	// describe index and check
    42  	indexes, _ := mc.DescribeIndex(ctx, collName, DefaultVectorFieldName)
    43  	expParams := map[string]string{
    44  		"metric_type": string(entity.IP),
    45  		"index_type":  string(entity.AUTOINDEX),
    46  	}
    47  	expIndex := entity.NewGenericIndex(DefaultVectorFieldName, entity.AUTOINDEX, expParams)
    48  	common.CheckIndexResult(t, indexes, expIndex)
    49  
    50  	// check collection is loaded
    51  	loadState, _ := mc.GetLoadState(ctx, collName, []string{})
    52  	require.Equal(t, entity.LoadStateLoaded, loadState)
    53  
    54  	// insert
    55  	pkColumn := common.GenColumnData(0, nb, entity.FieldTypeInt64, DefaultPkFieldName)
    56  	vecColumn := common.GenColumnData(0, nb, entity.FieldTypeFloatVector, DefaultVectorFieldName, common.WithVectorDim(common.DefaultDim))
    57  	_, err = mc.Insert(
    58  		ctx, collName, "",
    59  		pkColumn, vecColumn,
    60  	)
    61  	common.CheckErr(t, err, true)
    62  	//time.Sleep(10) // because consistence level
    63  
    64  	// get
    65  	queryResult, err := mc.Get(
    66  		ctx,
    67  		collName,
    68  		pkColumn.Slice(0, 10),
    69  	)
    70  	common.CheckErr(t, err, true)
    71  	common.CheckOutputFields(t, queryResult, []string{DefaultPkFieldName, DefaultVectorFieldName})
    72  	common.CheckQueryResult(t, queryResult, []entity.Column{
    73  		pkColumn.Slice(0, 10),
    74  		vecColumn.Slice(0, 10),
    75  	})
    76  
    77  	// search
    78  	sp, _ := entity.NewIndexAUTOINDEXSearchParam(2)
    79  	searchRes, errSearchEmpty := mc.Search(
    80  		ctx, collName,
    81  		[]string{},
    82  		"",
    83  		[]string{DefaultPkFieldName},
    84  		common.GenSearchVectors(common.DefaultNq, common.DefaultDim, entity.FieldTypeFloatVector),
    85  		DefaultVectorFieldName,
    86  		entity.IP,
    87  		common.DefaultTopK,
    88  		sp,
    89  	)
    90  	common.CheckErr(t, errSearchEmpty, true)
    91  	common.CheckSearchResult(t, searchRes, common.DefaultNq, common.DefaultTopK)
    92  }
    93  
    94  func TestNewCollectionCustomize(t *testing.T) {
    95  	ctx := createContext(t, time.Second*common.DefaultTimeout)
    96  	mc := createMilvusClient(ctx, t)
    97  
    98  	// new collection
    99  	collName := common.GenRandomString(5)
   100  	pkFieldName := "pk"
   101  	vectorFieldName := "vec"
   102  	err := mc.NewCollection(
   103  		ctx, collName, common.DefaultDim,
   104  		client.WithPKFieldName(pkFieldName), client.WithPKFieldType(entity.FieldTypeVarChar), client.WithPKMaxLength(2048),
   105  		client.WithVectorFieldName(vectorFieldName),
   106  		client.WithMetricsType(entity.L2), client.WithAutoID(false), client.WithEnableDynamicSchema(false),
   107  		client.WithConsistencyLevel(entity.ClStrong))
   108  	common.CheckErr(t, err, true)
   109  
   110  	// describe collection and check
   111  	collection, _ := mc.DescribeCollection(ctx, collName)
   112  
   113  	pkField := common.GenField(pkFieldName, entity.FieldTypeVarChar, common.WithIsPrimaryKey(true), common.WithMaxLength(2048))
   114  	vecField := common.GenField(vectorFieldName, entity.FieldTypeFloatVector, common.WithDim(common.DefaultDim))
   115  	expSchema := common.GenSchema(collName, false, []*entity.Field{pkField, vecField})
   116  	common.CheckCollection(t, collection, collName, 1, expSchema, entity.ClStrong)
   117  
   118  	// describe index and check
   119  	indexes, _ := mc.DescribeIndex(ctx, collName, vectorFieldName)
   120  	expParams := map[string]string{
   121  		"metric_type": string(entity.L2),
   122  		"index_type":  string(entity.AUTOINDEX),
   123  	}
   124  	expIndex := entity.NewGenericIndex(vectorFieldName, entity.AUTOINDEX, expParams)
   125  	common.CheckIndexResult(t, indexes, expIndex)
   126  
   127  	// check collection is loaded
   128  	loadState, _ := mc.GetLoadState(ctx, collName, []string{})
   129  	require.Equal(t, entity.LoadStateLoaded, loadState)
   130  
   131  	// insert
   132  	pkColumn := common.GenColumnData(0, nb, entity.FieldTypeVarChar, pkFieldName)
   133  	vecColumn := common.GenColumnData(0, nb, entity.FieldTypeFloatVector, vectorFieldName, common.WithVectorDim(common.DefaultDim))
   134  	_, err = mc.Insert(
   135  		ctx, collName, "",
   136  		pkColumn,
   137  		vecColumn,
   138  	)
   139  	common.CheckErr(t, err, true)
   140  
   141  	// get
   142  	queryResult, err := mc.Get(
   143  		ctx,
   144  		collName,
   145  		pkColumn.Slice(0, 10),
   146  	)
   147  	common.CheckErr(t, err, true)
   148  	common.CheckOutputFields(t, queryResult, []string{pkFieldName, vectorFieldName})
   149  	common.CheckQueryResult(t, queryResult, []entity.Column{
   150  		pkColumn.Slice(0, 10),
   151  		vecColumn.Slice(0, 10),
   152  	})
   153  
   154  	// search
   155  	sp, _ := entity.NewIndexAUTOINDEXSearchParam(2)
   156  	searchRes, errSearchEmpty := mc.Search(
   157  		ctx, collName,
   158  		[]string{},
   159  		"",
   160  		[]string{pkFieldName},
   161  		common.GenSearchVectors(common.DefaultNq, common.DefaultDim, entity.FieldTypeFloatVector),
   162  		vectorFieldName,
   163  		entity.L2,
   164  		common.DefaultTopK,
   165  		sp,
   166  	)
   167  	common.CheckErr(t, errSearchEmpty, true)
   168  	common.CheckSearchResult(t, searchRes, common.DefaultNq, common.DefaultTopK)
   169  }