Administrator / xlyAi

package com.xly.milvus.test;

import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.collection.response.GetCollectionStatsResp;
import io.milvus.v2.service.index.request.CreateIndexReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.QueryReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import io.milvus.v2.service.vector.response.QueryResp;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * Milvus Java SDK 2.6.15 完整测试代码
 *
 * 功能：
 * 1. 连接Milvus服务器
 * 2. 创建集合（包含多种字段类型）
 * 3. 创建向量索引
 * 4. 插入测试数据
 * 5. 查询数据（标量过滤）
 * 6. 向量相似性搜索
 * 7. 获取集合统计信息
 * 8. 清理资源
 *
 * @author xly
 * @version 1.0
 */
public class MilvusCompleteTest {

    // ==================== 配置参数 ====================
    private static final String MILVUS_HOST = "121.43.128.225";  // 您的Milvus服务器IP
    private static final int MILVUS_PORT = 19530;                // Milvus服务端口
    private static final String COLLECTION_NAME = "complete_test"; // 集合名称
    private static final int VECTOR_DIM = 128;                   // 向量维度
    private static final int INSERT_COUNT = 15;                  // 插入数据条数

    private MilvusClientV2 client;

    public static void main(String[] args) {
        MilvusCompleteTest test = new MilvusCompleteTest();
        try {
            // 执行测试流程
            test.run();
        } catch (Exception e) {
            System.err.println("❌ 测试过程中发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 执行完整的测试流程
     */
    private void run() {
        printHeader("Milvus Java SDK 2.6.15 完整测试");

        // 1. 连接Milvus
        connect();

        try {
            // 2. 清理旧集合
            cleanup();

            // 3. 创建集合
            createCollection();

            // 4. 创建索引
            createIndex();

            // 5. 插入数据
            insertData();

            // 6. 查询数据
            queryData();

            // 7. 向量搜索
            searchData();

            // 8. 获取集合统计
            getCollectionStats();

            printSuccess("所有测试完成！");

        } catch (Exception e) {
            System.err.println("❌ 测试失败: " + e.getMessage());
            throw e;
        } finally {
            // 9. 清理资源
            releaseAndClose();
        }
    }

    /**
     * 连接Milvus服务器
     */
    /**
     * 连接Milvus服务器
     */
    private void connect() {
        printSection("连接Milvus服务器");

        try {
            ConnectConfig connectConfig = ConnectConfig.builder()
                    .uri("http://" + MILVUS_HOST + ":" + MILVUS_PORT)
                    .connectTimeoutMs(TimeUnit.SECONDS.toMillis(30))      // 连接超时30秒
                    .keepAliveTimeMs(TimeUnit.MINUTES.toMillis(5))        // 保持连接5分钟
                    .keepAliveTimeoutMs(TimeUnit.SECONDS.toMillis(5))     // keep-alive超时5秒
                    .keepAliveWithoutCalls(true)                           // 无调用时保持连接
                    .rpcDeadlineMs(TimeUnit.SECONDS.toMillis(10))         // RPC超时10秒
                    .enablePrecheck(true)                                   // 启用预检查
                    .build();

            client = new MilvusClientV2(connectConfig);

            // 验证连接
            String serverVersion = client.getServerVersion();
            System.out.println("✅ 连接成功！");
            System.out.println("   - 服务器地址: " + MILVUS_HOST + ":" + MILVUS_PORT);
            System.out.println("   - 服务器版本: " + serverVersion);

        } catch (Exception e) {
            System.err.println("❌ 连接失败: " + e.getMessage());
            throw new RuntimeException("无法连接到Milvus服务器", e);
        }
    }

    /**
     * 清理已存在的集合
     */
    private void cleanup() {
        printSection("清理环境");

        HasCollectionReq hasCollectionReq = HasCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();

        boolean exists = client.hasCollection(hasCollectionReq);
        if (exists) {
            DropCollectionReq dropCollectionReq = DropCollectionReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .build();
            client.dropCollection(dropCollectionReq);
            System.out.println("✅ 已删除旧集合: " + COLLECTION_NAME);
        } else {
            System.out.println("⏭️ 无需清理，集合不存在");
        }
    }

    /**
     * 创建集合
     */
    private void createCollection() {
        printSection("创建集合");

        // 定义字段列表
        List<CreateCollectionReq.FieldSchema> fieldSchemas = Arrays.asList(
                // 1. 主键字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("id")
                        .dataType(DataType.Int64)
                        .isPrimaryKey(true)
                        .autoID(false)
                        .description("主键ID")
                        .build(),

                // 2. 向量字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("vector")
                        .dataType(DataType.FloatVector)
                        .dimension(VECTOR_DIM)
                        .description("向量字段，用于相似性搜索")
                        .build(),

                // 3. 标题字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("title")
                        .dataType(DataType.VarChar)
                        .maxLength(200)
                        .description("标题")
                        .build(),

                // 4. 内容字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("content")
                        .dataType(DataType.VarChar)
                        .maxLength(1000)
                        .description("详细内容")
                        .build(),

                // 5. 分类字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("category")
                        .dataType(DataType.VarChar)
                        .maxLength(50)
                        .description("分类标签")
                        .build(),

                // 6. 得分字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("score")
                        .dataType(DataType.Float)
                        .description("评分")
                        .build(),

                // 7. 时间戳字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("create_time")
                        .dataType(DataType.Int64)
                        .description("创建时间戳")
                        .build(),

                // 8. 标签数组字段
                CreateCollectionReq.FieldSchema.builder()
                        .name("tags")
                        .dataType(DataType.Array)
                        .elementType(DataType.VarChar)
                        .maxCapacity(10)
                        .description("标签数组")
                        .build()
        );

        // 创建集合schema
        CreateCollectionReq.CollectionSchema schema =
                CreateCollectionReq.CollectionSchema.builder()
                        .fieldSchemaList(fieldSchemas)
                        .enableDynamicField(true)  // 启用动态字段
                        .build();

        // 创建集合请求
        CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .collectionSchema(schema)
                .consistencyLevel(ConsistencyLevel.BOUNDED)  // 有界一致性
                .build();

        // 执行创建
        client.createCollection(createCollectionReq);
        System.out.println("✅ 集合创建成功: " + COLLECTION_NAME);
        System.out.println("   - 向量维度: " + VECTOR_DIM);
        System.out.println("   - 字段数量: " + fieldSchemas.size());
    }

    /**
     * 创建索引
     */
    /**
     * 创建索引
     */
    private void createIndex() {
        printSection("创建索引");

        // 创建索引参数
        Map<String, Object> extraParams = new HashMap<>();
        extraParams.put("nlist", 1024);

        IndexParam indexParam = IndexParam.builder()
                .fieldName("vector")
                .indexType(IndexParam.IndexType.IVF_FLAT)
                .metricType(IndexParam.MetricType.L2)
                .extraParams(extraParams)
                .build();

        // 创建索引请求
        CreateIndexReq createIndexReq = CreateIndexReq.builder()
                .collectionName(COLLECTION_NAME)
                .indexParams(Collections.singletonList(indexParam))
                .build();

        // 执行创建
        client.createIndex(createIndexReq);
        System.out.println("✅ 索引创建成功");
        System.out.println("   - 索引类型: IVF_FLAT");
        System.out.println("   - 度量方式: L2 (欧氏距离)");

        // 修正：使用 LoadCollectionReq 加载集合到内存
        LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();
        client.loadCollection(loadCollectionReq);
        System.out.println("✅ 集合已加载到内存");
    }

    /**
     * 插入测试数据
     */
    private void insertData() {
        printSection("插入测试数据");

        List<JsonObject> rows = new ArrayList<>();
        Random random = new Random(42);  // 固定种子，保证可重复性

        // 测试数据
        String[] categories = {"科技", "体育", "娱乐", "教育", "财经"};
        String[] titles = {
                "人工智能发展现状与未来趋势", "机器学习入门实战教程",
                "深度学习在图像识别中的应用", "NBA季后赛精彩回顾",
                "世界杯预选赛最新战况", "奥运会筹备工作进展",
                "热门电影推荐排行榜", "音乐榜单TOP10",
                "综艺节目收视率分析", "在线教育行业发展报告",
                "编程语言流行趋势", "高效学习方法分享",
                "股市投资策略分析", "区块链技术应用前景",
                "数字货币市场动态"
        };

        long currentTime = System.currentTimeMillis() / 1000;

        // 生成测试数据
        for (int i = 1; i <= INSERT_COUNT; i++) {
            JsonObject row = new JsonObject();

            // ID
            row.addProperty("id", (long) i);

            // 向量
            JsonArray vectorArray = new JsonArray();
            for (int j = 0; j < VECTOR_DIM; j++) {
                vectorArray.add(random.nextFloat());
            }
            row.add("vector", vectorArray);

            // 标题
            int index = (i - 1) % titles.length;
            row.addProperty("title", titles[index]);

            // 内容
            row.addProperty("content", "这是第" + i + "条测试数据的详细内容。用于测试Milvus的插入和查询功能。");

            // 分类
            row.addProperty("category", categories[random.nextInt(categories.length)]);

            // 得分
            row.addProperty("score", random.nextFloat() * 100);

            // 创建时间
            row.addProperty("create_time", currentTime - random.nextInt(86400) * i);

            // 标签数组
            JsonArray tagsArray = new JsonArray();
            tagsArray.add("tag_" + random.nextInt(5));
            tagsArray.add("tag_" + random.nextInt(5));
            row.add("tags", tagsArray);

            rows.add(row);
        }

        // 批量插入
        InsertReq insertReq = InsertReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(rows)
                .build();

        InsertResp insertResp = client.insert(insertReq);
        System.out.println("✅ 成功插入 " + insertResp.getInsertCnt() + " 条数据");
        System.out.println("   - 数据预览:");

        // 显示前3条数据的部分信息
        for (int i = 0; i < Math.min(3, rows.size()); i++) {
            JsonObject row = rows.get(i);
            System.out.printf("     [%d] ID: %d, 标题: %s, 分类: %s, 得分: %.2f%n",
                    i + 1,
                    row.get("id").getAsLong(),
                    row.get("title").getAsString(),
                    row.get("category").getAsString(),
                    row.get("score").getAsFloat());
        }
    }

    /**
     * 查询数据（标量过滤）
     */
    private void queryData() {
        printSection("标量查询");

        // 查询得分大于50的数据
        QueryReq queryReq = QueryReq.builder()
                .collectionName(COLLECTION_NAME)
                .filter("score > 50")  // 过滤条件
                .outputFields(Arrays.asList("id", "title", "category", "score"))
                .limit(10)
                .build();

        QueryResp queryResp = client.query(queryReq);

        List<QueryResp.QueryResult> results = queryResp.getQueryResults();
        System.out.println("查询条件: score > 50");
        System.out.println("查询结果: " + results.size() + " 条数据");

        if (!results.isEmpty()) {
            System.out.println("结果列表:");
            for (QueryResp.QueryResult result : results) {
                Map<String, Object> entity = result.getEntity();
                System.out.printf("  ID: %d | 标题: %s | 分类: %s | 得分: %.2f%n",
                        (Long) entity.get("id"),
                        entity.get("title"),
                        entity.get("category"),
                        (Float) entity.get("score"));
            }
        }
    }

    /**
     * 向量相似性搜索
     */
    private void searchData() {
        printSection("向量相似性搜索");

        // 创建查询向量
        List<Float> vectorList = new ArrayList<>();
        Random random = new Random();
        for (int j = 0; j < VECTOR_DIM; j++) {
            vectorList.add(random.nextFloat());
        }
        FloatVec queryVector = new FloatVec(vectorList);

        // 构建搜索参数
        Map<String, Object> searchParams = new HashMap<>();
        searchParams.put("nprobe", 10);

        // 构建搜索请求
        SearchReq searchReq = SearchReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(Collections.singletonList(queryVector))
                .annsField("vector")
                .limit(5)
                .outputFields(Arrays.asList("title", "category", "score", "tags"))
                .metricType(IndexParam.MetricType.L2)
                .searchParams(searchParams)
                .build();

        // 执行搜索
        SearchResp searchResp = client.search(searchReq);

        // 打印结果
        System.out.println("最相似的5条结果:");
        List<List<SearchResp.SearchResult>> results = searchResp.getSearchResults();

        if (!results.isEmpty() && !results.get(0).isEmpty()) {
            for (int i = 0; i < results.get(0).size(); i++) {
                SearchResp.SearchResult result = results.get(0).get(i);
                Map<String, Object> entity = result.getEntity();

                System.out.printf("  %d. ID: %d | 距离: %.4f%n",
                        i + 1,
                        (Long) result.getId(),
                        result.getScore());
                System.out.printf("     标题: %s | 分类: %s | 得分: %.2f%n",
                        entity.get("title"),
                        entity.get("category"),
                        (Float) entity.get("score"));
                System.out.printf("     标签: %s%n",
                        entity.get("tags"));
            }
        } else {
            System.out.println("  没有找到结果");
        }
    }

    /**
     * 获取集合统计信息
     */
    private void getCollectionStats() {
        printSection("集合统计信息");

        GetCollectionStatsReq statsReq = GetCollectionStatsReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();

        GetCollectionStatsResp statsResp = client.getCollectionStats(statsReq);

        System.out.println("集合名称: " + COLLECTION_NAME);
        System.out.println("实体数量: " + statsResp.getNumOfEntities());
    }

    /**
     * 释放集合并关闭连接
     */
    private void releaseAndClose() {
        printSection("清理资源");

        if (client != null) {
            try {
                // 释放集合
                ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                        .collectionName(COLLECTION_NAME)
                        .build();
                client.releaseCollection(releaseCollectionReq);
                System.out.println("✅ 集合已释放");

                // 关闭连接
                client.close();
                System.out.println("✅ 连接已关闭");

            } catch (Exception e) {
                System.err.println("⚠️ 清理资源时出错: " + e.getMessage());
            }
        }
    }

    /**
     * 打印标题
     */
    private void printHeader(String title) {
        System.out.println("\n" + "=".repeat(60));
        System.out.println("  " + title);
        System.out.println("=".repeat(60));
    }

    /**
     * 打印章节标题
     */
    private void printSection(String section) {
        System.out.println("\n" + "-".repeat(40));
        System.out.println("  " + section);
        System.out.println("-".repeat(40));
    }

    /**
     * 打印成功信息
     */
    private void printSuccess(String message) {
        System.out.println("\n" + "=".repeat(60));
        System.out.println("  ✅ " + message);
        System.out.println("=".repeat(60));
    }
}