diff --git a/src/main/java/com/xly/milvus/service/impl/MilvusServiceImpl.java b/src/main/java/com/xly/milvus/service/impl/MilvusServiceImpl.java index 4baa380..5c8b892 100644 --- a/src/main/java/com/xly/milvus/service/impl/MilvusServiceImpl.java +++ b/src/main/java/com/xly/milvus/service/impl/MilvusServiceImpl.java @@ -337,30 +337,122 @@ public class MilvusServiceImpl implements MilvusService { } @Override - public List> getDataToCollection(String collectionName, String milvusFilter,String searchText,Integer size,List fields,String vectorValue,String sceneName){ - log.info("开始相似度查询: collection={}, searchText={}", collectionName, searchText); - // 2. 设置范围搜索参数 - Map searchParams = new HashMap<>(); - searchParams.put("nprobe", 10); - // 对于 IP 度量，相似度范围在 [minScore, maxScore] - searchParams.put("radius", 0.9); // 最小相似度 - searchParams.put("range_filter", 1); // 最大相似度 - if(ObjectUtil.isEmpty(fields)){ + public List> getDataToCollection(String collectionName, + String milvusFilter, + String searchText, + Integer size, + List fields, + String vectorValue, + String sceneName) { + log.info("开始查询: collection={}, searchText={}, vectorValue={}", collectionName, searchText, vectorValue); + + // 设置输出字段 + if (ObjectUtil.isEmpty(fields)) { fields = new ArrayList<>(); } fields.add("sSlaveId"); fields.add("metadata"); - // 1. 构建查询（通用） + + // 根据 vectorValue 判断走哪个查询 + if (ObjectUtil.isNotEmpty(vectorValue)) { + // vectorValue 不为空：走向量相似度搜索（使用 SearchReq） + log.info("使用向量相似度搜索，vectorValue 存在"); + return similaritySearchWithVector(collectionName, milvusFilter, size, fields, vectorValue); +// } else if (ObjectUtil.isNotEmpty(searchText)) { +// // vectorValue 为空但 searchText 不为空：走文本向量化搜索（使用 SearchReq） +// log.info("使用文本向量化搜索，searchText={}", searchText); +// return similaritySearchWithText(collectionName, milvusFilter, searchText, size, fields); + } else { + // vectorValue 和 searchText 都为空：走全量数据查询（使用 QueryReq） + log.info("使用全量数据查询，无搜索条件"); + return queryAllData(collectionName, milvusFilter, fields,size); + } + } + + /** + * 全量数据查询（使用 Query API） + */ + private List> queryAllData(String collectionName, + String milvusFilter, + List fields, + Integer size) { + // 使用 QueryReqBuilder（不需要向量数据） + QueryReq.QueryReqBuilder builder = QueryReq.builder() + .collectionName(collectionName) + .outputFields(fields); + if (ObjectUtil.isNotEmpty(milvusFilter)) { + builder.filter(milvusFilter); + } + // 可选：如果 size 有值，可以设置 limit + if (ObjectUtil.isNotEmpty(size)) { + builder.limit(size); + } + QueryReq queryReq = builder.build(); + QueryResp queryResp = milvusClient.query(queryReq); + return processMilvusResults(queryResp); + } + + /** + * 使用文本向量化进行相似度搜索 + */ + private List> similaritySearchWithText(String collectionName, + String milvusFilter, + String searchText, + Integer size, + List fields) { + // 向量化搜索文本 + List vectorList = vectorizationService.textToVector(searchText); + if (vectorList == null || vectorList.isEmpty()) { + throw new RuntimeException("向量化失败"); + } + + // 转换为 float[] + float[] floatArray = new float[vectorList.size()]; + for (int i = 0; i < vectorList.size(); i++) { + floatArray[i] = vectorList.get(i); + } + + // 设置搜索参数 + Map searchParams = new HashMap<>(); + searchParams.put("nprobe", 10); + searchParams.put("radius", 0.9); + searchParams.put("range_filter", 1); + + // 构建搜索请求 + FloatVec floatVec = new FloatVec(floatArray); SearchReq.SearchReqBuilder builder = SearchReq.builder() .collectionName(collectionName) .topK(size) .metricType(IndexParam.MetricType.IP) .outputFields(fields) -// .filterType(FilterType.POST_FILTER) - .searchParams(searchParams); - if(ObjectUtil.isEmpty(vectorValue)){ - vectorValue = sceneName; + .searchParams(searchParams) + .data(Collections.singletonList(floatVec)) + .annsField("vector"); + + if (ObjectUtil.isNotEmpty(milvusFilter)) { + builder.filter(milvusFilter); } + + SearchReq searchReq = builder.build(); + SearchResp searchResp = milvusClient.search(searchReq); + + return processMilvusResults(searchResp); + } + + /** + * 使用已有向量进行相似度搜索 + */ + private List> similaritySearchWithVector(String collectionName, + String milvusFilter, + Integer size, + List fields, + String vectorValue) { + // 设置搜索参数 + Map searchParams = new HashMap<>(); + searchParams.put("nprobe", 10); + searchParams.put("radius", 0.9); + searchParams.put("range_filter", 1); + // 解析向量字符串为 float 数组 // 1. 向量化搜索文本 List vectorList = vectorizationService.textToVector(vectorValue); if (vectorList == null || vectorList.isEmpty()) { @@ -373,31 +465,72 @@ public class MilvusServiceImpl implements MilvusService { } // 3. 创建 Milvus FloatVec 对象 FloatVec floatVec = new FloatVec(floatArray); - builder.data(Collections.singletonList(floatVec)) - .annsField("vector"); // 向量字段名 + SearchReq.SearchReqBuilder builder = SearchReq.builder() + .collectionName(collectionName) + .topK(size) + .metricType(IndexParam.MetricType.IP) + .outputFields(fields) + .searchParams(searchParams) + .data(Collections.singletonList(floatVec)) + .annsField("vector"); - if(ObjectUtil.isNotEmpty(milvusFilter)){ + if (ObjectUtil.isNotEmpty(milvusFilter)) { builder.filter(milvusFilter); } - // 4. 构建搜索请求 -// SearchReq searchReq = SearchReq.builder() -// .collectionName(collectionName) -// .data(Collections.singletonList(floatVec)) -// .annsField("vector") // 向量字段名 -// .topK(size) // 返回最相似的10条 -// .metricType(IndexParam.MetricType.IP) // 内积相似度 -// .outputFields(fields) -// .searchParams(searchParams) -// .filter(milvusFilter) -// .build(); - // 5. 执行搜索 + SearchReq searchReq = builder.build(); SearchResp searchResp = milvusClient.search(searchReq); - - // 6. 处理结果 return processMilvusResults(searchResp); } +// public List> getDataToCollection(String collectionName, String milvusFilter,String searchText,Integer size,List fields,String vectorValue,String sceneName){ +// log.info("开始相似度查询: collection={}, searchText={}", collectionName, searchText); +// // 2. 设置范围搜索参数 +// Map searchParams = new HashMap<>(); +// searchParams.put("nprobe", 10); +// // 对于 IP 度量，相似度范围在 [minScore, maxScore] +// searchParams.put("radius", 0.9); // 最小相似度 +// searchParams.put("range_filter", 1); // 最大相似度 +// if(ObjectUtil.isEmpty(fields)){ +// fields = new ArrayList<>(); +// } +// fields.add("sSlaveId"); +// fields.add("metadata"); +// // 1. 构建查询（通用） +// SearchReq.SearchReqBuilder builder = SearchReq.builder() +// .collectionName(collectionName) +// .topK(size) +// .metricType(IndexParam.MetricType.IP) +// .outputFields(fields) +//// .filterType(FilterType.POST_FILTER) +// .searchParams(searchParams); +// if(ObjectUtil.isNotEmpty(vectorValue)){ +// // 1. 向量化搜索文本 +// List vectorList = vectorizationService.textToVector(vectorValue); +// if (vectorList == null || vectorList.isEmpty()) { +// throw new RuntimeException("向量化失败"); +// } +// // 2. 转换为 float[] +// float[] floatArray = new float[vectorList.size()]; +// for (int i = 0; i < vectorList.size(); i++) { +// floatArray[i] = vectorList.get(i); +// } +// // 3. 创建 Milvus FloatVec 对象 +// FloatVec floatVec = new FloatVec(floatArray); +// builder.data(Collections.singletonList(floatVec)) +// .annsField("vector"); // 向量字段名 +// } +// if(ObjectUtil.isNotEmpty(milvusFilter)){ +// builder.filter(milvusFilter); +// } +// // 5. 执行搜索 +// SearchReq searchReq = builder.build(); +// SearchResp searchResp = milvusClient.search(searchReq); +// +// // 6. 处理结果 +// return processMilvusResults(searchResp); +// } + /** * 判断 Milvus 过滤条件是否有效（支持 TEXT_MATCH 全文检索） @@ -559,7 +692,7 @@ public class MilvusServiceImpl implements MilvusService { } /** - * 处理 Milvus 查询结果 + * 处理 SearchResp 查询结果（向量相似度搜索） */ private List> processMilvusResults(SearchResp response) { List> results = new ArrayList<>(); @@ -572,32 +705,84 @@ public class MilvusServiceImpl implements MilvusService { log.warn("Milvus 搜索结果为空"); return results; } + // 遍历每个查询的结果集（通常只有一个查询） for (List resultList : searchResults) { // 遍历每个搜索结果 for (SearchResp.SearchResult result : resultList) { + Map resultMap = new HashMap<>(); + // 获取实体字段数据 Map entity = result.getEntity(); - Map metadata = new HashMap<>(); - if(ObjectUtil.isNotEmpty(entity.get("metadata"))){ - JsonObject obj = (JsonObject) entity.get("metadata"); - metadata.putAll( jsonObjectToMap(obj)); + if (entity != null && !entity.isEmpty()) { + resultMap.putAll(entity); + } + + // 处理 metadata 字段（如果是 JSON 对象） + if (resultMap.containsKey("metadata") && resultMap.get("metadata") instanceof JsonObject) { + JsonObject jsonObject = (JsonObject) resultMap.get("metadata"); + Map metadataMap = jsonObjectToMap(jsonObject); + // 可以选择合并或替换 + resultMap.putAll(metadataMap); + resultMap.remove("metadata"); // 移除原始的 metadata 对象 } + // 获取相似度分数 Float score = result.getScore(); if (score != null) { - metadata.put("score", score); + resultMap.put("score", score); } - // 将所有字段添加到结果中 -// item.putAll(entity); - results.add(metadata); + + results.add(resultMap); } } + log.info("处理完成，共 {} 条搜索结果", results.size()); return results; } /** + * 处理 QueryResp 查询结果（全量查询） + */ + private List> processMilvusResults(QueryResp response) { + List> results = new ArrayList<>(); + if (response == null) { + log.warn("Milvus 响应为空"); + return results; + } + + List queryResults = response.getQueryResults(); + if (queryResults == null || queryResults.isEmpty()) { + log.warn("Milvus 查询结果为空"); + return results; + } + + // 遍历每个查询结果 + for (QueryResp.QueryResult queryResult : queryResults) { + Map resultMap = new HashMap<>(); + + // 获取实体字段数据 + Map entity = queryResult.getEntity(); + if (entity != null && !entity.isEmpty()) { + resultMap.putAll(entity); + } + + // 处理 metadata 字段（如果是 JSON 对象） + if (resultMap.containsKey("metadata") && resultMap.get("metadata") instanceof JsonObject) { + JsonObject jsonObject = (JsonObject) resultMap.get("metadata"); + Map metadataMap = jsonObjectToMap(jsonObject); + resultMap.putAll(metadataMap); + resultMap.remove("metadata"); + } + + results.add(resultMap); + } + + log.info("处理完成，共 {} 条查询结果", results.size()); + return results; + } + + /** * JsonObject 转 Map */ public static Map jsonObjectToMap(JsonObject jsonObject) { @@ -693,7 +878,7 @@ public class MilvusServiceImpl implements MilvusService { } vectorText.append(" | ").append(fieldArr[0]).append(":").append(sText); } - vectorText.append(" ").append(sceneName); +// vectorText.append(" ").append(sceneName).append("数据"); }