本文介紹如何通過Python SDK在Collection中按分組進行相似性檢索。
前提條件
已創建Cluster：創建Cluster。
已獲得API-KEY：API-KEY管理。
已安裝最新版SDK：安裝DashVector SDK。
接口定義
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_group_by(
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_group_by(
vector=[0.1, 0.2, 0.3, 0.4], # 向量檢索
sparse_vector={1: 0.3, 20: 0.7},
group_by_field='name',
)入參描述
vector和id兩個入參需要二選一使用，并保證其中一個不為空。
參數 | 類型 | 默認值 | 說明 |
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 | 稀疏向量 |
vector_field (可選) | Optional[str] | 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] | 分組相似性檢索結果 |