search infra teamのmrkm4ntrです。我々の運用するElasticsearchにはFunction Score Queryを使ったリクエストが送られてきます。Function Score Queryはサブクエリのスコアに任意の関数を適用できるというもので、とても便利な機能ですが、同時にTop K(スコアが大きいものからK個を取得する場合)クエリ処理の最適化の恩恵を受けられなくなるという欠点もあります。この記事では、Function Score Queryに用いる関数の性質を利用し、Function Score QueryとTop Kのクエリ処理の最適化を両立させる方法について説明します。本記事は読者が検索エンジンの仕組みにある程度詳しいことを想定しています。
Top Kのクエリ処理の最適化
Elasticsearchの検索機能を提供しているライブラリLuceneには、Top Kを取得する際に、Top Kに入る見込みのないもののスコア計算をスキップすることで、パフォーマンスの最適化を図る機能が存在します。
例えば( ”search” OR “engine”)のようなクエリがあり、”search”というtermに対応するposting listの最大スコアが5.0、”engine”というtermに対応するposting listの最大スコアが3.0だとします。
BM25 (https://ja.wikipedia.org/wiki/Okapi_BM25) にて各termに対するドキュメントはスコア付けされるため、indexの構築時に最大スコアが決まります。両方を含む文書のスコアは5.0 + 3.0 = 8.0になります。
ここでスコアの高い順にTop 10を検索することを考えます。この時大きいものから10番目のスコアがmin competitive scoreと呼ばれるものになります。つまり、このmin competitive scoreよりも大きいスコアをとりえない文書はスコア計算する必要がありません。
仮にmin competitive scoreが3.0より大きい値とします。この場合”engine”のみを含む文書のスコアはmin competitive scoreより大きくはならないので”engine”のみを含む文書のスコア計算はスキップできます。本来ならばORはそれぞれのtermのposting listを全て走査する必要があるのですが、”search”のposting listに存在する文書のみのスコア計算をすれば良いことになります。このような手法によりクエリのパフォーマンスを向上させることができます。
Function Score Query
Function Score Queryは以下のようなクエリです。クエリを実行した結果はそのまま使うのではなく、サブクエリを実行した結果にfunctionsで指定された関数の戻り値を結合したスコアを最終スコアとして利用します。デフォルトの結合方法は乗算ですが、score_modeの値によって動作を変更できます。
{
"query": {
"function_score": {
"query": {
… // サブクエリ
},
"functions": […],
"score_mode": …
}
}
}上記のTop Kクエリ処理の最適化は最大スコアがindex構築時に決まることが前提でした。Function Score Queryを使った場合、サブクエリのスコアに任意の関数の戻り値を結合することができるため、文書の最終スコアがクエリの実行時に決まることになります。このような状況下では先ほど説明したTop Kクエリ処理の最適化を使うことができません。
Function Score Queryを使うとTop Kクエリ処理の最適化がされていないのは、コードを確認すると明白です。min competitive scoreをLuceneのScorerに伝えるには setMinCompetitiveScoreというメソッド(https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/search/Scorable.java#L48-L57) を使うのですが、ElasticsearchのFunction Score QueryのScorerであるFunctionFactorScorerにおいてはsetMinCompeitiveScoreを呼んでいません。これによりTop Kクエリ処理の最適化がされていないことがわかります。
https://github.com/elastic/elasticsearch/blob/v8.13.2/server/src/main/java/org/elasticsearch/common/lucene/search/function/FunctionScoreQuery.java#L371-L487
TopKクエリ処理の最適化が可能な関数
確かに任意の関数に対してTop Kクエリ処理の最適化を実現するのは不可能です。しかし、利用する関数によってはTop Kクエリ処理の最適化の恩恵を受けれるものも存在します。例えばよく使われる、作成日時からの時間経過でスコアを指数関数的に減衰させる以下のようなFunction Score Queryです。
{
"query": {
"function_score": {
"query": {
… // サブクエリ
},
"functions": [
{
"exp": {
"created_time": {
"scale": "10d",
"decay": 0.8
}
}
}
],
"score_mode": "multiply"
}
}
}このクエリでは、最終スコアは (指数関数的減衰 * サブクエリのスコア) となります。
重要なのはこの減衰は現在日時と作成日時の差において単調減少であるということです。
Luceneでは、posting listをあるフィールドの値で構築時にソートすることができます。posting listを作成日時の降順にソートすることで、上記の減衰関数を用いる際に、posting list内の前のドキュメントよりも必ず減衰値が大きくなることが保証できます。
これにより、サブクエリのみのmin competitive scoreが5.0、今の評価しているドキュメントの減衰が0.7だとすると、実質min competitive scoreを 5.0 / 0.7 = 7.14として扱うことができます。これは単にTop Kクエリ処理の最適化が使えるだけではなく、min competitive scoreがposting listを進むたびに増幅していくことになり、より多くのドキュメントの評価をスキップできる可能性が高まります。
PoCの実装とLuceneへの貢献
それをふまえて、前述の減衰関数を受け取りTop Kクエリ処理の最適化を実現するElasticsearchの新しいクエリをElasticsearchのpluginとして実装しました。基本的にはElasticsearchのFunction Score Queryと同じですが、サブクエリのScorerのsetMinCompetitiveScoreを適宜呼び出す部分が異なります。
実装自体は簡単でしたが、いざ動作確認すると全くパフォーマンスに変化がありませんでした。リモートデバッグで確認したところ、サブクエリ内のBoolean QueryにTop Kクエリ処理の最適化に必要なWANDScorer(https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/search/WANDScorer.java) やBlockMaxConjunctionScorer(https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/search/BlockMaxConjunctionScorer.java) が使われておらず、代わりに使われていたのは、それぞれDisjunctionSumScorer(https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/search/DisjunctionSumScorer.java) とConjunctionScorer(https://github.com/apache/lucene/blob/releases/lucene/9.10.0/lucene/core/src/java/org/apache/lucene/search/ConjunctionScorer.java) でした。
調べてみたところ、WANDScorerやBlockMaxConjunctionScorerはオーバーヘッドが大きいためトップレベルの句の場合でしか利用されないようです。つまり ((A AND B) OR (C AND D)) のようなクエリの場合はORはWANDScorerが使われますが、ANDにはConjunctionScorerが使われることになります。サブクエリをトップレベルの句と認識させるためには、新しく追加したクエリからサブクエリのScorerSupplierのsetTopLevelScoringClauseメソッド(https://github.com/apache/lucene/blob/695c0ac84508438302cd346a812cfa2fdc5a10df/lucene/core/src/java/org/apache/lucene/search/ScorerSupplier.java#L46-L54) を呼ぶ必要があります。そのように修正したところ、無事にWANDScorerとBlockMaxConjunctionScorerが使われるようになりました。
これでパフォーマンスが改善するかと思われましたが、相変わらず変化がありません。さらに調べると、WANDScorerの下のConjunctionScorer (ORの下のAND)とBlockMaxConjunctionScorerの下のDisunctionSumScorer (ANDの下のOR)が最大スコアとしてInfinityを返していました。
Luceneの実装を見ると確かにInfinityを返すようになっています。何故Infinityを返すのか全く意図が掴めずに頭を抱えましたが、トップレベルの句以外は最適化をしないという修正(https://github.com/apache/lucene/pull/12490) での漏れだということがわかりました。そこで、それぞれ最大スコアに正しい値を返すように修正したところ、ようやくパフォーマンスが大きく改善することが確認できました。同様にprofile APIにおいても修正漏れがあったため、以下のプルリクエストをLuceneのupstreamに送りました。今は全てmergeされています。
https://github.com/apache/lucene/pull/13031
https://github.com/apache/lucene/pull/13043
https://github.com/apache/lucene/pull/13066
上で実装したpluginを使って我々のワークロードを模したパフォーマンステストを実施したところ、既存のクエリの95pが約35%、99pが60%下がりました。さらにコストも約1/3削減できることがわかりました。
ただし、この最適化の恩恵を受けるためにはリクエストパラメータのtrack_total_hitsをfalseにする、もしくは十分小さい値(1,000以下)に設定する必要があります。というのも最低でもこの件数はヒットするドキュメントを検索する必要があるため、スキップ対象が少なくなるからです。歴史的経緯により、この値をすぐに小さくすることは難しく、また最近この最適化ができない形のクエリがテストされているため、この最適化を実際に我々の本番環境に適用できるかは検討中です。
さいごに
この記事ではElasticsearchのFunction Score Queryに使われる関数の単調減少性を利用してTopKスコア処理の最適化の恩恵を受ける方法について述べました。このタスクでLuceneのスコアリング周りの内部実装についての理解が深まりました。また、仮説を実装し、何度もうまくいかない原因を潰していくサイクルは大変でしたがエキサイティングでした。もし仮にこのような最適化を適用できるクエリを利用されている場合は、試してみていただけると幸いです。
