🧊

Cloudflare Workers + KVだけで、ベクトル検索を実現する

力ずくで。

やってやれんことはないし、データの件数によってはナシではないって感じか。

あらすじ

  • とあるメディアの記事を、Embeddingsのベクトルにして検索したい
  • Embeddingsのベクトルをどこに保存して検索するかが問題
  • いわゆるベクトルDBを用意するのが普通かもしれないが、もっとコンパクトにやりたい
  • Cloudflare Workers + KVでなんとかできないか?

という流れ。

有限ではあるがそれなりの数のベクトル集合から、いかに対象のベクトルに近しいものを見つけ出すのかが問題。

データが少ないなら

データが1000件未満って程度であれば、特に問題はない。

  • KVにベクトルをそのまま保存しておき
  • Workerから`KV.list()`して全件取得
  • 対象のベクトルとそれぞれを`KV.get(key)`して比較
  • 並べ替えして、topKを選んで返す

というような愚直な方法でも、まあなんとかやれる。

1000件未満としたのは、1WorkerリクエストでKVにアクセスできる回数が制限されていて、それが1000回だから。

https://developers.cloudflare.com/workers/platform/limits/#kv-limits

`Operations/worker invocation`ってやつ。Paidプランでも変わらず。

データが多いなら

データが1000件を超える場合、1Workerリクエスト内で全件を取得しきれない。

これは、Workerを多段構成にすれば回避できる。(ここが力ずくポイント)
メインのWorkerでリクエストを受けたら、サブのWorkerに対象のベクトルを投げて、それぞれにKVにあるベクトルと比較させる。

Main
 - Sub1 - KV: 1000件まで
 - Sub2 - KV: 1000件まで
 - ...
 - SubN - KV: 1000件まで

SubのWorkerへのサブリクエストを並列で行えば、パフォーマンスという意味ではさっきと変わらず処理できる。

ちなみに、このパターンで対応できるのは、データの件数が50000件未満 or 1000000件未満の場合。
というのも、1Workerリクエストで外部リソースに`fetch()`できる回数が制限されているから。

https://developers.cloudflare.com/workers/platform/limits

`Subrequest`ってやつ。Freeなら50回、Paid:Bundledも50回、Paid:Unboundなら1000回まで。

注意点としては、Workerから別のWorkerを呼ぶ必要があるが、

  • 同じゾーンに配置することはできない
    • Pages Functions + WorkersとかならOK
  • Service bindingsは使えない
    • あれはこういった制限をオフロードできるものではないから

うーむ、力ずくだわ。

気になる制限

効率的な検索アルゴリズムを使うわけではなく、気合で全走査してるわけなので。

KV

1検索のたびに、N000回もREADするとなると、そのコストはどうなるかも知っておく必要がある。

そもそもREADできる回数は、Freeプランだと100000READ/日だが、Paidでは無制限。神。

コストは、現状のPaidプランで`10 million/month, + $0.50/million`とのこと。

https://developers.cloudflare.com/workers/platform/pricing/#workers-kv

なので1検索で1000件を走査する場合は、10000回検索するまでは範囲内で、そこから追加で課金される。さすがの安さだわ。

もし可能であれば、1ベクトル/キーで保存するのではなく、Nベクトル/キーで保存するようにすれば、READの回数は減らせる。

Worker

1Workerリクエストで使えるメモリが128MBまで。

なので、KVから1000件を読み出せるにしても、それが溢れない範囲でやらないといけない。

OpenAIのEmbeddingsであれば、たとえばInt16の範囲で量子化したりすると、そのサイズはJSONでも19kbから7kbくらいに減るので、少しは楽になれるはず。

Embeddingsのベクトルと、そのベクトルが表すテキスト自体も、最初から別のKVに分けておくなども吉かと。

というわけで

ベクトルDBのサービスを用意しなくても、KVだけでなんとかベクトル検索もできなくはない・・ということをやってみた。

1000件ぽっちならともかく、やっぱある程度以上のデータ量になると、高尚なアルゴリズムやらインデックスを工夫する方面を頑張ったほうがいいのは間違いない。
ただそれをエッジまわりだけで手軽にやる方法を見つけられなかったので、こういう力ずくの方法を試してみたという話。

そのうち、KVでRedisのベクトル検索の拡張みたいなのが使えたり、D1でSQLiteのベクトル検索の拡張みたいなのが使えたりするようになりません?