为大型表生成和管理自动向量嵌入

选择文档版本：

您可以为整个表列生成和管理向量嵌入，从而获得可伸缩的大规模向量嵌入创建解决方案。此解决方案尤其有助于在文本内容上实现语义搜索和检索增强生成 (RAG)，包括以下内容：

  • 为新表创建初始向量嵌入
  • 在大量数据导入后生成嵌入
  • 在数据发生重大更改后刷新嵌入
  • 以增量方式维护嵌入

了解自动向量嵌入

AlloyDB Omni 中的自动向量嵌入提供了一种可伸缩的方式，可自动生成和维护数据的向量嵌入。您可以配置自动向量嵌入功能来处理此流程，而无需为每段新文本或更新后的文本手动生成嵌入。对于依赖最新嵌入来实现语义搜索、检索增强生成 (RAG) 和其他 AI 功能的应用，此功能尤为实用。

借助自动向量嵌入，您可以执行以下操作：

  • 初始化整个表的嵌入：使用单个命令为表列中的所有现有数据生成嵌入。
  • 使嵌入保持同步：在源数据发生更改时自动更新嵌入，确保 AI 应用始终使用最新信息。
  • 大规模生成嵌入：高效地为数百万行的大型表创建嵌入。
  • 通过为每个嵌入列调用管理函数，为同一表中的多个列配置和管理嵌入。

此功能通过抽象化向量嵌入创建和维护的复杂性，简化了 AI 应用的开发和维护。

准备工作

在为大型表生成和管理向量嵌入之前，请按照以下步骤操作：

  • 检查并增加 Vertex AI 配额：为确保最佳性能并避免操作缓慢、QPS 低或多次重试，请验证 Vertex AI 配额，并在需要时增加配额。

    如需检查和增加配额，请按以下步骤操作：

    1. 在 Google Cloud 控制台中，进入 配额 页面。
    2. 过滤与您的模型和区域相关的以下指标：
      • 对于每分钟请求数 (RPM)：Regional online prediction requests per base model per minute per region per base_model
      • 对于 Gemini token：Embed content input tokens per minute per region per base_model
    3. 如果当前限制不足以满足表大小的需求，请找到要更新的配额值，然后选中该配额旁边的复选框。

    4. 点击修改 。系统会显示配额更改 对话框。

      如需了解详情，请参阅申请配额调整

    如需了解 Vertex AI 嵌入模型限制，请参阅 Gemini 嵌入模型 token 限制。

    如需查看完整的配额列表，请参阅 Vertex AI 配额和限制。

  • 以 postgres 用户身份使用 psql 连接到数据库。

  • 验证是否已安装 google_ml_integration 扩展程序，且其版本为 1.5.6 或更高版本。

  • 验证 google_ml_integration.enable_model_support 标志和 google_ml_integration.enable_faster_embedding_generation 标志是否已设置为 on。

  • 您必须先将 AlloyDB Omni 配置为与 Vertex AI 搭配使用，然后才能通过 AlloyDB Omni 数据库生成嵌入。如需了解详情，请参阅将数据库与 Vertex AI 集成。

  • 如需管理和监控自动嵌入生成，用户默认拥有对 google_ml.embed_gen_progress 和 google_ml.embed_gen_settings 表的 Select 访问权限。

    如需允许用户管理自动嵌入生成，请向用户授予对 google_ml.embed_gen_progress 和 google_ml.embed_gen_settings 表的 INSERT、UPDATE 和 DELETE 权限：

    GRANT INSERT, UPDATE, DELETE ON google_ml.embed_gen_progress TO 'USER_NAME';

    替换以下内容：

    • USER_NAME：为其授予 权限的用户的名称。

  • 完成生成文本嵌入中所述的初始设置。

  • 验证您使用的 PostgreSQL 客户端中是否已将 AUTOCOMMIT 设置为 ON。

验证扩展程序版本

如需检查 google_ml_integration 扩展程序的版本，请运行以下命令：

SELECT extversion FROM pg_extension WHERE extname = 'google_ml_integration';

如果您需要更新扩展程序，请运行以下命令：

ALTER EXTENSION google_ml_integration UPDATE;

验证数据库标志是否已设置为 on

如需验证数据库标志是否已正确设置，请运行以下命令：

SHOW google_ml_integration.enable_model_support;
SHOW google_ml_integration.enable_faster_embedding_generation;

如果这些标志设置为 off，请参阅 配置实例的数据库标志。如需详细了解这些标志，请参阅支持的数据库标志。

准备表

在生成自动嵌入之前，您必须在表中创建一个列来存储生成的向量。此列通常使用 vector(DIMENSION) 类型，并且必须具有 DEFAULT NULL 值。

例如，如需向名为 user_reviews 的表中添加一个用于存储 768 维嵌入的列，请运行以下命令：

ALTER TABLE user_reviews ADD COLUMN IF NOT EXISTS content_embeddings vector(768) DEFAULT NULL;

初始化表的嵌入

用于管理自动向量嵌入的函数在 ai 和 google_ml 架构中均可用。ai 架构为 AlloyDB Omni 中的最新 AI 功能提供了一个简化的界面。

使用 ai.initialize_embeddings() SQL 函数为表的内容列生成嵌入。这是一个阻塞调用，这意味着数据库会话会等待操作完成，然后返回结果，并允许您在该会话中发出新命令。不过，与数据库的其他连接不会被阻塞，并且可以继续使用该表：

  • 如果该函数返回成功，则向量嵌入创建完成。
  • 该函数会自动尝试从模型配额错误等暂时性问题中恢复。只有在这些恢复尝试失败时，才会向您返回失败。对于持久性问题（例如，配置错误的 batch_size 导致请求超出大小限制，或者手动取消了操作），您必须手动重新发出调用。

此函数支持 Google 提供的模型（例如 Vertex AI 的 text-embedding-005）以及您注册的自定义模型。

执行批量生成

默认情况下，AlloyDB Omni 使用批处理在单个请求中为多个文本输入生成嵌入，从而提高效率。如果您未提供具体的批次大小，AlloyDB Omni 会应用自动确定的默认值。

提示批次大小

ai.initialize_embeddings 中的 batch_size 参数可让您通过为直接支持的模型建议首选批次大小来指导 AlloyDB Omni 的查询优化器。AlloyDB Omni 可能会根据模型限制或配额动态减小此大小，但您的提示有助于影响查询执行计划。

CALL ai.initialize_embeddings(
    model_id => 'text-embedding-005',
    table_name => 'user_reviews',
    content_column => 'content',
    embedding_column => 'content_embeddings',
    batch_size => 50
);

使用支持批处理的自定义嵌入模型

如果您想使用支持批处理的自定义模型或外部支持的模型，请定义批处理转换函数，并在创建模型时将其指定为 model_batch_in_transform_fn 和 model_batch_out_transform_fn。您还可以在 initialize_embeddings 调用中指定 batch_size。对于支持批处理的模型，我们建议您使用大于 1 的 batch_size 以获得更好的性能。

  1. 定义自定义模型的输入、输出和批处理转换函数。

    -- Scalar input transform functions
CREATE OR REPLACE FUNCTION acme_text_input_transform(model_id TEXT, input TEXT) RETURNS JSON;
CREATE OR REPLACE FUNCTION acme_text_output_transform(model_id TEXT, model_output JSON) RETURNS real[];
CREATE OR REPLACE FUNCTION acme_generate_headers(model_id TEXT, input TEXT) RETURNS JSON;
-- Batch input transform functions
CREATE OR REPLACE FUNCTION acme_text_batch_input_transform(model_id TEXT, input TEXT[]) RETURNS JSON;
CREATE OR REPLACE FUNCTION acme_text_batch_output_transform(model_id TEXT, model_output JSON) RETURNS real[][];

  2. 如需创建模型，请指定批处理转换函数。

    CALL
  ai.create_model(
    model_id => 'custom-embedding-model',
    model_request_url => 'https://acme.com/models/text/embeddings/v1',
    model_type => 'text_embedding',
    model_in_transform_fn => 'acme_text_input_transform',
    model_out_transform_fn => 'acme_text_output_transform',
    generate_headers_fn => 'acme_generate_headers',
    model_batch_in_transform_fn => 'acme_text_batch_input_transform',
    model_batch_out_transform_fn => 'acme_text_batch_output_transform'
  );

  3. 使用自定义模型生成向量嵌入。

    CALL
  ai.initialize_embeddings(
    model_id => 'custom-embedding-model',
    table_name => 'user_reviews',
    content_column => 'content',
    embedding_column => 'content_embeddings',
    batch_size => 10
);

您还可以将自动嵌入功能与本身不支持批处理的自定义模型搭配使用。为此，您仍必须定义批处理转换函数 model_batch_in_transform_fn 和 model_batch_out_transform_fn。对于非批处理模型，请定义这些函数以一次处理输入数组中的单个输入。为此模型调用 ai.initialize_embeddings 时，请将 batch_size 设置为 1。

使用支持批处理的自定义维度

如需将自定义批处理输入转换函数与 OUTPUT_DIMENSIONALITY 参数搭配使用，您可以定义一个函数，该函数在请求参数中指定所选维度。这有助于优化支持可变输出大小的模型的嵌入生成。

例如，以下函数为需要输出维度为 768 的模型定义了自定义批处理输入转换：

CREATE OR REPLACE FUNCTION google_ml.vertexai_text_embedding_batch_input_transform_with_768_dims(model_id VARCHAR(100), input_list TEXT[])
RETURNS JSON
LANGUAGE SQL
AS $$
  SELECT pg_catalog.json_build_object(
    'instances',
    pg_catalog.json_agg(pg_catalog.json_build_object('content', content)),
    'parameters',
    pg_catalog.json_build_object('outputDimensionality', 768)
  ) FROM unnest(input_list) AS content;
$$;

将 JSONB 优化与自定义模型搭配使用

您可以在自定义输出转换函数中使用 JSONB 数据类型来提高性能。当您创建支持批处理的自定义模型时，google_ml_integration 扩展程序会自动查找并使用转换函数的 JSONB 变体。

此优化可以显著提高自动嵌入生成的性能，因为 JSONB 是一种更高效的二进制格式，用于在 PostgreSQL 中存储和处理 JSON 数据。

如需利用此功能，您需要提供转换函数的另一个版本，该版本接受 JSONB 实参而不是 JSON。

例如，如果您有一个具有以下签名的批处理输出转换函数：

CREATE OR REPLACE FUNCTION my_batch_output_transform(model_id TEXT, model_output JSON) RETURNS real[][];

您可以创建如下所示的 JSONB 变体：

CREATE OR REPLACE FUNCTION my_batch_output_transform(model_id TEXT, model_output JSONB) RETURNS real[][];

该扩展程序会自动检测到该函数的 JSONB 版本，并将其用于批处理。如果您已创建使用 JSON 转换函数的模型注册，则无需更新 ai.create_model 调用。只要新的 JSONB 函数使用与现有 JSON 函数完全相同的名称，该扩展程序就会透明地检测并使用 JSONB 变体进行批处理。

以增量方式刷新嵌入

刷新嵌入时，系统会根据输入内容列中的最新值重新生成嵌入。

为了让您能够控制一致性和性能，AlloyDB Omni 支持各种增量嵌入刷新模式。您可以使用 ai.initialize_embeddings() 中的 incremental_refresh_mode 枚举实参选择模式。下面列出了可能的模式：

  • transactional：在更新内容列的事务中刷新嵌入。此过程通常使用类似于数据库触发器的机制，在内容列更新时自动生成嵌入，但可能会产生开销并减慢更新操作的速度。引入的开销是为了维护事务语义并确保嵌入与内容同步而做出的权衡。此模式依赖于模型的标量转换函数，因此您必须在创建模型时定义 model_in_transform_fn 和 model_out_transform_fn。如需使用 transactional 模式，您必须拥有对表的所有者 角色。

    CALL
  ai.initialize_embeddings(
    model_id => 'text-embedding-005',
    table_name => 'user_reviews',
    content_column => 'content',
    embedding_column => 'content_embeddings',
    batch_size => 10,
    incremental_refresh_mode => 'transactional'
);

    在 transactional 模式下，ai.refresh_embeddings() 函数处于停用状态，因为嵌入是使用触发器自动保持同步的。如需在此模式下为整个表重新生成嵌入，或从中断的 ai.initialize_embeddings() 调用中恢复，您必须先使用 ai.drop_embedding_config() 函数删除配置，然后重新发出 ai.initialize_embeddings() 调用。

    对于演示版、小型数据集或插入和更新量与初始加载量相比很低的表（例如，一个百万行表每天更新几百次），此自动化模式非常有用。当维护即时数据一致性比更新操作期间引入的额外延迟更重要时，此模式非常适用。

  • manual：此为默认模式。在此模式下，系统会在表中添加一个新的布尔值跟踪列，以跟踪过时的嵌入。调用 ai.refresh_embeddings() 函数会定期执行增量刷新，仅为新行或更新后的行生成嵌入。我们建议需要更好地控制性能的用户使用此模式，尤其是在初始加载后处理大量插入或更新时。此模式适用于以下场景：最大限度地减少写入延迟是首要任务，并且在触发使用 ai.refresh_embeddings 的定期增量刷新之前，暂时存在过时或 null 嵌入是可以接受的。如需更新过时或新行，请按照刷新表的所有嵌入部分中的说明使用 ai.refresh_embeddings() 函数。

选择刷新模式

下表比较了两种增量刷新模式，以帮助您为应用选择最佳方法。

模式 说明 一致性和性能 适用场景
transactional 嵌入会使用类似于触发器的机制作为数据库事务的一部分进行更新。请注意，在此模式下，ai.refresh_embeddings() 处于停用状态。 即时一致性 。写入延迟开销较高，因为每个 INSERT 或 UPDATE 都会触发嵌入生成。 演示、小型数据集或数据一致性至关重要的表。
manual（默认） 使用布尔值跟踪列来监控过时或新行。您必须调用 ai.refresh_embeddings() 才能触发定期增量刷新。 最终一致性 。写入延迟极低，因为嵌入是在您选择的时间批量生成的。 大型数据集、写入频率高的生产环境或对性能敏感的应用。

刷新表的所有嵌入

在使用 manual 增量刷新模式成功为表运行 ai.initialize_embeddings() 后，您可以定期执行嵌入的增量刷新，该刷新使用 ai.refresh_embeddings 触发。您可以使用刷新操作来更新在初始 initialize_embeddings 调用期间并发修改的行的嵌入，或执行定期增量刷新。

如果在手动模式下嵌入创建过程被中断（例如，被 pg_cancel 中断），请调用 ai.refresh_embeddings() 函数以完成剩余行的生成。

刷新函数会重复使用初始调用中的设置，因此您只需指定表和嵌入列。您还可以提供可选的 batch_size 来替换默认值。

CALL ai.refresh_embeddings(
    table_name => 'user_reviews',
    embedding_column => 'content_embeddings',
    batch_size => 50  -- Optional override
);

在创建向量嵌入期间使用表数据

虽然 ai.initialize_embeddings() 是其运行所在会话的阻塞调用，但其他连接可以继续使用该表。自动向量嵌入过程使用标准行级锁定按批更新行。这意味着，如果其他连接尝试修改活动嵌入任务所针对的相同行，则来自其他连接的并发数据修改语言 (DML) 操作（例如 UPDATE 或 DELETE）只会短暂阻塞。非修改型 SELECT 查询不会被阻塞。

删除自动向量嵌入设置

如果您需要移除特定表和嵌入列组合的自动向量嵌入配置，请使用 ai.drop_embedding_config() 函数。此函数可用于清理，或在您为列重新配置嵌入管理时使用。

CALL
  ai.drop_embedding_config(
    table_name => 'user_reviews',
    embedding_column => 'content_embeddings');

使用分区表

自动向量嵌入功能支持分区表。这样，您就可以高效地管理大型分区数据集的嵌入。以下是使用分区表的常见应用场景。

在分区表上初始化嵌入

您只能在表的根分区上初始化嵌入。对于整个分区表，此操作只需执行一次。

CALL ai.initialize_embeddings(
    model_id => 'text-embeddings-005',
    table_name => 'documents', -- This is the root partitioned table
    content_column => 'content',
    embedding_column => 'content_embeddings'
);

刷新分区表上的嵌入

初始化后，您可以刷新任何分区上的嵌入，包括根分区、子分区或各个叶分区。对于大型数据集，您可以从不同的数据库连接并行刷新不同分区的嵌入，从而提高性能：

  • 如需刷新整个表，请运行以下命令：
CALL ai.refresh_embeddings(
    table_name => 'documents', -- This is the root partitioned table
    embedding_column => 'content_embeddings'
);
  • 如需刷新单个分区，请运行以下命令：
CALL ai.refresh_embeddings(
    table_name => 'documents_eu',
    embedding_column => 'content_embeddings'
);

刷新新添加或附加的分区的嵌入

自动嵌入功能支持为在初始设置后并入主表的各个分区生成嵌入。具体步骤取决于您是添加全新分区还是附加预先存在的表。

  • 新添加的分区：如果您向表中添加了新分区，则可以通过对新分区调用 ai.refresh_embeddings 为其生成嵌入。
-- Add a new partition
CREATE TABLE documents_africa PARTITION OF documents
    FOR VALUES IN ('africa');

-- Refresh embeddings for the new partition
CALL ai.refresh_embeddings(
    table_name => 'documents_africa',
    embedding_column => 'content_embeddings'
);
  • 新附加的分区：如需将现有表附加为分区，请先使用 ai.embedding_prepare_partition 过程确保其架构与分区表兼容。自动嵌入功能支持在分层分区设置的任何级别添加或附加分区。ai.embedding_prepare_partition 过程可确保架构与层次结构中的任何父表兼容。

如需将现有表附加为分区，请先使用 ai.embedding_prepare_partition 过程确保其架构与分区表兼容：

-- Prepare the table to be attached
CALL ai.embedding_prepare_partition(
    parent_table => 'documents',
    child_table => 'documents_misc'
);

-- Attach the partition
ALTER TABLE documents ATTACH partition documents_misc DEFAULT;

-- Refresh embeddings for the newly attached partition
CALL ai.refresh_embeddings(
    table_name => 'documents_misc',
    embedding_column => 'content_embeddings'
);

自动嵌入功能支持在分层分区设置的任何级别添加或附加分区。ai.embedding_prepare_partition 过程可确保架构与层次结构中的任何父表兼容：

-- Prepare a sub-partition for a non-root parent table
CALL ai.embedding_prepare_partition(
    parent_table => 'documents_eu', -- An existing partition
    child_table => 'documents_eu_germany'
);

-- Attach the new sub-partition
ALTER TABLE documents_eu ATTACH PARTITION documents_eu_germany
    FOR VALUES IN ('germany');

-- Refresh embeddings for the new sub-partition
CALL ai.refresh_embeddings(
    table_name => 'documents_eu_germany',
    embedding_column => 'content_embeddings'
);

监控嵌入生成进度

您可以通过查询 ai.embedding_progress_view 来监控活跃 initialize_embeddings 和 refresh_embeddings 调用的实时状态。此视图提供了有关操作进度的详细信息，包括完成百分比、已用时间和剩余估计时间。

如需检查进度，请运行以下查询：

SELECT
  table_name,
  content_column,
  embedding_column,
  model_id,
  percent_progress,
  status,
  elapsed_time,
  rows_processed,
  partition_root
FROM
  ai.embedding_progress_view;

该视图提供以下信息：

列 说明
table_name 正在处理的表或分区的名称。
content_column 包含嵌入的源内容的列。
embedding_column 用于存储嵌入的列。
model_id 用于生成的模型。
percent_progress 已完成的操作百分比。
status 操作的当前状态（例如，正在运行、成功）。
elapsed_time 自操作开始以来经过的时间。
rows_processed 到目前为止已处理的行数。
partition_root 根分区表的名称。

自动生成嵌入的示例

本部分提供了使用已注册模型端点自动生成嵌入的示例。

OpenAI 嵌入模型

如需使用 OpenAI 提供的已注册 text-embedding-3-small 模型端点生成嵌入，请运行以下语句：

CALL ai.initialize_embeddings(
    model_id => 'text-embedding-3-small',
    table_name => 'user_reviews',
    content_column => 'content',
    embedding_column => 'content_embeddings'
);

自定义嵌入模型

对于您自己的模型或外部支持的模型，您必须定义输入和输出转换函数，并使用 ai.create_model 注册它们。如果您计划使用自动嵌入功能，则必须同时指定标量转换函数（例如 acme_text_input_transform、acme_text_output_transform）和批处理转换函数（例如 acme_text_batch_input_transform、acme_text_batch_output_transform）。

后续步骤