SpringBoot整合Milvus

5,743 阅读3分钟

什么是Milvus?

  • Milvus,一个开源的高性能向量数据库,它在各种应用场景中展现出强大的性能和灵活性。 在许多现代应用中,处理和分析大规模向量数据变得越来越重要。例如,在图像和视频搜索、推荐系统、自然语言处理和生物信息学等领域,向量数据被广泛应用。

项目背景

  • 在公司推荐系统中,我们需要根据用户的历史行为和兴趣,为其推荐相关的内容。于是将用户和内容表示为向量,并使用 Milvus 进行相似度匹配。通过将用户向量和内容向量存储在 Milvus 中,并利用其高效的相似度查询功能,我们可以快速找到与用户兴趣最匹配的内容,并进行个性化推荐。

  • 向量的生成由spark任务生成数据并写入,本文只写SpringBoot集成Milvus实现数据查询部分,面向C端,性能已测

Maven依赖引入

  • 开始使用的是1.x版本,后来由于2.x新增了过滤筛选功能,升级了版本为2.2.3,1版本和2版本查询还是有一些区别,建议采用2版本
<dependency>
    <groupId>io.milvus</groupId>
    <artifactId>milvus-sdk-java</artifactId>
    <version>2.2.3</version>
</dependency>

自动配置

@Configuration
public class MilvusConfiguration {

    /**
     *  milvus ip addr
     */
    @Value("${milvus.config.ipAddr}")
    private String ipAddr;

    /**
     * milvus   port
     */
    @Value("${milvus.config.port}")
    private Integer  port;

    @Bean
    @Scope("singleton")
    public MilvusServiceClient getMilvusClient() {
        return getMilvusFactory().getMilvusClient();
    }

    @Bean(initMethod = "init", destroyMethod = "close")
    public MilvusRestClientFactory getMilvusFactory() {
        return  MilvusRestClientFactory.build(ipAddr, port);
    }
}

milvus Rest client 封装

public class MilvusRestClientFactory {

    private static String  IP_ADDR;

    private static Integer PORT ;

    private MilvusServiceClient milvusServiceClient;

    private ConnectParam.Builder  connectParamBuilder;


    private static MilvusRestClientFactory milvusRestClientFactory = new MilvusRestClientFactory();

    private MilvusRestClientFactory(){

    }

    public static MilvusRestClientFactory build(String ipAddr, Integer  port) {
        IP_ADDR = ipAddr;
        PORT = port;
        return milvusRestClientFactory;
    }

    private ConnectParam.Builder connectParamBuilder(String host, int port) {
        return  ConnectParam.newBuilder().withHost(host).withPort(port);
    }



    public void init() {
        connectParamBuilder =  connectParamBuilder(IP_ADDR,PORT);
        ConnectParam connectParam = connectParamBuilder.build();
        milvusServiceClient =new MilvusServiceClient(connectParam);
    }


    public MilvusServiceClient getMilvusClient() {
        return milvusServiceClient;
    }


    public void close() {
        if (milvusServiceClient != null) {
            try {
                milvusServiceClient.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}

查询

写入数据不同,获取结果不同,我这里最后获取的是Long类型的数据集合,仅供参考

  • 同步搜索milvus
/**
 * 同步搜索milvus
 * @param collectionName 表名
 * @param vectors 查询向量
 * @param topK 最相似的向量个数
 * @return
 */
public List<Long> search(String collectionName, List<List<Float>> vectors, Integer topK) {

    Assert.notNull(collectionName, "collectionName  is null");
    Assert.notNull(vectors, "vectors is null");
    Assert.notEmpty(vectors, "vectors is empty");
    Assert.notNull(topK, "topK is null");
    int nprobeVectorSize = vectors.get(0).size();
    String paramsInJson = "{"nprobe": " + nprobeVectorSize + "}";
    SearchParam searchParam =
            SearchParam.newBuilder().withCollectionName(collectionName)
                    .withParams(paramsInJson)
                    .withMetricType(MetricType.IP)
                    .withVectors(vectors)
                    .withVectorFieldName("embedding")
                    .withTopK(topK)
                    .build();

    R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam);
    SearchResults searchResultsRData = searchResultsR.getData();
    List<Long> topksList = searchResultsRData.getResults().getIds().getIntId().getDataList();
    return topksList;
}
  • 同步搜索milvus,增加过滤条件搜索
/**
 * 同步搜索milvus,增加过滤条件搜索
 *
 * @param collectionName 表名
 * @param vectors 查询向量
 * @param topK 最相似的向量个数
 * @param exp 过滤条件:status=1
 * @return
 */
public List<Long> search(String collectionName, List<List<Float>> vectors, Integer topK, String exp) {
    Assert.notNull(collectionName, "collectionName  is null");
    Assert.notNull(vectors, "vectors is null");
    Assert.notEmpty(vectors, "vectors is empty");
    Assert.notNull(topK, "topK is null");
    Assert.notNull(exp, "exp is null");
    int nprobeVectorSize = vectors.get(0).size();
    String paramsInJson = "{"nprobe": " + nprobeVectorSize + "}";
    SearchParam searchParam =
            SearchParam.newBuilder().withCollectionName(collectionName)
                    .withParams(paramsInJson)
                    .withMetricType(MetricType.IP)
                    .withVectors(vectors)
                    .withExpr(exp)
                    .withVectorFieldName("embedding")
                    .withTopK(topK)
                    .build();

    R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam);
    SearchResults searchResultsRData = searchResultsR.getData();
    List<Long> topksList = searchResultsRData.getResults().getIds().getIntId().getDataList();
    return topksList;
}
  • 异步搜索milvus:针对实时结果要求不高的场景
/**
 * 异步搜索milvus
 *
 * @param collectionName 表名
 * @param vectors 查询向量
 * @param partitionList 最相似的向量个数
 * @param topK
 * @return
 */
public List<Long> searchAsync(String collectionName, List<List<Float>> vectors,
                              List<String> partitionList, Integer topK) throws ExecutionException, InterruptedException {

    Assert.notNull(collectionName, "collectionName  is null");
    Assert.notNull(vectors, "vectors is null");
    Assert.notEmpty(vectors, "vectors is empty");
    Assert.notNull(partitionList, "partitionList is null");
    Assert.notEmpty(partitionList, "partitionList is empty");
    Assert.notNull(topK, "topK is null");
    int nprobeVectorSize = vectors.get(0).size();
    String paramsInJson = "{"nprobe": " + nprobeVectorSize + "}";
    SearchParam searchParam =
            SearchParam.newBuilder().withCollectionName(collectionName)
                    .withParams(paramsInJson)
                    .withVectors(vectors)
                    .withTopK(topK)
                    .withPartitionNames(partitionList)
                    .build();
    ListenableFuture<R<SearchResults>> listenableFuture = milvusServiceClient.searchAsync(searchParam);

    List<Long> resultIdsList = listenableFuture.get().getData().getResults().getTopksList();

    return resultIdsList;
}
  • 获取分区集合
/**
 * 获取分区集合
 * @param collectionName 表名
 * @return
 */
public List<String> getPartitionsList(String collectionName) {
    Assert.notNull(collectionName, "collectionName  is null");
    ShowPartitionsParam searchParam = ShowPartitionsParam.newBuilder().withCollectionName(collectionName).build();
    List<ByteString> byteStrings = milvusServiceClient.showPartitions(searchParam).getData().getPartitionNamesList().asByteStringList();
    List<String> partitionList = Lists.newLinkedList();
    byteStrings.forEach(s -> {
        partitionList.add(s.toStringUtf8());
    });
    return partitionList;
}

yml配置数据

milvus:
  config:
    ipAddr: xxx.xxx.xxx.xxx
    port: 19531