ちょっとテクい検索ランキングをVespaで実現する

この記事は何?

情報検索・検索技術 Advent Calendar 2021 - Adventarの4日目の記事です。

検索システムを開発・運用していると、検索ランキングの精度を改善したくなります。(なりますよね?)
ElasticsearchやSolrでは、検索ランキングで使われるデフォルトのスコアとして、検索クエリとドキュメントのフィールドとのマッチ度(BM25)が使われます。
マッチ度だけでなく、ドキュメントの新しさや人気度を考慮することで、よりユーザーが求めているドキュメントを検索結果の上位に掲出できるでしょう。
マッチ度・新しさ・人気度・…を考慮したスコアリングロジックを人手でコネコネ作ることもできますが、訓練データを用意することで、各要素を良い感じに考慮した機械学習によるランキングも実現できます。

さて、検索ランキングの改善案を考案・実装していくと、
「あれっ、これってElasticsearchやSolrでどうやって実現するんだろう?」
「実現はできそうだけど、ゴリ押し実装だな…」
「複雑怪奇な検索クエリになっちゃった🥺」
という沼に陥っていきます、つらいですね。。


…話は変わりますが、OSSの検索エンジンであるVespaを紹介します。
VespaはもともとYahoo! Inc.が内製していた検索エンジンですが、2017年にOath Inc.(Verizon社で、Yahoo, AOL等50以上のテクノロジーとメディアを運営する子会社)からOSSとして公開された…という経緯があります。
Vespaの特徴として「検索パフォーマンスがとても良い」という点があげられます。
また、検索ランキング改善の視点からすると、「検索ランキングロジックを非常に柔軟に記述できる」というメリットがあります。


さてこの記事では、Vespaを使うとちょっとテクい検索ランキングを簡単に実現できますよ、ということを紹介したいと思います。
検索ランキング改善に携わっている方に楽しんでいただけると嬉しいです。

Vespaの準備

Vespaは豊富なチュートリアルがあります むしろ、ありすぎてどこから手をつけたら良いのか混乱する が、この記事ではQuick Startを利用して解説します。

Quick Startでは、以下の5つのアルバムに関するドキュメントをフィードします。

{
    "put": "id:mynamespace:music::a-head-full-of-dreams",
    "fields": {
        "album": "A Head Full of Dreams",
        "artist": "Coldplay",
        "year": 2015,
        "category_scores": {
            "cells": [
                { "address" : { "cat" : "pop" },  "value": 1 },
                { "address" : { "cat" : "rock" }, "value": 0.2 },
                { "address" : { "cat" : "jazz" }, "value": 0 }
            ]
        }
    }
}
{
    "put": "id:mynamespace:music::hardwired-to-self-destruct",
    "fields": {
        "album": "Hardwired...To Self-Destruct",
        "artist": "Metallica",
        "year": 2016,
        "category_scores": {
            "cells": [
                { "address" : { "cat" : "pop" },  "value": 0 },
                { "address" : { "cat" : "rock" }, "value": 1 },
                { "address" : { "cat" : "jazz" }, "value": 0 }
            ]
        }
     }
}

{
    "put": "id:mynamespace:music::liebe-ist-für-alle-da",
    "fields": {
        "album": "Liebe ist für alle da",
        "artist": "Rammstein",
        "year": 2009,
        "category_scores": {
            "cells": [
                { "address" : { "cat" : "pop" },  "value": 0.1 },
                { "address" : { "cat" : "rock" }, "value": 1.0 },
                { "address" : { "cat" : "jazz" }, "value": 0 }
            ]
        }
    }
}
{
    "put": "id:mynamespace:music::love-id-here-to-stay",
    "fields": {
        "album": "Love Is Here To Stay",
        "artist": "Diana Krall",
        "year": 2018,
        "category_scores": {
            "cells": [
                { "address" : { "cat" : "pop" },  "value": 0.4 },
                { "address" : { "cat" : "rock" }, "value": 0   },
                { "address" : { "cat" : "jazz" }, "value": 0.8 }
            ]
        }
     }
}

{
    "put": "id:mynamespace:music::when-we-all-fall-asleep-where-do-we-go",
    "fields": {
        "album": "When We All Fall Asleep, Where Do We Go?",
        "artist": "Billie Eilish",
        "year": 2019,
        "category_scores": {
            "cells": [
                { "address" : { "cat" : "pop" },  "value": 1.0 },
                { "address" : { "cat" : "rock" }, "value": 0 },
                { "address" : { "cat" : "jazz" }, "value": 0.1 }
            ]
        }
    }
}

各ドキュメント(アルバム)には、

  • album:タイトル
  • artist:アーティスト
  • year:発表年月
  • category_scores:アルバムが各カテゴリ(pop, rock, jazz)に属するスコア

が紐付いています。

Vespaでは以下のような形式でスキーマを定義します。

    document music {
        field artist type string {
            indexing: summary | index
        }
        field album type string {
            indexing: summary | index
        }
        field year type int {
            indexing: summary | attribute
        }
        field category_scores type tensor<float>(cat{}) {
            indexing: summary | attribute
        }
    }

本筋と離れてしまうため詳しくは説明しませんが、スキーマの詳細な説明についてはSchemasをご覧ください。

こちらのチュートリアルに従うと、以下のように検索を行うことができます。

例1「全ドキュメントを取得」

$ vespa query "yql=select id from music where true;"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 5
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 0.0,
                "source": "music"
            },
            {
                "id": "index:music/0/82705986116f101ee965c5c5",
                "relevance": 0.0,
                "source": "music"
            },
            {
                "id": "index:music/0/7b5f5c18d5fc145be152c0aa",
                "relevance": 0.0,
                "source": "music"
            },
            {
                "id": "index:music/0/e2ccdf7ee8e87d8e7b95b485",
                "relevance": 0.0,
                "source": "music"
            },
            {
                "id": "index:music/0/22fbf0ae2f3e60a661aa740a",
                "relevance": 0.0,
                "source": "music"
            }
        ]
    }
}

例2「albumフィールドに'head'を含むドキュメントを取得」

$ vespa query "yql=select id, album from music where album contains 'head';"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 1
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 0.16343879032006287,
                "source": "music",
                "fields": {
                    "album": "A Head Full of Dreams"
                }
            }
        ]
    }
}

Vespaにおけるランキングロジックの実装(rank-profile)

Vespaにドキュメントをフィードできたので、次はランキングロジックを実装してみましょう。
Vespaでは、rank-profileという形式でランキングロジックを実装します。

例えば、year順にドキュメントをランキングするには、以下のように記述します。

# rank_yearというランキングを定義

rank-profile rank_year inherits default {
    # Vespaでは2段階のランキングを行える。今回は1段階のランキングで実装。
    first-phase {
          # フィールド: yearの値をスコアとする
          expression: attribute(year)
    }
}

rank-profileの詳細な説明はRanking Expressions and Featuresをご覧ください。
上記で定義したrank_yearでランキングするには、以下のようにリクエストします。

$ vespa query \
"yql=select id, year from music where true;" \
"ranking=rank_year"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 5
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/22fbf0ae2f3e60a661aa740a",
                "relevance": 2019.0,
                "source": "music",
                "fields": {
                    "year": 2019
                }
            },
            {
                "id": "index:music/0/82705986116f101ee965c5c5",
                "relevance": 2018.0,
                "source": "music",
                "fields": {
                    "year": 2018
                }
            },
            {
                "id": "index:music/0/7b5f5c18d5fc145be152c0aa",
                "relevance": 2016.0,
                "source": "music",
                "fields": {
                    "year": 2016
                }
            },
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 2015.0,
                "source": "music",
                "fields": {
                    "year": 2015
                }
            },
            {
                "id": "index:music/0/e2ccdf7ee8e87d8e7b95b485",
                "relevance": 2009.0,
                "source": "music",
                "fields": {
                    "year": 2009
                }
            }
        ]
    }
}

relevanceの値がyearの値と等しくなっており、yearの大きい順にドキュメントがランキングされていることが分かります。

Vespaにおける、検索クエリとフィールド間のマッチ度計算には、nativeRankという計算方法がデフォルトで利用されます。
前節で「albumフィールドに'head'を含むドキュメントを取得」の例を紹介しましたが、検索クエリ'head'とalbumフィールドの間で計算されたnativeRankの値がrelevance=0.16343879032006287として出力されています。
nativeRankの詳細な説明はnativeRank Referenceをご覧ください。


rank-profileは四則演算、cos, log, sqrtをはじめとした数学関数、さらにif文なども利用できます。
例えば、yearが2016以上ならyear+10を、2016未満ならsqrt(year)をスコアとして利用する、という無茶苦茶なランキングロジックは、以下のように記述できます。

    rank-profile rank_complicated inherits default {
        first-phase {
            expression: if(attribute(year) >= 2016, attribute(year) + 10, sqrt(attribute(year)))
        }
    }
$ vespa query \
"yql=select id, year from music where true;" \
"ranking=rank_complicated"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 5
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/22fbf0ae2f3e60a661aa740a",
                "relevance": 2029.0,
                "source": "music",
                "fields": {
                    "year": 2019
                }
            },
            {
                "id": "index:music/0/82705986116f101ee965c5c5",
                "relevance": 2028.0,
                "source": "music",
                "fields": {
                    "year": 2018
                }
            },
            {
                "id": "index:music/0/7b5f5c18d5fc145be152c0aa",
                "relevance": 2026.0,
                "source": "music",
                "fields": {
                    "year": 2016
                }
            },
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 44.88875137492688,
                "source": "music",
                "fields": {
                    "year": 2015
                }
            },
            {
                "id": "index:music/0/e2ccdf7ee8e87d8e7b95b485",
                "relevance": 44.82186966202994,
                "source": "music",
                "fields": {
                    "year": 2009
                }
            }
        ]
    }
}

ちょっとテクい検索ランキング集

ここからは、Vespaを使ってちょっとテクい検索ランキングをいくつか紹介していきます。

検索結果をランダムに並び替える

検索クエリにマッチしたドキュメントから、ランダムにドキュメントを掲出したい…というケースがあります。 あんまりないです、レコメンドならあるかな…
Vespaではrandomという記述で、ドキュメントごとにランダム値を得ることができます。

    rank-profile rank_random inherits default {
        first-phase {
            expression: random
        }
    }

もちろん、ランダム値をif文の条件などに使うことも可能です。

検索時の時刻を利用して並び替える

ドキュメントの新しさを考慮してランキングするために、検索時の時刻を利用したい…というケースがあります。
例えば、フィールドに投稿時刻を格納しておき、検索時の時刻との差分を取ることで、ドキュメントの新しさを考慮できます。

Vespaではnowという記述で、検索時の時刻をunixtimeで取得できます。

    rank-profile rank_now inherits default {
        # 以下の例では全てのドキュメントのスコアが検索時の時刻(unixtime)となる。
        # フィールド値との差分を取りたい場合は、例えば now - attribute(indexed_time) といった指定をすればOK
        first-phase {
            expression: now
        }
    }

検索リクエスト時にランキングで利用する値を渡す

ドキュメントに付与されている情報だけでなく、検索リクエストに紐づく情報もランキングで利用したい…というケースがあります。
例えば、

  • ユーザーの年齢をランキングで考慮したい
  • 新規ユーザー・既存ユーザーでランキングを変更したい

といったケースが考えられます。

Vespaでは検索リクエスト時のパラメータをrank-profile内で利用する機能が実装されています。
例えば、yearからユーザーの年齢を引いた値をスコアとして利用したい際には、以下のようにrank-profileを定義します。

    rank-profile rank_with_user_age inherits default {
        # 検索リクエスト時に受け取るパラメータおよびデフォルト値を指定する
        rank-properties {
            query(user_age): 0
        }

        first-phase {
            expression: attribute(year) - query(user_age)
        }
    }

検索時は、ranking.features.query(key)=valueを付与してリクエストします。

$ vespa query \
"yql=select id, year from music where true;" \
"ranking=rank_with_user_age" \
"ranking.features.query(user_age)=30"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 5
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/22fbf0ae2f3e60a661aa740a",
                "relevance": 1989.0,
                "source": "music",
                "fields": {
                    "year": 2019
                }
            },
            {
                "id": "index:music/0/82705986116f101ee965c5c5",
                "relevance": 1988.0,
                "source": "music",
                "fields": {
                    "year": 2018
                }
            },
            {
                "id": "index:music/0/7b5f5c18d5fc145be152c0aa",
                "relevance": 1986.0,
                "source": "music",
                "fields": {
                    "year": 2016
                }
            },
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 1985.0,
                "source": "music",
                "fields": {
                    "year": 2015
                }
            },
            {
                "id": "index:music/0/e2ccdf7ee8e87d8e7b95b485",
                "relevance": 1979.0,
                "source": "music",
                "fields": {
                    "year": 2009
                }
            }
        ]
    }
}

もちろんif文の条件にも使えるので、ユーザーの属性によるランキングを出し分けなどにも応用できます。

クエリとドキュメントのベクトル表現の内積をランキングに使う

世はNeural Ranking時代です。
クエリベクトルとドキュメントベクトルの間で計算されるスコアをランキングで使いたくなります。

Vespaではベクトルにまつわる計算を行うために、tensor型をサポートしています。
実は、この記事で紹介したチュートリアルでも、アルバムが各カテゴリ(pop, rock, jazz)に属するスコアを表すために、tensor型を使っています *1

        field category_scores type tensor<float>(cat{}) {
            indexing: summary | attribute
        }

以下のドキュメントは、popに属するスコア=1、rockに属するスコア=0.2、jazzに属するスコア=0を表します。

{
    "put": "id:mynamespace:music::a-head-full-of-dreams",
    "fields": {
        "album": "A Head Full of Dreams",
        "artist": "Coldplay",
        "year": 2015,
        "category_scores": {
            "cells": [
                { "address" : { "cat" : "pop" },  "value": 1 },
                { "address" : { "cat" : "rock" }, "value": 0.2 },
                { "address" : { "cat" : "jazz" }, "value": 0 }
            ]
        }
    }
}

さて、ユーザーの各カテゴリへの興味度合いをランキングに反映してみましょう。
ユーザーのpopへの興味が0.5、rockへの興味が0.3、jazzへの興味が0.2だとします。
ドキュメントのスコアを、ユーザーの各カテゴリへの興味度合いと、アルバムが各カテゴリに属するスコアとの内積で定義します。
上のドキュメントの例では、1 * 0.5 + 0.2 * 0.3 + 0 * 0.2 = 0.56となります。

rank-profileは以下のように定義します。

    rank-profile rank_albums inherits default {
        first-phase {
            expression: sum(query(user_profile) * attribute(category_scores))
        }
    }

上で、検索リクエストにranking.features.query(key)=valueを付与することで、ランキングで利用する値を渡せることを紹介しました。
同じように、検索リクエストにユーザーの各カテゴリへの興味度合いを渡してリクエストしてみます。

$ vespa query \
"yql=select id, category_scores from music where true;" \
"ranking=rank_albums" \
"ranking.features.query(user_profile)={{cat:pop}:0.5,{cat:rock}:0.3,{cat:jazz}:0.2}"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 5
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 0.5600000023841858,
                "source": "music",
                "fields": {
                    "category_scores": {
                        "cells": [
                            {
                                "address": {
                                    "cat": "pop"
                                },
                                "value": 1.0
                            },
                            {
                                "address": {
                                    "cat": "rock"
                                },
                                "value": 0.20000000298023224
                            },
                            {
                                "address": {
                                    "cat": "jazz"
                                },
                                "value": 0.0
                            }
                        ]
                    }
                }
            },
            {
                "id": "index:music/0/22fbf0ae2f3e60a661aa740a",
                "relevance": 0.5200000014156103,
                "source": "music",
                "fields": {
                    "category_scores": {
                        "cells": [
                            {
                                "address": {
                                    "cat": "pop"
                                },
                                "value": 1.0
                            },
                            {
                                "address": {
                                    "cat": "rock"
                                },
                                "value": 0.0
                            },
                            {
                                "address": {
                                    "cat": "jazz"
                                },
                                "value": 0.10000000149011612
                            }
                        ]
                    }
                }
            },
            {
                "id": "index:music/0/82705986116f101ee965c5c5",
                "relevance": 0.36000001430511475,
                "source": "music",
                "fields": {
                    "category_scores": {
                        "cells": [
                            {
                                "address": {
                                    "cat": "pop"
                                },
                                "value": 0.4000000059604645
                            },
                            {
                                "address": {
                                    "cat": "rock"
                                },
                                "value": 0.0
                            },
                            {
                                "address": {
                                    "cat": "jazz"
                                },
                                "value": 0.800000011920929
                            }
                        ]
                    }
                }
            },
            {
                "id": "index:music/0/e2ccdf7ee8e87d8e7b95b485",
                "relevance": 0.350000012665987,
                "source": "music",
                "fields": {
                    "category_scores": {
                        "cells": [
                            {
                                "address": {
                                    "cat": "pop"
                                },
                                "value": 0.10000000149011612
                            },
                            {
                                "address": {
                                    "cat": "rock"
                                },
                                "value": 1.0
                            },
                            {
                                "address": {
                                    "cat": "jazz"
                                },
                                "value": 0.0
                            }
                        ]
                    }
                }
            },
            {
                "id": "index:music/0/7b5f5c18d5fc145be152c0aa",
                "relevance": 0.30000001192092896,
                "source": "music",
                "fields": {
                    "category_scores": {
                        "cells": [
                            {
                                "address": {
                                    "cat": "pop"
                                },
                                "value": 0.0
                            },
                            {
                                "address": {
                                    "cat": "rock"
                                },
                                "value": 1.0
                            },
                            {
                                "address": {
                                    "cat": "jazz"
                                },
                                "value": 0.0
                            }
                        ]
                    }
                }
            }
        ]
    }
}

tensor型を使いこなすと、様々なランキングロジックを実装できます。

例えば、

  • 検索クエリから推定されたカテゴリ(Query Understanding)と、ドキュメントのカテゴリスコアとの内積を計算
  • ユーザー×ドキュメントの閲覧行列からMatrix Factorizationなどでユーザーベクトル・ドキュメントベクトルを作成し、内積を計算

なんかも実現できます。

また、内積値をそのまま使うだけでなく、検索クエリとのマッチ度との和など、他スコアと組み合わせて使うこともできます。
ランキングアイデアが膨らみますね!

機械学習モデルでランキングする

Vespaは四則演算やif文をサポートしているので、その気になれば線形モデル・NeuralNetモデル・GBDTモデルを記述していくことができます…が、かなり面倒です。
なんとVespaでは、機械学習モデルをそのまま読み込むことができます!
具体的には、TensorFlow・XGBoost・LightGBMで構築したモデルおよび、ONNX形式のモデルに対応しています *2

ここでは、LightGBMで構築したモデルをVespaのランキングで使う方法を紹介します。
rank-profileは以下のように記述します。

    rank-profile rank_with_lightgbm inherits default {
        # モデルで利用する特徴量を明示的に定義した
        function f_year() {
            expression: attribute(year)
        }

        function f_nativerank_score() {
            expression: nativeRank(album)
        }

        function f_inner_product() {
            expression: sum(query(user_profile) * attribute(category_scores))
        }
        # lightgbm_model.jsonを利用
        first-phase {
            expression: lightgbm("lightgbm_model.json")
        }
    }

LightGBMモデルは、以下のコードで適当に作ってみました。

import json
import lightgbm as lgb
import pandas as pd
import numpy as np

N = 10000
df = pd.DataFrame()
df["label"] = np.random.choice([0, 1], N)
df["f_year"] = np.random.choice([2009, 2015, 2016, 2018, 2019], N)
df["f_nativerank_score"] = np.random.rand(N)
df["f_inner_product"] = np.random.rand(N)

train_data = lgb.Dataset(
    df[["f_year", "f_nativerank_score", "f_inner_product"]],
    df["label"]
)
params = {
    "objective": "binary",
    "metric": "binary_logloss",
    "num_leaves": 3
}
model = lgb.train(params, train_data, num_boost_round=5)

with open("lightgbm_model.json", "w") as f:
    json.dump(model.dump_model(), f)

リクエストしてみましょう!

$ vespa query \
"yql=select id from music where album contains 'head';" \
"ranking=rank_with_lightgbm" \
"ranking.features.query(user_profile)={{cat:pop}:0.5,{cat:rock}:0.3,{cat:jazz}:0.2}"
{
    "root": {
        "id": "toplevel",
        "relevance": 1.0,
        "fields": {
            "totalCount": 1
        },
        "coverage": {
            "coverage": 100,
            "documents": 5,
            "full": true,
            "nodes": 1,
            "results": 1,
            "resultsFull": 1
        },
        "children": [
            {
                "id": "index:music/0/de97c3f0cf0d1122b3494a44",
                "relevance": 0.5075908067816924,
                "source": "music"
            }
        ]
    }
}

というわけで、LightGBMモデルを使ってスコア計算できました。
LightGBMモデルのスコアをそのままランキングに使っても良いですが、ビジネスロジックなどを混ぜ込むことも可能です!

まとめ

この記事ではVespaの機能および、ちょっとテクい検索ランキングの実現方法を紹介しました。
普段ElasticsearchやSolrを利用している方は、Vespaではとても柔軟にランキングロジックを記述できることに驚かれたのではないでしょうか。 驚いてください
Vespaの機能はとても豊富なので、この記事では扱いきれていない機能もあります。
ランキング改善に興味がある方は、ぜひVespaを試してみてください!

*1: この記事の例ではcatをキーとしたmap型のように定義していますが、配列型のように定義することも可能です。

*2:これらのフレームワークに対応していれば、実用上はほぼ不便無いと思います。