--- title: "给 Postgres 写一个向量插件 - 介绍" date: 2024-12-18 lastmod: 2024-12-18 description: "本文探讨了如何从零为 Postgres 实现向量支持，以构建一个简单的向量数据库。通过 Rust 和 `pgrx` 库，作者创建了一个扩展项目 `pg_vector_ext`，初步实现了基本功能（如 `hello_pg_vector_ext()` 函数），并计划后续实现向量类型和相似度查询操作符 `<=>`。" tags: ["vector database", "Postgres"] summary: "为什么以及什么 # 向量数据库现在是非常热门的话题。我一直对它们是什么以及它们是如何在背后工作的感到好奇，所以我们自己来构建一个。从头开始构建一个全新的数据库并不现实，我们需要一些构建块，或者，直接使用一个真正的数据库系统。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 来表示，它可能看起来像这样：" --- # 为什么以及什么 向量数据库现在是非常热门的话题。我一直对它们是什么以及它们是如何在背后工作的感到好奇，所以我们自己来构建一个。从头开始构建一个全新的数据库并不现实，我们需要一些构建块，或者，直接使用一个真正的数据库系统。Postgres 因其扩展性而享有长期的声誉，这使它成为我们需求的完美选择，像 [pgvector] 这样的项目已经证明，将向量支持作为扩展添加到 Postgres 是可行的。 我们将为 Postgres 实现向量支持，但需要实现哪些详细功能呢？这个问题并不难，维基百科对[向量数据库][vector_db_wikipedia] 的定义为我们指明了正确的方向： > 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 来表示，它可能看起来像这样： ```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 非常简单，只需按照[此处][install_rust] 的说明进行操作，然后运行以下命令来设置 `cargo-pgrx`，这是一个用于管理所有与 pgrx 相关事宜的 cargo 子命令： ```sh $ cargo install --locked cargo-pgrx $ cargo pgrx --version # 验证安装是否成功 ``` 现在我们需要一个 Postgres 服务器来运行和测试我们的项目，我将让 pgrx 为我安装一个全新的 Postgres 版本以简化操作。在撰写本文时，[Postgres 17][pg17_release] 是最新版本，所以我将使用它。 pgrx 从源代码构建 Postgres，因此你需要确保满足这些[构建要求][build_pg_requirements]。pgrx 还提供了一个关于[系统要求][pgrx_system_requiremens] 的页面，但 Postgres 文档非常完善，值得一读。设置好一切后，运行： ```sh $ cargo pgrx init --pg17 download ``` # 初始提交 现在让我们开始编写代码，cargo pgrx 和 cargo 一样提供了一个 new 子命令来创建新项目，假设我们将项目命名为 pg_vector_ext，运行： ```sh $ 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 文件。 ```sh $ 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* 特性： ```sh $ cargo pgrx run ``` 它将启动 Postgres 17 实例并通过 psql 连接到它，我们可以安装扩展并运行该函数： ```sql 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 可以存储向量。 [pgvector]: https://github.com/pgvector/pgvector [vector_db_wikipedia]: https://en.wikipedia.org/wiki/Vector_database [pgrx]: https://github.com/pgcentralfoundation/pgrx [install_rust]: https://www.rust-lang.org/tools/install [pg17_release]: https://www.postgresql.org/about/news/postgresql-17-released-2936/ [build_pg_requirements]: https://www.postgresql.org/docs/current/install-requirements.html [pgrx_system_requiremens]: https://github.com/pgcentralfoundation/pgrx/?tab=readme-ov-file#system-requirements