创建索引和查询向量

本文档介绍了如何使用存储的嵌入生成索引和查询嵌入。如需详细了解如何存储嵌入，请参阅存储向量嵌入。

您可以使用 AlloyDB 创建 ScaNN、IVF、IVFFlat 和 HNSW 索引。

准备工作

在开始创建索引之前，您必须完成以下前提条件。

  • 嵌入向量已添加到 AlloyDB 数据库中的表。

  • 安装了基于 pgvector 的 vector 扩展程序版本 0.5.0 或更高版本（由 Google 针对 AlloyDB 进行扩展）。

    CREATE EXTENSION IF NOT EXISTS vector;

  • 如需生成 ScaNN 索引，请安装 alloydb_scann 扩展程序以及 vector 扩展程序。

    CREATE EXTENSION IF NOT EXISTS alloydb_scann;

创建索引

您可以为数据库中的表创建以下索引类型之一。

创建 ScaNN 索引

AlloyDB alloydb_scann 是由 Google 开发的 PostgreSQL 扩展程序，可实现依托 ScaNN 算法的高效最近邻索引。

ScaNN 索引是一种基于树的量化索引，用于近似最近邻搜索。与 HNSW 相比，它可以缩短索引构建时间并减少内存占用。此外，与 HNSW 相比，它可根据工作负载提供更快的 QPS。

二级树 ScaNN 索引

如需将使用 ScaNN 算法的二级树索引应用于包含存储的向量嵌入的列，请运行以下 DDL 查询：

CREATE INDEX INDEX_NAME ON TABLE
  USING scann (EMBEDDING_COLUMN DISTANCE_FUNCTION)
  WITH (num_leaves=NUM_LEAVES_VALUE);

替换以下内容：

  • INDEX_NAME：要创建的索引的名称，例如 my-scann-index。索引名称会在整个数据库中共享。确保每个索引名称对数据库中的每个表中都是唯一的。

  • TABLE：要向其中添加索引的表。

  • EMBEDDING_COLUMN：用于存储 vector 数据的列。

  • DISTANCE_FUNCTION：要与此索引一起使用的距离函数。请按以下方式之一操作：

    • L2 距离：l2

    • 点积：dot_product

    • 余弦距离：cosine

  • NUM_LEAVES_VALUE：要应用于此索引的分区数量。设置为介于 1 到 1048576 之间的任意值。如需详细了解如何确定此值，请参阅对 ScaNN 索引进行调优。

三级树 ScaNN 索引

如需对包含存储的向量嵌入的列创建使用 ScaNN 算法的三级树索引，请运行以下 DDL 查询：

CREATE INDEX INDEX_NAME ON TABLE
  USING scann (EMBEDDING_COLUMN DISTANCE_FUNCTION)
  WITH (num_leaves=NUM_LEAVES_VALUE, max_num_levels = MAX_NUM_LEVELS);

替换以下内容：

  • MAX_NUM_LEVELS：K-means 聚类树的级别数量上限。对于基于二级树的量化，设置为 1（默认值）；对于基于三级树的量化，设置为 2。

创建索引后，您可以按照使用给定文本执行最近邻查询中的说明，运行利用该索引的最近邻搜索查询。

索引参数必须设置为可在 QPS 和召回率之间取得适当的平衡。如需详细了解如何对 ScaNN 索引进行调优，请参阅对 ScaNN 索引进行调优。

如需在使用 real[] 数据类型（而非 vector）的嵌入列上创建此索引，请将该列转换为 vector 数据类型：

CREATE INDEX INDEX_NAME ON TABLE
  USING scann (CAST(EMBEDDING_COLUMN AS vector(DIMENSIONS)) DISTANCE_FUNCTION)
  WITH (num_leaves=NUM_LEAVES_VALUE, max_num_levels = MAX_NUM_LEVELS);

将 DIMENSIONS 替换为嵌入列的维度宽度。如需详细了解如何查找维度，请参阅向量函数中的 vector_dims 函数。

如需查看索引编制进度，请使用 pg_stat_progress_create_index 视图：

SELECT * FROM pg_stat_progress_create_index;

phase 列会显示索引创建的当前状态，building index: tree training 阶段会在索引创建后消失。

如需对索引进行调优以实现目标召回率和 QPS 平衡，请参阅对 ScaNN 索引进行调优。

分析已编制索引的表

创建 ScaNN 索引后，您必须运行 ANALYZE 命令来更新有关数据的统计信息。

ANALYZE TABLE;

运行查询

在将嵌入存储到数据库中并为其编制索引后，您可以开始使用 pgvector 查询功能进行查询。您无法使用 alloydb_scann 扩展程序运行批量搜索查询。

如需查找嵌入向量的语义最近邻，您可以运行以下示例查询，其中设置了在创建索引期间使用的相同距离函数。

  SELECT * FROM TABLE
    ORDER BY EMBEDDING_COLUMN DISTANCE_FUNCTION_QUERY ['EMBEDDING']
    LIMIT ROW_COUNT

替换以下内容：

  • TABLE：包含要与文本进行比较的嵌入的表。

  • INDEX_NAME：要使用的索引的名称，例如 my-scann-index。

  • EMBEDDING_COLUMN：包含存储的嵌入的列。

  • DISTANCE_FUNCTION_QUERY：要用于此查询的距离函数。根据创建索引时使用的距离函数，选择以下各项之一：

    • L2 距离：<->

    • 内积：<#>

    • 余弦距离：<=>

  • EMBEDDING：您要为其查找存储的语义最近邻的嵌入向量。

  • ROW_COUNT：要返回的行数。

    如果您只想获得单个最佳匹配项，请指定 1。

如需详细了解其他查询示例，请参阅查询。

您还可以使用 embedding() 函数将文本转换为向量。您可以将向量应用于某个 pgvector 最近邻运算符（适用于 L2 距离的 <->），以查找具有语义最相似的嵌入的数据库行。

由于 embedding() 会返回 real 数组，因此您必须明确将 embedding() 调用转换为 vector，以便将这些值与 pgvector 运算符结合使用。

  CREATE EXTENSION IF NOT EXISTS google_ml_integration;
  CREATE EXTENSION IF NOT EXISTS vector;

  SELECT * FROM TABLE
    ORDER BY EMBEDDING_COLUMN::vector
    <-> embedding('MODEL_IDVERSION_TAG', 'TEXT')
    LIMIT ROW_COUNT

替换以下内容：

  • MODEL_ID：要查询的模型的 ID。

    如果您使用的是 Vertex AI Model Garden，请将 text-embedding-005 指定为模型 ID。这些是 AlloyDB 可用于文本嵌入的云端模型。如需了解详情，请参阅文本嵌入。

  • 可选：VERSION_TAG：要查询的模型的版本标记。在标记前面添加 @。

    如果您将某个 text-embedding 英语模型与 Vertex AI 搭配使用，请指定模型版本中列出的某个版本标记，例如 text-embedding-005。

    Google 强烈建议您始终指定版本标记。如果您未指定版本标记，AlloyDB 会始终使用最新的模型版本，这可能会导致意外结果。

  • TEXT：要转换为向量嵌入的文本。

后续步骤