为什么以及什么

向量数据库现在是非常热门的话题。我一直对它们是什么以及它们是如何在背后工作的感到好奇，所以我们自己来构建一个。从头开始构建一个全新的数据库并不现实，我们需要一些构建块，或者，直接使用一个真正的数据库系统。Postgres 因其扩展性而享有长期的声誉，这使它成为我们需求的完美选择，像 pgvector 这样的项目已经证明，将向量支持作为扩展添加到 Postgres 是可行的。

我们将为 Postgres 实现向量支持，但需要实现哪些详细功能呢？这个问题并不难，维基百科对 向量数据库 的定义为我们指明了正确的方向：

A vector database, vector store or vector search engine is a database that can store vectors (fixed-length lists of numbers) along with other data items. Vector databases typically implement one or more Approximate Nearest Neighbor algorithms so that one can search the database with a query vector to retrieve the closest matching database records

好的，那么我们需要使 Postgres 能够存储向量，并能够执行 Top-K 查询。即对于给定的输入向量，Postgres 应返回与之最相似（或最近）的 K 个向量。如果用 SQL 来表示，它可能看起来像这样：

-- 创建一个表，其中有一个 `vector(3)` 类型的列，3 是向量的维度
CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

-- 插入向量，Postgres 应该能够存储它们！
INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]');

-- 现在，Postgres 应该返回与 [3, 1, 2] 最相似的 Top-5 向量
SELECT * FROM items ORDER BY embedding <=> '[3,1,2]' LIMIT 5;

上述 SQL 示例使事情变得清晰，简而言之，我们需要：

1. 为 Postgres 实现一个 vector 类型，它应该接受维度参数。
2. 实现 <=> 二元操作符，它应该计算两个向量的相似度并返回结果。

设置环境

我将使用 Rust 语言和一个名为 [pgrx] 的库。安装 Rust 非常简单，只需按照 此处 的说明进行操作，然后运行以下命令来设置 cargo-pgrx，这是一个用于管理所有与 pgrx 相关事宜的 cargo 子命令：

$ cargo install --locked cargo-pgrx
$ cargo pgrx --version # 验证安装是否成功

现在我们需要一个 Postgres 服务器来运行和测试我们的项目，我将让 pgrx 为我安装一个全新的 Postgres 版本以简化操作。在撰写本文时，Postgres 17 是最新版本，所以我将使用它。

pgrx 从源代码构建 Postgres，因此你需要确保满足这些 构建要求。pgrx 还提供了一个关于 系统要求 的页面，但 Postgres 文档非常完善，值得一读。设置好一切后，运行：

$ cargo pgrx init --pg17 download

初始提交

现在让我们开始编写代码，cargo pgrx 和 cargo 一样提供了一个 new 子命令来创建新项目，假设我们将项目命名为 pg_vector_ext，运行：

$ cargo pgrx new pg_vector_ext

$ cd pg_vector_ext
$ tree .
pg_vector_ext/
├── Cargo.toml
├── pg_vector_ext.control
├── sql
└── src
    ├── bin
    │   └── pgrx_embed.rs
    └── lib.rs

4 directories, 4 files

从这些文件中可以看到，pgrx 为我们创建了一些模板文件。现在我们只关心 src/lib.rs 文件。

$ bat src/lib.rs
───────┬───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
       │ 文件: src/lib.rs
───────┼───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ use pgrx::prelude::*;
   2   │
   3   │ ::pgrx::pg_module_magic!();
   4   │
   5   │ #[pg_extern]
   6   │ fn hello_pg_vector_ext() -> &'static str {
   7   │     "Hello, pg_vector_ext"
   8   │ }
   9   │
  10   │ #[cfg(any(test, feature = "pg_test"))]
  11   │ #[pg_schema]
  12   │ mod tests {
  13   │     use pgrx::prelude::*;
  14   │
  15   │     #[pg_test]
  16   │     fn test_hello_pg_vector_ext() {
  17   │         assert_eq!("Hello, pg_vector_ext", crate::hello_pg_vector_ext());
  18   │     }
  19   │
  20   │ }
  21   │
  22   │ /// This module is required by `cargo pgrx test` invocations.
  23   │ /// It must be visible at the root of your extension crate.
  24   │ #[cfg(test)]
  25   │ pub mod pg_test {
  26   │     pub fn setup(_options: Vec<&str>) {
  27   │         // perform one-off initialization when the pg_test framework starts
  28   │     }
  29   │
  30   │     #[must_use]
  31   │     pub fn postgresql_conf_options() -> Vec<&'static str> {
  32   │         // return any postgresql.conf settings that are required for your tests
  33   │         vec![]
  34   │     }
  35   │ }
───────┴───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────

忽略 tests 模块（它用于测试），我们可以看到 pgrx 创建了一个函数 hello_pg_vector_ext()，这是一个可以在 SQL 中调用的函数。如果我们运行项目：

在运行之前，你需要编辑 Cargo.toml 文件，在 features 部分将默认特性改为 pg17，并且可以选择删除其他不使用的 pg* 特性：

$ cargo pgrx run

它将启动 Postgres 17 实例并通过 psql 连接到它，我们可以安装扩展并运行该函数：

pg_vector_ext=# CREATE EXTENSION pg_vector_ext;
CREATE EXTENSION
pg_vector_ext=# SELECT hello_pg_vector_ext();
 hello_pg_vector_ext
----------------------
 Hello, pg_vector_ext
(1 row)

这是我们使用 pgrx 的第一次尝试，也是我们对项目的第一次提交。在下一篇文章中，我将实现 vector 类型，以便 Postgres 可以存储向量。