# 下載Milvus 2.4.x 配置文件
wget https://github.com/milvus-io/milvus/releases/download/v2.4.4/milvus-standalone-docker-compose.yml -O docker-compose.yml
# 啟動服務(wù)
docker-compose up -d

pip install pymilvus>=2.4.0  # 安裝官方Python SDK
pip install numpy           # 用于生成模擬向量

# 1. 導(dǎo)入依賴包
from pymilvus import (
    connections,  # 連接Milvus
    Collection,   # 集合操作
    FieldSchema,  # 字段定義
    CollectionSchema,  # 集合結(jié)構(gòu)
    DataType,     # 數(shù)據(jù)類型
    utility       # 工具函數(shù)（檢查集合、刪除集合等）
)
import numpy as np
# ====================== 2. 連接Milvus服務(wù) ======================
# 默認連接本地Milvus，參數(shù)：host(地址)、port(端口)、alias(連接別名)
connections.connect(
    alias="default", 
    host="localhost", 
    port="19530"
)
print("? 成功連接Milvus服務(wù)")
# ====================== 3. 定義集合結(jié)構(gòu)（Schema） ======================
# 集合 = 表，字段 = 列，必須包含：主鍵字段 + 向量字段
# 字段1：主鍵ID（自增整數(shù)，唯一標識）
primary_field = FieldSchema(
    name="id", 
    dtype=DataType.INT64, 
    is_primary=True,  # 主鍵
    auto_id=True      # 自動生成ID，無需手動插入
)
# 字段2：向量字段（核心！存儲特征向量，512維）
vector_field = FieldSchema(
    name="embedding", 
    dtype=DataType.FLOAT_VECTOR, 
    dim=512  # 向量維度：必須和你的數(shù)據(jù)一致
)
# 字段3：標量字段（輔助過濾，比如分類標簽）
tag_field = FieldSchema(
    name="tag", 
    dtype=DataType.INT64
)
# 定義集合整體結(jié)構(gòu)
schema = CollectionSchema(
    fields=[primary_field, vector_field, tag_field],
    description="Milvus基礎(chǔ)案例：特征向量存儲與檢索"
)
# ====================== 4. 創(chuàng)建集合 ======================
collection_name = "milvus_demo"
# 如果集合已存在，先刪除（避免重復(fù)創(chuàng)建報錯）
if utility.has_collection(collection_name):
    utility.drop_collection(collection_name)
# 創(chuàng)建集合
collection = Collection(
    name=collection_name,
    schema=schema,
    using="default"  # 使用默認連接
)
print(f"? 成功創(chuàng)建集合：{collection_name}")
# ====================== 5. 創(chuàng)建向量索引（必須?。?======================
# Milvus必須建索引才能檢索，索引決定檢索速度和精度
# 入門推薦：IVF_FLAT（平衡速度和精度）
index_params = {
    "index_type": "IVF_FLAT",      # 索引類型
    "metric_type": "L2",           # 相似度計算方式：L2(歐氏距離)、IP(內(nèi)積)、COSINE(余弦)
    "params": {"nlist": 128}       # 索引參數(shù)（聚類數(shù)量）
}
# 為向量字段創(chuàng)建索引
collection.create_index(
    field_name="embedding",
    index_params=index_params
)
print("? 成功創(chuàng)建向量索引")
# ====================== 6. 加載集合到內(nèi)存（必須?。?======================
# 集合創(chuàng)建后默認在磁盤，必須加載到內(nèi)存才能讀寫
collection.load()
print("? 集合已加載到內(nèi)存")
# ====================== 7. 插入模擬數(shù)據(jù) ======================
# 生成1000條512維隨機向量（模擬文本/圖像特征）
num_entities = 1000
vectors = np.random.rand(num_entities, 512).tolist()  # 轉(zhuǎn)列表格式
# 生成標量標簽（0-9隨機分類）
tags = np.random.randint(0, 10, num_entities).tolist()
# 插入數(shù)據(jù)：字段順序和Schema定義一致
insert_result = collection.insert(
    data=[vectors, tags]  # 自動填充主鍵（auto_id=True）
)
print(f"? 成功插入 {num_entities} 條數(shù)據(jù)，主鍵ID：{insert_result.ids[:5]}...")
# 刷新數(shù)據(jù)（確保插入后立即可檢索）
collection.flush()
# ====================== 8. 核心功能：向量相似度檢索 ======================
# 生成1條查詢向量（和插入數(shù)據(jù)同維度）
search_vector = np.random.rand(1, 512).tolist()
# 檢索參數(shù)
search_params = {
    "metric_type": "L2",
    "params": {"nprobe": 10}  # 檢索精度：值越大越準，速度越慢
}
# 執(zhí)行檢索
results = collection.search(
    data=search_vector,        # 查詢向量
    anns_field="embedding",    # 要檢索的向量字段
    param=search_params,       # 檢索參數(shù)
    limit=5,                   # 返回Top5最相似的結(jié)果
    output_fields=["id", "tag"]# 同時返回的標量字段
)
# 打印檢索結(jié)果
print("\n?? 向量相似度檢索結(jié)果（Top5）：")
for hits in results:
    for hit in hits:
        print(f"ID: {hit.id}, 相似度距離: {hit.distance:.4f}, 標簽: {hit.entity.get('tag')}")
# ====================== 9. 進階：標量過濾 + 向量檢索 ======================
# 場景：只檢索【標簽=3】的向量中，最相似的Top3
filtered_results = collection.search(
    data=search_vector,
    anns_field="embedding",
    param=search_params,
    limit=3,
    expr="tag == 3",  # 標量過濾條件（SQL風格）
    output_fields=["id", "tag"]
)
print("\n?? 過濾后檢索結(jié)果（tag=3，Top3）：")
for hits in filtered_results:
    for hit in hits:
        print(f"ID: {hit.id}, 相似度距離: {hit.distance:.4f}, 標簽: {hit.entity.get('tag')}")
# ====================== 10. 數(shù)據(jù)管理 ======================
# 1. 刪除數(shù)據(jù)（根據(jù)主鍵ID）
if len(insert_result.ids) > 0:
    del_id = insert_result.ids[0]
    collection.delete(expr=f"id == {del_id}")
    print(f"\n? 成功刪除ID={del_id}的數(shù)據(jù)")
# 2. 釋放集合（從內(nèi)存卸載）
collection.release()
print("? 集合已從內(nèi)存釋放")
# 3. 刪除集合（清空所有數(shù)據(jù)）
# utility.drop_collection(collection_name)
# print("? 已刪除集合")
# ====================== 11. 斷開連接 ======================
connections.disconnect("default")
print("\n? 已斷開Milvus連接")