検索エンジンとtop-k検索

検索エンジンの基本的な動作の一つに、クエリに対してマッチするドキュメントを全てスコアリングし、スコア順で上位から$k$件を返却するtop-k検索がある。これは単純には、マッチしたドキュメントをリストに入れ、全てスコアリングした後、スコアの降順にソートを行い、リストの先頭$k$件を取得することで実現できる。

優先度付きキュー

しかし、優先度付きキューという抽象データ構造を用いれば、top-k検索をもっと効率的に行うことができる。これは優先度付きの要素を優先度が高い順に所定の個数まで保持するデータ構造で、例えば以下の操作を効率的にサポートする。

キューの中で優先度が最低の要素の参照および取り出し。
要素の挿入。入り切らない場合は：
- まず追加しようとしている要素とキューの中で優先度が最低の要素を比較。
- 前者の優先度のほうが高い場合のみ、後者を削除し前者を挿入する。

サイズ$k$の優先度付きキューを用意し、ドキュメントごとにスコアリングするそばからスコアを優先度としてキューに挿入していけば、常にスコア順で上位の$k$件を保持しておくことができる。全てのドキュメントをスコアリングし終わったら、キューの全ての要素を優先度が低い順に取り出し、逆順に返却すればよい。

二分ヒープによる優先度付きキューの実装

ドキュメントを要素とする優先度付きキューは、より抽象度の低いデータ構造としては、ドキュメントをノードとする二分ヒープによって実現できる。二分ヒープについて詳細は省くが、これは平衡二分木であり、例えば常に親ドキュメントの優先度よりも子ドキュメントの優先度のほうが大きくなるように制約される。

二分ヒープの配列への詰め込み

そして二分ヒープは、具体的なデータ構造としては配列によって実現できる。

あるドキュメントが配列の$n$番目の要素に配置されているとすると：

その子は$2n$と$2n+1$番目の要素
（言い換えれば）その親は$n/2$番目の要素（切り捨て）

と、一意に決めて詰め込むことができるからである。

Apache Solrにおける実装

例えばApache Solrは、これまで述べてきたアイデアを全て実装している（正確には、Apache Luceneが実装して、Solrはそこに依存している）。

以下、Solr 8.6.3を想定して話を進める。

優先度付きキューの実装はorg.apache.lucene.util.PriorityQueueにある。

実際に要素を保持しているのは配列である。

  private final T[] heap;

そして、このクラスには以下のメソッドが実装されている。

top: 優先度が最低の要素の参照

二分ヒープを前述のように制約すると、キューの中で優先度が最低の要素は二分ヒープのルートに来る。二分ヒープを配列に詰め込んだ場合、二分ヒープのルートは常に配列の1番目の要素なので、以下のように定数時間で参照することができる。

  public final T top() {
    // We don't need to check size here: if maxSize is 0,
    // then heap is length 2 array with both entries null.
    // If size is 0 then heap[1] is already null.
    return heap[1];
  }

pop: 優先度が最低の要素の取り出し

参照と同様、優先度が最低の要素は配列の1番目の要素である。取り出すということは参照したあとに削除するということになるが、このうち削除は配列中で最後の要素で1番目の要素を上書きしたあと（こうすることで配列をコンパクトに保つ）二分ヒープの制約が維持されるように要素を並び替えることで行う。

  public final T pop() {
    if (size > 0) {
      T result = heap[1];       // save first value
      heap[1] = heap[size];     // move last to first
      heap[size] = null;        // permit GC of objects
      size--;
      downHeap(1);              // adjust heap
      return result;
    } else {
      return null;
    }
  }

insertWithOverflow: オーバーフローありの挿入

引数で与えられた要素を挿入する。
入り切らない場合（オーバーフロー）、入り切らない要素を返す。
そうでない場合はnullを返す。

  public T insertWithOverflow(T element) {
    if (size < maxSize) {  // 入り切るのであれば、
      add(element);  // 単に挿入を行い、二分ヒープの制約が維持されるように要素を並び替える。
      return null;
    } else if (size > 0 && !lessThan(element, heap[1])) {
      // 入り切らないのであれば、挿入しようとしている要素とキューの中で優先度が最低の要素を比較。
      T ret = heap[1];
      heap[1] = element;  // 前者の優先度のほうが高ければ、後者を前者で上書きし、
      updateTop();  // 二分ヒープの制約が維持されるように要素を並び替える。
      return ret;
    } else {  // 前者の優先度のほうが低ければ、何もしない。
      return element;
    }
  }

（コメントは筆者による）

（追記）Sentinel Object

優先度つきキューの状態としては空、フル、どちらでもない、が考えられ、状態によって処理を分岐する必要がある。ここで登場するのがsentinel objectで、優先度つきキューの初期化時にsentinel objectでフルの状態にしておくことで、分岐を減らす役割がある。たとえば検索エンジンの動作においては、(a) 全てのドキュメントを挿入してから、(b) キューの全ての要素を取り出す、ことになるが、sentinel objectにより (a) の処理において常にキューがフルの状態になる。

なお、sentinel objectの優先度は最低にする必要がある。そうしないと本来の（sentinel objectでない）要素が捨てられる可能性があるためである。また、(b) の処理までsentinel objectが残っていた場合は捨てる。

Solrでは TopScoreDocCollector（スコアの降順の検索において、優先度つきキューを内包するオブジェクト）で利用されている。

HitQueue（PriorityQueue の実装）のコンストラクタで、引数によってsentinel objectの生成を指示する。なお、スコアの降順の検索において優先度が最低であるとは、スコアが負の無限大であるという実装となっている。

  HitQueue(int size, boolean prePopulate) {
    super(size, () -> {
      if (prePopulate) {
        // Always set the doc Id to MAX_VALUE so that it won't be favored by
        // lessThan. This generally should not happen since if score is not NEG_INF,
        // TopScoreDocCollector will always add the object to the queue.
        return new ScoreDoc(Integer.MAX_VALUE, Float.NEGATIVE_INFINITY);
      } else {
        return null;
      }
    });
  }

HitQueue は常にフルの状態になるので、そのサイズとは別に挿入の回数　collectedHits　を保存しておき、それを用いて (b) の処理まで残ったsentinel objectを捨てる。

    @Override
    protected int topDocsSize() {
      return collectedHits < pq.size() ? collectedHits : pq.size();
    }

    // pq's pop() returns the 'least' element in the queue, therefore need
    // to discard the first ones, until we reach the requested range.
    // Note that this loop will usually not be executed, since the common usage
    // should be that the caller asks for the last howMany results. However it's
    // needed here for completeness.
    for (int i = pq.size() - start - howMany; i > 0; i--) { pq.pop(); }

ちなみに、mutableなオブジェクトでフルの優先度つきキューはオブジェクトプールと見なすこともでき、実際、ここではオブジェクトを再利用する実装となっている。

          if (score <= pqTop.score) {
            ...
            return;
          }
          collectedHits++;
          pqTop.doc = doc + docBase;
          pqTop.score = score;
          pqTop = pq.updateTop();