PostgreSQL pgvector：入门与扩展

PostgreSQL pgvector：入门与扩展

摘要： 本文提供了 pgvector 扩展在 PostgreSQL 中的完整入门指南，涵盖安装、将文本嵌入（embedding）存储为向量、使用距离函数（余弦距离、欧几里得距离、内积）执行相似性搜索、预过滤策略，以及 HNSW 和 IVFFlat 索引的对比。文章还探讨了如何使用基于 PostgreSQL 构建的分布式 SQL 数据库 YugabyteDB 来扩展基于 pgvector 的应用。

原文链接

开始使用 Pgvector

要开始使用 pgvector，需要在 PostgreSQL 中安装该扩展。

1. 安装带有 pgvector 扩展的 PostgreSQL Docker 镜像。

docker pull pgvector/pgvector:pg16

2. 运行该 Docker 镜像

docker run --name postgres -p 5432:5432 -e POSTGRES_PASSWORD=mysecretpassword -d pgvector/pgvector:pg16

3. 连接到 Docker 容器并在 PostgreSQL 中启用 pgvector。

docker exec -it postgres ./bin/psql -U postgres

postgres=# CREATE EXTENSION IF NOT EXISTS vector;
postgres=# SELECT extname from pg_extension;
 extname
----------
 plpgsql
 vector
(2 rows)

pgvector 安装完成后，就可以开始在 PostgreSQL 中存储嵌入向量了。

存储文本嵌入

文本嵌入（text embedding）是文本的数值化表示，以固定维度向量的形式存储。例如，回想一下代数课上在二维空间中绘制点 (x,y) 坐标。

在现代文本嵌入模型中，维度数量要大得多，通常是 768、1536 甚至更多。pgvector 扩展目前支持最多 2000 维的向量，最近还新增了对最多 4000 维的半向量（half vector）存储支持。这些高维向量表示允许通过比较向量每个维度的值来进行相似性搜索。

假设我们有一个流行超级英雄的数据集，并且可以访问一个能为每个描述生成二维文本嵌入的模型。

以下是使用 pgvector 存储嵌入的方法：

1. 创建存储超级英雄及其关联嵌入的表，使用 vector 数据类型。这里，我们在 description_embedding 列中指定嵌入将有两个维度。

CREATE TABLE superheroes (
  id int4 NULL,
  hero_name varchar(50) NULL,
  gender varchar(50) NULL,
  eye_color varchar(50) NULL,
  race varchar(50) NULL,
  hair_color varchar(50) NULL,
  height float4 NULL,
  publisher varchar(50) NULL,
  skin_color varchar(50) NULL,
  alignment varchar(50) NULL,
  weight float4 null,
  description text,
description_embedding vector(2),
);

2. 将 superheroes 插入数据库。

INSERT INTO superheroes (superhero_name, superhero_description) VALUES 
('superman', 'possesses super strength, flight, and invulnerability, which he uses to fight for truth and justice on Earth'),
('batman', 'relies on his intellect, physical prowess, and an array of technological gadgets to combat crime and corruption'),
('the incredible hulk', 'transforms into a massive green creature with immense strength whenever he becomes angry');

3. 为每个超级英雄描述分配 description_embeddings。注意：这只是一个示例。在真实场景中，会使用文本嵌入模型从 description 列中存储的文本生成向量。

UPDATE superheroes 
SET description_embedding = '[1, 3]' 
WHERE id = 1; 

UPDATE superheroes 
SET description_embedding = '[1, 4]'
WHERE id = 2;

UPDATE superheroes 
SET description_embedding = '[4, 6]'
WHERE id = 3;

在这个例子中，我们可以绘制为每个超级英雄描述生成的嵌入，以可视化它们的相似程度。

在下一节中，我们将介绍如何使用 pgvector 中的距离函数来查询这些嵌入。

查询嵌入

Pgvector 允许我们使用熟悉的 PostgreSQL 语法查询记录。查询向量最常见的方式是通过距离来排序。

pgvector 提供了多种距离函数，适用于不同的实际应用场景。

1. 余弦距离（Cosine Distance，<=>）： 衡量两个向量之间夹角的余弦值。在进行角度比较时，相似性与向量幅度无关。向量之间的余弦距离越小，它们越相似。该函数常用于文本比较，当向量的方向比幅度更重要时尤为适用。

   余弦距离的公式可以从余弦相似度推导得出：

• a ⋅ b：这是两个向量的点积（dot product）。它计算两个向量对应元素乘积的总和。

• ∥a∥：表示向量 a 的幅度（或范数），计算公式为

• ∥b∥：类似地，表示向量 b 的幅度，计算公式为

余弦相似度的范围是 -1 到 1，其中 1 表示最大相似度，0 表示无相似度，-1 表示最大不相似度。

以下是余弦距离从余弦相似度推导的方法：

余弦距离的范围是 0 到 2，其中 0 表示无距离或最大相似度，2 表示最大距离和最大不相似度。

以下是使用 pgvector 进行余弦距离查询的示例：

SELECT hero_name, description_embedding FROM superheroes ORDER BY description_embedding <=> ['1,5]' LIMIT 5;
 
 hero_name         | description_embedding 
Batman               [1,4]
Superman             [1,3]
The Incredible Hulk  [4,6]

虽然没有内置的直接按余弦相似度查询的函数，但这可以通过代数方式实现。如果 Distance = 1 - Similarity，那么 Similarity = 1 - Distance。

因此，我们可以按余弦相似度排序：

SELECT 1 -  (description_embedding <=> '[3,1]') AS similarity FROM superheroes ORDER BY similarity;

2. 欧几里得距离 / L2 距离（Euclidean / L2 distance，<->）： 衡量两点之间的直线距离。通过比较距离，假设较短的距离代表更高的相似度。

SELECT * FROM superheroes ORDER BY description_embedding <-> '[3,1]';

3. 负内积（Negative Inner Product，<#>）： 衡量两个向量之间的幅度（距离）和方向（角度）。Pgvector 计算两个向量的内积并取其负值，值越低表示相似度越高。

SELECT (description_embedding <#> '[3,1]') * -1 AS inner_product FROM superheroes;
...

接下来，让我们探讨如何通过预过滤数据来提高查询效率。

预过滤数据

使用提供的距离函数通过文本嵌入查询数据很简单，但在大数据集上可能会成为高延迟操作。为了提高应用效率，在查询中预过滤数据是一种常见策略。

这在向量数据库中很常见，数据被分成"块"（chunk）。每个块都有一个关联的元数据对象，允许数据库在执行向量搜索前高效地过滤数据。

对于使用 SQL 数据库的人来说，这是一种熟悉的做法。例如，我们可以先限制数据集为 Marvel 超级英雄来预过滤。

SELECT COUNT(*) from superheroes;
 count 
-------
   734

SELECT COUNT(*) from superheroes where publisher = 'Marvel Comics';
 count 
-------
   388

我们还可以进一步缩小范围，搜索蓝眼睛的 Marvel 超级英雄。

SELECT COUNT(*) from superheroes where publisher = 'Marvel Comics' AND eye_color = 'blue';
 count 
-------
   125

我使用 1536 维的 OpenAI 嵌入模型为每个角色的描述生成了文本嵌入，并通过 pgvector 将其存储在数据库中。通过使用相同的模型为文本"has the power of invisibility"生成嵌入，我们可以按余弦距离排序。

SELECT hero_name from superheroes ORDER BY description_embedding <=> '[0.2452, 0.62356, ...] LIMIT 5';
    
hero_name    
-----------------
 Phantom
 Invisible Woman
 Vanisher
 Cloak
 Hollow
(5 rows)


Sort Method: top-N heapsort  Memory: 25kB
         ->  Seq Scan on superheroes  (cost=0.00..67.17 rows=734 width=18) (actual time=0.951..21.068 rows=734 loops=1)
Planning Time: 0.730 ms
Execution Time: 21.494 ms

接下来，我们在缩减后的数据集上通过预过滤来执行查询。在下面的示例中，我们搜索 DC Comics 中蓝眼睛且具有隐身能力的超级英雄。

首先，我们在 publisher 列上创建一个索引以提高查询效率。

CREATE INDEX publisher_idx on superheroes(publisher);

现在，执行查询并观察查询计划。

SELECT hero_name from superheroes where publisher = 'DC Comics' AND eye_color = 'blue' ORDER BY description_embedding <=> '[0.2452, 0.62356, ...] LIMIT 5';

hero_name
----------------
Rorschach
Misfit
Deadman
Scarecrow
John Constantine
(5 rows)

 Limit  (cost=68.27..68.28 rows=5 width=1095) (actual time=1.019..1.021 rows=5 loops=1)
   ->  Sort  (cost=68.27..68.43 rows=66 width=1095) (actual time=1.017..1.018 rows=5 loops=1)
         Sort Key: ((description_embedding <=> '[0.014116188,-0.00544635,...]'::vector)
         Sort Method: top-N heapsort  Memory: 33kB
         ->  Bitmap Heap Scan on superheroes  (cost=5.78..67.17 rows=66 width=1095) (actual time=0.075..0.936 rows=83 loops=1)
               Recheck Cond: ((publisher)::text = 'DC Comics'::text)
               Filter: ((eye_color)::text = 'blue'::text)
               Rows Removed by Filter: 132
               Heap Blocks: exact=57
               ->  Bitmap Index Scan on publisher_idx  (cost=0.00..5.76 rows=215 width=0) (actual time=0.032..0.032 rows=215 loops=1)
                     Index Cond: ((publisher)::text = 'DC Comics'::text)
 Planning Time: 0.146 ms
 Execution Time: 1.060 ms

这里，我们展示了通过创建索引和预过滤数据集，延迟从约 22ms 显著下降到约 5ms。

现在，让我们看看如何通过 pgvector 支持的专业向量索引进一步提高查询性能。

比较向量索引

默认情况下，pgvector 会对数据集执行精确最近邻查询，即对数据进行全顺序扫描。如上所示，在大数据集上执行时，这可能导致查询缓慢。

不过，我们可以使用索引通过近似最近邻（ANN）搜索来更高效地查询数据。ANN 算法允许我们牺牲召回率来换取速度，这通常是应用中值得的权衡。目前，pgvector 支持 HNSW 和 IVFFlat 索引，其他索引正在开发中。

大多数用例的首选索引是分层导航小世界（HNSW）索引。

分层导航小世界（HNSW）

HNSW 是一种多层基于图的索引，兼具高性能和高准确度。

顾名思义，HNSW 是导航小世界（Navigable Small Worlds）的扩展，这是一种基于图的最近邻查找算法。通过使这些图具有层次结构，可以减少时间和计算资源。

以下是 HNSW 从入口点 E 到查询向量的 ANN（顶点 3）的简化遍历示意图：

每一层（图）都比上一层更详细，通过将向量组织成"邻域"（neighborhood）来实现高查询性能。这类似于使用数字地图的体验——放大时细节级别会增加。

HNSW 索引中的插入是迭代的，索引时间随图中数据量的增加而增加。这使得索引易于管理，但需要付出索引时间的代价。

CREATE INDEX superhero_description_hnsw_idx ON superheroes USING hnsw (description_embedding vector_cosine_ops) WITH (m = 4, ef_construction = 10);

Limit  (cost=100.18..101.82 rows=5 width=1095) (actual time=0.849..0.912 rows=5 loops=1)
   ->  Index Scan using superheroes_description_hnsw_idx on superheroes  (cost=100.18..341.35 rows=734 width=1095) (actual time=0.846..0.909 rows=5 loops=1)
         Order By: (description_embedding <=> '[0.014116188,-0.00544635,...]'::vector)
 Planning Time: 0.185 ms
 Execution Time: 0.649 ms

倒排文件平面索引（IVFFlat）

IVFFlat 索引与 HNSW 不同，它在构建时将数据聚类成列表。通过这种方式，可以将搜索向量与每个聚类的质心（也称为列表）进行比较，以确定最相关的聚类。

通过配置查询期间要访问的列表数量，我们可以调整召回率和速度之间的权衡。这通常称为设置"探针"（probe）数量。在查询中设置更多的探针将搜索更多的列表，从而获得更好的结果。

例如，如果您希望访问 10 个列表以提高查询的召回率，可以设置以下参数：

SET ivfflat.probes = 10;

类似地，可以设置列表的数量，更多的列表意味着每个列表中的向量更少。这将通过减少搜索空间中的向量数量来提高查询速度，但可能因排除有效数据点而导致召回率误差。

列表数量可以在创建索引时设置。例如，以下语句将数据分割为 50 个列表：

CREATE INDEX superhero_description_ivfflat_idx ON superheroes USING ivfflat (description_embedding vector_cosine_ops) WITH (lists = 50);


SELECT hero_name FROM superheroes ORDER BY description_embedding <=> '[0.014116188,-0.00544635,... ]' LIMIT 5;

Limit  (cost=40.11..40.51 rows=5 width=18) (actual time=0.477..0.538 rows=5 loops=1)
  ->  Index Scan using superhero_description_ivfflat_idx on superheroes  (cost=40.11..98.10 rows=734 width=18) (actual time=0.476..0.535 rows=5 loops=1)
        Order By: (description_embedding <=> '[0.014116188,-0.00544635,...]'::vector)
Planning Time: 0.142 ms
Execution Time: 0.567 ms

对比

让我们快速对比两种索引：

索引	构建速度	查询速度	更新后是否需要重建
HNSW	最慢	快	否
IVFFlat	最快	最慢	是

HNSW

• 基于图（Graph）
• 将向量组织成邻域
• 迭代插入
• 插入时间随图中数据量增加而增加

适用场景：

• 数据频繁变更
• 需要高查询性能和召回率

IVFFlat

• 基于 K-means
• 将向量组织成列表
• 需要预填充数据
• 插入时间受列表数量限制

适用场景：

• 数据基本静态
• 需要快速索引

使用 YugabyteDB 进行扩展

让我们探讨如何利用分布式 SQL 使我们的应用更具可扩展性和弹性。

以下是使用 pgvector 的 AI 应用受益于分布式 PostgreSQL 数据库（如 YugabyteDB）的一些关键原因：

1. 嵌入向量消耗大量存储和内存。 例如，一个具有 1536 维的 OpenAI 模型在 1000 万条记录下将占用约 57GB 空间。水平扩展提供了存储向量所需的空间。
2. 向量相似性搜索非常消耗计算资源。 通过扩展到多个节点，应用可以访问无限制的 CPU 和 GPU 资源。
3. 服务中断不再是问题。 数据库将对节点、数据中心或区域故障具有弹性，这意味着 AI 应用永远不会因数据库层而经历停机。

YugabyteDB 是一个基于 PostgreSQL 构建的分布式 SQL 数据库。它与 Postgres 功能和运行时兼容，允许您重用为标准版 Postgres 创建的库、驱动程序、工具和框架。

YugabyteDB 与 pgvector 兼容，并提供原生 PostgreSQL 中的所有功能。这使其成为希望提升 AI 应用水平的开发者的理想选择。

以下是如何在本地 Docker 中运行 3 节点 YugabyteDB 集群的方法：

1. 创建本地存储数据的目录。

mkdir ~/yb_docker_data

2. 创建 Docker 网络

docker network create custom-network

3. 在此网络上部署 3 节点集群。

docker run -d --name yugabytedb-node1 --net custom-network \
      -p 15433:15433 -p 7001:7000 -p 9001:9000 -p 5433:5433 \
      -v ~/yb_docker_data/node1:/home/yugabyte/yb_data --restart unless-stopped \
      yugabytedb/yugabyte:latest \
      bin/yugabyted start \
      --base_dir=/home/yugabyte/yb_data --background=false

docker run -d --name yugabytedb-node2 --net custom-network \
      -p 15434:15433 -p 7002:7000 -p 9002:9000 -p 5434:5433 \
      -v ~/yb_docker_data/node2:/home/yugabyte/yb_data --restart unless-stopped \
      yugabytedb/yugabyte:latest \
      bin/yugabyted start --join=yugabytedb-node1 \
      --base_dir=/home/yugabyte/yb_data --background=false

docker run -d --name yugabytedb-node3 --net custom-network \
      -p 15435:15433 -p 7003:7000 -p 9003:9000 -p 5435:5433 \
      -v ~/yb_docker_data/node3:/home/yugabyte/yb_data --restart unless-stopped \
      yugabytedb/yugabyte:latest \
      bin/yugabyted start --join=yugabytedb-node1 \
      --base_dir=/home/yugabyte/yb_data --background=false

结论

pgvector 扩展是开源社区为 PostgreSQL 增加功能的又一实例，使其成为现代应用的理想数据库选择。

随着更多距离函数、索引等功能的开发，pgvector 持续为那些在其应用中添加 AI 功能的开发者创造价值。希望在 AI 应用中增加弹性和可扩展性的开发者可以转向 YugabyteDB，在不牺牲性能的情况下水平分布数据。

有兴趣了解更多关于使用 YugabyteDB 构建生成式 AI 应用的信息？请查看以下资源：

• Follow an AI tutorial backed by YugabyteDB
• Watch the getting started with pgvector video series
• View the AI travel agent with pgvector on GitHub