|
@@ -1,279 +0,0 @@
|
|
|
-package com.tzld.videoVector.util;
|
|
|
|
|
-
|
|
|
|
|
-import com.tzld.videoVector.config.MilvusConfig;
|
|
|
|
|
-import io.milvus.client.MilvusClient;
|
|
|
|
|
-import io.milvus.client.MilvusServiceClient;
|
|
|
|
|
-import io.milvus.grpc.DataType;
|
|
|
|
|
-import io.milvus.grpc.MutationResult;
|
|
|
|
|
-import io.milvus.grpc.QueryResults;
|
|
|
|
|
-import io.milvus.grpc.SearchResults;
|
|
|
|
|
-import io.milvus.param.ConnectParam;
|
|
|
|
|
-import io.milvus.param.IndexType;
|
|
|
|
|
-import io.milvus.param.MetricType;
|
|
|
|
|
-import io.milvus.param.R;
|
|
|
|
|
-import io.milvus.param.collection.*;
|
|
|
|
|
-import io.milvus.param.dml.InsertParam;
|
|
|
|
|
-import io.milvus.param.dml.QueryParam;
|
|
|
|
|
-import io.milvus.param.dml.SearchParam;
|
|
|
|
|
-import io.milvus.param.index.CreateIndexParam;
|
|
|
|
|
-import io.milvus.response.MutationResultWrapper;
|
|
|
|
|
-import io.milvus.response.QueryResultsWrapper;
|
|
|
|
|
-import io.milvus.response.SearchResultsWrapper;
|
|
|
|
|
-import lombok.extern.slf4j.Slf4j;
|
|
|
|
|
-import org.springframework.stereotype.Component;
|
|
|
|
|
-
|
|
|
|
|
-import javax.annotation.PostConstruct;
|
|
|
|
|
-import javax.annotation.PreDestroy;
|
|
|
|
|
-import java.util.ArrayList;
|
|
|
|
|
-import java.util.Collections;
|
|
|
|
|
-import java.util.List;
|
|
|
|
|
-
|
|
|
|
|
-/**
|
|
|
|
|
- * Milvus 工具类 - 封装常用操作(使用 MilvusClient V1 版本)
|
|
|
|
|
- */
|
|
|
|
|
-@Slf4j
|
|
|
|
|
-@Component
|
|
|
|
|
-public class MilvusUtil {
|
|
|
|
|
-
|
|
|
|
|
- private final MilvusConfig milvusConfig;
|
|
|
|
|
-
|
|
|
|
|
- private MilvusClient milvusClient;
|
|
|
|
|
-
|
|
|
|
|
- public MilvusUtil(MilvusConfig milvusConfig) {
|
|
|
|
|
- this.milvusConfig = milvusConfig;
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- @PostConstruct
|
|
|
|
|
- public void init() {
|
|
|
|
|
- try {
|
|
|
|
|
- // 初始化 Milvus 客户端 V1
|
|
|
|
|
- ConnectParam connectParam = milvusConfig.buildConnectParam();
|
|
|
|
|
- this.milvusClient = new MilvusServiceClient(connectParam);
|
|
|
|
|
- log.info("Milvus 客户端初始化成功");
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("Milvus 客户端初始化失败:{}", e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- @PreDestroy
|
|
|
|
|
- public void destroy() {
|
|
|
|
|
- if (milvusClient != null) {
|
|
|
|
|
- try {
|
|
|
|
|
- milvusClient.close();
|
|
|
|
|
- log.info("Milvus 客户端已关闭");
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("关闭 Milvus 客户端时出错:{}", e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 检查集合是否存在
|
|
|
|
|
- */
|
|
|
|
|
- public boolean hasCollection(String collectionName) {
|
|
|
|
|
- try {
|
|
|
|
|
- R<Boolean> response = milvusClient.hasCollection(HasCollectionParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .build());
|
|
|
|
|
- return response.getData();
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("检查集合 {} 是否存在失败:{}", collectionName, e.getMessage(), e);
|
|
|
|
|
- return false;
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 创建集合
|
|
|
|
|
- */
|
|
|
|
|
- public void createCollection(String collectionName, int dimension, String indexType) {
|
|
|
|
|
- try {
|
|
|
|
|
- if (hasCollection(collectionName)) {
|
|
|
|
|
- log.warn("集合 {} 已存在", collectionName);
|
|
|
|
|
- return;
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- // 定义字段
|
|
|
|
|
- List<FieldType> fieldsSchema = new ArrayList<>();
|
|
|
|
|
-
|
|
|
|
|
- // ID 字段
|
|
|
|
|
- fieldsSchema.add(FieldType.newBuilder()
|
|
|
|
|
- .withName("id")
|
|
|
|
|
- .withDataType(DataType.Int64)
|
|
|
|
|
- .withPrimaryKey(true)
|
|
|
|
|
- .withAutoID(true)
|
|
|
|
|
- .build());
|
|
|
|
|
-
|
|
|
|
|
- // 向量字段
|
|
|
|
|
- fieldsSchema.add(FieldType.newBuilder()
|
|
|
|
|
- .withName("vector")
|
|
|
|
|
- .withDataType(DataType.FloatVector)
|
|
|
|
|
- .withDimension(dimension)
|
|
|
|
|
- .build());
|
|
|
|
|
-
|
|
|
|
|
- // 创建集合
|
|
|
|
|
- milvusClient.createCollection(CreateCollectionParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .withFieldTypes(fieldsSchema)
|
|
|
|
|
- .build());
|
|
|
|
|
-
|
|
|
|
|
- log.info("集合 {} 创建成功", collectionName);
|
|
|
|
|
-
|
|
|
|
|
- // 创建索引
|
|
|
|
|
- createIndex(collectionName, "vector", indexType, dimension);
|
|
|
|
|
-
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("创建集合 {} 失败:{}", collectionName, e.getMessage(), e);
|
|
|
|
|
- throw new RuntimeException("创建集合失败:" + e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 创建索引
|
|
|
|
|
- */
|
|
|
|
|
- public void createIndex(String collectionName, String fieldName, String indexType, int dimension) {
|
|
|
|
|
- try {
|
|
|
|
|
- CreateIndexParam indexParam = CreateIndexParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .withFieldName(fieldName)
|
|
|
|
|
- .withIndexType(IndexType.valueOf(indexType))
|
|
|
|
|
- .withMetricType(MetricType.IP) // 内积相似度
|
|
|
|
|
- .withExtraParam("{\"nlist\":1024}")
|
|
|
|
|
- .build();
|
|
|
|
|
-
|
|
|
|
|
- milvusClient.createIndex(indexParam);
|
|
|
|
|
-
|
|
|
|
|
- log.info("集合 {} 的字段 {} 索引创建成功,类型:{}", collectionName, fieldName, indexType);
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("创建索引失败:{}", e.getMessage(), e);
|
|
|
|
|
- throw new RuntimeException("创建索引失败:" + e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 加载集合到内存
|
|
|
|
|
- */
|
|
|
|
|
- public void loadCollection(String collectionName) {
|
|
|
|
|
- try {
|
|
|
|
|
- milvusClient.loadCollection(LoadCollectionParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .build());
|
|
|
|
|
- log.info("集合 {} 加载到内存成功", collectionName);
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("加载集合 {} 失败:{}", collectionName, e.getMessage(), e);
|
|
|
|
|
- throw new RuntimeException("加载集合失败:" + e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 插入向量数据
|
|
|
|
|
- */
|
|
|
|
|
- public List<Long> insertVectors(String collectionName, List<List<Float>> vectors) {
|
|
|
|
|
- try {
|
|
|
|
|
- // 构建插入数据
|
|
|
|
|
- List<InsertParam.Field> fields = new ArrayList<>();
|
|
|
|
|
- fields.add(new InsertParam.Field("vector", vectors));
|
|
|
|
|
-
|
|
|
|
|
- R<MutationResult> response = milvusClient.insert(InsertParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .withFields(fields)
|
|
|
|
|
- .build());
|
|
|
|
|
-
|
|
|
|
|
- MutationResultWrapper wrapper = new MutationResultWrapper(response.getData());
|
|
|
|
|
- List<Long> ids = wrapper.getLongIDs();
|
|
|
|
|
- log.info("成功插入 {} 条向量数据到集合 {}", vectors.size(), collectionName);
|
|
|
|
|
- return ids;
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("插入向量数据失败:{}", e.getMessage(), e);
|
|
|
|
|
- throw new RuntimeException("插入向量数据失败:" + e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 向量搜索
|
|
|
|
|
- */
|
|
|
|
|
- public List<SearchResultsWrapper.IDScore> searchVectors(String collectionName, List<Float> queryVector, int topK) {
|
|
|
|
|
- try {
|
|
|
|
|
- // 执行搜索
|
|
|
|
|
- R<SearchResults> response = milvusClient.search(SearchParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .withVectors(Collections.singletonList(queryVector))
|
|
|
|
|
- .withVectorFieldName("vector")
|
|
|
|
|
- .withTopK(topK)
|
|
|
|
|
- .withMetricType(MetricType.IP)
|
|
|
|
|
- .withOutFields(Collections.singletonList("id"))
|
|
|
|
|
- .withParams("{\"nprobe\":10}")
|
|
|
|
|
- .build());
|
|
|
|
|
-
|
|
|
|
|
- SearchResultsWrapper wrapper = new SearchResultsWrapper(response.getData().getResults());
|
|
|
|
|
- List<SearchResultsWrapper.IDScore> records = wrapper.getIDScore(0);
|
|
|
|
|
-
|
|
|
|
|
- log.info("搜索完成,找到 {} 条相似向量", records.size());
|
|
|
|
|
- return records;
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("搜索向量失败:{}", e.getMessage(), e);
|
|
|
|
|
- throw new RuntimeException("搜索向量失败:" + e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 删除集合
|
|
|
|
|
- */
|
|
|
|
|
- public void dropCollection(String collectionName) {
|
|
|
|
|
- try {
|
|
|
|
|
- milvusClient.dropCollection(DropCollectionParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .build());
|
|
|
|
|
- log.info("集合 {} 删除成功", collectionName);
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("删除集合 {} 失败:{}", collectionName, e.getMessage(), e);
|
|
|
|
|
- throw new RuntimeException("删除集合失败:" + e.getMessage(), e);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 根据 videoId 列表查询在 Milvus 中存在的 ID
|
|
|
|
|
- * @param collectionName 集合名称
|
|
|
|
|
- * @param videoIds 要查询的 videoId 列表
|
|
|
|
|
- * @return 存在的 videoId 集合
|
|
|
|
|
- */
|
|
|
|
|
- public java.util.Set<Long> existsByIds(String collectionName, List<Long> videoIds) {
|
|
|
|
|
- java.util.Set<Long> existingIds = new java.util.HashSet<>();
|
|
|
|
|
- if (videoIds == null || videoIds.isEmpty()) {
|
|
|
|
|
- return existingIds;
|
|
|
|
|
- }
|
|
|
|
|
- try {
|
|
|
|
|
- // 构建 IN 查询表达式
|
|
|
|
|
- String expr = "video_id in [" + videoIds.stream()
|
|
|
|
|
- .map(String::valueOf)
|
|
|
|
|
- .collect(java.util.stream.Collectors.joining(",")) + "]";
|
|
|
|
|
-
|
|
|
|
|
- R<QueryResults> response = milvusClient.query(QueryParam.newBuilder()
|
|
|
|
|
- .withCollectionName(collectionName)
|
|
|
|
|
- .withExpr(expr)
|
|
|
|
|
- .withOutFields(Collections.singletonList("video_id"))
|
|
|
|
|
- .build());
|
|
|
|
|
-
|
|
|
|
|
- if (response.getData() != null) {
|
|
|
|
|
- QueryResultsWrapper wrapper = new QueryResultsWrapper(response.getData());
|
|
|
|
|
- List<?> fieldData = wrapper.getFieldWrapper("video_id").getFieldData();
|
|
|
|
|
- for (Object id : fieldData) {
|
|
|
|
|
- if (id instanceof Long) {
|
|
|
|
|
- existingIds.add((Long) id);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
- log.info("查询完成,找到 {} 条已存在的记录", existingIds.size());
|
|
|
|
|
- return existingIds;
|
|
|
|
|
- } catch (Exception e) {
|
|
|
|
|
- log.error("查询 videoId 是否存在失败:{}", e.getMessage(), e);
|
|
|
|
|
- return existingIds;
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /**
|
|
|
|
|
- * 获取客户端实例
|
|
|
|
|
- */
|
|
|
|
|
- public MilvusClient getClient() {
|
|
|
|
|
- return milvusClient;
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|