Spanner emulator で vector search をしてみます

Oct 22, 2024 · AI/ML Google Cloud Spanner Python ·

Google Cloud の提供する Cloud Spanner (以下 Spanner) はスケーラブルなクラウドの database で、最近では graph database の機能が追加されたりなど、開発も盛んに行われています。

Google Cloud の database 製品では vector search の機能を持つものも多く、Big Query や Cloud SQL などその特性に合わせて選択することができます。Spanner に機能追加されたのもしばらく前だったと記憶していますが、試したことがなかったので今回はその使い勝手を試してみようと思います。

※ ただし Vector functions は記事作成時点ではまだ preview 機能です。また、Spanner Standard では利用できず、Enterprise または Enterprise Plus が必要なようです。Spanner の Edition ごとの違いはこちら。

Spanner emulator を使います

1/10 instance から起動できるようになったとはいえ割とお金がかかるともよく言われる Spanner です。本番環境で使うには頼もしいことこの上ないんですけどね。というわけで、今回は Spanner の emulator を使用します。

この emulator、あまり気にしたことがなかったのですがこちらでソースが公開されており、見てみると (ハリボテではなく) 結構しっかりと実装されています。これだけメンテするのもなかなかたいへんそうです。ちなみにこの blog でもよくとりあげている Dev Container にも対応しており、 .devcontainer/devcontainer.json が用意されています。単に pre-built の image でかんたんに試したいときにはよいかもしれません。

emulator では動作しない機能もいくらかあり、性能が出るものではないとも言われていますが、閉じた環境での自動テストなどでは非常に使い勝手がよさそうです。

Dev Container を使います

いつもながら Visual Studio Code の Dev Container 機能を使用します。他の環境をご利用の方はお手数ですが、以下の構築手順を参考にして環境構築してください。

このブログでは他にも Dev Container に関する記事があります。Dev Container って何? と思われたかたにもご覧いただけたら幸いです。

Dev Container 用の設定をつくります。

{
	"name": "Dev",
	"dockerComposeFile": "compose.yml",
	"service": "dev",
	"workspaceFolder": "/workspace",
	"features": {
		"ghcr.io/devcontainers/features/go:1": {},
		"ghcr.io/devcontainers/features/python:1": {}
	}
}

services:

  dev:
    image: mcr.microsoft.com/devcontainers/base:ubuntu
    tty: true
    volumes:
      - ..:/workspace:cached

  spanner:
    image: gcr.io/cloud-spanner-emulator/emulator
    ports:
      - "127.0.0.1:9010:9010"
      - "127.0.0.1:9020:9020"

これらを置いた folder を VS Code で開き、Dev Container を起動すると、Spanner emulator の container と、Go と Python が入った container が起動します。後者には workspace (開いた folder) が mount されて VS Code の explorer に表示されますので、そこから file を開いて編集できてとても便利です。

準備ができましたら次に進みましょう。

database をつくります

Spanner emulator では Google Cloud の web UI が使えません (と思う) ので、gcloud CLI を使用します。

一度だけの操作なので Dev Container 側に入れず、host 側にあるものを使用します (無ければ install してください…)。Docker container として spanner emulator が起動していれば、 localhost:9020 が expose されており host 側からも接続可能になっていると思います。

まずは、gcloud CLI に対し、emulator 用の configuration を追加します。

gcloud config configurations create emulator
gcloud config set auth/disable_credentials true
gcloud config set project <project_id>
gcloud config set api_endpoint_overrides/spanner http://localhost:9020/

emulator なので、<project_id> はたぶん何でもいいのだと思います。(後で使用するものと合わせる必要はあります)

以下のようにして元の設定と切り替えることができます。

gcloud config configurations activate [emulator | default]

Spanner emulator に対し、instance と database を作成します。

❯ gcloud spanner instances create test \
   --config=emulator-config --description="Test Instance" --nodes=1
Creating instance...done.

❯ gcloud spanner databases create db1 \
  --instance=test \
  --database-dialect=GOOGLE_STANDARD_SQL
Creating database...done.

emulator に接続されていれば、いずれもあっという間に完了すると思います。
実物の Spanner instance がクラウド側にできちゃったりしていないか心配なようでしたら、web console 等から確認しておけばよいでしょう。

spanner-cli を使います (option)

gcloud CLI でも以下のようにして SQL を実行することができます。

gcloud spanner databases execute-sql db1 --instance=test \
  --sql='SELECT model FROM Bikes LIMIT 3'

が、OSS でメンテされている spanner-cli もあり、 mysql 等の CLI と同様の使用感で便利に使用できます。以下のようにして install できます。

go install github.com/cloudspannerecosystem/spanner-cli@latest

以下のように環境変数で emulator container に接続するように設定して使用します。

$ export export SPANNER_EMULATOR_HOST=spanner:9010
$ spanner-cli -p <project_id> -i test -d db1
Connected.
spanner> show tables ;
Empty set (0.02 sec)

やってみます (本題)

以前実質無料で気安くベクトル検索を体験するという記事で、Redis を使用して vector search をする例を書きました。今回は同様のことを Spanner (emulator) に対してやってみようと思います。

すこし長くなったので別ファイル main.py にしました。コード内容はほとんど Microsoft Copilot に生成してもらったもので、今回の用途に合わせて調整してあります。

前の記事同様、json に整形されたサンプルデータを Spanner (emulator) に load して使用します。

処理の流れ (main.py) は以下のとおりです。

SentenceTransformer で model を load (実行時に自動で download されます)
Spanner emulator に接続
json data、embeddings を格納する table を作成
json file を load
embeddings を生成しながら json data とあわせて table へ insert
検索文字列を定義、その embedding を生成
Spanner の COSINE_DISTANCE を使用して近傍検索を実行
結果を出力

なんというかことさらここに引用して説明するような部分も無いのですが、table 定義において

    CREATE TABLE Bikes (
        model STRING(1024),
        brand STRING(1024),
        price INT64,
        type STRING(1024),
        material STRING(1024),
        weight STRING(1024),
        description STRING(MAX),
        desc_embs ARRAY<FLOAT64>, # ←これ
    ) PRIMARY KEY (model)

という ARRAY 型の column を以下のようにして検索できるのが特徴的なのだと思います。(main.py からの引用)

    SELECT COSINE_DISTANCE(desc_embs, ARRAY[{query_embedding_str}]) as distance,
    brand, model, description
    FROM Bikes
    ORDER BY distance ASC
    LIMIT {k}

main.py を実行すると以下のようになります。

$ python main.py 
  ... (Python/library version 差に起因するらしい warning を省略)
Data inserted successfully.
Nearest neighbors:
COS Dist.: 0.4626744858732623, Brand: Nord, Model: Chook air 5
COS Dist.: 0.49474421824658277, Brand: nHill, Model: Summit
COS Dist.: 0.5413842678724139, Brand: Velorim, Model: Jigger
COS Dist.: 0.5493878815438571, Brand: Bicyk, Model: Hillcraft
COS Dist.: 0.5676556307656708, Brand: ScramBikes, Model: WattBike

以前の記事に貼ったものと同様の結果 (結果の順) が得られていることがわかります。

まとめ

emulator でも問題なく、Spanner の vector search を体験することができました。お金がかからない方法というのはそれだけで尊いものです。実装については直感的で迷うこともなかったので、特に「もともと Spanner を使っているがちょっとした用途に vector search を足したい」というときに便利に使えそうです。実物でやるとまた何かあったりするかもしれませんが、今日は忘れてここまでにしましょう。

Spanner の graph 機能はどうなったんだという心残りがずっとありますが、もしかしたらまた今度ということで… 👋

参考

Building a Search App with Spanner, Vector Search, and Gemini 1.0 Pro
- Spanner の Vertex AI 連携を使用して embeddings の管理を楽にしています。
- CREATE MODEL や ML.PREDICT を使用しています。(emulator で動くのかな… 🤔 ↓ 補足へ)

補足: Spanner emulator で ML.PREDICT を実行してみる

spanner> CREATE MODEL test_embs INPUT (
      ->   content STRING(MAX)
      -> ) OUTPUT (
      ->   embeddings STRUCT<statistics STRUCT<truncated BOOL, token_count FLOAT64>, values ARRAY<FLOAT64>>
      -> )
      -> REMOTE OPTIONS (
      ->   endpoint = '//aiplatform.googleapis.com/projects/<project_id>/locations/us-central1/publishers/google/models/textembedding-gecko@003'
      -> );
Query OK, 0 rows affected (0.00 sec)

spanner> SELECT embeddings.values 
      -> FROM ML.PREDICT(MODEL test_embs, (
      ->     SELECT description as content
      ->     FROM Bikes
      ->     LIMIT 3
      ->   )
      -> );
+-------------------------------+
| values                        |
+-------------------------------+
| [6397185019836432384.000000]  |
| [14047431176493074432.000000] |
| [5524949853424734208.000000]  |
+-------------------------------+
3 rows in set (2.17825ms)

Vertex AI 連携無し、というか Google Cloud 認証なしでも動作するようになっていました。偉い！

今回は前の記事に合わせて SentenceTransformer を使用しましたが、こちらの方法でサンプルを組んだほうがシンプルだったかもしれません。

補足: 検索用の index を使う

Spanner の vector search で index の使用はできるのでしょうか。

❯ gcloud spanner databases ddl update db1 --instance=test \
    --ddl='CREATE INDEX BikesDesc ON Bikes(desc_embs)'
ERROR: (gcloud.spanner.databases.ddl.update) HTTPError 400: {"code":9, "message":"Cannot reference ARRAY desc_embs in the creation of index BikesDesc."}

残念ながら、ARRAY 型には直接 index が作成できないようです。

こちらの例では、embeddings とは別の column に対して index を作成し、 STORING 句で embedding を含めることで、index-only scan はできるよと言っています。(が、embeddings そのものに対して index を利用しているわけではなさそうです)

すこし調べたのですが、index を使わなくても効率に問題がないか、データ量が増えても性能に問題が出ないか、残念ながら今回調査した範囲ではわかりませんでした。

ちなみに蛇足として、一般的には以下のようにして index を作成することができます。

❯ gcloud spanner databases ddl update db1 --instance=test \
    --ddl='CREATE INDEX BikesDesc ON Bikes(model)'
Schema updating...done.

(spanner-cli で確認)
...
spanner> SHOW INDEX FROM Bikes ;
+-------+--------------+-------------+-------------+-----------+------------------+-------------+
| Table | Parent_table | Index_name  | Index_type  | Is_unique | Is_null_filtered | Index_state |
+-------+--------------+-------------+-------------+-----------+------------------+-------------+
| Bikes |              | PRIMARY_KEY | PRIMARY_KEY | true      | false            | NULL        |
| Bikes |              | BikesDesc   | INDEX       | false     | false            | READ_WRITE  |
+-------+--------------+-------------+-------------+-----------+------------------+-------------+
2 rows in set (0.03 sec)