本章介绍Go 如何快速上手Milvus 向量数据库的读写操作。

安装Go SDK

提示: go sdk要求go 1.15以后的版本。

go get -u github.com/milvus-io/milvus-sdk-go/v2

导入依赖包

import (
    "context"
    "fmt"
    "log"
    "math/rand"
    "time"

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

连接Milvus数据库

    ctx := context.Background()

    log.Printf(msgFmt, "开始连接 Milvus")
    c, err := client.NewClient(ctx, client.Config{
        Address: "localhost:19530",
    })
    if err != nil {
        log.Fatal("连接 milvus失败, err: ", err.Error())
    }
    defer c.Close()

创建集合

创建集合的过程，类似创建MYSQL表结构，需要描述各个字段的类型

   collectionName = "hello_tizi365"
    log.Printf(msgFmt, fmt.Sprintf("创建集合, `%s`", collectionName))

    // 定义集合的字段类型, 这里定义ID、random和embeddings三个字段
    schema := entity.NewSchema().WithName(collectionName).WithDescription("集合描述").
        WithField(entity.NewField().WithName("ID").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(false)).
        WithField(entity.NewField().WithName("random").WithDataType(entity.FieldTypeDouble)).
        WithField(entity.NewField().WithName("embeddings").WithDataType(entity.FieldTypeFloatVector).WithDim(dim))

    // 根据集合字段类型定义创建集合
    if err := c.CreateCollection(ctx, schema, entity.DefaultShardNumber); err != nil { // 使用默认分区，每个集合可以定义数据分区
        log.Fatalf("创建集合失败, err: %v", err)
    }

向集合插入数据

    idList, randomList := make([]int64, 0, nEntities), make([]float64, 0, nEntities)
    embeddingList := make([][]float32, 0, nEntities)

    rand.Seed(time.Now().UnixNano())

    // 为了测试，这里随机生成一些测试数据
    // 生成id
    for i := 0; i < nEntities; i++ {
        idList = append(idList, int64(i))
    }
    // 生成random值
    for i := 0; i < nEntities; i++ {
        randomList = append(randomList, rand.Float64())
    }
    // 生成向量值
    for i := 0; i < nEntities; i++ {
        vec := make([]float32, 0, dim)
        for j := 0; j < dim; j++ {
            vec = append(vec, rand.Float32())
        }
        embeddingList = append(embeddingList, vec)
    }
    // 这里插入多条数据，按字段列组织多条数据
    idColData := entity.NewColumnInt64("ID", idList)
    randomColData := entity.NewColumnDouble("random", randomList)
    embeddingColData := entity.NewColumnFloatVector("embeddings", dim, embeddingList)

    // 插入数据
    if _, err := c.Insert(ctx, collectionName, "", idColData, randomColData, embeddingColData); err != nil {
        log.Fatalf("数据写入失败 , err: %v", err)
    }

    // 刷新数据，让数据存盘
    if err := c.Flush(ctx, collectionName, false); err != nil {
        log.Fatalf("failed to flush data, err: %v", err)
    }

创建索引

    // 定义索引类型
    idx, err := entity.NewIndexIvfFlat(entity.L2, 128)
    if err != nil {
        log.Fatalf("定义索引失败, err: %v", err)
    }
    // 为指定集合创建索引
    if err := c.CreateIndex(ctx, collectionName, embeddingCol, idx, false); err != nil {
        log.Fatalf("索引创建失败, err: %v", err)
    }

    //  将指定集合加载到内存中提高查询效率
    err = c.LoadCollection(ctx, collectionName, false)
    if err != nil {
        log.Fatalf("failed to load collection, err: %v", err)
    }

向量相似搜索

   // 准备需要搜索的向量，这里从前面的生成的数据提取部分向量数据用于测试，下面会查询跟这些向量相似的结果
    vec2search := []entity.Vector{
        entity.FloatVector(embeddingList[len(embeddingList)-2]),
        entity.FloatVector(embeddingList[len(embeddingList)-1]),
    }
    begin := time.Now()
    sp, _ := entity.NewIndexIvfFlatSearchParam(16)
    // 执行向量相似搜索
    sRet, err := c.Search(ctx, collectionName, nil, "", []string{"random"}, vec2search,
        "embeddings", entity.L2, topK, sp)
    end := time.Now()
    if err != nil {
        log.Fatalf("搜索失败, err: %v", err)
    }

    log.Println("打印结果:")
    for _, res := range sRet {
        printResult(&res)
    }

printResult函数定义

func printResult(sRet *client.SearchResult) {
    randoms := make([]float64, 0, sRet.ResultCount)
    scores := make([]float32, 0, sRet.ResultCount)

    var randCol *entity.ColumnDouble
    for _, field := range sRet.Fields {
        if field.Name() == randomCol {
            c, ok := field.(*entity.ColumnDouble)
            if ok {
                randCol = c
            }
        }
    }
    for i := 0; i < sRet.ResultCount; i++ {
        val, err := randCol.ValueByIdx(i)
        if err != nil {
            log.Fatal(err)
        }
        randoms = append(randoms, val)
        scores = append(scores, sRet.Scores[i])
    }
    log.Printf("\trandoms: %v, scores: %v\n", randoms, scores)
}

完整的的演示源码，Go Milvus入门源码