パーティション化された 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 の値をサポートしていることを確認します。クエリヒント @{remote_udf_max_rows_per_rpc=NEW_NUMBER} を使用して default_batch_size をオーバーライドできます。各リージョンの 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 の割り当てを増やします。GoogleSql の使用中に、ステートメント レベルのヒント @{pdml_max_parallelism=DESIRED_NUMBER} を使用して並列処理を調整することもできます。次の例では、並列処理を「5」に設定しています。

GoogleSQL 

@{pdml_max_parallelism=5} UPDATE products
SET products.desc_embed =(
  SELECT embeddings.values
  FROM SAFE.ML.PREDICT(MODEL gecko_model, (
        SELECT products.value AS CONTENT
        )
  )
      @{remote_udf_max_rows_per_rpc=200}
),
products.desc_embed_model_version = 003
WHERE products.desc_embed IS NULL;

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

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 に送信されるリクエスト数、レイテンシ、ネットワーク バイト数をモニタリングできます。

次のステップ