在Postgres中实现高效向量搜索：使用PGEmbedding与HNSW

原创于 2025-07-02 16:13:03 发布 · 266 阅读

CC 4.0 BY-SA版权

文章标签：

Postgres Embedding是一种开源解决方案，利用Hierarchical Navigable Small Worlds (HNSW)算法为Postgres数据库提供向量相似性搜索支持。该解决方案能够执行精准和近似的邻居搜索，特别是基于L2距离的搜索能力。今天，我们将深入探讨如何在Postgres中使用PGEmbedding进行向量搜索，并提供可运行的代码示例。

技术背景介绍

在处理复杂数据集时，向量搜索是一种关键的技术手段。传统数据库通常不擅长处理向量数据，而PGEmbedding通过引入HNSW算法使得Postgres能够高效处理向量相似性搜索。HNSW是一种基于图的搜索算法，适用于大规模的向量数据集。

核心原理解析

PGEmbedding利用Postgres的扩展功能，引入HNSW近似最近邻搜索算法。此算法通过构建一个导航小世界图，从而在大规模数据集中快速找到与查询向量相似的向量。

代码实现演示

下面，我们将演示如何使用PGEmbedding进行向量搜索。首先，我们需要安装必要的Python包，并将OpenAI API密钥和数据库URL保存到环境变量中：

# 安装所需的Python包
%pip install --upgrade --quiet langchain-openai langchain-community psycopg2-binary tiktoken

# 加载环境变量
import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
os.environ["DATABASE_URL"] = getpass.getpass("Database Url:")

# 导入所需模块
from typing import List, Tuple
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import PGEmbedding
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 加载文档并进行文本拆分
loader = TextLoader("state_of_the_union.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# 创建OpenAI嵌入对象
embeddings = OpenAIEmbeddings()
connection_string = os.environ.get("DATABASE_URL")
collection_name = "state_of_the_union"

# 使用从文档中创建PGEmbedding实例
db = PGEmbedding.from_documents(
    embedding=embeddings,
    documents=docs,
    collection_name=collection_name,
    connection_string=connection_string,
)

# 执行相似性搜索
query = "What did the president say about Ketanji Brown Jackson"
docs_with_score: List[Tuple[Document, float]] = db.similarity_search_with_score(query)

# 输出搜索结果
for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)

应用场景分析

PGEmbedding适用于需要高效执行向量相似性搜索的应用场景，例如推荐系统、图像检索和自然语言处理。通过创建HNSW索引，可以显著提升搜索速度，对于大规模数据集尤为重要。

实践建议

配置HNSW索引：在大规模数据集上创建HNSW索引可以大幅提高搜索效率。
调整参数：根据数据集规模和维度，合理设置maxelements、dims等参数。
环境整合：确保API密钥和数据库URL正确配置，避免连接问题。

存储与检索

创建HNSW索引以便更快地检索数据，示例如下：

PGEmbedding.create_hnsw_index(
    max_elements=10000, dims=1536, m=8, ef_construction=16, ef_search=16
)

如果遇到问题欢迎在评论区交流。

—END—