你当前正在访问 Microsoft Azure Global Edition 技术文档网站。 如果需要访问由世纪互联运营的 Microsoft Azure 中国技术文档网站,请访问 https://docs.azure.cn

在 Azure Cosmos DB for PostgreSQL 上使用 pgvector 时如何优化性能

适用对象:PostgreSQL 的 Azure Cosmos DB (由 PostgreSQL 的 Citus 数据库扩展提供支持)

pgvector 扩展将开源向量相似性搜索添加到 PostgreSQL。

本文探讨 pgvector 的局限性和利弊,并介绍如何使用分区、索引和搜索设置来提高性能。

有关该扩展本身的详细信息,请参阅 pgvector 基础知识。 建议另外参阅该项目的官方自述文件

性能

始终从调查查询计划开始。 如果查询以相当快的速度终止,请运行 EXPLAIN (ANALYZE,VERBOSE, BUFFERS)

EXPLAIN (ANALYZE, VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

对于执行时间过长的查询,请考虑删除 ANALYZE 关键字。 结果包含更少的详细信息,但可以立即提供。

EXPLAIN (VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

第三方站点(例如 explain.depesz.com)有助于理解查询计划。 应该尝试回答的一些问题包括:

向量是否已规范化为长度 1,例如 OpenAI 嵌入项。 应考虑使用内积 (<#>) 来获得最佳性能。

并行执行

在解释计划的输出中,查找 Workers PlannedWorkers Launched(后者仅在使用 ANALYZE 关键字时适用)。 max_parallel_workers_per_gather PostgreSQL 参数定义数据库可为每个 GatherGather Merge 计划节点启动多少个后台工作器。 增大此值可以加快精确搜索查询,而无需创建索引。 但请注意,即使此值较高,数据库也可能不会决定并行运行计划。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 3;
                                        QUERY PLAN
------------------------------------------------------------------------------------------
 Limit  (cost=4214.82..4215.16 rows=3 width=33)
   ->  Gather Merge  (cost=4214.82..13961.30 rows=84752 width=33)
         Workers Planned: 1
         ->  Sort  (cost=3214.81..3426.69 rows=84752 width=33)
               Sort Key: ((embedding <-> '[1,2,3]'::vector))
               ->  Parallel Seq Scan on t_test  (cost=0.00..2119.40 rows=84752 width=33)
(6 rows)

索引

在没有索引的情况下,该扩展会执行精确搜索,以牺牲性能为代价提供完美召回率。

为了执行近似最近的邻域搜索,可以基于数据创建索引,以召回率换取执行性能。

如果可能,请始终在为数据编制索引之前加载数据。 以这种方式创建索引速度更快,生成的布局更佳。

有两种支持的索引类型:

IVFFlat 索引的生成时间更快,使用的内存比 HNSW 少,但查询性能较低(在速度-召回率权衡方面)。

限制

  • 若要为某个列编制索引,该列必须已定义维度。 尝试为定义为 col vector 的列编制索引会导致错误:ERROR: column does not have dimensions
  • 只能为最多包含 2000 个维度的列编制索引。 尝试为包含更多维度的列编制索引会导致错误:ERROR: column cannot have more than 2000 dimensions for INDEX_TYPE index,其中 INDEX_TYPEivfflathnsw

虽然可以存储维度数超过 2000 的向量,但无法为这些向量编制索引。 可以使用降维来避免超出限制。 或者,依赖 Azure Cosmos DB for PostgreSQL 的分区和/或分片,在不编制索引的情况下实现可接受的性能。

反向文件与平面压缩 (IVVFlat)

ivfflat 是近似最近的邻域 (ANN) 搜索的索引。 此方法使用倒排文件索引将数据集分区成多个列表。 probes 参数控制搜索的列表数,能够以搜索速度减慢为代价提高搜索结果的准确度。

如果将 probes 参数设置为索引中的列表数,则会搜索所有列表,并且搜索将变成精确最近的邻域搜索。 在这种情况下,规划器不会使用索引,因为搜索所有列表等同于对整个数据集执行暴力搜索。

索引编制方法使用 k 平均值聚类算法将数据集分区成多个列表。 每个列表包含最接近特定聚类中心的向量。 在搜索过程中,会将查询向量与聚类中心进行比较,以确定哪些列表最有可能包含最近的邻域。 如果 probes 参数设置为 1,则只搜索与最近的聚类中心对应的列表。

索引选项

为要执行的探测数量和列表大小选择正确的值可能会影响搜索性能。 良好的起点是:

  1. 对于包含最多 100 万行的表,使用等于 rows / 1000lists;对于更大的数据集,使用 sqrt(rows)
  2. 对于 probes,首先为包含最多 100 万行的表使用 lists / 10,并为更大的数据集使用 sqrt(lists)

lists 的数量是在使用 lists 选项创建索引时定义的:

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 5000);

可为整个连接或每个事务设置探测(在事务块中使用 SET LOCAL):

SET ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes
BEGIN;

SET LOCAL ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, one probe

索引进度

在 PostgreSQL 12 及更新版本中,可使用 pg_stat_progress_create_index 检查索引进度。

SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

生成 IVFFlat 索引的阶段包括:

  1. initializing
  2. performing k-means
  3. assigning tuples
  4. loading tuples

注意

进度百分比 (%) 仅在“loading tuples”阶段填充。

分层可导航小世界 (HNSW)

hnsw 是使用分层可导航小世界算法进行的近似最近邻域 (ANN) 搜索的一个索引。 它的工作原理是:围绕随机选择的入口点创建一个图形,查找它们最近的邻域,然后用多个层扩展该图形,每个下层包含更多点。 当搜索从顶部开始时,此多层图形会向下缩小,直到达到包含查询的最近邻域的最低层为止。

构建此索引所需的时间和内存比 IVFFlat 更多,但它具有更好的速度-召回率权衡。 此外,没有类似于 IVFFlat 的训练步骤,因此可以在空表中创建索引。

索引选项

创建索引时,可以优化两个参数:

  1. m - 每层最大连接数(默认值为 16)
  2. ef_construction - 用于图形构造的动态候选列表的大小(默认值为 64)
CREATE INDEX t_test_hnsw_l2_idx ON t_test USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);

在查询期间,可以指定搜索的动态候选列表(默认值为 40)。

可以针对整个连接或每个事务(在事务块中使用 SET LOCAL)设置搜索的动态候选列表:

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates
BEGIN;

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, 40 candidates

索引进度

在 PostgreSQL 12 及更新版本中,可使用 pg_stat_progress_create_index 检查索引进度。

SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

生成 HNSW 索引的阶段包括:

  1. initializing
  2. loading tuples

选择索引访问函数

vector 类型允许对存储的向量执行三种类型的搜索。 需要为索引选择正确的访问函数,以便数据库在执行查询时考虑索引。 这些示例演示了 ivfflat 索引类型,但也可以为 hnsw 索引执行相同的操作。 lists 选项仅适用于 ivfflat 索引。

余弦距离

对于余弦相似性搜索,请使用 vector_cosine_ops 访问方法。

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

若要使用上述索引,查询需要使用 <=> 运算符执行余弦相似性搜索。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5;
                                              QUERY PLAN
------------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_cosine_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <=> '[1,2,3]'::vector)
(3 rows)

L2 距离

对于 L2 距离(也称为欧氏距离),请使用 vector_l2_ops 访问方法。

CREATE INDEX t_test_embedding_l2_idx ON t_test USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

若要使用上述索引,查询需要使用 <-> 运算符执行 L2 距离搜索。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;
                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_l2_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <-> '[1,2,3]'::vector)
(3 rows)

inner_product

对于内积相似性,请使用 vector_ip_ops 访问方法。

CREATE INDEX t_test_embedding_ip_idx ON t_test USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);

若要使用上述索引,查询需要使用 <#> 运算符执行内积相似性搜索。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <#> '[1,2,3]' LIMIT 5;
                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_ip_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <#> '[1,2,3]'::vector)
(3 rows)

部分索引

在某些情况下,使用仅涵盖部分数据集的索引是有利的。 例如,我们可以只为高级用户生成一个索引:

CREATE INDEX t_premium ON t_test USING ivfflat (vec vector_ip_ops) WITH (lists = 100) WHERE tier = 'premium';

可以看到,高级层现在使用了该索引:

explain select * from t_test where tier = 'premium' order by vec <#> '[2,2,2]';
                                     QUERY PLAN
------------------------------------------------------------------------------------
 Index Scan using t_premium on t_test  (cost=65.57..25638.05 rows=245478 width=39)
   Order By: (vec <#> '[2,2,2]'::vector)
(2 rows)

而免费层用户则缺乏这种优势:

explain select * from t_test where tier = 'free' order by vec <#> '[2,2,2]';
                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44019.01..44631.37 rows=244941 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15395.25 rows=244941 width=39)
         Filter: (tier = 'free'::text)
(4 rows)

只为一部分数据编制索引意味着索引占用的磁盘空间更少,且搜索速度更快。

如果部分索引定义的 WHERE 子句中使用的格式与查询中使用的格式不匹配,则 PostgreSQL 可能无法确定使用索引是否安全。 在我们的示例数据集中,只有确切值 'free''test''premium' 是层列的非重复值。 即使对于使用 tier LIKE 'premium' 的查询,PostgreSQL 也不会使用索引。

explain select * from t_test where tier like 'premium' order by vec <#> '[2,2,2]';
                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44086.30..44700.00 rows=245478 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15396.59 rows=245478 width=39)
         Filter: (tier ~~ 'premium'::text)
(4 rows)

分区

提高性能的一种方法是将数据集拆分为多个分区。 假设某个系统可以自然地引用当年或过去两年的数据。 在此类系统中,可以按日期范围将数据分区,然后在系统能够仅读取按查询年份定义的相关分区时受益于改进的性能。

让我们定义一个已分区表:

CREATE TABLE t_test_partitioned(vec vector(3), vec_date date default now()) partition by range (vec_date);

可以手动为每个年份创建分区,或使用 Citus 实用工具函数(适用于 Cosmos DB for PostgreSQL)。

    select create_time_partitions(
      table_name         := 't_test_partitioned',
      partition_interval := '1 year',
      start_from         := '2020-01-01'::timestamptz,
      end_at             := '2024-01-01'::timestamptz
    );

检查创建的分区:

\d+ t_test_partitioned
                                Partitioned table "public.t_test_partitioned"
  Column  |   Type    | Collation | Nullable | Default | Storage  | Compression | Stats target | Description
----------+-----------+-----------+----------+---------+----------+-------------+--------------+-------------
 vec      | vector(3) |           |          |         | extended |             |              |
 vec_date | date      |           |          | now()   | plain    |             |              |
Partition key: RANGE (vec_date)
Partitions: t_test_partitioned_p2020 FOR VALUES FROM ('2020-01-01') TO ('2021-01-01'),
            t_test_partitioned_p2021 FOR VALUES FROM ('2021-01-01') TO ('2022-01-01'),
            t_test_partitioned_p2022 FOR VALUES FROM ('2022-01-01') TO ('2023-01-01'),
            t_test_partitioned_p2023 FOR VALUES FROM ('2023-01-01') TO ('2024-01-01')

若要手动创建分区,请运行:

CREATE TABLE t_test_partitioned_p2019 PARTITION OF t_test_partitioned FOR VALUES FROM ('2019-01-01') TO ('2020-01-01');

然后,确保查询实际向下筛选到可用分区的子集。 例如,在下面的查询中,我们向下筛选到两个分区:

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01';
                                                                  QUERY PLAN
-----------------------------------------------------------------------------------------------------------------------------------------------
 Append  (cost=0.00..58.16 rows=12 width=36) (actual time=0.014..0.018 rows=3 loops=1)
   ->  Seq Scan on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.00..29.05 rows=6 width=36) (actual time=0.013..0.014 rows=1 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
   ->  Seq Scan on t_test_partitioned_p2023 t_test_partitioned_2  (cost=0.00..29.05 rows=6 width=36) (actual time=0.002..0.003 rows=2 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.125 ms
 Execution Time: 0.036 ms

可为已分区表编制索引。

CREATE INDEX ON t_test_partitioned USING ivfflat (vec vector_cosine_ops) WITH (lists = 100);
explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01' order by vec <=> '[1,2,3]' limit 5;
                                                                                         QUERY PLAN
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Limit  (cost=4.13..12.20 rows=2 width=44) (actual time=0.040..0.042 rows=1 loops=1)
   ->  Merge Append  (cost=4.13..12.20 rows=2 width=44) (actual time=0.039..0.040 rows=1 loops=1)
         Sort Key: ((t_test_partitioned.vec <=> '[1,2,3]'::vector))
         ->  Index Scan using t_test_partitioned_p2022_vec_idx on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.04..4.06 rows=1 width=44) (actual time=0.022..0.023 rows=0 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
         ->  Index Scan using t_test_partitioned_p2023_vec_idx on t_test_partitioned_p2023 t_test_partitioned_2  (cost=4.08..8.11 rows=1 width=44) (actual time=0.015..0.016 rows=1 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.167 ms
 Execution Time: 0.139 ms
(11 rows)

结论

祝贺你,你已了解 pgvector 的利弊和局限性,以及实现其最佳性能的最佳做法。