文档

分组检索Doc

更新时间:

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

前提条件

接口定义

Collection.query_group_by(
        self,
        vector: Optional[Union[List[Union[int, float]], np.ndarray]] = None,
        *,
        group_by_field: str,
        group_count: int = 10,
        group_topk: int = 10,
        id: Optional[str] = None,
        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,
        async_req: bool = False,
    ) -> DashVectorResponse:

使用示例

说明

需要使用您的api-key替换示例中的YOUR_API_KEY、您的Cluster Endpoint替换示例中的YOUR_CLUSTER_ENDPOINT,代码才能正常运行。

import dashvector
import numpy as np

client = dashvector.Client(
    api_key='YOUR_API_KEY',
    endpoint='YOUR_CLUSTER_ENDPOINT'
)
ret = client.create(
    name='group_by_demo',
    dimension=4,
    fields_schema={'document_id': str, 'chunk_id': int}
)
assert ret

collection = client.get(name='group_by_demo')

ret = collection.insert([
    ('1', np.random.rand(4), {'document_id': 'paper-01', 'chunk_id': 1, 'content': 'xxxA'}),
    ('2', np.random.rand(4), {'document_id': 'paper-01', 'chunk_id': 2, 'content': 'xxxB'}),
    ('3', np.random.rand(4), {'document_id': 'paper-02', 'chunk_id': 1, 'content': 'xxxC'}),
    ('4', np.random.rand(4), {'document_id': 'paper-02', 'chunk_id': 2, 'content': 'xxxD'}),
    ('5', np.random.rand(4), {'document_id': 'paper-02', 'chunk_id': 3, 'content': 'xxxE'}),
    ('6', np.random.rand(4), {'document_id': 'paper-03', 'chunk_id': 1, 'content': 'xxxF'}),
])
assert ret

根据向量进行分组相似性检索

ret = collection.query_group_by(
    vector=[0.1, 0.2, 0.3, 0.4],
    group_by_field='document_id',  # 按document_id字段的值分组
    group_count=2,  # 返回2个分组
    group_topk=2,   # 每个分组最多返回2个doc
)
# 判断是否成功
if ret:
    print('query_group_by success')
    print(len(ret))
    print('------------------------')
    for group in ret:
        print('group key:', group.group_id)
        for doc in group.docs:
            prefix = ' -'
            print(prefix, doc)

参考输出如下

query_group_by success
4
------------------------
group key: paper-01
 - {"id": "2", "fields": {"document_id": "paper-01", "chunk_id": 2, "content": "xxxB"}, "score": 0.6807}
 - {"id": "1", "fields": {"document_id": "paper-01", "chunk_id": 1, "content": "xxxA"}, "score": 0.4289}
group key: paper-02
 - {"id": "3", "fields": {"document_id": "paper-02", "chunk_id": 1, "content": "xxxC"}, "score": 0.6553}
 - {"id": "5", "fields": {"document_id": "paper-02", "chunk_id": 3, "content": "xxxE"}, "score": 0.4401}

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

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

带过滤条件的分组相似性检索

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

带有Sparse Vector的分组向量检索

说明

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

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

入参描述

说明

vectorid两个入参需要二选一使用,并保证其中一个不为空。

参数

类型

默认值

说明

group_by_field(必填)

str

按指定字段的值来分组检索,目前不支持schema-free字段。

vector(可选)

Optional[Union[List[Union[int, float]], np.ndarray]]

None

向量数据

id(可选)

Optional[str]

None

主键,表示根据主键对应的向量进行相似性检索

group_count(可选)

int

10

最多返回的分组个数,尽力而为参数,一般可以返回group_count个分组。

group_topk(可选)

int

10

每个分组返回group_topk条相似性结果,尽力而为参数,优先级低于group_count。

filter(可选)

Optional[str]

None

过滤条件,需满足SQL where子句规范,详见

include_vector(可选)

bool

False

是否返回向量数据

partition(可选)

Optional[str]

None

Partition名称

output_fields(可选)

Optional[List[str]]

None

默认返回所有Fields

sparse_vector(可选)

Optional[Dict[int, float]]

None

稀疏向量

async_req(可选)

bool

False

是否异步

出参描述

说明

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

字段

类型

描述

示例

code

int

返回值,参考返回状态码说明

0

message

str

返回消息

success

request_id

str

请求唯一id

19215409-ea66-4db9-8764-26ce2eb5bb99

output

List[Group]

分组相似性检索结果