パーティション化 DML を使用してテキストデータのベクトル エンベディングを一括で生成する

このドキュメントでは、SQL と Vertex AI textembedding-gecko モデルを使用して、Spanner に保存されているテキストデータ（STRING または JSON）のベクトル エンベディングを一括で生成してバックフィルする方法について説明します。

前提条件

Spanner データベースには、テキストデータ（STRING または JSON）を含むテーブルが必要です。データのインポートの詳細については、Spanner のインポートとエクスポートの概要をご覧ください。

使用例

Spanner に次のスキーマを持つテーブルがあるとします。 このテーブルには数百万のレコードが含まれています。

GoogleSQL

CREATE TABLE Products (
  product_id INT64 NOT NULL,
  name STRING(MAX),
  description STRING(MAX)
) PRIMARY KEY(product_id);

PostgreSQL

CREATE TABLE Products (
  product_id INT8 NOT NULL,
  name TEXT,
  description TEXT,
  PRIMARY KEY(product_id)
);

目標は、このテーブルの description 列のベクトル エンベディングを生成し、ベクトル検索を使用して、顧客におすすめできる類似アイテムを見つけてショッピング エクスペリエンスを向上させることです。

エンべディング モデルを登録する

GoogleSQL

Spanner データベースの Vertex AI textembedding-gecko エンドポイントにエンべディング モデルを登録します。

CREATE MODEL MODEL_NAME
INPUT(
  content STRING(MAX)
)
OUTPUT(
  embeddings STRUCT<values ARRAY<FLOAT32>>
)
REMOTE OPTIONS(
    endpoint = '//aiplatform.googleapis.com/projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION',
  default_batch_size = 5
)

次のように置き換えます。

  • MODEL_NAME: エンべディング　モデルの名前
  • PROJECT: Vertex AI エンドポイントをホストするプロジェクト
  • LOCATION: Vertex AI エンドポイントのロケーション
  • MODEL_VERSION: textembedding-gecko エンべディング モデルのバージョン

PostgreSQL

PostgreSQL 言語では、モデルを登録する必要はありません。エンドポイント名を spanner.ML_PREDICT_ROW 関数呼び出しに直接渡します。

ベスト プラクティスとして、次の点を考慮してください。

  • 割り当ての分離を維持するには、別のプロジェクトのエンドポイントを使用して、本番環境のエンドポイントとは異なるエンベディングを生成してバックフィルします。本番環境のトラフィックの処理用に、本番環境のエンドポイントを予約します。
  • モデル エンドポイントが default_batch_size の値をサポートしていることを確認します。default_batch_size は、クエリヒント @{remote_udf_max_rows_per_rpc=NEW_NUMBER} でオーバーライドできます。各リージョンの default_batch_size の上限については、テキストのスニペットのテキスト エンベディングを取得するをご覧ください。
  • @latest ではなく、特定のモデル バージョン（@003 など）でエンドポイントを定義します。これは、同じテキストに対して生成されるエンべディング ベクトルが、使用するモデルのバージョンによって異なる場合があるためです。そのため、同じデータセットでエンベディングを生成するために異なるモデル バージョンを使用しないようにします。また、モデル定義ステートメントでモデル バージョンを更新しても、このモデルですでに生成されたエンベディングは更新されません。エンベディング用のモデル バージョンを管理する方法の 1 つは、モデル バージョンを格納するテーブルに追加の列を作成することです。
  • カスタム チューニングされた textembedding-gecko モデルは、GoogleSQL の ML.PREDICT 関数と PostgreSQL spanner.ML_PREDICT_ROW 関数ではサポートされていません。

エンべディング モデルのエンドツーエンドの統合をテストする

クエリを実行して、エンベディング モデルが正常に構成され、エンベディングが取得されていることをテストできます。たとえば、次のクエリを実行します。

GoogleSQL

SELECT embeddings.values
FROM SAFE.ML.PREDICT(
  MODEL MODEL_NAME,
  (SELECT description AS content FROM products LIMIT 10)
);

次のように置き換えます。

  • MODEL_NAME: エンべディング　モデルの名前

PostgreSQL

SELECT spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION',
    JSONB_BUILD_OBJECT('instances', JSONB_BUILD_ARRAY(JSONB_BUILD_OBJECT('content', description))))
FROM Products
LIMIT 10;

次のように置き換えます。

  • PROJECT: Vertex AI エンドポイントをホストするプロジェクト
  • LOCATION: Vertex AI エンドポイントのロケーション
  • MODEL_VERSION: textembedding-gecko エンべディング モデルのバージョン

ソーステーブルを更新して、エンベディングを保存する列を追加する

次に、ソーステーブル スキーマを更新して、生成されたエンベディングを格納するためのデータ型 ARRAY<FLOAT32> の追加列を含めます。

GoogleSQL

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME ARRAY<FLOAT32>;

次のように置き換えます。

  • TABLE_NAME: ソーステーブルの名前
  • EMBEDDING_COLUMN_NAME: 生成されたエンべディングを追加する列の名前

PostgreSQL

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME real[];

次のように置き換えます。

  • TABLE_NAME: ソーステーブルの名前
  • EMBEDDING_COLUMN_NAME: 生成されたエンべディングを追加する列の名前

たとえば、products テーブルの例を使用して次のコマンドを実行します。

GoogleSQL

ALTER TABLE Products
ADD COLUMN desc_embed ARRAY<FLOAT32>;

PostgreSQL

ALTER TABLE Products
ADD COLUMN desc_embed real[];

別の列を追加して、エンべディング モデルのバージョンを管理できます。

GoogleSQL

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT64;

PostgreSQL

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT8;

Vertex AI の割り当てを増やす

モデルを使用するリージョンで、textembedding-gecko に対する Vertex AI API の割り当てを増やす必要があります。割り当ての増加をリクエストするには、Vertex AI の割り当ての増加をご覧ください。

詳細については、Vertex AI の割り当てと上限をご覧ください。

バックフィル エンベディング

最後に、パーティション化 DML を使用して次の UPDATE ステートメントを実行し、テキストデータ列のエンベディングを生成し、そのエンベディングをデータベースに保存します。モデル バージョンをエンベディングとともに保存できます。このクエリは、データベースのトラフィックが少ない期間に実行することをおすすめします。

GoogleSQL

UPDATE TABLE_NAME
SET
  TABLE_NAME.EMBEDDING_COLUMN_NAME = (
    SELECT embeddings.values
    FROM SAFE.ML.PREDICT(
      MODEL MODEL_NAME,
      (SELECT TABLE_NAME.DATA_COLUMN_NAME AS content)
    ) @{remote_udf_max_rows_per_rpc=MAX_ROWS}
  ),
  TABLE_NAME.EMBEDDING_VERSION_COLUMN = MODEL_VERSION
WHERE FILTER_CONDITION;

次のように置き換えます。

  • TABLE_NAME: テキストデータを含むテーブルの名前
  • EMBEDDING_COLUMN_NAME: 生成されたエンべディングを追加する列の名前
  • DATA_COLUMN_NAME: テキストデータを含む列の名前
  • MODEL_NAME: エンべディング　モデルの名前
  • MAX_ROWS: RPC あたりの最大行数
  • EMBEDDING_VERSION_COLUMN: エンベディングのバックフィルに使用する textembedding-gecko エンベディング モデルのバージョンを管理する列
  • MODEL_VERSION: textembedding-gecko エンべディング モデルのバージョン
  • FILTER_CONDITION: 適用する分割可能なフィルタ条件

SAFE.ML.PREDICT を使用すると、失敗したリクエストに対して NULL が返されます。SAFE.ML.PREDICT を WHERE embedding_column IS NULL フィルタと組み合わせて使用すると、すでに計算されたフィールドのエンベディングを計算せずにクエリを再実行できます。

PostgreSQL

UPDATE TABLE_NAME
SET
  EMBEDDING_COLUMN_NAME = spanner.FLOAT32_ARRAY(spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION',
    JSONB_BUILD_OBJECT('instances', JSONB_BUILD_ARRAY(JSONB_BUILD_OBJECT('content', DATA_COLUMN_NAME)))
  ) /*@ remote_udf_max_rows_per_rpc=MAX_ROWS */ ->'predictions'->0->'embeddings'->'values'),
  EMBEDDING_VERSION_COLUMN = MODEL_VERSION
WHERE FILTER_CONDITION;

次のように置き換えます。

  • TABLE_NAME: テキストデータを含むテーブルの名前
  • EMBEDDING_COLUMN_NAME: 生成されたエンべディングを追加する列の名前
  • DATA_COLUMN_NAME: テキストデータを含む列の名前
  • PROJECT: Vertex AI エンドポイントをホストするプロジェクト
  • LOCATION: Vertex AI エンドポイントのロケーション
  • MODEL_VERSION: textembedding-gecko エンべディング モデルのバージョン
  • MAX_ROWS: RPC あたりの最大行数
  • EMBEDDING_VERSION_COLUMN: エンベディングのバックフィルに使用する textembedding-gecko エンベディング モデルのバージョンを管理する列
  • FILTER_CONDITION: 適用する分割可能なフィルタ条件

products テーブルのバックフィル クエリの例:

GoogleSQL

UPDATE products
SET
  products.desc_embed = (
    SELECT embeddings.values
    FROM SAFE.ML.PREDICT(
      MODEL gecko_model,
      (SELECT products.description AS content)
    ) @{remote_udf_max_rows_per_rpc=200}
  ),
  products.desc_embed_model_version = 3
WHERE products.desc_embed IS NULL;

PostgreSQL

UPDATE products
SET
  desc_embed = spanner.FLOAT32_ARRAY(spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko@003',
    JSONB_BUILD_OBJECT('instances', JSONB_BUILD_ARRAY(JSONB_BUILD_OBJECT('content', description)))
  ) /*@ remote_udf_max_rows_per_rpc=200 */ ->'predictions'->0->'embeddings'->'values'),
  desc_embed_model_version = 3
WHERE desc_embed IS NULL;

ベスト プラクティスとして、次の点を考慮してください。

  • Spanner API のデフォルトの gRPC タイムアウトは 1 時間です。 バックフィルする埋め込みの量によっては、UPDATE パーティション化 DML を完了するのに十分な時間を確保するために、このタイムアウトを長くする必要が生じる場合があります。詳細については、カスタム タイムアウトと再試行を構成するをご覧ください。

パフォーマンスとその他の考慮事項

埋め込みデータをバックフィルするときにパフォーマンスを最適化するために、次の点を考慮してください。

ノード数

パーティション化 DML は、異なるパーティションで特定の DML ステートメントを並行して実行します。ノード数の多いインスタンスの場合、パーティション化 DML の実行時に割り当てエラーが発生することがあります。Vertex AI API の割り当て上限により Vertex AI API リクエストがスロットルされている場合、Spanner はこのような失敗をパーティション化 DML トランザクション モードで最大 20 回再試行します。Vertex AI で割り当てエラーの発生率が高い場合は、Vertex AI の割り当てを増やします。

データ列内のテキストのサイズ

Vertex AI エンべディング モデルでは、テキスト入力ごとのトークンの最大数に制限があります。モデル バージョンごとに、トークンの上限が異なります。各 Vertex AI リクエストには複数の入力テキスト フィールドを含めることができますが、1 つのリクエストに含まれるトークンの最大数には制限があります。GoogleSQL データベースの場合、INVALID_ARGUMENT エラーが発生して「リクエストが大きすぎます」というメッセージが表示された場合は、エラーを避けるためにバッチサイズを小さくしてみてください。これを行うには、default_batch_size を構成するか、モデルの登録時に @{remote_udf_max_outstanding_rpcs} クエリヒントを使用します。

Vertex AI に送信された API リクエストの数

クエリヒント @{remote_udf_max_outstanding_rpcs} を使用すると、Spanner から Vertex AI に送信されるリクエストの数を増減できます。この上限を増やすと、Spanner インスタンスの CPU とメモリの使用量が増加する可能性があるので注意してください。GoogleSQL データベースの場合、このクエリヒントを使用すると、モデル用に構成された default_batch_size がオーバーライドされます。

バックフィルの進捗状況をモニタリングする

システム分析情報ダッシュボードを使用して、Spanner から Vertex AI に送信されるリクエスト数、レイテンシ、ネットワーク バイト数をモニタリングできます。

次のステップ