告别低效查询！用Python轻松玩转pgvector向量数据库

【免费下载链接】pgvector Open-source vector similarity search for Postgres 项目地址: https://gitcode.com/GitHub_Trending/pg/pgvector

你是否还在为向量数据查询效率低下而烦恼？是否想在PostgreSQL中轻松实现向量相似度搜索？本文将带你一步步掌握如何使用Python操作pgvector扩展，从环境搭建到高级查询，让你彻底摆脱向量数据管理的困境。读完本文，你将能够：搭建完整的pgvector环境、创建和管理向量数据表、执行高效的相似度查询、优化查询性能，以及处理实际应用中的常见问题。

什么是pgvector？

pgvector是一个开源的PostgreSQL扩展，专为向量相似度搜索设计。它允许你将向量数据与其他数据一起存储在PostgreSQL中，并支持精确和近似最近邻搜索。pgvector支持多种向量类型，包括单精度、半精度、二进制和稀疏向量，以及多种距离度量方式，如L2距离、内积、余弦距离等。

pgvector的核心优势在于与PostgreSQL的无缝集成，你可以利用PostgreSQL的强大功能来管理和查询向量数据，而无需使用专门的向量数据库。这不仅简化了架构，还能充分利用PostgreSQL的事务支持、索引功能和安全特性。

环境准备

安装pgvector

在开始之前，我们需要先安装pgvector扩展。以下是在Linux和Mac系统上的安装步骤：

cd /tmp
git clone --branch v0.8.1 https://gitcode.com/GitHub_Trending/pg/pgvector.git
cd pgvector
make
make install # 可能需要sudo权限

对于Windows用户，可以使用以下命令：

set "PGROOT=C:\Program Files\PostgreSQL\18"
cd %TEMP%
git clone --branch v0.8.1 https://gitcode.com/GitHub_Trending/pg/pgvector.git
cd pgvector
nmake /F Makefile.win
nmake /F Makefile.win install

更多安装选项，如Docker、Homebrew、PGXN等，可以参考官方文档。

安装Python依赖

接下来，我们需要安装Python的PostgreSQL客户端和pgvector的Python库：

pip install psycopg2-binary pgvector

快速入门

启用pgvector扩展

首先，我们需要在PostgreSQL数据库中启用pgvector扩展。这一步只需在每个数据库中执行一次：

import psycopg2
from pgvector.psycopg2 import register_vector

# 连接到PostgreSQL数据库
conn = psycopg2.connect(
    dbname="your_db",
    user="your_user",
    password="your_password",
    host="your_host",
    port="your_port"
)
conn.autocommit = True

# 创建游标
cur = conn.cursor()

# 启用pgvector扩展
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")

# 注册向量类型
register_vector(conn)

创建向量数据表

接下来，我们创建一个包含向量列的数据表。例如，创建一个存储3维向量的表：

# 创建数据表
cur.execute("CREATE TABLE IF NOT EXISTS items (id bigserial PRIMARY KEY, embedding vector(3))")

这里，vector(3)表示我们将存储3维向量。你可以根据需要调整维度，pgvector支持最高2000维的向量。

插入向量数据

现在，我们可以向表中插入向量数据了：

# 插入向量数据
cur.execute("INSERT INTO items (embedding) VALUES (%s), (%s)", (
    [1, 2, 3],
    [4, 5, 6]
))

执行相似度查询

pgvector支持多种相似度查询方式。以下是一些常用的示例：

L2距离查询

L2距离（欧氏距离）是最常用的相似度度量之一：

# 查询与[3, 1, 2]最相似的5个向量
cur.execute("SELECT id, embedding <-> %s AS distance FROM items ORDER BY distance LIMIT 5", ([3, 1, 2],))
results = cur.fetchall()

for row in results:
    print(f"id: {row[0]}, distance: {row[1]}")

内积查询

内积也是常用的相似度度量，特别是当向量已经归一化时：

# 查询与[3, 1, 2]内积最大的5个向量
# 注意：pgvector使用<#>操作符返回负内积，因此需要取反
cur.execute("SELECT id, (embedding <#> %s) * -1 AS inner_product FROM items ORDER BY embedding <#> %s LIMIT 5", ([3, 1, 2], [3, 1, 2]))
results = cur.fetchall()

for row in results:
    print(f"id: {row[0]}, inner_product: {row[1]}")

余弦相似度查询

余弦相似度是衡量两个向量方向相似度的常用指标：

# 查询与[3, 1, 2]余弦相似度最高的5个向量
# 余弦相似度 = 1 - 余弦距离
cur.execute("SELECT id, 1 - (embedding <=> %s) AS cosine_similarity FROM items ORDER BY embedding <=> %s LIMIT 5", ([3, 1, 2], [3, 1, 2]))
results = cur.fetchall()

for row in results:
    print(f"id: {row[0]}, cosine_similarity: {row[1]}")

高级功能

创建索引提升查询性能

对于大规模向量数据，创建索引可以显著提高查询性能。pgvector支持两种主要的索引类型：HNSW和IVFFlat。

HNSW索引

HNSW（Hierarchical Navigable Small World）是一种近似最近邻搜索算法，具有较好的查询性能，但构建时间较长，内存占用较大。

# 创建L2距离的HNSW索引
cur.execute("CREATE INDEX IF NOT EXISTS items_embedding_idx ON items USING hnsw (embedding vector_l2_ops)")

你还可以为其他距离类型创建HNSW索引：

# 内积
cur.execute("CREATE INDEX IF NOT EXISTS items_embedding_ip_idx ON items USING hnsw (embedding vector_ip_ops)")

# 余弦距离
cur.execute("CREATE INDEX IF NOT EXISTS items_embedding_cosine_idx ON items USING hnsw (embedding vector_cosine_ops)")

IVFFlat索引

IVFFlat（Inverted File with Flat Compression）是另一种近似最近邻搜索算法，构建时间较短，内存占用较小，但查询性能略低于HNSW。

# 创建L2距离的IVFFlat索引
cur.execute("CREATE INDEX IF NOT EXISTS items_embedding_ivf_idx ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")

这里的lists参数指定了倒排表的数量，一般建议设置为数据量的平方根。

向量聚合操作

pgvector支持向量的聚合操作，如计算平均向量：

# 计算所有向量的平均值
cur.execute("SELECT AVG(embedding) FROM items")
avg_embedding = cur.fetchone()[0]
print(f"Average embedding: {avg_embedding}")

# 按类别计算平均向量
cur.execute("SELECT category_id, AVG(embedding) FROM items GROUP BY category_id")
category_avgs = cur.fetchall()
for category_id, avg_emb in category_avgs:
    print(f"Category {category_id} average embedding: {avg_emb}")

向量过滤查询

在实际应用中，我们经常需要结合过滤条件进行向量查询。pgvector支持多种过滤方式：

# 结合WHERE子句的向量查询
cur.execute("SELECT * FROM items WHERE category_id = 123 ORDER BY embedding <-> %s LIMIT 5", ([3, 1, 2],))
results = cur.fetchall()

# 使用CTE优化带距离过滤的查询
cur.execute("""
WITH nearest_results AS MATERIALIZED (
    SELECT id, embedding <-> %s AS distance FROM items ORDER BY distance LIMIT 5
) SELECT * FROM nearest_results WHERE distance < 5 ORDER BY distance
""", ([3, 1, 2],))
results = cur.fetchall()

实战案例：图像相似性搜索

让我们通过一个实际案例来展示pgvector的强大功能。假设我们要实现一个图像相似性搜索系统，步骤如下：

1. 提取图像特征向量
2. 将向量存储到PostgreSQL中
3. 查询相似图像

import psycopg2
from pgvector.psycopg2 import register_vector
from PIL import Image
import numpy as np

# 连接数据库
conn = psycopg2.connect(dbname="images_db", user="postgres", password="password")
register_vector(conn)
cur = conn.cursor()

# 创建图像表
cur.execute("CREATE TABLE IF NOT EXISTS images (id bigserial PRIMARY KEY, name text, embedding vector(2048))")

# 假设我们有一个函数可以从图像中提取特征向量
def extract_features(image_path):
    # 这里使用假数据模拟，实际应用中可以使用预训练的CNN模型
    return np.random.rand(2048).tolist()

# 插入图像特征
image_paths = ["image1.jpg", "image2.jpg", "image3.jpg"]
for path in image_paths:
    features = extract_features(path)
    cur.execute("INSERT INTO images (name, embedding) VALUES (%s, %s)", (path, features))

# 查询相似图像
query_image = "query_image.jpg"
query_features = extract_features(query_image)
cur.execute("SELECT name, embedding <-> %s AS distance FROM images ORDER BY distance LIMIT 5", (query_features,))
similar_images = cur.fetchall()

print("Similar images to", query_image)
for name, distance in similar_images:
    print(f"{name}: distance {distance}")

conn.commit()
cur.close()
conn.close()

性能优化技巧

1. 使用COPY命令批量加载数据，效率远高于逐条插入：

   cur.copy_expert("COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)", binary_file)
   

2. 对于大规模数据，创建索引时可以增加并行工作进程：

   SET max_parallel_maintenance_workers = 7;
   

3. 使用EXPLAIN ANALYZE分析查询性能：

   EXPLAIN ANALYZE SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;
   

4. 对于近似搜索，调整索引参数平衡速度和召回率：

   -- HNSW
   SET hnsw.ef_search = 100;
   
   -- IVFFlat
   SET ivfflat.probes = 10;
   

常见问题解答

Q: pgvector支持多少维度的向量？

A: pgvector支持多种向量类型，vector类型最多支持2000维，halfvec类型支持4000维，bit类型支持64000维。

Q: 如何在已有的表中添加向量列？

A: 可以使用ALTER TABLE命令添加向量列：

ALTER TABLE existing_table ADD COLUMN embedding vector(3);

Q: pgvector是否支持向量的更新和删除？

A: 是的，pgvector完全支持向量数据的更新和删除操作：

UPDATE items SET embedding = '[1,2,3]' WHERE id = 1;
DELETE FROM items WHERE id = 1;

Q: 如何监控pgvector的性能？

A: 可以使用PostgreSQL的pg_stat_statements扩展来监控查询性能：

CREATE EXTENSION pg_stat_statements;
SELECT query, calls, total_time FROM pg_stat_statements ORDER BY total_time DESC LIMIT 10;

总结

通过本文的介绍，你已经掌握了使用Python操作pgvector的基本方法和高级技巧。从环境搭建到性能优化，pgvector提供了一套完整的解决方案，让你能够在PostgreSQL中轻松实现高效的向量相似度搜索。无论是图像识别、自然语言处理，还是推荐系统，pgvector都能为你的应用提供强大的向量数据支持。

现在，是时候将这些知识应用到你的项目中了。开始尝试使用pgvector，体验向量数据库的强大功能吧！如果你有任何问题或建议，欢迎在评论区留言讨论。别忘了点赞、收藏本文，关注我们获取更多pgvector和PostgreSQL的实用教程！

下一篇，我们将深入探讨pgvector在大规模数据集上的性能优化策略，敬请期待！

【免费下载链接】pgvector Open-source vector similarity search for Postgres 项目地址: https://gitcode.com/GitHub_Trending/pg/pgvector