Springboot整合向量数据库Milvus

package com.cdn.milvus.component; import cn.hutool.core.lang.Dict; import cn.hutool.json.JSONUtil; import com.google.common.collect.Lists; import com.google.common.util.concurrent.ListenableFuture; import com.google.protobuf.ByteString; import io.milvus.client.MilvusServiceClient; import io.milvus.grpc.DataType; import io.milvus.grpc.MutationResult; import io.milvus.grpc.SearchResultData; import io.milvus.grpc.SearchResults; import io.milvus.param.IndexType; import io.milvus.param.MetricType; import io.milvus.param.R; import io.milvus.param.RpcStatus; import io.milvus.param.collection.CreateCollectionParam; import io.milvus.param.collection.FieldType; import io.milvus.param.collection.HasCollectionParam; import io.milvus.param.dml.InsertParam; import io.milvus.param.dml.SearchParam; import io.milvus.param.highlevel.dml.DeleteIdsParam; import io.milvus.param.highlevel.dml.QuerySimpleParam; import io.milvus.param.highlevel.dml.response.DeleteResponse; import io.milvus.param.highlevel.dml.response.QueryResponse; import io.milvus.param.index.CreateIndexParam; import io.milvus.param.partition.ShowPartitionsParam; import io.milvus.response.QueryResultsWrapper; import org.springframework.stereotype.Component; import org.springframework.util.Assert; import javax.annotation.Resource; import java.util.ArrayList; import java.util.List; import java.util.concurrent.ExecutionException; /** * @author 蔡定努 * @date 2024/05/11 11:56 */ @Component public class MilvusSearch { @Resource private MilvusServiceClient milvusServiceClient; /** * 同步搜索milvus * 欧氏距离而言，score 越接近于 0 表示两个向量之间的距离越近，相似度越高； * 而 score 越大表示两个向量之间的距离越远，相似度越低。因此，对于欧氏距离，score 越小表示相似度越高 * * @param collectionName 表名 * @param vectors 查询向量 * @param topK 最相似的向量个数 * @return */ public SearchResultData 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.size(); String paramsInJson = "{\"nprobe\": " + nprobeVectorSize + "}"; // 使用查询表达式构建标量字段过滤条件 String filterExpr = "user_id==" + 27337; SearchParam searchParam = SearchParam.newBuilder().withCollectionName(collectionName) .withParams(paramsInJson) .withMetricType(MetricType.L2) .withFloatVectors(vectors) .withVectorFieldName("lan")// 向量字段名称 .withTopK(topK) //返回条数 .withRoundDecimal(8) //保留的小数位数 .withExpr(filterExpr) //标量字段查询条件 a==1 and b==2 .build(); R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam); SearchResults searchResultsRData = searchResultsR.getData(); // List<Long> topksList = searchResultsRData.getResults().getIds().getIntId().getDataList(); return searchResultsRData.getResults(); } /** * 同步搜索milvus，增加过滤条件搜索 * * @param collectionName 表名 * @param vectors 查询向量 * @param topK 最相似的向量个数 * @param exp 过滤条件：status==1 and user_id = '4' * @return */ public SearchResultData search(String collectionName, List<List<Float>> vectors, Integer topK, String exp, Double radius) { 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(); /** * Milvus 会根据查询向量的 hash 值或者最近邻居信息来确定与查询向量最相似的桶。然后，Milvus 将在这些桶中搜索最相似的向量，并返回给定数量的最相似结果。 * * 增加 nprobe 可以提高搜索速度，因为它减少了搜索的候选桶的数量，但可能会降低搜索的准确度。减少 nprobe 可以提高搜索的准确度，但可能会增加搜索的时间 * * radius：必要参数。决定搜索请求将执行 Range Search 还是 KNN Search。 * * range_filter：可选参数。如果设置该参数，函数将对结果进行二次过滤。 */ Dict dict = Dict.create(); dict.put("nprobe", nprobeVectorSize); // dict.put("range_filter", 0); dict.put("radius", radius); //范围查询 [range_filter,radius) SearchParam searchParam = SearchParam.newBuilder().withCollectionName(collectionName) .withParams(JSONUtil.toJsonStr(dict)) .withMetricType(MetricType.L2) .withFloatVectors(vectors) .withExpr(exp) .withVectorFieldName("lan") // 向量字段名称 .withTopK(topK) .addOutField("app_id").addOutField("user_id") // 添加返回的字段 .build(); R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam); SearchResults searchResultsRData = searchResultsR.getData(); return searchResultsRData.getResults(); } /** * 异步搜索milvus * @param nprobe 参数指定了在搜索时要访问的索引桶（buckets）的数量， * 当你执行搜索操作时，Milvus 会根据查询向量的 hash 值或者最近邻居信息来确定与查询向量最相似的桶。然后，Milvus 将在这些桶中搜索最相似的向量，并返回给定数量的最相似结果。 * 增加 nprobe 可以提高搜索速度，因为它减少了搜索的候选桶的数量，但可能会降低搜索的准确度。减少 nprobe 可以提高搜索的准确度，但可能会增加搜索的时间。 * 通常情况下，你可以根据实际情况调整 nprobe 的值，通过权衡搜索速度和准确度来选择合适的值。在大多数情况下，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(); SearchParam searchParam = SearchParam.newBuilder().withCollectionName(collectionName) .withParams("{\"nprobe\": " + nprobeVectorSize + "}") // .withFloatVectors(vectors) .withTopK(topK) .withPartitionNames(partitionList) .build(); Listenable