Elasticsearch 8.X で類似ベクトル検索

Elasticsearch 8.X では "knn" で類似ベクトル検索を行うことができます。
類似度の尺度としてここでは "cosine" を使いますが、"l2_norm" を指定することもできます。
それぞれ以下に基づく検索となります。
・cosine: ２ベクトルのコサイン
・l2_norm: ２ベクトル間の距離の二乗
cosine、l2_norm の両方を使いたい場合には、"vector_cos" と "vector_l2" に同じベクトルを登録しておき、用途に応じて検索対象を切り替えて使うということもできます。
■スキーマ（create_knn_test1.json）

{
  "mappings": {
    "dynamic": "strict",
    "properties": {
      "id": {"type": "keyword", "store": "true"},
      "vector": {"type": "dense_vector", "dims": 2, "index": true, "similarity": "cosine"}
    }
  }
}

インデックス作成

$ curl 'http://localhost:9200/knn_test1/?pretty' \
    -X PUT \
    -H 'Content-Type: application/json' \
    -T create_knn_test1.json

■データ登録例（bulk_knn_test1.jsonl）

{"index": {"_index": "knn_test1", "_id": "id_01"}}
{"id": "id_01", "vector": [1.0, 0.50]}
{"index": {"_index": "knn_test1", "_id": "id_02"}}
{"id": "id_02", "vector": [1.0, 0.60]}
{"index": {"_index": "knn_test1", "_id": "id_03"}}
{"id": "id_03", "vector": [1.0, 0.70]}
{"index": {"_index": "knn_test1", "_id": "id_04"}}
{"id": "id_04", "vector": [1.0, 0.80]}
{"index": {"_index": "knn_test1", "_id": "id_05"}}
{"id": "id_05", "vector": [1.0, 0.90]}
{"index": {"_index": "knn_test1", "_id": "id_06"}}
{"id": "id_06", "vector": [1.0, 1.00]}
{"index": {"_index": "knn_test1", "_id": "id_07"}}
{"id": "id_07", "vector": [1.0, 1.15]}
{"index": {"_index": "knn_test1", "_id": "id_08"}}
{"id": "id_08", "vector": [1.0, 1.25]}
{"index": {"_index": "knn_test1", "_id": "id_09"}}
{"id": "id_09", "vector": [1.0, 1.35]}
{"index": {"_index": "knn_test1", "_id": "id_10"}}
{"id": "id_10", "vector": [1.0, 1.45]}

登録

$ curl "http://localhost:9200/knn_test1/_bulk?pretty" \
     -X POST \
     -H 'Content-Type: application/x-ndjson' \
     -T bulk_knn_test1.jsonl

■検索プログラム（k=3, num_candidates=5）

import sys
import json
from elasticsearch import Elasticsearch

url = 'http://localhost:9200'
es = Elasticsearch(url)

q = {
    'knn': {
        'field': 'vector',
        'query_vector': [1.0, 1.0],
        'k': 3,
        'num_candidates': 5,
    },
    'fields': ['id']
}

res = es.knn_search(index='knn_test1', body=q)
print(json.dumps(res['hits']['hits'], indent=2))

■実行結果（k=3, num_candidates=5）
k=3, num_candidates=5 の場合、3レコードのみが返却されます。

[
  {
    "id": "id_06",
    "score": 1.0
  },
  {
    "id": "id_05",
    "score": 0.99930894
  },
  {
    "id": "id_07",
    "score": 0.9987876
  }
]

■パラメータ変更（k=5, num_candidates=5）
k=5, num_candidates=5 にプログラムを変更します。

q = {
    'knn': {
        'field': 'vector',
        'query_vector': [1.0, 1.0],
        'k': 5,
        'num_candidates': 5,
    },
    'fields': ['id']
}

res = es.knn_search(index='knn_test1', body=q)
items = [
    {'id': item['_id'], 'score': item['_score']}
    for item in res['hits']['hits']
]
print(json.dumps(items, indent=2))

■実行結果（k=5, num_candidates=5）
k=5, num_candidates=5 に変更すると、検索結果は 5件になります。

[
  {
    "id": "id_06",
    "score": 1.0
  },
  {
    "id": "id_05",
    "score": 0.99930894
  },
  {
    "id": "id_07",
    "score": 0.9987876
  },
  {
    "id": "id_08",
    "score": 0.99694186
  },
  {
    "id": "id_04",
    "score": 0.9969418
  }
]