Milvus 基本操作

潇锐killer

已于 2025-01-17 14:40:36 修改

阅读量808

点赞数 2

CC 4.0 BY-SA版权

分类专栏：大模型应用文章标签： milvus 开发语言 java

于 2024-06-03 15:39:01 首次发布

本文链接：https://blog.youkuaiyun.com/wangqiaowq/article/details/139416294

1、maven 依赖

<dependency>
            <groupId>io.milvus</groupId>
            <artifactId>milvus-sdk-java</artifactId>
            <version>2.3.3</version>
            <exclusions>
                <exclusion>
                    <groupId>org.slf4j</groupId>
                    <artifactId>slf4j-api</artifactId>
                </exclusion>
                <exclusion>
                    <groupId>org.apache.logging.log4j</groupId>
                    <artifactId>log4j-slf4j-impl</artifactId>
                </exclusion>
            </exclusions>
        </dependency>

2、MivusService 封装了基本操作

@Service
@Slf4j
public class MivusService {


    @Autowired
    MilvusServiceClient milvusClient;

    private String clientId;

    /**
     * 同步搜索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.L2)
                        .withVectors(vectors)
                        .withVectorFieldName("embeddings")
                        .withTopK(topK)
                        .build();

        R<SearchResults> searchResultsR = milvusClient.search(searchParam);
        SearchResults searchResultsRData = searchResultsR.getData();
        List<Long> topksList = searchResultsRData.getResults().getIds().getIntId().getDataList();
        return topksList;
    }

    /**
     * 同步搜索milvus
     * @param collectionName 表名
     * @param vectors 查询向量
     * @param topK 最相似的向量个数
     * @return
     */
    public List<Long> search1(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 = milvusClient.search(searchParam);
        SearchResults searchResultsRData = searchResultsR.getData();
        List<Long> topksList = searchResultsRData.getResults().getIds().getIntId().getDataList();
        return topksList;
    }


    /**
     * 同步搜索milvus，增加过滤条件搜索
     *
     * @param collectionName 表名
     * @param vectors 查询向量
     * @param topK 最相似的向量个数
     * @param exp 过滤条件：status=1
     * @return
     */
    public List<Long> search2(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 = milvusClient.search(searchParam);
        SearchResults searchResultsRData = searchResultsR.getData();
        List<Long> topksList = searchResultsRData.getResults().getIds().getIntId().getDataList();
        return topksList;
    }

    /**
     * 异步搜索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 = milvusClient.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 = milvusClient.showPartitions(searchParam).getData().getPartitionNamesList().asByteStringList();
        List<String> partitionList = Lists.newLinkedList();
        byteStrings.forEach(s -> {
            partitionList.add(s.toStringUtf8());
        });
        return partitionList;
    }



    public void loadCollection(String collectionName) {
        LoadCollectionParam loadCollectionParam = LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        R<RpcStatus> response = milvusClient.loadCollection(loadCollectionParam);
        log.info("loadCollection {} is {}", collectionName, response.getData().getMsg());
    }


    public void releaseCollection(String collectionName) {
        ReleaseCollectionParam param = ReleaseCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        R<RpcStatus> response = milvusClient.releaseCollection(param);
        log.info("releaseCollection {} is {}", collectionName, response.getData().getMsg());
    }


    public void loadPartitions(String collectionName, List<String> partitionsName) {
        LoadPartitionsParam build = LoadPartitionsParam.newBuilder()
                .withCollectionName(collectionName)
                .withPartitionNames(partitionsName)
                .build();
        R<RpcStatus> rpcStatusR = milvusClient.loadPartitions(build);
        log.info("loadPartitions {} is {}", partitionsName, rpcStatusR.getData().getMsg());
    }


    public void releasePartitions(String collectionName, List<String> partitionsName) {
        ReleasePartitionsParam build = ReleasePartitionsParam.newBuilder()
                .withCollectionName(collectionName)
                .withPartitionNames(partitionsName)
                .build();
        R<RpcStatus> rpcStatusR = milvusClient.releasePartitions(build);
        log.info("releasePartition {} is {}", collectionName, rpcStatusR.getData().getMsg());
    }


    public boolean isExitCollection(String collectionName) {
        HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        R<Boolean> response = milvusClient.hasCollection(hasCollectionParam);
        Boolean isExists = response.getData();
        log.info("collection {} is exists: {}", collectionName, isExists);
        return isExists;
    }


    public Boolean creatCollection(String collectionName) {
        // 主键字段
        FieldType fieldType1 = FieldType.newBuilder()
                .withName(Content.Field.ID)
                .withDescription("primary key")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build();
        // 文本字段
        FieldType fieldType2 = FieldType.newBuilder()
                .withName(Content.Field.CONTENT)
                .withDataType(DataType.VarChar)
                .withMaxLength(Content.MAX_LENGTH)
                .build();
        // 向量字段
        FieldType fieldType3 = FieldType.newBuilder()
                .withName(Content.Field.CONTENT_VECTOR)
                .withDataType(DataType.FloatVector)
                .withDimension(Content.FEATURE_DIM)
                .build();
        // 创建collection
        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription("Schema of Content")
                .withShardsNum(Content.SHARDS_NUM)
                .addFieldType(fieldType1)
                .addFieldType(fieldType2)
                .addFieldType(fieldType3)
                .build();
        R<RpcStatus> response = milvusClient.createCollection(createCollectionReq);

        log.info("collection: {} is created ? status: = {}", collectionName, response.getData().getMsg());
        return response.getData().getMsg().equals("Success");
    }


    public Boolean dropCollection(String collectionN