github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/examples/partitionkey/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"log"
     6  	"math/rand"
     7  	"time"
     8  
     9  	"github.com/milvus-io/milvus-sdk-go/v2/client"
    10  	"github.com/milvus-io/milvus-sdk-go/v2/entity"
    11  )
    12  
    13  const (
    14  	milvusAddr     = `localhost:19530`
    15  	nEntities, dim = 10000, 128
    16  	collectionName = "hello_partition_key"
    17  
    18  	idCol, keyCol, embeddingCol = "ID", "key", "embeddings"
    19  	topK                        = 3
    20  )
    21  
    22  func main() {
    23  	ctx := context.Background()
    24  
    25  	log.Println("start connecting to Milvus")
    26  	c, err := client.NewClient(ctx, client.Config{
    27  		Address: milvusAddr,
    28  	})
    29  	if err != nil {
    30  		log.Fatalf("failed to connect to milvus, err: %v", err)
    31  	}
    32  	defer c.Close()
    33  
    34  	// delete collection if exists
    35  	has, err := c.HasCollection(ctx, collectionName)
    36  	if err != nil {
    37  		log.Fatalf("failed to check collection exists, err: %v", err)
    38  	}
    39  	if has {
    40  		c.DropCollection(ctx, collectionName)
    41  	}
    42  
    43  	// create collection
    44  	log.Printf("create collection `%s`\n", collectionName)
    45  	schema := entity.NewSchema().WithName(collectionName).WithDescription("hello_partition_key is the a demo to introduce the partition key related APIs").
    46  		WithField(entity.NewField().WithName(idCol).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)).
    47  		WithField(entity.NewField().WithName(keyCol).WithDataType(entity.FieldTypeInt64).WithIsPartitionKey(true)).
    48  		WithField(entity.NewField().WithName(embeddingCol).WithDataType(entity.FieldTypeFloatVector).WithDim(dim))
    49  
    50  	if err := c.CreateCollection(ctx, schema, entity.DefaultShardNumber, client.WithPartitionNum(32)); err != nil { // use default shard number
    51  		log.Fatalf("create collection failed, err: %v", err)
    52  	}
    53  
    54  	var keyList []int64
    55  	var embeddingList [][]float32
    56  	keyList = make([]int64, 0, nEntities)
    57  	embeddingList = make([][]float32, 0, nEntities)
    58  	for i := 0; i < nEntities; i++ {
    59  		keyList = append(keyList, rand.Int63()%512)
    60  	}
    61  	for i := 0; i < nEntities; i++ {
    62  		vec := make([]float32, 0, dim)
    63  		for j := 0; j < dim; j++ {
    64  			vec = append(vec, rand.Float32())
    65  		}
    66  		embeddingList = append(embeddingList, vec)
    67  	}
    68  	keyColData := entity.NewColumnInt64(keyCol, keyList)
    69  	embeddingColData := entity.NewColumnFloatVector(embeddingCol, dim, embeddingList)
    70  
    71  	log.Println("start to insert data into collection")
    72  
    73  	if _, err := c.Insert(ctx, collectionName, "", keyColData, embeddingColData); err != nil {
    74  		log.Fatalf("failed to insert random data into `%s`, err: %v", collectionName, err)
    75  	}
    76  
    77  	log.Println("insert data done, start to flush")
    78  
    79  	if err := c.Flush(ctx, collectionName, false); err != nil {
    80  		log.Fatalf("failed to flush data, err: %v", err)
    81  	}
    82  	log.Println("flush data done")
    83  
    84  	// build index
    85  	log.Println("start creating index HNSW")
    86  	idx, err := entity.NewIndexHNSW(entity.L2, 16, 256)
    87  	if err != nil {
    88  		log.Fatalf("failed to create ivf flat index, err: %v", err)
    89  	}
    90  	if err := c.CreateIndex(ctx, collectionName, embeddingCol, idx, false); err != nil {
    91  		log.Fatalf("failed to create index, err: %v", err)
    92  	}
    93  
    94  	log.Printf("build HNSW index done for collection `%s`\n", collectionName)
    95  	log.Printf("start to load collection `%s`\n", collectionName)
    96  
    97  	// load collection
    98  	if err := c.LoadCollection(ctx, collectionName, false); err != nil {
    99  		log.Fatalf("failed to load collection, err: %v", err)
   100  	}
   101  
   102  	log.Println("load collection done")
   103  
   104  	vec2search := []entity.Vector{
   105  		entity.FloatVector(embeddingList[len(embeddingList)-2]),
   106  		entity.FloatVector(embeddingList[len(embeddingList)-1]),
   107  	}
   108  	begin := time.Now()
   109  	sp, _ := entity.NewIndexHNSWSearchParam(30)
   110  	_, err = c.Search(ctx, collectionName, nil, "", []string{keyCol}, vec2search,
   111  		embeddingCol, entity.L2, topK, sp)
   112  	if err != nil {
   113  		log.Fatalf("failed to search collection, err: %v", err)
   114  	}
   115  
   116  	log.Printf("search `%s` done, latency %v\n", collectionName, time.Since(begin))
   117  
   118  	c.DropCollection(ctx, collectionName)
   119  }