Milvus向量检索Demo初探

Milvus 是什么

Milvus 于 2019 年开源，致力于存储、索引和管理由深度神经网络学习与其他机器学习模型生成的海量 Embedding 向量。

Milvus 向量数据库专为向量查询与检索设计，能够为万亿级向量数据建立索引。与现有的主要用作处理结构化数据的关系型数据库不同，Milvus 在底层设计上就是为了处理由各种非结构化数据转换而来的 Embedding 向量而生。

Milvus 2.0 是一款云原生向量数据库，采用存储与计算分离的架构设计，所有组件均为无状态组件，极大地增强了系统弹性和灵活性。

Demo

from pymilvus import connections
# 1 创建连接
connections.connect(
	alias="default", 
	host='localhost', 
	port='19530'
	)
# 2 创建 collection，相当于一个表，类似 mongo 中 collection 概念
from pymilvus import CollectionSchema, FieldSchema, DataType
book_id = FieldSchema(
    name="book_id", 
    dtype=DataType.INT64, 
    is_primary=True, 
    )
word_count = FieldSchema(
    name="word_count", 
    dtype=DataType.INT64,  
    )
book_intro = FieldSchema(
    name="book_intro", 
    dtype=DataType.FLOAT_VECTOR, 
    dim=2
    )
schema = CollectionSchema(
    fields=[book_id, word_count, book_intro], 
    description="Test book search"
    )
collection_name = "book" #类似于表名
from pymilvus import Collection
collection = Collection(
    name=collection_name, 
    schema=schema, 
    using='default', # milvus server name
    shards_num=2,
    consistency_level="Strong"
    )


# 3 为 collections 创建分区，默认为 collection创建两个分区
collection = Collection("book")      # Get an existing collection.
collection.create_partition("novel")

# 4 插入数据到 collection
import random
# data[0] book_id, data[1] word_count, data[2]向量，总共 2000个 item 
data = [[i for i in range(2000)],
		[i for i in range(10000, 12000)],
    	[[random.random() for _ in range(2)] for _ in range(2000)],
		]
mr = collection.insert(data)

# 5 创建索引
# 索引参数
index_params = {
        "metric_type":"L2",
        "index_type":"IVF_FLAT", #加聚类的倒排索引
        "params":{"nlist":1024} # 聚成 1024 个类
    }
# 向 collections 中插入索引
collection.create_index(
    field_name="book_intro", # collection中的存储向量的字段
    index_params=index_params
    )
# 6 向量搜索，进行搜索之前要先将 collection加载到内存
collection = Collection("book")      # Get an existing collection.
collection.load()
import time

# 准备搜索参数
search_params = {"metric_type": "L2", "params": {"nprobe": 100}} # 搜索最相似的 100 个item
start = time.time()
results = collection.search(
	data=[[0.1, 0.2]], # query 向量，这个列子中向量维度是2
	anns_field="book_intro", 
	param=search_params, 
	limit=10, # 返回 topK
	expr=None,
	consistency_level="Strong"
)
end = time.time()
# 搜索结果
print(results[0].ids, results[0].distances, end - start)
'''
([1073, 448, 489, 23, 752, 1496, 10, 1965, 1534, 107],
 [3.3897744287969545e-05, 0.000469639606308192, 0.0006169890984892845, 0.0006427020998671651, 0.000682430574670434, 0.000686873565427959, 0.0007882573409006, 0.0013403998455032706, 0.0016471443232148886, 0.0016562778037041426],
 0.2413792610168457)
'''


# 7 混合向量检索，在向量检索基础上，支持标量的指定范围查询
search_param = {
    "data": [[0.1, 0.2]],
    "anns_field": "book_intro",
    "param": {"metric_type": "L2", "params": {"nprobe": 10}},
    "limit": 2,
    "expr": "word_count <= 11000", # 指定查询表达式
}
res = collection.search(**search_param)
print(res[0])
'''
['(distance: 0.000469639606308192, id: 448)', '(distance: 0.0006169890984892845, id: 489)']
'''