milvus分批写入测试数据

场景:

假如现在需要向milvus写入500W数据，调用milvus的insert api一次性写入500W数据是不现实的，内存也吃不消。

这时候需要分批写入，例如每次写入1000行，直至写完500W数据。

实现代码如下:

package mainimport ("context""fmt""github.com/milvus-io/milvus-sdk-go/v2/client""github.com/milvus-io/milvus-sdk-go/v2/entity""log""math/rand""strconv"
)const (milvusAddr     = `192.168.230.71:19530`dim            = 512collectionName = "hello_iterator"msgFmt                                     = "==== %s ====\n"idCol, randomCol, addressCol, embeddingCol = "ID", "random", "address", "embeddings"totalRows                                  = 100005batchSize                                  = 1000
)func main() {ctx := context.Background()log.Printf(msgFmt, "start connecting to Milvus")c, err := client.NewClient(ctx, client.Config{Address: milvusAddr,})if err != nil {log.Fatal("failed to connect to milvus, err: ", err.Error())}defer c.Close()// delete collection if existshas, err := c.HasCollection(ctx, collectionName)if err != nil {log.Fatalf("failed to check collection exists, err: %v", err)}if has {//c.DropCollection(ctx, collectionName)log.Println("collection exists")return}// create collectionlog.Printf(msgFmt, fmt.Sprintf("create collection, `%s`", collectionName))schema := entity.NewSchema().WithName(collectionName).WithDescription("hello_milvus is the simplest demo to introduce the APIs").WithField(entity.NewField().WithName(idCol).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(false)).WithField(entity.NewField().WithName(randomCol).WithDataType(entity.FieldTypeDouble)).WithField(entity.NewField().WithName(addressCol).WithDataType(entity.FieldTypeVarChar).WithTypeParams(entity.TypeParamMaxLength, "50")).WithField(entity.NewField().WithName(embeddingCol).WithDataType(entity.FieldTypeFloatVector).WithDim(dim))if err := c.CreateCollection(ctx, schema, 1); err != nil {log.Fatalf("create collection failed, err: %v", err)}// 分批数groupNum := totalRows / batchSizeif totalRows%batchSize > 0 {groupNum++ // 如果有余数，则需要多写一批}// 分批写入数据for i := 0; i < groupNum; i++ {// 计算当前批次的起始和结束索引start := i * batchSizeend := start + batchSizeif end > totalRows {end = totalRows // 确保不会超出总数}// insert dataidList := make([]int64, 0, batchSize)randomList := make([]float64, 0, batchSize)addressList := make([]string, 0, batchSize)embeddingList := make([][]float32, 0, batchSize)// generate datafor i := start; i < end; i++ {idList = append(idList, int64(i))}for i := start; i < end; i++ {randomList = append(randomList, rand.Float64())}for i := start; i < end; i++ {addressList = append(addressList, "wuhan"+strconv.Itoa(i))}for i := start; i < end; i++ {vec := make([]float32, 0, dim)for j := 0; j < dim; j++ {vec = append(vec, rand.Float32())}embeddingList = append(embeddingList, vec)}idColData := entity.NewColumnInt64(idCol, idList)randomColData := entity.NewColumnDouble(randomCol, randomList)addressColData := entity.NewColumnVarChar(addressCol, addressList)embeddingColData := entity.NewColumnFloatVector(embeddingCol, dim, embeddingList)if _, err := c.Insert(ctx, collectionName, "", idColData, randomColData, addressColData, embeddingColData); err != nil {log.Fatalf("failed to insert random data into `hello_milvus, err: %v", err)}log.Printf("inserted:%d\n", end)}if err := c.Flush(ctx, collectionName, false); err != nil {log.Fatalf("failed to flush data, err: %v", err)}// build indexlog.Printf(msgFmt, "start creating index HNSW")idx, err := entity.NewIndexHNSW(entity.COSINE, 15, 50)if err != nil {log.Fatalf("failed to create ivf flat index, err: %v", err)}if err := c.CreateIndex(ctx, collectionName, embeddingCol, idx, false); err != nil {log.Fatalf("failed to create index, err: %v", err)}log.Printf(msgFmt, "start loading collection")err = c.LoadCollection(ctx, collectionName, false)if err != nil {log.Fatalf("failed to load collection, err: %v", err)}
}