分散top-k検索の仕組み

ドキュメント検索において、サイズの大きいドキュメントや多数のドキュメントを扱う場合、サーバごとのストレージのサイズとの兼ね合いによっては、ドキュメントコレクションを複数のサーバに分割して保持する必要がある。ここで、複数のサーバに分割されたコレクション上の検索を分散検索と呼ぶ。具体的な検索処理としては、コレクション中の全てのドキュメントをスコアの降順にソートし、その上位$k$件を取得するtop-k検索が考えられる。本記事では、分散検索で単一のサーバ上のtop-k検索と同じ結果を実現する分散top-k検索の仕組みと、そのSolrにおける実装について述べる。

シャードとサーチヘッド

本記事では、コレクションを漏れも重なりもなくドキュメント単位で分割した単位をシャードと呼ぶ。例えばシャード0, 1とドキュメントA, B, C, Dがあったら、シャード0にはAとCを、シャード1にはBとDを振り分けるということである。

通常のサービスでは、冗長化のために各シャードをそれぞれ複数のサーバにコピーする必要があり、こちらの単位はレプリカと呼ぶ。上の例で言えば、ドキュメントA, Cを2組、B, Dを2組、それぞれ別のサーバ（計4台）に保持しておくということである。レプリカの概念は本記事ではあまり使わないが、シャードとの混同を避けるために明記した。

ユーザからの検索リクエストを直接受けるサーバをサーチヘッドと呼ぶ。Solrのように、サーチヘッドがシャード（レプリカ）を兼ねる場合もあるが、本記事では説明の単純化のため、サーチヘッドは別のサーバとする。

フィールド値の取得

もしユーザがtop-kドキュメントのIDとスコアのみをリクエストしていれば、上述の処理のあと、すぐサーチヘッドからユーザに結果を返すことができる。しかし一般には、top-kドキュメントの他のフィールド値（例えば、商品検索システムにおける商品名や価格）もユーザに返す必要がある。フィールド値の取得は、ユーザに検索結果を返す前に、サーチヘッドを起点として、以下のように行われる。

フィールド値のリクエスト。サーチヘッドは、コレクション全体のtop-kドキュメントのうち1件以上を保持している全てのシャードに宛てて、そのドキュメントIDと必要なフィールドを指定してリクエストする。
ドキュメントの引き当て。各シャードは、IDに対応するドキュメントを引き当て、サーチヘッドに返す。
マージとソート（2回目）。サーチヘッドは、返ってきた$k$件すべてのドキュメントを一つのリストにマージし、スコアの降順にソートし直す。これが検索結果なので、ユーザに返す。

Apache Solrにおける実装

ここからは、分散top-k検索のApache Solrにおける実装を見ていく。以下、Solr 8.6.3を想定して話を進める。
https://github.com/apache/solr/releases/tag/releases%2Flucene-solr%2F8.6.3

本記事で述べた処理を扱うクラスはSearchHandlerである。
https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/handler/component/SearchHandler.java

サーチヘッドにおける処理

サーチヘッドにおいては、SearchHandlerがステージ（実体はEnumに近いint）を進めつつ、各ステージで各SearchComponentのdistributedProcess, handleResponses（これはシャードごとのレスポンスが返るごとに呼び出される）, finishStageの各メソッドを順に呼び出す。また、実際にリクエストの送受信を行うのもSearchHandlerである。

        do {
          rb.stage = nextStage;
          nextStage = ResponseBuilder.STAGE_DONE;

          // call all components
          for (SearchComponent c : components) {
            // the next stage is the minimum of what all components report
            nextStage = Math.min(nextStage, c.distributedProcess(rb));
          }

          // check the outgoing queue and send requests
          while (rb.outgoing.size() > 0) {

// （中略。シャードごとのリクエスト送信の処理などが入る）

            while (rb.outgoing.size() == 0) {

// （中略。シャードごとのレスポンス受信の処理などが入る）

              // let the components see the responses to the request
              for (SearchComponent c : components) {
                c.handleResponses(rb, srsp.getShardRequest());
              }
            }
          }

          for (SearchComponent c : components) {
            c.finishStage(rb);
          }

          // we are done when the next stage is MAX_VALUE
        } while (nextStage != Integer.MAX_VALUE);

ランキングの生成がSTAGE_EXECUTE_QUERYに、フィールド値の取得がSTAGE_GET_FIELDSに対応する。

シャードごとの処理

シャードごとの処理においては、単に各SearchComponentのprocessメソッドを呼び出す。

            for (SearchComponent c : components) {
              c.process(rb);
            }

仕組みと実装の対応表

Top-k検索を行うSearchComponentはQueryComponentである。
https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java

仕組みで説明した各ステップを、SearchHandlerのステージとQueryComponentのメソッドとに対応させると、以下のようになる。より具体的な処理の記述されているメソッドをカッコ内に示したが、グルーピングを行う場合など、リクエストによっては異なるコードパスを通ることがある。

処理するサーバ	本稿におけるステップ	ステージ	メソッド
サーチヘッド	検索リクエストの転送	`STAGE_EXECUTE_QUERY`	`distributedProcess` (`createMainQuery`)
各シャード	シャードごとのtop-k検索	`STAGE_EXECUTE_QUERY`	`process` (`doProcessUngroupedSearch`)
サーチヘッド	マージとソート（1回目）	`STAGE_EXECUTE_QUERY`	`handleResponses` (`mergeIds`)
サーチヘッド	フィールド値のリクエスト	`STAGE_GET_FIELDS`	`distributedProcess` (`createRetrieveDocs`)
各シャード	ドキュメントの引き当て	`STAGE_GET_FIELDS`	`process` (`doProcessUngroupedSearch`)
サーチヘッド	マージとソート（2回目）	`STAGE_GET_FIELDS`	`handleResponses` (`mergeIds`)

分散top-k検索の仕組みとSolrにおける実装