golang-cookbook

Appearance

pgvector 是一个 PostgreSQL 扩展，用于在数据库中存储和查询高维向量。它非常适合机器学习、推荐系统和其他需要相似度搜索的场景。使用 pgvector，你可以存储向量数据（如嵌入向量）并执行高效的相似度查询，例如 余弦相似度 或 欧氏距离。

1. 安装 pgvector 扩展 ​

首先，需要确保在 PostgreSQL 中安装了 pgvector 扩展。你可以通过以下步骤安装和配置：

1.1 安装 pgvector 扩展 ​

sql
-- 创建扩展
CREATE EXTENSION IF NOT EXISTS vector;

如果你使用的是 Docker 或其他容器化环境，也可以通过以下命令安装扩展：

bash
docker exec -it your_postgres_container_name psql -U postgres
CREATE EXTENSION IF NOT EXISTS vector;

2. 创建包含向量数据的表 ​

你可以使用 vector 数据类型在 PostgreSQL 中创建表来存储向量数据。向量字段可以存储任意维度的向量。

2.1 创建表并存储向量数据 ​

假设我们要存储包含文本内容和相应嵌入向量的文档：

sql
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    content TEXT,
    embedding vector(1536)  -- 1536 是向量的维度，根据实际情况调整
);

这里，embedding 字段是 vector(1536) 类型，表示存储 1536 维的浮点向量。

2.2 插入向量数据 ​

插入向量数据时，可以通过 INSERT 语句将文本内容和向量插入表中。向量通常是通过模型（如 BERT、GPT、Word2Vec 等）计算得到的。

sql
-- 插入文档和其嵌入向量
INSERT INTO documents (content, embedding)
VALUES
('This is a sample document', '[0.1, 0.2, 0.3, ..., 0.1536]');

embedding 字段接收一个包含浮动值的数组，表示向量数据。通常，在实际应用中，你会使用程序（例如 Python 或 Go）将嵌入向量存入数据库。

3. 查询与向量相似的文档 ​

PostgreSQL 中的 pgvector 扩展支持使用 余弦相似度（cosine similarity）和 欧氏距离（Euclidean distance）来查询向量相似度。

3.1 查询与给定向量最相似的文档 ​

假设你有一个查询向量，想要查找与该向量最相似的文档，可以使用 <=> 操作符计算相似度，并按距离排序。

sql
-- 查询与给定向量最相似的 5 个文档（按余弦相似度排序）
SELECT id, content, embedding <=> '[0.1, 0.2, 0.3, ..., 0.1536]' AS distance
FROM documents
ORDER BY distance
LIMIT 5;

<=> 操作符用于计算与查询向量的 余弦相似度（或者根据实际情况选择欧氏距离）。

3.2 查询与给定向量的欧氏距离最小的文档 ​

如果你想使用 欧氏距离（Euclidean distance）来度量相似度，可以使用 vector <-> vector 操作符：

sql
-- 查询与给定向量最接近的 5 个文档（按欧氏距离排序）
SELECT id, content, embedding <-> '[0.1, 0.2, 0.3, ..., 0.1536]' AS distance
FROM documents
ORDER BY distance
LIMIT 5;

这里，<-> 操作符用于计算与查询向量的欧氏距离。

4. 向量索引 ​

为了加速向量查询，尤其是在有大量文档时，pgvector 支持使用 倒排索引 和 IVF（倒排文件）索引，这对于高效的向量检索至关重要。

4.1 创建向量索引 ​

你可以创建基于向量字段的索引，以加速查询。例如：

sql
-- 创建索引以加速相似度搜索
CREATE INDEX ON documents USING ivfflat (embedding vector_l2_ops);

ivfflat 是 pgvector 提供的基于倒排文件的索引类型，适合欧氏距离（L2 距离）。你也可以使用 cosine_ops 来进行余弦相似度的索引。

4.2 查询时使用索引 ​

创建索引后，查询时 PostgreSQL 会自动使用索引来加速相似度查询。如果使用 ORDER BY distance，数据库会优先选择更快的索引查询方式。

5. 从 Go 程序中操作 pgvector ​

如果你使用 Go 来与 PostgreSQL 和 pgvector 交互，可以通过 pgx 或 gorm 等库来实现。

5.1 使用 pgx 插入和查询向量数据 ​

首先，你需要安装 pgx 库：

bash
go get github.com/jackc/pgx/v4

然后，使用以下代码插入和查询向量数据：

go
package main

import (
	"context"
	"fmt"
	"log"
	"github.com/jackc/pgx/v4"
	"github.com/lib/pq"
)

func main() {
	// 设置数据库连接
	conn, err := pgx.Connect(context.Background(), "postgres://user:password@localhost:5432/mydb")
	if err != nil {
		log.Fatal("Unable to connect to database:", err)
	}
	defer conn.Close(context.Background())

	// 插入向量数据
	embedding := pq.Array([]float32{0.1, 0.2, 0.3, 0.4, 0.5}) // 示例向量
	_, err = conn.Exec(context.Background(), `
		INSERT INTO documents (content, embedding) 
		VALUES ($1, $2)
	`, "Example document", embedding)
	if err != nil {
		log.Fatal("Failed to insert data:", err)
	}

	// 查询与给定向量最相似的文档
	queryVector := pq.Array([]float32{0.1, 0.2, 0.3, 0.4, 0.5}) // 查询向量
	rows, err := conn.Query(context.Background(), `
		SELECT id, content, embedding <=> $1 AS distance
		FROM documents
		ORDER BY distance
		LIMIT 5
	`, queryVector)
	if err != nil {
		log.Fatal("Failed to query data:", err)
	}

	// 输出查询结果
	for rows.Next() {
		var id int
		var content string
		var distance float32
		err := rows.Scan(&id, &content, &distance)
		if err != nil {
			log.Fatal("Failed to scan row:", err)
		}
		fmt.Printf("ID: %d, Content: %s, Distance: %f\n", id, content, distance)
	}
}

6. 总结 ​

  • pgvector 是 PostgreSQL 的一个扩展，支持高效地存储和查询向量数据。
  • 可以使用 余弦相似度 或 欧氏距离 来查询与给定向量最相似的文档。
  • 使用 IVF索引 可以加速大规模向量数据的查询。
  • 通过 Go 程序，可以轻松地与 PostgreSQL 和 pgvector 进行交互，执行插入、查询和索引操作。

pgvector 为 PostgreSQL 提供了强大的向量处理能力，使得在数据库中处理大规模高维向量数据变得更加高效。