如何通过Python SDK在Collection中进行相似性检索

77 阅读4分钟

本文介绍如何通过Python SDK在Collection中进行相似性检索。


前提条件

接口定义

Python示例:

Collection.query(
    vector: Optional[Union[List[Union[int, float]], np.ndarray, VectorQuery, Dict[str, VectorQuery]]] = None,   
    id: Optional[str] = None,
    topk: int = 10,
    filter: Optional[str] = None,
    include_vector: bool = False,
    partition: Optional[str] = None,
    output_fields: Optional[List[str]] = None,
    sparse_vector: Optional[Dict[int, float]] = None,
    rerank: Optional[BaseRanker] = None,
    async_req: False
) -> DashVectorResponse

使用示例

说明

  1. 需要使用您的api-key替换示例中的YOUR_API_KEY、您的Cluster Endpoint替换示例中的YOUR_CLUSTER_ENDPOINT,代码才能正常运行。
  2. 本示例需要参考新建Collection-使用示例提前创建好名称为quickstart的Collection,并参考插入Doc提前插入部分数据。

Python示例:

import dashvector
import numpy as np

client = dashvector.Client(
    api_key='YOUR_API_KEY',
    endpoint='YOUR_CLUSTER_ENDPOINT'
)
collection = client.get(name='quickstart')

根据向量进行相似性检索

Python示例:

ret = collection.query(
    vector=[0.1, 0.2, 0.3, 0.4]
)
# 判断query接口是否成功
if ret:
    print('query success')
    print(len(ret))
    for doc in ret:
        print(doc)
        print(doc.id)
        print(doc.vector)
        print(doc.fields)

ret = collection.query(
    vector=[0.1, 0.2, 0.3, 0.4],
    topk=100,
    output_fields=['name', 'age'], # 仅返回name、age这2个Field
    include_vector=True
)

根据主键(对应的向量)进行相似性检索

Python示例:

ret = collection.query(
    id='1'
)
# 判断query接口是否成功
if ret:
    print('query success')
    print(len(ret))
    for doc in ret:
        print(doc)
        print(doc.id)
        print(doc.vector)
        print(doc.fields)

ret = collection.query(
    id='1',
    topk=100,
    output_fields=['name', 'age'], # 仅返回name、age这2个Field
    include_vector=True
)

带过滤条件的相似性检索

Python示例:

# 根据向量或者主键进行相似性检索 + 条件过滤
ret = collection.query(
    vector=[0.1, 0.2, 0.3, 0.4],   # 向量检索,也可设置主键检索
    topk=100,
    filter='age > 18',             # 条件过滤,仅对age > 18的Doc进行相似性检索
    output_fields=['name', 'age'], # 仅返回name、age这2个Field
    include_vector=True
)

带有Sparse Vector的向量检索

说明

Sparse Vector(稀疏向量)可用于关键词权重表示,实现带关键词感知能力的向量检索

Python示例:

# 根据向量进行相似性检索 + 稀疏向量
ret = collection.query(
    vector=[0.1, 0.2, 0.3, 0.4],   # 向量检索
    sparse_vector={1: 0.3, 20: 0.7}
)

通过过滤条件进行匹配查询

Python示例:

# 支持向量和主键都不传入,那么只进行条件过滤
ret = collection.query(
    topk=100,
    filter='age > 18',             # 条件过滤,仅对age > 18的Doc进行匹配查询
    output_fields=['name', 'age'], # 仅返回name、age这2个Field
    include_vector=True
)

向量检索高级参数

说明

详情可参考 向量检索高级参数

ret = collection.query(vector=VectorQuery(vector=[0.1,0.2,0.3,0.4], is_linear=False, ef=1000, radius=0.53))

多向量检索

说明

详情可参考 多向量检索

RrfRanker 示例:

title_vector = [0.1, 0.2, 0.3, 0.4]
content_vector = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
vectors = {
    "title": VectorQuery(vector=title_vector, num_candidates=10),
    "content": VectorQuery(vector=content_vector),
}
ret = collection.query(
    vector=vectors,
    include_vector=False,
    rerank=RrfRanker(rank_constant=100),
)
assert ret
print(ret.output)

WeightedRanker 示例:

title_vector = [0.1, 0.2, 0.3, 0.4]
content_vector = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
vectors = {
    "title": VectorQuery(vector=title_vector, num_candidates=10),
    "content": VectorQuery(vector=content_vector),
}
ret = collection.query(
    vector=vectors,
    include_vector=False,
    rerank=WeightedRanker(weights={"title": 1.0, "content": 1.0}),
)
assert ret
print(ret.output)

使用多向量的一个向量执行检索:

title_vector = [0.1, 0.2, 0.3, 0.4]
ret = collection.query(vector={"title": VectorQuery(vector=title_vector, is_linear=False, ef=1000, radius=0.53)})

入参描述

说明

vectorid两个入参需要二选一使用,如都不传入,则仅完成条件过滤。

参数类型默认值说明
vector(可选)Optional[Union[List[Union[int, float]], np.ndarray, VectorQuery, Dict[str, VectorQuery]]]None一个或多个向量数据,或者使用VectorQuery执行高级检索,详情参考向量检索高级参数
id(可选)Optional[str]None主键,表示根据主键对应的向量进行相似性检索
topk(可选)int10返回topk相似性结果
filter(可选)Optional[str]None过滤条件,需满足SQL where子句规范,详见条件过滤检索
include_vector(可选)boolFalse是否返回向量数据
partition(可选)Optional[str]NonePartition名称
output_fields(可选)Optional[List[str]]None默认返回所有Fields
sparse_vector(可选)Optional[Dict[int, float]]None稀疏向量
rerankOptional[BaseRanker]None融合排序参数,详见多向量检索
async_req(可选)boolFalse是否异步

出参描述

说明

返回结果为DashVectorResponse对象,DashVectorResponse对象中可获取本次操作结果信息,如下表所示。

字段类型描述示例
codeint返回值,参考返回状态码说明0
messagestr返回消息success
request_idstr请求唯一id19215409-ea66-4db9-8764-26ce2eb5bb99
outputList[Doc]返回相似性检索结果
usageRequestUsage对Serverless实例(按量付费)集合的Doc检索请求,成功后返回实际消耗的读请求单元数