本記事でやること
ElasticsearchのLearning to Rank(LTR)プラグインを使用した検索順位のリスコアにおいて、go-elasticsearchライブラリを使ったRescoreクエリの構築実装パターンを紹介します。
今回実装したコードは以下のレポジトリで公開しています。
対象読者
- go-elasticsearchでより高度な検索クエリを構築したい方
使用言語・技術
- Go 1.21
- Elasticsearch 8.14.1
- go-elasticsearch v8
- Docker/Docker Compose
背景
go-elasticsearchレポジトリの以下issueにもあるように、LTRのRescoreクエリを直接サポートする型定義が存在しません。
そのため、LTRを使用した検索クエリを実行するためには独自にSearchRequestを実装する必要があります。本記事では、この課題に対する実践的な解決策を提示します。
Elasticsearchクラスターの準備
本記事では、https://github.com/o19s/hello-ltr レポジトリを使用してLTRが利用できる環境を構築します。このレポジトリには、LTRプラグインがプリインストールされたElasticsearchコンテナとモデル学習に必要なコードが含まれています。
git clone https://github.com/o19s/hello-ltr.git
cd hello-ltr/notebooks/elasticsearch
docker compose up -d
Elasticsearchクラスタのコンテナをローカルに起動した後は、TMDB(The Movie DataBase)データセットのインデックス作成、FeatureSetの定義、モデルの学習とアップロードを、hello-ltrのJupyter Notebook(hello-ltr(ES).ipynb)を使用して行います。
docker compose
したコンテナにはkibanaのコンテナもあるためlocalhost:5601
にアクセスし、コンソールから以下を実行しLTRプラグインがインストールされていることを確認します。
GET _cat/plugins?v
name component version
eaf520a41cb6 ltr 1.5.8-es8.8.2
LTRモデルの準備
hello-ltrレポジトリのNotebook(hello-ltr(ES).ipynb)を使用してLTRモデルを準備します。
cd hello-ltr/notebooks/elasticsearch/tmdb
hello-ltr(ES).ipynbを開き、上から順にセルを実行します。大まかな処理内容は以下になります。
- ライブラリのインポートとデータセットのダウンロード
- Elasticsearchへのインデックス作成
- 約28,000件の映画データをtmdbインデックスに投入
- Feature Setの定義
- release_year(映画の公開年)を特徴量として定義
- 訓練データの生成とモデル学習
- 古い映画を好むclassicモデル
- 新しい映画を好むlatestモデル
- 学習済みモデルのElasticsearchへのアップロード
Feature Setとモデルが正常にアップロードされたかを確認するため、Kibanaの開発ツールコンソールで以下のコマンドを実行します。
# Feature Setの確認
GET _ltr/_featureset/release
{
"_index": ".ltrstore",
"_id": "featureset-release",
"_version": 1,
"_seq_no": 0,
"_primary_term": 1,
"found": true,
"_source": {
"name": "release",
"type": "featureset",
"featureset": {
"name": "release",
"features": [
{
"name": "release_year",
"params": [],
"template_language": "mustache",
"template": {
"function_score": {
"field_value_factor": {
"field": "release_year",
"missing": 2000
},
"query": {
"match_all": {}
}
}
}
}
]
}
}
}
GET _ltr/_model/latest
{
"_index": ".ltrstore",
"_id": "model-latest",
"_version": 2,
"_seq_no": 2,
"_primary_term": 1,
"found": true,
"_source": {
"name": "latest",
"type": "model",
"model": {
"name": "latest",
"feature_set": {
"name": "release",
"features": [
{
"name": "release_year",
"params": [],
"template_language": "mustache",
"template": {
"function_score": {
"field_value_factor": {
"field": "release_year",
"missing": 2000
},
"query": {
"match_all": {}
}
}
}
}
]
},
"model": {
"type": "model/ranklib",
"definition": """## LambdaMART
...省略
これでElasticsearchでLTRを使用する準備が整いました。
go-elasticsearchでのLTRクエリ実装パターン
go-elasticsearchのtyped-apiではLTRのRescoreを表現する型定義がないので、独自にクエリを組み立てる実装が必要があります。ここでは2つのパターンを紹介します。
シンプルな文字列ベースの実装
このパターンはJSONクエリを文字列として直接構築する最もシンプルなアプローチです。クエリの構造を文字列テンプレートとして定義し、必要なパラメータを文字列連結して埋め込みます。
type StringLTRQueryBuilder struct {
keyword string
model string
}
func NewStringLTRQueryBuilder(keyword, model string) *StringLTRQueryBuilder {
return &StringLTRQueryBuilder{
keyword: keyword,
model: model,
}
}
func (s *StringLTRQueryBuilder) Build() ([]byte, error) {
query := `{
"query": {
"bool": {
"must": {"match_all": {}},
"filter": {"match": {"title": "` + s.keyword + `"}}
}
},
"rescore": {
"window_size": 1000,
"query": {
"rescore_query": {
"sltr": {
"params": {},
"model": "` + s.model + `"
}
}
}
}
}`
return []byte(query), nil
}
構造体を活用したタイプセーフな実装
このパターンでは、Rescoreクエリの各要素を構造体として定義し、JSONタグを使用してマーシャリングします。型安全性を保ちながら、メソッドチェーンで設定を追加できる柔軟な実装です。
type SimpleLTRQueryBuilder struct {
BaseQuery *types.Query
RescorePart *Rescore
}
type SLTR struct {
Params map[string]interface{} `json:"params"`
Model string `json:"model"`
}
type RescoreQuery struct {
RescoreQuery map[string]SLTR `json:"rescore_query"`
}
type Rescore struct {
WindowSize int `json:"window_size"`
Query RescoreQuery `json:"query"`
}
func (r SimpleLTRQueryBuilder) Build() ([]byte, error) {
req := &search.Request{
Query: r.BaseQuery,
}
base, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("failed to marshal base request: %w", err)
}
var m map[string]interface{}
if err := json.Unmarshal(base, &m); err != nil {
return nil, fmt.Errorf("failed to unmarshal base request: %w", err)
}
if r.RescorePart != nil {
m["rescore"] = r.RescorePart
}
return json.Marshal(m)
func NewSimpleLTRQueryBuilder(baseQuery *types.Query, model_name string) SimpleLTRQueryBuilder {
rescore := Rescore{
WindowSize: 1000,
Query: RescoreQuery{
RescoreQuery: map[string]SLTR{
"sltr": {
Params: map[string]interface{}{},
Model: model_name,
},
},
},
}
return SimpleLTRQueryBuilder{
BaseQuery: baseQuery,
RescorePart: &rescore,
}
}
func (s *SimpleLTRQueryBuilder) WithWindowSize(size int) *SimpleLTRQueryBuilder {
s.RescorePart.WindowSize = size
return s
}
func (s *SimpleLTRQueryBuilder) WithParams(params map[string]interface{}) *SimpleLTRQueryBuilder {
sltr := s.RescorePart.Query.RescoreQuery["sltr"]
if sltr.Params == nil {
sltr.Params = make(map[string]interface{})
}
for k, v := range params {
sltr.Params[k] = v
}
s.RescorePart.Query.RescoreQuery["sltr"] = sltr
return s
}
各構造体の役割としては
-
SimpleLTRQueryBuilder構造体
-
BaseQuery
: 基本的な検索クエリ(types.Query
型) -
RescorePart
リスコアの設定
-
-
SLTR構造体
-
Params
: モデルに渡すパラメータ -
Model
: 使用するモデル名
-
-
RescoreQuery / Rescore構造体
-
WindowSize
: リスコア対象のドキュメント数 -
Query
: リスコアクエリの詳細
-
実装の流れは
-
Build() メソッド
- 基本的な検索クエリを
search.Request
でマーシャル - JSONをmap(
map[string]interface{}
)型にアンマーシャル - mapに直接定義した
rescore
キーのバリューにRescorePart
型で定義した値を追加 - 最終的にJSONを生成
- 基本的な検索クエリを
-
NewSimpleLTRQueryBuilder() 関数
- デフォルト設定(WindowSize:1000)でBuilderを初期化
- 指定されたモデル名でSLT構造体を生成
-
チェーンメソッド
-
WithWindowSize()
: ウィンドウサイズの設定 -
WithParams()
: LTRモデルのパラメータ追加
-
この実装により、typed-api構造体と生JSONを組み合わせてLTRクエリを柔軟に構築できます。
呼び出し側の例を以下の記載します。
queryBuilder := NewSimpleLTRQueryBuilder(baseQuery, "latest")
queryBuilder.WithWindowSize(500).
WithParams(map[string]interface{}{
"parameter1": "value1",
"parameter2": "value2",
})
res, err := client.Search(ctx, "tmdb", queryBuilder)
if err != nil {
fmt.Printf("Error executing search: %v\n", err)
return
}
for i, r := range res {
if i >= 3 {
break
}
fmt.Printf("Found result: ID=%s, Title=%s, Release Year=%s, Score=%.2f\n",
r.Id, r.Title, r.ReleaseYear, r.Score)
}
Found result: ID=242643, Title=Batman: Assault on Arkham, Release Year=2014, Score=6.40
Found result: ID=251519, Title=Son of Batman, Release Year=2014, Score=6.40
Found result: ID=366924, Title=Batman: Bad Blood, Release Year=2016, Score=1.37
まとめ
本記事ではLTRを使用できる環境の構築からgo-elasticsearchでLTRのRescoreクエリを構築する2つの実装パターンを紹介しました。
go-elasticsearchはドキュメントが充実しているわけではないので、今回のようなプラグイン固有の機能を実装する際は、Elasticsearchの公式ドキュメントとgo-elasticsearchのソースコードを併せて参照することが重要です。また、typed-apiでサポートされていない機能については、JSONと組み合わせることで柔軟なアプローチが有効であることがわかりました。
参考リソース