使用分区 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 限制，请参阅获取文本片段的文本嵌入。
  • 使用特定模型版本（例如 @003，而不是 @latest）定义端点。这是因为针对同一段文本生成的嵌入向量可能会因所用模型的版本而异；因此，您应避免在同一数据集中使用不同的模型版本生成嵌入。此外，更新模型定义语句中的模型版本不会更新已使用此模型生成的嵌入。为嵌入管理模型版本的一种方法是在表中创建一个用于存储模型版本的额外列。
  • GoogleSQL ML.PREDICT 和 PostgreSQL spanner.ML_PREDICT_ROW 函数不支持自定义调优 textembedding-gecko 模型。

测试嵌入模型的端到端集成

您可以执行查询来测试嵌入模型是否已成功配置，以及是否检索到嵌入。例如运行以下查询：

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 嵌入模型会对每个文本输入的最大 token 数量施加限制。不同模型版本的 token 限制各不相同。每个 Vertex AI 请求可以包含多个输入文本字段，但系统会向单个请求中存在的最大 token 数量施加限制。对于 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 的请求数量、延迟时间和网络字节数量。

后续步骤