milvus分批写入测试数据

milvus分批写入测试数据

场景:

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

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

实现代码如下:

package main
import (
	"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            = 512
	collectionName = "hello_iterator"
	msgFmt                                     = "==== %s ====\n"
	idCol, randomCol, addressCol, embeddingCol = "ID", "random", "address", "embeddings"
	totalRows                                  = 100005
	batchSize                                  = 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 exists
	has, 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 collection
	log.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 / batchSize
	if totalRows%batchSize > 0 {
		groupNum++ // 如果有余数，则需要多写一批
	}
	// 分批写入数据
	for i := 0; i  totalRows {
			end = totalRows // 确保不会超出总数
		}
		// insert data
		idList := make([]int64, 0, batchSize)
		randomList := make([]float64, 0, batchSize)
		addressList := make([]string, 0, batchSize)
		embeddingList := make([][]float32, 0, batchSize)
		// generate data
		for i := start; i