ElasticsearchのKNNクエリーができることを一通りまとめたほうが良いなーと思っていたので、今回まとめてみることにする。ANNとしての性能や精度は一旦忘れて、クエリーとしてできることに注目する。
KNNクエリーを試すにあたって、わかりやすいデータでないと良いのかどうかがわからないので、都道府県の県庁所在地的な緯度経度のデータを用いることにする。（ジオ・サーチで良いじゃんとか、地球は球面だよねとか、は置いておいて、結果の理解のしやすさだけに注目しているので、そのへんの細かいことは気にしないものとする）
環境準備
今回は、Elasticsearch 8.12で確認するので、別途起動して置いてください。Dockereとか使えるなら、codelibs/docker-elasticsearchでも簡単に起動できます。
まずは、インデックスマッピングを指定して、インデクスを作成します。今回は、japanインデックスを作成します。
$ curl -XPUT -H "Content-Type: application/json" localhost:9200/japan -d '
{
"mappings": {
"properties": {
"region": {
"type": "keyword"
},
"prefecture": {
"type": "keyword"
},
"location": {
"type": "dense_vector",
"dims": 2,
"similarity": "l2_norm"
}
}
},
"settings": {
"index": {
"number_of_shards": 1,
"number_of_replicas": 0
}
}
}
'
regionに地方名、prefectureに都道府県名、locationに緯度経度を入れてます。japanインデックスが正しく作成されているかを以下を実行して確認しておきます。
$ curl -XGET "localhost:9200/japan?pretty"
次に、データをバルクで入れます。
$ curl -XPOST -H "Content-Type: application/json" "localhost:9200/japan/_bulk?refresh=true" -d '
{"index": {"_index": "japan", "_id": 1}}
{"region": "北海道地方", "prefecture": "北海道", "location": [43.064359, 141.347449]}
{"index": {"_index": "japan", "_id": 2}}
{"region": "東北地方", "prefecture": "青森県", "location": [40.824294, 140.740054]}
{"index": {"_index": "japan", "_id": 3}}
{"region": "東北地方", "prefecture": "岩手県", "location": [39.70353, 141.152667]}
{"index": {"_index": "japan", "_id": 4}}
{"region": "東北地方", "prefecture": "宮城県", "location": [38.268737, 140.872183]}
{"index": {"_index": "japan", "_id": 5}}
{"region": "東北地方", "prefecture": "秋田県", "location": [39.718175, 140.103356]}
{"index": {"_index": "japan", "_id": 6}}
{"region": "東北地方", "prefecture": "山形県", "location": [38.240127, 140.362533]}
{"index": {"_index": "japan", "_id": 7}}
{"region": "東北地方", "prefecture": "福島県", "location": [37.750146, 140.466754]}
{"index": {"_index": "japan", "_id": 8}}
{"region": "関東地方", "prefecture": "茨城県", "location": [36.341817, 140.446796]}
{"index": {"_index": "japan", "_id": 9}}
{"region": "関東地方", "prefecture": "栃木県", "location": [36.56575, 139.883526]}
{"index": {"_index": "japan", "_id": 10}}
{"region": "関東地方", "prefecture": "群馬県", "location": [36.391205, 139.060917]}
{"index": {"_index": "japan", "_id": 11}}
{"region": "関東地方", "prefecture": "埼玉県", "location": [35.857771, 139.647804]}
{"index": {"_index": "japan", "_id": 12}}
{"region": "関東地方", "prefecture": "千葉県", "location": [35.604563, 140.123179]}
{"index": {"_index": "japan", "_id": 13}}
{"region": "関東地方", "prefecture": "東京都", "location": [35.689185, 139.691648]}
{"index": {"_index": "japan", "_id": 14}}
{"region": "関東地方", "prefecture": "神奈川県", "location": [35.447505, 139.642347]}
{"index": {"_index": "japan", "_id": 15}}
{"region": "中部地方", "prefecture": "新潟県", "location": [37.901699, 139.022728]}
{"index": {"_index": "japan", "_id": 16}}
{"region": "中部地方", "prefecture": "富山県", "location": [36.695274, 137.211302]}
{"index": {"_index": "japan", "_id": 17}}
{"region": "中部地方", "prefecture": "石川県", "location": [36.594729, 136.62555]}
{"index": {"_index": "japan", "_id": 18}}
{"region": "中部地方", "prefecture": "福井県", "location": [36.06522, 136.221641]}
{"index": {"_index": "japan", "_id": 19}}
{"region": "中部地方", "prefecture": "山梨県", "location": [35.665102, 138.568985]}
{"index": {"_index": "japan", "_id": 20}}
{"region": "中部地方", "prefecture": "長野県", "location": [36.651282, 138.180972]}
{"index": {"_index": "japan", "_id": 21}}
{"region": "中部地方", "prefecture": "岐阜県", "location": [35.39116, 136.722204]}
{"index": {"_index": "japan", "_id": 22}}
{"region": "中部地方", "prefecture": "静岡県", "location": [34.976987, 138.383057]}
{"index": {"_index": "japan", "_id": 23}}
{"region": "中部地方", "prefecture": "愛知県", "location": [35.180247, 136.906698]}
{"index": {"_index": "japan", "_id": 24}}
{"region": "近畿地方", "prefecture": "三重県", "location": [34.730547, 136.50861]}
{"index": {"_index": "japan", "_id": 25}}
{"region": "近畿地方", "prefecture": "滋賀県", "location": [35.004532, 135.868588]}
{"index": {"_index": "japan", "_id": 26}}
{"region": "近畿地方", "prefecture": "京都県", "location": [35.0209962, 135.7531135]}
{"index": {"_index": "japan", "_id": 27}}
{"region": "近畿地方", "prefecture": "大阪府", "location": [34.686492, 135.518992]}
{"index": {"_index": "japan", "_id": 28}}
{"region": "近畿地方", "prefecture": "兵庫県", "location": [34.69128, 135.183087]}
{"index": {"_index": "japan", "_id": 29}}
{"region": "近畿地方", "prefecture": "奈良県", "location": [34.685296, 135.832745]}
{"index": {"_index": "japan", "_id": 30}}
{"region": "近畿地方", "prefecture": "和歌山県", "location": [34.224806, 135.16795]}
{"index": {"_index": "japan", "_id": 31}}
{"region": "中国地方", "prefecture": "鳥取県", "location": [35.503463, 134.238258]}
{"index": {"_index": "japan", "_id": 32}}
{"region": "中国地方", "prefecture": "島根県", "location": [35.472248, 133.05083]}
{"index": {"_index": "japan", "_id": 33}}
{"region": "中国地方", "prefecture": "岡山県", "location": [34.66132, 133.934414]}
{"index": {"_index": "japan", "_id": 34}}
{"region": "中国地方", "prefecture": "広島県", "location": [34.396033, 132.459595]}
{"index": {"_index": "japan", "_id": 35}}
{"region": "中国地方", "prefecture": "山口県", "location": [34.185648, 131.470755]}
{"index": {"_index": "japan", "_id": 36}}
{"region": "四国地方", "prefecture": "徳島県", "location": [34.065732, 134.559293]}
{"index": {"_index": "japan", "_id": 37}}
{"region": "四国地方", "prefecture": "香川県", "location": [34.34014, 134.04297]}
{"index": {"_index": "japan", "_id": 38}}
{"region": "四国地方", "prefecture": "愛媛県", "location": [33.841649, 132.76585]}
{"index": {"_index": "japan", "_id": 39}}
{"region": "四国地方", "prefecture": "高知県", "location": [33.55969, 133.530887]}
{"index": {"_index": "japan", "_id": 40}}
{"region": "九州地方", "prefecture": "福岡県", "location": [33.606767, 130.418228]}
{"index": {"_index": "japan", "_id": 41}}
{"region": "九州地方", "prefecture": "佐賀県", "location": [33.249367, 130.298822]}
{"index": {"_index": "japan", "_id": 42}}
{"region": "九州地方", "prefecture": "長崎県", "location": [32.744542, 129.873037]}
{"index": {"_index": "japan", "_id": 43}}
{"region": "九州地方", "prefecture": "熊本県", "location": [32.790385, 130.742345]}
{"index": {"_index": "japan", "_id": 44}}
{"region": "九州地方", "prefecture": "大分県", "location": [33.2382, 131.612674]}
{"index": {"_index": "japan", "_id": 45}}
{"region": "九州地方", "prefecture": "宮崎県", "location": [31.91109, 131.423855]}
{"index": {"_index": "japan", "_id": 46}}
{"region": "九州地方", "prefecture": "鹿児島県", "location": [31.560219, 130.557906]}
{"index": {"_index": "japan", "_id": 47}}
{"region": "九州地方", "prefecture": "沖縄県", "location": [26.211538, 127.681115]}
'
47件のデータが登録されているか確認しておきます。
$ curl localhost:9200/_cat/indices (git)-[master]
green open japan U0RsVDw-S72OuDf-PgBtCx 1 0 47 0 9.4kb 9.4kb 9.4kb
データが登録されたら、早速、KNNクエリーを投げていきましょう。KNNクエリーを投げるには、
- knnオプションで指定する
- knnクエリーで指定する
の２パターンがあります。どちらを利用するかは、検索要件により選んでください。
以前は_knn_searchエンドポイントに投げる方法もありましたが、Elasticsearch 8.4でdeprecatedになっています。
KNNオプション
まずは、knnオプションでのクエリーを投げてみます。東京の緯度経度から近い３件を取得する検索です。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"k": 3,
"num_candidates": 10
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 11,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 3,
"relation" : "eq"
},
"max_score" : 1.0,
"hits" : [
{
"_index" : "japan",
"_id" : "13",
"_score" : 1.0,
"fields" : {
"prefecture" : [
"東京都"
]
}
},
{
"_index" : "japan",
"_id" : "11",
"_score" : 0.97054905,
"fields" : {
"prefecture" : [
"埼玉県"
]
}
},
{
"_index" : "japan",
"_id" : "14",
"_score" : 0.9426493,
"fields" : {
"prefecture" : [
"神奈川県"
]
}
}
]
}
}
東京で探しているので、東京が返ってくるのは期待通りで、それ以外も近いところが返ってきているようなので、たぶん、良さそうです。実際にベクトル検索する際には、ベクトルの部分がそこそこの量の数値配列になるので、_sourceはfalseで投げるのが良いです。必要な値は、fieldsで指定して取得します。_scoreはl2_normの値をベースに1 / (1 + l2_norm(query, vector)^2)という感じで、近いものが1になるように算出されます。
取得する件数は、kで指定しています。今回は、num_candidatesが10、kが3なので、各シャードから10件取って、3件を返すことになります。
今回は、距離を考えているので、l2_normを利用していますが、ベクトルの種類によって、dot_productやcosineを適切に選んでください。
KNNクエリー
次に、KNNクエリーの場合です。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"query": {
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"num_candidates": 10
}
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 2,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 10,
"relation" : "eq"
},
"max_score" : 1.0,
"hits" : [
{
"_index" : "japan",
"_id" : "13",
"_score" : 1.0,
"fields" : {
"prefecture" : [
"東京都"
]
}
},
{
"_index" : "japan",
"_id" : "11",
"_score" : 0.97054905,
"fields" : {
"prefecture" : [
"埼玉県"
]
}
},
{
"_index" : "japan",
"_id" : "14",
"_score" : 0.9426493,
"fields" : {
"prefecture" : [
"神奈川県"
]
}
}
]
}
}
同じように返ってきます。KNNクエリーでは、queryの配下に普通のクエリーと同様に指定してます。KNNクエリーではkがありません。なので、今回は、num_candidatesが10なので、各シャードから10件取得して、sizeが3なので、3件が検索結果として返ります。
rescoreクエリーとの組み合わせ
rescoreクエリーと組み合わせて使ってみます。ベクトル検索をして、取得したものを更にベクトル計算して並び替えるような場合です。
今回は、東京から近い１０件を取得して、取得したものをさらに北海道から近い順に並び替えて、３件取得するクエリーです。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"k": 10,
"num_candidates": 10
},
"rescore": {
"window_size": 10,
"query": {
"rescore_query": {
"script_score": {
"query": {
"match_all": {}
},
"script": {
"source": "1 / (1 + l2norm(params.query_vector, '\''location'\''))",
"params": {
"query_vector": [43.064359, 141.347449]
}
}
}
},
"query_weight" : 0,
"rescore_query_weight" : 1.0
}
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 4,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 10,
"relation" : "eq"
},
"max_score" : 0.13052356,
"hits" : [
{
"_index" : "japan",
"_id" : "9",
"_score" : 0.13052356,
"fields" : {
"prefecture" : [
"栃木県"
]
}
},
{
"_index" : "japan",
"_id" : "8",
"_score" : 0.12849167,
"fields" : {
"prefecture" : [
"茨城県"
]
}
},
{
"_index" : "japan",
"_id" : "10",
"_score" : 0.12416161,
"fields" : {
"prefecture" : [
"群馬県"
]
}
}
]
}
}
KNNオプションで10件を取得しています（KNNクエリーでも同じことができます）。rescoreのクエリーで北海道の緯度経度を指定して、スクリプトで距離計算をして、rescoreの計算結果を最終スコアとして利用しています。北関東の結果が返ってきているので、良さそうな感じですね。
このようにrescoreでスクリプトを利用して、ベクトル計算をして並び替えるというのも可能です。しかし、KNNオプションの方はHNSWで近似最近傍探索になりますが、rescoreの方はスクリプトで計算する普通に距離計算されるので、window_sizeが大きい場合やベクトルの次元数が高い場合は、計算性能を考慮して利用を考えるのが良いです。
他のクエリーとの組み合わせ
他のクエリーとの組み合わせて利用することを考えてみます。たとえば、KNNクエリーをtermクエリーとORで使う場合を考えてみます。東京から近いものを１０件取得して、中部地方のデータをORで取得する検索を投げてみます。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"query": {
"bool": {
"should": [
{
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"num_candidates": 10
}
},
{
"term": {
"region": "中部地方"
}
}
]
}
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 10,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 16,
"relation" : "eq"
},
"max_score" : 2.0621996,
"hits" : [
{
"_index" : "japan",
"_id" : "19",
"_score" : 2.0621996,
"fields" : {
"prefecture" : [
"山梨県"
]
}
},
{
"_index" : "japan",
"_id" : "22",
"_score" : 1.9305023,
"fields" : {
"prefecture" : [
"静岡県"
]
}
},
{
"_index" : "japan",
"_id" : "20",
"_score" : 1.8575637,
"fields" : {
"prefecture" : [
"長野県"
]
}
}
]
}
}
それっぽい３件が返ってきました。スコアに注目すると、2を超えるものもあるので、termクエリーとKNNクエリーの合算スコアになっています。今回のKNNクエリーのスコアは1以下になりますが、termクエリーの方は1.6…くらいの１以上の値が返ってくるので、このような結果になります。ですので、KNNクエリーを他のクエリーと組み合わせて利用する場合は、スコアの大きさの違いなどにも注意が必要です。
KNNオプションでも同様に投げてみます。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"query": {
"term": {
"region": "中部地方"
}
},
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"k": 10,
"num_candidates": 10
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 4,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 16,
"relation" : "eq"
},
"max_score" : 2.0621996,
"hits" : [
{
"_index" : "japan",
"_id" : "19",
"_score" : 2.0621996,
"fields" : {
"prefecture" : [
"山梨県"
]
}
},
{
"_index" : "japan",
"_id" : "22",
"_score" : 1.9305023,
"fields" : {
"prefecture" : [
"静岡県"
]
}
},
{
"_index" : "japan",
"_id" : "20",
"_score" : 1.8575637,
"fields" : {
"prefecture" : [
"長野県"
]
}
}
]
}
}
同じ感じの結果になりました。
スコアのスケールが異なり、termクエリーのほうが強いため、boostを追加して、調整してみます。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"query": {
"term": {
"region": {
"value": "中部地方",
"boost": 0.4
}
}
},
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"k": 10,
"num_candidates": 10,
"boost": 1.0
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 3,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 16,
"relation" : "eq"
},
"max_score" : 1.0902542,
"hits" : [
{
"_index" : "japan",
"_id" : "19",
"_score" : 1.0902542,
"fields" : {
"prefecture" : [
"山梨県"
]
}
},
{
"_index" : "japan",
"_id" : "13",
"_score" : 1.0,
"fields" : {
"prefecture" : [
"東京都"
]
}
},
{
"_index" : "japan",
"_id" : "11",
"_score" : 0.97054905,
"fields" : {
"prefecture" : [
"埼玉県"
]
}
}
]
}
}
termクエリーを0.4にして、KNNの方は1のままにすると、結果が変わることを確認できました。
しかし、実際の環境では、スコアが最大いくつになるのか？を考えるのは難しいので、boostでバランス調整するのも困難な場合が多いと思います。データの分布（ばらつき）的な話は一旦置いておいて、クエリー側のスコアを1以下に調整する方法として、function_scoreを使って以下のように調整するのもあると思います。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"query": {
"function_score": {
"query": {
"term": {
"region": "中部地方"
}
},
"script_score": {
"script": {
"source": "sigmoid(_score, params.k, params.a)",
"params": {
"k": 1,
"a": 1
},
"lang": "painless"
}
},
"boost_mode": "replace"
}
},
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"k": 10,
"num_candidates": 10
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 4,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 16,
"relation" : "eq"
},
"max_score" : 1.0605978,
"hits" : [
{
"_index" : "japan",
"_id" : "19",
"_score" : 1.0605978,
"fields" : {
"prefecture" : [
"山梨県"
]
}
},
{
"_index" : "japan",
"_id" : "13",
"_score" : 1.0,
"fields" : {
"prefecture" : [
"東京都"
]
}
},
{
"_index" : "japan",
"_id" : "11",
"_score" : 0.97054905,
"fields" : {
"prefecture" : [
"埼玉県"
]
}
}
]
}
}
sigmoid関数を使って、スコアを1以下になるようにしています。sigmoid(value, k, a) = value^a/ (k^a + value^a)で計算されていますが、kとaを調整してバランスを取る的な感じになると思います。別に、sigmoidでなくても、データに適した計算式に変えるのでも良いと思います。
この辺の話として、ハイブリッドサーチで検索結果をマージする方法が悩ましいので、Elastic社は有償でRRF(Reciprocal Rank Fusion)でベクトル検索の結果をマージする方法を提供してます。
フィルター検索
ベクトル検索をするとしても実際の場面においては、近傍探索以外にも絞り込み条件がある場合がほとんどだと思います。KNNの検索においてもフィルタする方法が提供されています。フィルタ条件の適用の仕方として
- KNNの前にフィルタする: pre filter
- KNNの後にフィルタする: post filter
のパターンが存在します。KNNクエリーで他の検索条件と合わせて検索する場合がpost filterになります。今までのKNNクエリーの例ではbool-shouldを使っているので、フィルタされていませんが、bool-mustなどで検索条件を組み立てれば、KNNで取得したものに対して検索結果が絞り込まれることになります。
post filterで検索する場合、たとえば、KNNで10件取得したい場合、絞り込み条件が厳しすぎると、10件取得できないなどもありえます。なので、KNNで取得したい件数が決まっている場合は注意必要です。
pre filterで検索するには、KNNオプションまたはKNNクエリーでfilterとして指定します。
$ curl -XPOST -H 'Content-Type: application/json' "localhost:9200/japan/_search?pretty" -d'
{
"size": 3,
"knn": {
"field": "location",
"query_vector": [35.689185, 139.691648],
"k": 10,
"num_candidates": 10,
"filter": {
"term": {
"region": "中部地方"
}
}
},
"fields": ["prefecture"],
"_source": false
}
'
{
"took" : 5,
"timed_out" : false,
"_shards" : {
"total" : 1,
"successful" : 1,
"skipped" : 0,
"failed" : 0
},
"hits" : {
"total" : {
"value" : 9,
"relation" : "eq"
},
"max_score" : 0.44229046,
"hits" : [
{
"_index" : "japan",
"_id" : "19",
"_score" : 0.44229046,
"fields" : {
"prefecture" : [
"山梨県"
]
}
},
{
"_index" : "japan",
"_id" : "22",
"_score" : 0.3105931,
"fields" : {
"prefecture" : [
"静岡県"
]
}
},
{
"_index" : "japan",
"_id" : "20",
"_score" : 0.23765454,
"fields" : {
"prefecture" : [
"長野県"
]
}
}
]
}
}
KNNクエリーでも同様にfilterで指定することができます。絞り込まれた対象に対して、KNNで取得したい件数を取りたい場合に利用することができます。
pre filterを利用する場合、絞り込みの対象が多い場合の性能問題など、pre filter利用時に検討すべき課題はあるので、利用する際には注意が必要です。
まとめ
他には、knnオプションで複数のフィールドを指定する話とか、Quantizationの話とかもありますが、KNNクエリーの基本的なところ話は一通り書いたような気がします。たぶん、上記の話を考えて使えば、ベクトル検索もいい感じにできるようになるのではないかなと思います。
