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

# 分散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検索を実現できる。

1. サーチヘッドがユーザからの検索リクエストを受ける。
1. **検索リクエストの転送**。サーチヘッドは、全$n$シャード（複数レプリカがある場合は、そのうち各シャードごとに1つずつ、計$n$個。以下、単にシャードと呼ぶ）に宛てて検索リクエストを転送する。
1. **シャードごとのtop-k検索**。各シャードは、シャードごとのtop-k検索を行い、その結果のドキュメントIDとスコアのペアのリストをサーチヘッドに返す。
1. **マージとソート（1回目）**。サーチヘッドは、返ってきた$nk$件すべてのペアを一つのリストにマージし、スコアの降順にソートし直す。そのtop-kを切り出すと、それがコレクション全体のtop-kになっている。

![分散top-k検索の仕組みとSolrにおける実装_ランキングの生成.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/155352/007a3bb9-745a-b611-27f6-1523f55eedab.png)

$k$件のドキュメントをリクエストされているのに対して、サーチヘッドには$kn$件のペアが集まることになり、非効率に思える。しかし、コレクション全体のtop-kすべてが特定のシャードに偏っている最悪ケースを考えると、この数のペアが必要である。

## フィールド値の取得

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

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

![分散top-k検索の仕組みとSolrにおける実装_フィールド値の取得.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/155352/b0974474-384a-de11-9e94-9e2fd0903986.png)

# 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`である。

```java
        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`メソッドを呼び出す。

```java
            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`)|