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

このドキュメントでは、SQL と Vertex AI テキストエンベディングモデルを使用して、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 モデルエンドポイントに、テキストエンベディングモデルを登録します。

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/$MODEL_NAME',
  default_batch_size = 5
)

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

MODEL_NAME: Vertex AI テキストエンベディングモデルの名前
PROJECT: Vertex AI エンドポイントをホストするプロジェクト
LOCATION: Vertex AI エンドポイントのロケーション

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 など）を使用してエンドポイントを定義します。これは、同じテキストに対して生成されたエンベディングベクトルが、使用するモデルのバージョンによって異なる可能性があるためです。そのため、同じデータセットで異なるバージョンのモデルを使用してエンベディングを生成することは避けてください。また、モデル定義ステートメントでモデルのバージョンを更新しても、このモデルですでに生成されているエンベディングは更新されません。エンベディングのモデルのバージョンを管理する方法の一つとして、モデルのバージョンを格納する列をテーブルに追加する方法もあります。
カスタムチューニングされたテキストエンベディングモデルは、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: Vertex AI テキストエンベディングモデルの名前

PostgreSQL

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

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

PROJECT: Vertex AI エンドポイントをホストするプロジェクト
LOCATION: Vertex AI エンドポイントのロケーション
MODEL_NAME: Vertex AI テキストエンベディングモデルの名前

エンベディング格納のための列を追加するようソーステーブルを更新する

次に、ソーステーブルのスキーマを更新して、生成されたエンベディングを保存するデータ型 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 の割り当てを増やす

テキストエンベディングモデルを使用するリージョンで、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: Vertex AI エンベディングモデルの名前
MAX_ROWS: RPC あたりの行の最大数
EMBEDDING_VERSION_COLUMN: エンベディングのバックフィルに使用されるエンベディングモデルのバージョンを管理する列
MODEL_VERSION: テキストエンベディングモデルのバージョン
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/$MODEL_NAME',
    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_NAME: Vertex AI エンベディングモデルの名前
MODEL_VERSION: Vertex AI エンベディングモデルのバージョン
MAX_ROWS: RPC あたりの行の最大数
EMBEDDING_VERSION_COLUMN: エンベディングのバックフィルに使用されるテキストエンベディングモデルのバージョンを管理する列
FILTER_CONDITION: 適用するパーティショニング可能なフィルタの条件

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

GoogleSQL

UPDATE products
SET
  products.desc_embed = (
    SELECT embeddings.values
    FROM SAFE.ML.PREDICT(
      MODEL embedding_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/$MODEL_NAME',
    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 embedding_model, (
        SELECT products.value AS CONTENT
        )
  )
      @{remote_udf_max_rows_per_rpc=200}
),
products.desc_embed_model_version = MODEL_VERSION
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 に送信されるリクエスト数、レイテンシ、ネットワークバイト数をモニタリングできます。

次のステップ

K 最近傍探索を行って類似ベクトル検索を行う方法を確認する。
エンベディングに関する集中講座で ML とエンベディングの詳細を確認する。
Vertex AI テキストエンベディングモデルの詳細を確認する。

パーティション化された DML を使用してテキストデータのベクター エンベディングを一括生成する コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

前提条件

ユースケースの例

GoogleSQL

PostgreSQL

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

GoogleSQL

PostgreSQL

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

GoogleSQL

PostgreSQL

エンベディング格納のための列を追加するようソーステーブルを更新する

GoogleSQL

PostgreSQL

GoogleSQL

PostgreSQL

GoogleSQL

PostgreSQL

Vertex AI の割り当てを増やす

エンベディングをバックフィルする

GoogleSQL

PostgreSQL

GoogleSQL

PostgreSQL

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

ノード数

GoogleSQL

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

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

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

次のステップ

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

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

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