python milvus 案例

**"""
Milvus 向量数据库学习 Demo
涵盖:

  1. 连接 Milvus(MilvusClient)
  2. 创建 Collection(schema + 索引)
  3. 插入数据(含向量 + 标量字段)
  4. 向量相似度搜索(search)
  5. 标量字段过滤查询(query)
  6. 按 ID 查询(get)
  7. 更新数据(vector + 标量字段,upsert 方式)
  8. 删除数据
  9. 删除 Collection
    """**
python 复制代码
"""
安裝依賴
pip install numpy pymilvus -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn --timeout 120
"""

"""
Milvus 向量数据库学习 Demo
涵盖:
  1. 连接 Milvus(MilvusClient)
  2. 创建 Collection(schema + 索引)
  3. 插入数据(含向量 + 标量字段)
  4. 向量相似度搜索(search)
  5. 标量字段过滤查询(query)
  6. 按 ID 查询(get)
  7. 更新数据(vector + 标量字段,upsert 方式)
  8. 删除数据
  9. 删除 Collection
"""
import time
from typing import Optional, List

"""#向量归一化"""
import numpy as np
from pymilvus import MilvusClient, DataType
def normalize(v):
    norm = np.linalg.norm(v)
    return [x / norm for x in v]


"""
# ===========================
# 配置(按项目实际地址修改)
# ===========================
"""
MILVUS_HOST = "127.0.0.1"
MILVUS_PORT = 19530
"""库"""
MILVUS_DB = "study"
"""表"""
COLLECTION_NAME = "t_employee"
"""维度信息"""
VECTOR_DIM = 4

"""
# ===========================
# 1. 连接 Milvus
# ===========================
"""
def connect_milvus() -> Optional[MilvusClient]:
    """
    创建 MilvusClient 连接
    项目模式:MilvusClient(uri="http://host:port", db_name="...")
    """
    try:
        client = MilvusClient(
            uri=f"http://{MILVUS_HOST}:{MILVUS_PORT}",
            db_name=MILVUS_DB,
        )
        # ping 验证连接
        client.list_collections()
        print(f"[OK] 连接 Milvus 成功: {MILVUS_HOST}:{MILVUS_PORT}, db={MILVUS_DB}")
        return client
    except Exception as e:
        print(f"[ERROR] 连接 Milvus 失败: {e}")
        return None


"""
# ===========================
# 2. 创建 Collection
# ===========================
"""
def create_collection(client: MilvusClient):
    """
    创建 t_employee 集合
    字段设计:
      - id: VARCHAR 主键
      - name: VARCHAR 员工姓名
      - department: VARCHAR 部门
      - salary: INT64 薪资
      - vector: FLOAT_VECTOR 向量(4维,demo 用)
      - metadata: JSON 额外信息
    """
    # 如果已存在,先删除(demo 需要)
    if client.has_collection(collection_name=COLLECTION_NAME):
        print(f"[INFO] 集合已存在,先删除: {COLLECTION_NAME}")
        client.drop_collection(collection_name=COLLECTION_NAME)

    """
    # 创建 Schema
    # 手动指定主键
    # 允许动态字段(项目用 True)
    """
    schema = client.create_schema(
        auto_id=False,
        enable_dynamic_field=True,
    )

    """
    # 添加字段(对标项目 add_field())
    """
    schema.add_field(
        field_name="id",
        datatype=DataType.VARCHAR,
        max_length=100,
        is_primary=True,
        nullable=False,
    )
    schema.add_field(
        field_name="name",
        datatype=DataType.VARCHAR,
        max_length=200,
        nullable=True,
    )
    schema.add_field(
        field_name="department",
        datatype=DataType.VARCHAR,
        max_length=200,
        nullable=True,
    )
    schema.add_field(
        field_name="salary",
        datatype=DataType.INT64,
        nullable=True,
    )
    """# 项目里可用大模型解析成1024 维"""
    schema.add_field(
        field_name="vector",
        datatype=DataType.FLOAT_VECTOR,
        dim=VECTOR_DIM,
    )
    """# 代表扩展字段"""
    schema.add_field(
        field_name="metadata",
        datatype=DataType.JSON,
        nullable=True,
    )

    """# 创建集合"""
    client.create_collection(
        collection_name=COLLECTION_NAME,
        schema=schema,
    )
    print(f"[OK] 创建集合成功: {COLLECTION_NAME}")

    """# 创建索引"""
    index_params = client.prepare_index_params()

    """# 标量字段索引(用于过滤加速)"""
    index_params.add_index(
        field_name="department",
        index_type="INVERTED",
        index_name="department_index",
        params={},
    )

    """
    # 向量字段索引(用于相似度搜索)
    # 配置:AUTOINDEX + IP(内积),AUTOINDEX 自动选择最优索引算法
    """
    index_params.add_index(
        field_name="vector",
        index_type="AUTOINDEX",
        index_name="vector_index",
        metric_type="IP",  # 内积(Inner Product)
    )

    client.create_index(
        collection_name=COLLECTION_NAME,
        index_params=index_params,
    )




def refresh_memory_data(client):
    """
    # 加载集合到内存(查询前必须加载)
    """
    client.release_collection(collection_name=COLLECTION_NAME)
    client.load_collection(collection_name=COLLECTION_NAME)
    """
    # 新增:阻塞等待加载完成
    """
    while True:
        state = client.get_load_state(collection_name=COLLECTION_NAME)
        if state["state"] == 3:  # 3 = Loaded
            break
        print("等待集合加载中...")
        time.sleep(0.2)
    print(f"[OK] 索引创建并加载完成")


"""
# ===========================
# 3. 插入数据
# ===========================
"""
def insert_employees(client: MilvusClient):
    """
    插入员工数据
    项目模式:client.insert(collection_name=..., data=[dict1, dict2, ...])
    """
    employees = [
        {
            "id": "emp_001",
            "name": "张三",
            "department": "技术部",
            "salary": 25000,
            "vector": normalize([0.12, 0.34, 0.56, 0.78]),  # 模拟向量
            "metadata": {"level": "高级", "skills": ["Python", "ML"]},
        },
        {
            "id": "emp_002",
            "name": "李四",
            "department": "技术部",
            "salary": 20000,
            "vector": normalize([0.15, 0.32, 0.54, 0.80]),
            "metadata": {"level": "中级", "skills": ["Java", "Spring"]},
        },
        {
            "id": "emp_003",
            "name": "王五",
            "department": "市场部",
            "salary": 18000,
            "vector": normalize([0.90, 0.78, 0.12, 0.34]),
            "metadata": {"level": "高级", "skills": ["营销", "品牌"]},
        },
        {
            "id": "emp_004",
            "name": "赵六",
            "department": "市场部",
            "salary": 15000,
            "vector": normalize([0.85, 0.80, 0.15, 0.30]),
            "metadata": {"level": "中级", "skills": ["广告", "策划"]},
        },
        {
            "id": "emp_005",
            "name": "陈七",
            "department": "人事部",
            "salary": 16000,
            "vector": normalize([0.45, 0.55, 0.65, 0.75]),
            "metadata": {"level": "初级", "skills": ["招聘", "培训"]},
        },
    ]

    result = client.insert(
        collection_name=COLLECTION_NAME,
        data=employees,
    )
    print(f"[OK] 插入 {len(employees)} 条员工数据, insert_count={result['insert_count']}")


"""
# ===========================
# 4. 向量相似度搜索
# ===========================
"""
def search_similar_employees(client: MilvusClient, query_vector: List[float], top_k: int = 3):
    """
    向量相似度搜索(核心功能!)
    项目模式:client.search(collection_name=..., data=[vectors], ...)
    搜索逻辑:
      1. 用 query_vector(用户问题向量化后的结果)去 Milvus 中找最相似的向量
      2. 返回 top_k 条最匹配的结果

    data  # 查询向量(支持批量,传入列表套列表)
    anns_field # 在哪个向量字段上搜索
    search_params # 搜索参数
    output_fields # 返回的标量字段
    """
    result = client.search(
        collection_name=COLLECTION_NAME,
        data=[query_vector],  # 查询向量(支持批量,传入列表套列表)
        anns_field="vector",  # 在哪个向量字段上搜索
        search_params={"nprobe": 10},  # 搜索参数
        limit=top_k,  # 返回 top-K
        output_fields=["name", "department", "salary", "metadata"],  # 返回的标量字段
    )

    print(f"\n--- 向量搜索 (top_k={top_k}) ---")
    print(f"  查询向量: {query_vector}")
    for i, hit in enumerate(result[0]):
        print(f"  [{i + 1}] id={hit['id']}, 距离={hit['distance']:.4f}")
        print(f"      姓名: {hit['entity']['name']}, 部门: {hit['entity']['department']}, "
              f"薪资: {hit['entity']['salary']}")
    return result


"""
# ===========================
# 5. 标量字段过滤查询
# ===========================
"""
def query_by_department(client: MilvusClient, department: str):
    """
    按部门过滤查询(标量字段过滤)
    项目模式:client.query(collection_name=..., filter=..., output_fields=...)
    """
    result = client.query(
        collection_name=COLLECTION_NAME,
        filter=f'department == "{department}"',
        output_fields=["id", "name", "department", "salary", "metadata"],
    )

    print(f"\n--- 按部门查询: {department} ---")
    print(f"  找到 {len(result)} 人:")
    for item in result:
        print(f"  - {item['name']} | 薪资 {item['salary']} | {item['metadata']}")
    return result


def query_by_salary_range(client: MilvusClient, min_salary: int, max_salary: int):
    """
    按薪资范围过滤查询
    """
    result = client.query(
        collection_name=COLLECTION_NAME,
        filter=f"salary >= {min_salary} && salary <= {max_salary}",
        output_fields=["id", "name", "department", "salary"],
    )

    print(f"\n--- 按薪资范围查询: {min_salary} ~ {max_salary} ---")
    print(f"  找到 {len(result)} 人:")
    for item in result:
        print(f"  - {item['name']} ({item['department']}) 薪资 {item['salary']}")
    return result


"""
# ===========================
# 6. 按 ID 查询
# ===========================
"""
def get_by_ids(client: MilvusClient, ids: List[str]):
    """
    按主键 ID 批量查询
    """
    result = client.get(
        collection_name=COLLECTION_NAME,
        ids=ids,
        output_fields=["id", "name", "department", "salary", "vector"],
    )

    print(f"\n--- 按 ID 查询: {ids} ---")
    for item in result:
        print(f"  - {item['id']}: {item['name']} | {item['department']} | 向量前3维: {item['vector'][:3]}")
    return result


"""
# ===========================
# 7. 删除数据
# ===========================
"""
def delete_by_filter(client: MilvusClient, filter_expr: str):
    """
    按条件删除(对标项目中的 filter 删除)
    项目模式:client.delete(collection_name=..., filter=...)
    """
    result = client.delete(
        collection_name=COLLECTION_NAME,
        filter=filter_expr,
    )
    print(f"\n--- 删除: filter='{filter_expr}' ---")
    print(f"  删除数量: {result['delete_count']}")
    return result


"""
# ===========================
# 8. 更新数据(含向量字段,upsert 方式)
# ===========================
"""
def update_employee(client: MilvusClient, emp_id: str, update_data: dict):
    """
    更新员工数据(支持向量字段 + 标量字段同时更新)

    Milvus 没有直接 update API,用 upsert 实现:
      主键相同则覆盖更新,主键不存在则新增。
    """
    # 1. 查出当前数据(保留未被覆盖的字段)
    existing = client.get(
        collection_name=COLLECTION_NAME,
        ids=[emp_id],
        output_fields=["id", "name", "department", "salary", "vector", "metadata"],
    )
    if not existing:
        print(f"[WARN] 未找到员工: {emp_id},跳过更新")
        return

    # 2. 合并更新字段
    merged = dict(existing[0])
    merged.update(update_data)

    # 3. upsert 覆盖写入(主键相同 = 更新)
    result = client.upsert(
        collection_name=COLLECTION_NAME,
        data=[merged],
    )
    print(f"\n--- 更新员工: {emp_id} ---")
    for k, v in update_data.items():
        if k == "vector":
            print(f"  vector 更新为: {[round(x, 4) for x in v]}")
        else:
            print(f"  {k} 更新为: {v}")
    print(f"  upsert_count={result['upsert_count']}")
    return result


def update_employee_salary_and_vector(client: MilvusClient, emp_id: str,
                                      new_salary: int, new_vector: List[float]):
    """同时更新薪资 + 向量字段的便捷方法"""
    return update_employee(client, emp_id, {
        "salary": new_salary,
        "vector": normalize(new_vector),
    })


"""
# ===========================
# 9. 主程序
# ===========================
"""
def main():
    print("=" * 60)
    print("Milvus 向量数据库 Demo --- 基于项目模式")
    print("=" * 60)

    """# 1. 连接"""
    client = connect_milvus()
    if not client:
        return

    """# 2. 创建 Collection"""
    create_collection(client)
    refresh_memory_data(client)
    """# 3. 插入数据"""
    insert_employees(client)
    refresh_memory_data(client)
    """等待数据加载到内存"""
    """
    # 4. 向量搜索:找"最像"张三的人
    #    张三的向量是 [0.12, 0.34, 0.56, 0.78]
    #    李四和他最像(同技术部),王五(市场部)最不像
    """
    search_similar_employees(client, query_vector=normalize([0.12, 0.34, 0.56, 0.78]), top_k=3)

    """# 5. 向量搜索:找"最像"王五的人"""
    search_similar_employees(client, query_vector=normalize([0.90, 0.78, 0.12, 0.34]), top_k=3)

    """# 6. 按部门过滤"""
    query_by_department(client, "技术部")

    """# 7. 按薪资范围过滤"""
    query_by_salary_range(client, 18000, 30000)

    """# 8. 按 ID 查询"""
    get_by_ids(client, ["emp_001", "emp_003", "emp_005"])

    """# 9. 更新数据(含向量字段)
    #    张三(emp_001)调薪到 30000,同时向量改为 [0.88,0.82,0.14,0.32](接近市场部方向)
    #    更新前:张三向量 [0.12,0.34,0.56,0.78] → 搜索最像"张三"时李四排第一
    #    更新后:张三向量 [0.88,0.82,0.14,0.32] → 搜索最像"张三"时王五排第一
    """
    update_employee_salary_and_vector(
        client,
        emp_id="emp_001",
        new_salary=30000,
        new_vector=[0.88, 0.82, 0.14, 0.32],
    )
    """"""
    """
    # upsert后手动刷盘,增量更新内存,无需全量重载
    # client.load_collection(collection_name=COLLECTION_NAME) 
    load_collection:磁盘全量数据 → 内存重建完整索引(清空旧内存数据)
    """
    refresh_memory_data(client)  # 重新加载最新数据到内存
    """# 验证更新"""
    get_by_ids(client, ["emp_001"])
    search_similar_employees(client, query_vector=normalize([0.12,0.34,0.56,0.78]), top_k=3)

    """# 10. 删除数据"""
    delete_by_filter(client, 'department == "人事部"')

    """# 验证删除"""
    remaining = client.query(
        collection_name=COLLECTION_NAME,
        filter="id != ''",
        output_fields=["id", "name"],
    )
    print(f"\n当前剩余数据: {len(remaining)} 条")
    for item in remaining:
        print(f"  - {item['id']}: {item['name']}")

    # 11. 清理
    print("\n--- 清理 ---")
    client.drop_collection(collection_name=COLLECTION_NAME)
    print(f"[OK] 删除集合: {COLLECTION_NAME}")

    print("\n" + "=" * 60)
    print("Demo 运行完毕!")
    print("=" * 60)


if __name__ == "__main__":
    main()
相关推荐
Gauss松鼠会2 小时前
【GaussDB】GaussDB锁阻塞源头查询
java·开发语言·前端·数据库·算法·gaussdb·经验总结
mingo_敏2 小时前
DeepAgents : 后端(Backends)
java·开发语言
mabing9932 小时前
Qt QMessageBox、QDialogButtonBox中英文翻译动态切换
开发语言·qt
BerryS3N2 小时前
Java 后端转型大模型:Demo 能跑不等于能上线
java·人工智能·python·java后端·spring ai·langchain4j·大模型转型
Sagittarius_A*2 小时前
【RCELABS】Level 17~18 —— PHP命令执行函数与环境变量注入
开发语言·安全·web安全·靶场·php·rce
冻柠檬飞冰走茶2 小时前
PTA基础编程题目集 7-19 支票面额(C语言实现)
c语言·开发语言·数据结构·算法
制造数据与AI践行者老蒋2 小时前
智联工坊实战:从“金鱼记忆”到“记住了”:给制造Agent装上记忆芯片的完整指南
人工智能·python·langchain·制造
某林2122 小时前
ros从底层硬件到 Web 端部署
python·机器人·硬件架构·ros2
破碎的南瓜3 小时前
靶场中使用到的php函数以及每种漏洞的防御方式
开发语言·php