SolrのRankQueryによるドキュメントの並べ替え

# はじめに

本記事は、以下の既存の記事でも述べられているようなことを、別の切り口から述べたものですので、以下の記事も読んでいただけると理解が深まると思います。

https://qiita.com/yuyano/items/789cd9e0accd32fba412

https://qiita.com/ft28/items/47d7add1341293d3ab35

以下、細かい説明はSolr 8.6.3に基づきます。
https://github.com/apache/solr/releases/tag/releases%2Flucene-solr%2F8.6.3


# RankQueryでできること

`RankQuery`は抽象クラスで、以下の抽象メソッドが実装されています。

```java
  public abstract TopDocsCollector getTopDocsCollector(int len, QueryCommand cmd, IndexSearcher searcher) throws IOException;
  public abstract MergeStrategy getMergeStrategy();
  public abstract RankQuery wrap(Query mainQuery);
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/search/RankQuery.java

それぞれ、以下の動作に対応します。

|抽象メソッド名|動作|
|----|----|
|getTopDocsCollector|優先度付きキューを内包するオブジェクト`TopDocsCollector`を生成する。優先度付きキューの使われ方については以下を参照。https://qiita.com/tomanabe/items/34fbb3b1386f4df6dcd9|
|getMergeStrategy|シャードごとのレスポンスのマージを責務とするオブジェクト`MergeStrategy`を返す。このマージが必要な理由とタイミングについては以下を参照。https://qiita.com/tomanabe/items/7ea6aaba4dda85d12ca2|
|wrap|元のクエリをラップする。|

一言でまとめると、`RankQuery`とは、**元のクエリをラップし、優先度付きキューの動作とシャードごとのレスポンスのマージの動作を元のクエリとは変えるもの**です。そしてこれらは、ドキュメントの順位に関わる動作を過不足なく網羅しています。


# RankQueryがセットされていない場合のデフォルトの動作

それぞれのメソッドについて、`RankQuery`がセットされていない場合のデフォルトの動作を見ていきます。これは、動作をより具体的にイメージするためと、単にデフォルトの動作の場所が分かりにくいためです。

## デフォルトのTopDocsCollectorの生成

`SolrIndexSearcher#buildTopDocsCollector`で行います。前半で`Query`が`RankQuery`かどうかを調べ、そうなら`RankQuery#getTopDocsCollector`を呼び、そうでないなら後半でデフォルトの`TopDocsCollector`の生成を行います。

```java
  private TopDocsCollector buildTopDocsCollector(int len, QueryCommand cmd) throws IOException {
    int minNumFound = cmd.getMinExactCount();
    Query q = cmd.getQuery();
    if (q instanceof RankQuery) {
      RankQuery rq = (RankQuery) q;
      return rq.getTopDocsCollector(len, cmd, this);
    }

    if (null == cmd.getSort()) {
      assert null == cmd.getCursorMark() : "have cursor but no sort";
      return TopScoreDocCollector.create(len, minNumFound);
    } else {
      // we have a sort
      final Sort weightedSort = weightSort(cmd.getSort());
      final CursorMark cursor = cmd.getCursorMark();

      final FieldDoc searchAfter = (null != cursor ? cursor.getSearchAfterFieldDoc() : null);
      return TopFieldCollector.create(weightedSort, len, searchAfter, minNumFound);
    }
  }
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java#L1495

具体的なデフォルトの動作としては、ソート条件の指定がない場合は高速な`TopScoreDocCollector`を、そうでない場合は`TopFieldCollector`を生成します。


## デフォルトのMergeStrategy

`MergeStrategy`はシャードごとのレスポンスのマージを責務とするオブジェクトです。実はデフォルトの`MergeStrategy`がまとまっているわけではなく、`MergeStrategy`インタフェースで定義されている`merge`と`handleMergeFields`の各メソッドについて、別々の箇所にデフォルトの動作が書かれています。

### デフォルトのMergeStrategy#merge

サーチヘッドにおいて、シャードごとのtop-kドキュメントのIDとスコアを収集する動作です。通常、全体のtop-kドキュメントのIDとスコアのリストへのマージも行います。

デフォルトの動作は`QueryComponent#mergeIds`に記述されています。メソッドの冒頭で`MergeStrategy`が存在するかどうかを調べ、存在するなら`MergeStrategy#merge`を呼び（複数可で、複数の場合はコストの低い順に呼ぶ）、存在しないかどの`MergeStrategy`もマージを行わない場合は、メソッドの残りの部分でデフォルトのマージを行います。複数の`MergeStrategy`がマージを行うと衝突しそうですが、そこはケアしていないようです。

```java
      List<MergeStrategy> mergeStrategies = rb.getMergeStrategies();
      if(mergeStrategies != null) {
        Collections.sort(mergeStrategies, MergeStrategy.MERGE_COMP);
        boolean idsMerged = false;
        for(MergeStrategy mergeStrategy : mergeStrategies) {
          mergeStrategy.merge(rb, sreq);
          if(mergeStrategy.mergesIds()) {
            idsMerged = true;
          }
        }

        if(idsMerged) {
          return; //ids were merged above so return.
        }
      }
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java#L810

### デフォルトのMergeStategy#handleMergeFields

サーチヘッドで全体のtop-kドキュメントのIDとスコアのリストを生成するとき、各ドキュメントのIDとスコアに加えて、追加の情報が必要なことがあります。例えば、スコアではなく特定のフィールドの値でソートしようと思ったら、各ドキュメントについてそのフィールドの値が必要です。これはシャードが持っている情報で、マージを行うのはサーチヘッドですので、シャードからサーチヘッドに情報を送る必要があります。`MergeStrategy#handleMergeFields`は、シャードごとのレスポンスの生成において、ここでいうフィールド値のような追加の情報をレスポンスに追加する動作です。

デフォルトの動作は`QueryComponent#doProcessUngroupedSearch`で行います。例によって`MergeStrategy`が存在するかどうかを調べ（こちらは複数不可）、存在するなら`MergeStrategy#handleMergeFields`を呼び、存在しないならデフォルトの処理である`QueryComponent#doFieldSortValues`を呼びます。

```java
    if(rb.mergeFieldHandler != null) {
      rb.mergeFieldHandler.handleMergeFields(rb, searcher);
    } else {
      doFieldSortValues(rb, searcher);
    }
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java#L1505

`QueryComponent#doFieldSortValues`では、まさに特定のフィールドの値でソートする場合に、各ドキュメントのそのフィールドの値をレスポンスに追加する動作を行います。
https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java#L414


# 具体的なRankQueryの実装

では、これらデフォルトの動作を上書きする具体的な`RankQuery`の実装を見ていきます。


## AbstractReRankQuery#getTopDocsCollectorの実装

例えば`AbstractReRankQuery`は、`getTopDocsCollector`メソッドを実装しています。内容としては以下の通り、独自実装した`ReRankCollector`を返すというものです。

```java
  public TopDocsCollector getTopDocsCollector(int len, QueryCommand cmd, IndexSearcher searcher) throws IOException {
    if(this.boostedPriority == null) {
      SolrRequestInfo info = SolrRequestInfo.getRequestInfo();
      if(info != null) {
        Map context = info.getReq().getContext();
        this.boostedPriority = (Set<BytesRef>)context.get(QueryElevationComponent.BOOSTED);
      }
    }

    return new ReRankCollector(reRankDocs, len, reRankQueryRescorer, cmd, searcher, boostedPriority);
  }
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/search/AbstractReRankQuery.java#L60

やや本記事の主題からズレますが、`ReRankCollector`は以下の動作を行います。

- 内包する優先度付きキューはデフォルトと同じだが、そのサイズはデフォルトよりも大きくすることができる。
- 優先度付きキューからドキュメントを取り出すとき、キューに詰めるときの1度目のスコアリングとは別の2度目のスコアリングを行い、その順序にドキュメントを並べ替える（リランキング）。

クエリにマッチした全ドキュメントには軽いスコアリングを行い、優先度付きキューに残ったドキュメントには重くて高精度なスコアリングを行うことができます。


## 具体的なRankQuery#getMergeStrategyの実装

なんとSolr 8.6.3にバンドルされている`RankQuery`の実装で、まともに`getMergeStrategy`を実装しているものが無さそうなため、説明は省略します。


# QParserPluginでRankQueryをセットする

`RankQuery`をセットするには、`QParserPlugin`を実装し、クラスロードし、`solrconfig.xml`の`queryParser`要素で参照します。`ReRankQParserPlugin`はデフォルトで有効なので、`LTRQParserPlugin`を例に取ると、サンプルの`solrconfig.xml`に以下の記述があります。

```xml
  <queryParser enable="${solr.ltr.enabled:false}" name="ltr" class="org.apache.solr.ltr.search.LTRQParserPlugin"/>
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/server/solr/configsets/sample_techproducts_configs/conf/solrconfig.xml#L1590

（この例ではJVMのコマンドライン引数でも有効化する必要、具体的には`-Dsolr.ltr.enabled=true`とする必要があるので、やや特殊な例ですが）こうしておくと、その要素の`name`属性値を`rq`リクエストパラメータに指定して`QParserPlugin#createParser`（`QParser`を返す）を、続いて`QParser#parse`を呼び出すことができます。あとはそこで`RankQuery`を生成すれば良いです。

前述の`AbstractReRankQuery`の実装である`ReRankQuery`の例を挙げます。コメント部分もリクエストの例になっていますので、読んでいただけると良いと思います。

```java
/*
*
*  Syntax: q=*:*&rq={!rerank reRankQuery=$rqq reRankDocs=300 reRankWeight=3}
*
*/

public class ReRankQParserPlugin extends QParserPlugin {

  public static final String NAME = "rerank";
// （中略）
  public QParser createParser(String query, SolrParams localParams, SolrParams params, SolrQueryRequest req) {
    return new ReRankQParser(query, localParams, params, req);
  }

  private class ReRankQParser extends QParser  {
// （中略）
    public Query parse() throws SyntaxError {
      String reRankQueryString = localParams.get(RERANK_QUERY);
      if (StringUtils.isBlank(reRankQueryString)) {
        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, RERANK_QUERY+" parameter is mandatory");
      }
      QParser reRankParser = QParser.getParser(reRankQueryString, req);
      Query reRankQuery = reRankParser.parse();

      int reRankDocs  = localParams.getInt(RERANK_DOCS, RERANK_DOCS_DEFAULT);
      reRankDocs = Math.max(1, reRankDocs); //

      double reRankWeight = localParams.getDouble(RERANK_WEIGHT, RERANK_WEIGHT_DEFAULT);

      return new ReRankQuery(reRankQuery, reRankDocs, reRankWeight);
    }
  }
```
引用元：https://github.com/apache/solr/blob/e001c2221812a0ba9e9378855040ce72f93eced4/solr/core/src/java/org/apache/solr/search/ReRankQParserPlugin.java