検索イテレータ

ANN Search には、1 回のクエリで呼び出すことができるエンティティの数に上限があり、基本的な ANN Search を使用するだけでは大規模検索の要求を満たせない場合がある。topKが16,384を超えるようなANN Searchリクエストでは、SearchIteratorの使用を検討することが望ましい。このセクションでは、SearchIterator の使用方法とそれに関連する考察を紹介する。

概要

Search リクエストは検索結果を返すが、SearchIterator はイテレーターを返す。このイテレータのnext()メソッドを呼び出すことで、検索結果を取得することができます。

具体的には、SearchIterator を次のように使用します。

1. SearchIterator を作成し、検索要求ごとに返すエンティティの数と返すエンティティの総数を設定します。

2. 検索結果をページ分割して取得するために、ループ内で SearchIterator のnext()メソッドを呼び出します。

3. next()メソッドが空の結果を返した場合は、イテレータのclose()メソッドを呼び出してループを終了します。

SearchIterator の作成

次のコード・スニペットはSearchIteratorの作成方法を示しています。

from pymilvus import connections, Collection

connections.connect(
    uri="http://localhost:19530",
    token="root:Milvus"
)

# create iterator
query_vectors = [
    [0.3580376395471989, -0.6023495712049978, 0.18414012509913835, -0.26286205330961354, 0.9029438446296592]]

collection = Collection("iterator_collection")

iterator = collection.search_iterator(
    data=query_vectors,
    anns_field="vector",
    param={"metric_type": "L2", "params": {"nprobe": 16}},
    # highlight-next-line
    batch_size=50,
    output_fields=["color"],
    # highlight-next-line
    limit=20000
)

import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.orm.iterator.SearchIterator;
import io.milvus.v2.common.IndexParam.MetricType;
import io.milvus.v2.service.vector.request.data.FloatVec;

import java.util.*;

MilvusClientV2 client = new MilvusClientV2(ConnectConfig.builder()
        .uri("http://localhost:19530")
        .token("root:Milvus")
        .build());

FloatVec queryVector = new FloatVec(new float[]{0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f});
SearchIterator searchIterator = client.searchIterator(SearchIteratorReq.builder()
        .collectionName("iterator_collection")
        .vectors(Collections.singletonList(queryVector))
        .vectorFieldName("vector")
        .batchSize(500L)
        .outputFields(Lists.newArrayList("color"))
        .topK(20000)
        .metricType(IndexParam.MetricType.COSINE)
        .build());

上記の例では、1回の検索で返すエンティティの数（batch_size/batchSize）を50、返すエンティティの総数（topK）を20,000に設定しています。

SearchIterator の使用

SearchIteratorの準備ができたら、そのnext()メソッドを呼び出して、ページ分割された検索結果を得ることができる。

results = []

while True:
    # highlight-next-line
    result = iterator.next()
    if not result:
        # highlight-next-line
        iterator.close()
        break
    
    for hit in result:
        results.append(hit.to_dict())

import io.milvus.response.QueryResultsWrapper;

while (true) {
    List<QueryResultsWrapper.RowRecord> res = searchIterator.next();
    if (res.isEmpty()) {
        searchIterator.close();
        break;
    }

    for (QueryResultsWrapper.RowRecord record : res) {
        System.out.println(record);
    }
}

上記のコード例では、無限ループを作成し、ループ内でnext()メソッドを呼び出して検索結果を変数に格納し、next()が何も返さなかったときにイテレータを閉じています。