Python RAG 实战手册——向量数据库与相似性搜索

向量数据库的受欢迎程度迅速增长，尤其是在 LLMs 兴起之后。图 6-1 展示了一些最流行的向量数据库，以及它们的发布时间。

FAISS 于 2017 年发布，随后 Milvus 和 Weaviate 于 2019 年发布，Vald 于 2020 年发布，Pinecone 于 2021 年发布，Chroma 于 2023 年发布。与此同时，PostgreSQL 和 Elasticsearch 等传统数据库，也已经在其现有平台中加入了 vector search 功能。这意味着，如果你现有的 SQL 或 NoSQL 数据库已经支持 vector search，你不一定需要在技术栈中额外加入一个专用向量数据库。

图 6-1：Vector stores 的演进

本章介绍支持 vector operations 的流行 libraries 和 databases，包括 FAISS、Chroma 和 PostgreSQL。你将学习如何根据需求选择合适的 vector store、实现 similarity searches，并通过 indexing techniques 优化性能。

NOTE

本章中，你会同时看到 “similarity search” 和 “semantic search” 两个说法。Similarity search 指的是寻找 vector space 中彼此接近 vectors 的技术操作。Semantic search 描述的是面向用户的能力，也就是按含义而不是精确关键词进行搜索。Semantic search 是以 similarity search 作为底层机制实现的。

你可以在本书 GitHub repository 中找到本章所有代码示例。

6.1 选择合适的向量数据库

Problem

你需要选择一个适合 RAG workload、团队和基础设施的向量数据库，同时避免过度工程化或被锁定在错误技术栈中。

Solution

使用以下决策路径，将选项缩小到一到两个现实选择。按照这些步骤评估并选择 RAG 应用的向量数据库：

Step 1：定义需求

Vector search 是只在 preprocessing、experiments 或 offline pipelines 中使用？还是 live application 中服务用户的一部分？

如果 vector search 不是 live system 的一部分，可以使用 vector library 或 in-file vector store，例如：

• FAISS 或 Annoy，用于 pipelines 和 experiments
• Chroma，用于本地 RAG prototypes

如果 vector search 是 live application 的一部分，继续 Step 2。

Step 2：考虑运维需求

Embeddings 是否需要与 business data、users、documents 或 permissions 放在一起？

如果是，使用你已经在运维的数据库：

• PostgreSQL with pgvector
• MongoDB with vector search
• Redis 或 Cassandra with vector support

如果不是，继续 Step 3。

Step 3：决定所需规模和运维能力

系统是否会增长到数百万或数千万 vectors？是否需要 filtering、access control 和 predictable latency？

如果是，使用 purpose-built 或 managed vector database：

• Weaviate、Qdrant、Milvus，用于 self-hosted systems
• Pinecone，用于 fully managed systems

如果不是，Chroma 这样的 in-file store 可能仍然足够。

Step 4：根据实践约束验证

问自己表 6-1 中的问题。

表 6-1：向量数据库选择的实践问题

Question	If yes, lean toward
是否允许将数据发送到第三方 managed cloud service？	Pinecone 或 cloud-managed PostgreSQL
是否已经在运维 PostgreSQL 或 MongoDB？	pgvector 或 MongoDB vector search
是否需要 keyword search 与 vectors 结合？	Elasticsearch、OpenSearch 或 Weaviate
是否需要小团队也能最小化设置成本？	Chroma 或 managed Pinecone

到这一步，你应该已经得到一到两个现实候选项。表 6-2 提供了比较，用于作出最终决定。

表 6-2：常见 vector stores 的优缺点

Database	Type	Pros	Cons
FAISS	Vector library	Similarity search 非常快支持 GPU acceleration非常适合 pipelines	无 persistence无 filtering 或 access control意味着其他能力都需要自己构建
Chroma	Embedded vector store	即时设置非常适合 notebooks 和 prototypesPython-first API	非 production gradeScaling 和 durability 有限
PostgreSQL（pgvector）	Relational database with vectors	Vectors 与 business data 放在一起使用现有 backups 和 security易于 joins 和 filters	比专用引擎慢大规模时需要 tuning
Weaviate	Vector database	为大规模 RAG 构建支持 hybrid search丰富 filtering	新的 operational stack复杂度更高
Pinecone	Managed vector database	无需管理 infrastructure强 performance guarantees易于扩展	仅 cloudVendor lock-in成本随用量增长

Discussion

向量数据库存在的原因，是传统数据库并不针对 high-dimensional vectors 上的 nearest neighbor search 做优化。Similarity search 的工作方式，是计算 high-dimensional space 中 vectors 之间的距离。现代 embedding models 会产生数百或数千维 vectors，这使 brute-force search 太慢。彼此接近的 vectors 表示语义相似内容，而距离较远的 vectors 表示不同概念。

RAG 应用中对向量数据库的主要需求，是在数千或数百万 vectors 上进行快速 similarity search。Store 还必须支持 high-dimensional vectors，因为现代 embedding models 通常会生成数千维 vectors。

很多团队犯的错误，是根据 benchmarks 选择，而不是根据 workflow 选择。一个只有几千个 chunks 的 RAG prototype 不需要 distributed vector database；一个服务大量用户的 production assistant 则需要。

图 6-2 展示了生态系统如何聚类为 libraries、embedded stores、带 vector extensions 的 databases，以及 purpose-built vector engines。

图 6-2：按类型和 use case 组织的 vector stores

Pinecone、Weaviate 和 Qdrant 等 purpose-built vector databases 针对 vector search 优化，同时也提供 persistence 和 user management 等核心数据库能力。

Elasticsearch 和 OpenSearch 等面向文本搜索的数据库，在 RAG 应用中也很流行。一个搜索 semantic text chunks 的 retriever，本质上像是一个专门 search engine，因此当 keyword search 已经扮演重要角色时，这些系统自然适合。

PostgreSQL、MongoDB、Redis 和 Cassandra 等 SQL 与 NoSQL 系统，通过 extensions 或 native features 支持 vector search。这使你可以将 application data、metadata 和 embeddings 保存在同一个地方，并使用熟悉的 query languages 结合 filtering 与 semantic search。

带 vector support 的 SQL databases，或 Pinecone 和 Milvus 这样的 dedicated vector databases，最适合 production systems。它们支持 hybrid search、metadata filtering，以及大型 RAG 和 agentic applications 中常见的 workflows。对于较小项目和 prototypes，vector libraries 或 Chroma 这样的 lightweight vector stores 设置成本最低，并且在 standalone vector operations 上表现优秀。

主要取舍是 simplicity 与 capability。Vector libraries 设置简单，但缺乏 persistence 和 access control。Purpose-built vector databases 提供全面功能，但需要 infrastructure management 或 vendor costs。带 vector extensions 的 SQL databases，如果你已经使用它们，则提供最佳集成；但在超大规模下，原始性能可能不如专用方案。

如果你从简单 vector library 开始，之后需要完整 database，迁移通常很直接。你必须重建 index，但不需要重新生成 embeddings，而 embedding 才是成本最高的部分。

See Also

ANN Benchmarks repository 比较 vector search engines 的性能。

Chroma documentation 展示了如何运行 in-file vector store。

FAISS repository 解释了 vector libraries 如何实现 similarity search。

6.2 使用 FAISS 存储和搜索 Embeddings

Problem

你需要一个轻量方案来存储 embeddings 并执行 similarity searches。

Solution

FAISS 非常适合那些不需要 persistence 或 metadata filtering 的 RAG workflows。图 6-3 展示了 FAISS 如何融入用于临时分析任务的 RAG workflow。

图 6-3：在 RAG workflow 中使用 FAISS

这个 recipe 使用 CPU 版本的 FAISS，因此几乎可以在任何笔记本上运行。安装用于 CPU 执行的 FAISS library，以及用于生成和处理 embeddings 的 Python libraries NumPy 和 OpenAI：

pip install faiss-cpu openai numpy

接下来，你将生成一个 FAISS index，并用示例文本 embeddings 填充它。Index 是 FAISS 用来存储 vectors 并执行搜索的数据结构。

你还会使用 OpenAI embedding models 生成 vectors，然后将其加入 FAISS index。示例使用 IndexFlatL2，它通过 L2 distance 衡量 similarity。L2 distance 和 cosine distance 是最常用的 similarity metrics。代码如下：

import faiss
import numpy as np
from openai import OpenAI

# Example list of sample strings
text_chunks = [
    "The sky is blue.",
    "The sun is shining.",
    "I love chocolate.",
    "Ice cream is delicious.",
    "Roses are red.",
    "Violets are blue.",
]

# Initialize the OpenAI embeddings model
client = OpenAI()
model = "text-embedding-3-small"

# Generate embeddings for the sample strings
def get_embedding(text):
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

document_embeddings = np.array(
    [get_embedding(text) for text in text_chunks]
)

# Convert embeddings to float32 (FAISS requires float32 type)
document_embeddings = document_embeddings.astype("float32")

# Create a FAISS index (using L2 distance)
index = faiss.IndexFlatL2(document_embeddings.shape[1])

# Add embeddings to the index
index.add(document_embeddings)

下一步是使用刚创建的 index。当用户提出问题时，你会为 query 生成 embedding，然后通过计算 L2 distance 搜索 index。

下面代码用一个示例 query 演示，问题是紫罗兰是什么颜色：

# Generate a query embedding for the user query
query = "What color are violets?"
query_embedding = np.array(get_embedding(query)).astype("float32")

# Perform the search: k = number of closest documents you want to retrieve
k = 5
distances, indices = index.search(query_embedding.reshape(1, -1), k)

# Retrieve the documents corresponding to the indices
retrieved_documents = [text_chunks[i] for i in indices[0]]

图 6-4 展示了搜索结果，包括 distance、chunk location 和 retrieved text chunks。说明紫罗兰是蓝色的句子具有最小 distance，这意味着它具有最高 semantic similarity。

图 6-4：FAISS query result

Discussion

FAISS 是一个用于 dense vectors 快速 similarity search 的 library。它由 Meta 于 2017 年发布，至今仍被广泛使用，因为它轻量、支持 CPU 和 GPU，并且在 in-memory vector operations 上性能优秀。

FAISS 的工作方式，是将 vectors 组织进高效 index structures，以支持 fast nearest neighbor search。L2（Euclidean distance）和 cosine similarity metrics 用于衡量 vectors 在 high-dimensional space 中有多接近。L2 衡量两个点之间的直线距离，而 cosine similarity 衡量 vectors 之间的角度。两者都能有效识别语义相似内容。

当你的 workflow 不需要跨运行 persistence、metadata filtering 或 user-specific access controls 时，使用 FAISS。FAISS 擅长临时、一次性分析任务：你构建 index，查询它，然后丢弃它。例如，分析合同以提取 parties 和 dates，只在该处理任务中需要 vector search。类似地，对语料库处理一次并保存结果的 research scripts，也非常适合 FAISS。

当你需要数据在 application restarts 后仍然存在、多个用户需要不同 access permissions，或必须在搜索前按 metadata fields 过滤时，不要使用 FAISS。FAISS 将所有内容存储在内存或简单文件格式中，没有内置 query language 或 access control。

FAISS 的主要取舍是 simplicity versus features。FAISS 用极少代码提供最快的纯 vector search performance，但不提供数据库能力。你获得速度和简单性，但必须在 application layer 自己处理 persistence、access control 和 metadata filtering。

FAISS 是一个 in-memory vector index，可以序列化到磁盘，但它不提供数据库级 persistence、metadata 或 concurrency control。Chroma 增加了轻量 persistence 和基本 metadata filtering。PostgreSQL with pgvector 则增加完整 SQL capabilities、transactions 和 production-grade access control，但需要数据库管理。Batch jobs 和 offline indexing 选择 FAISS；需要简单 persistence 的 prototypes 选择 Chroma；production systems 选择 PostgreSQL。

对于需要 persistence、metadata filtering 或 multiuser access 的生产系统，可以使用 PostgreSQL with pgvector，如果你已经使用 PostgreSQL；或者使用 Weaviate、Pinecone 等 dedicated vector databases。这些方案支持 role-based access control，并且可作为云平台上的 managed services 使用。

See Also

FAISS repository 包含官方文档和使用示例。

FAISS wiki 提供了 index types 和 performance tuning 的详细指南。

原始 FAISS paper 解释了该 library 背后的算法和设计原则。

6.3 在 Chroma Vector Database 中存储和处理 Embeddings

Problem

你需要一个用于较小项目的轻量级向量数据库，并且可以在几分钟内设置完成。

Solution

Chroma 是一个为快速 prototyping 和小型项目设计的轻量级向量数据库。它提供 persistence 和 metadata handling，同时保持设置和使用简单。

Chroma 可以用两种方式存储数据：

Ephemeral client

将数据存储在内存中，这意味着应用停止时数据会丢失。

Persistent client

将数据保存到磁盘，因此重启 app 后可以重新加载。

Chroma 内置函数可以自动处理 text processing 和 embedding generation。对于 prototypes，这些内置功能能加速开发。对于计划未来迁移的 production systems，应在自己的代码中生成 embeddings，只用 Chroma 存储和搜索它们。这样切换到 PostgreSQL 或 Pinecone 会更直接。

图 6-5 中可以看到两个主要步骤：

1. 加载文本、chunk 它，并使用 embedding model 生成 embeddings。
2. 将这些 embeddings 插入 Chroma。如果启用 persistence，Chroma 也会将它们存储到磁盘。

当用户提出新问题时，你会用同一个 embedding model 为问题生成 embedding，然后在 Chroma 中执行 vector search。

图 6-5：从 Chroma 存储和检索 vector embeddings

要使用 Chroma，只需要 Chroma Python package。因为这个示例使用 OpenAI embedding model，所以还需要 OpenAI SDK。使用以下命令安装两者：

pip install chromadb openai

开始使用 Chroma 时，你需要创建一个 Chroma client 和一个新 collection。因为这个示例使用 persistent client，所以必须提供一个 directory path，让 Chroma 存储其文件：

import chromadb

# Vector store settings
VECTOR_STORE_PATH = r"../02_Data/00_Vector_Store"
COLLECTION_NAME = "my_collection"

# Get/create a Chroma client and collection
chroma_client = chromadb.PersistentClient(path=VECTOR_STORE_PATH)
collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME)

下一步，将 text chunks 添加到 collection。Chroma 可以替你处理 tokenization、embedding 和 indexing，并默认使用 Sentence Transformer 模型 all-MiniLM-L6-v2 创建 embeddings。不过，Chroma 也允许你提供在其他地方生成的 embeddings。

在这个示例中，embeddings 是在你的 application code 中用 OpenAI text embedding model 生成的，而 Chroma 只负责存储和搜索。使用 OpenAI embedding model，为你想插入 Chroma 的 text chunks 创建 embeddings 列表，然后通过 add 函数将该列表传给 vector store：

from openai import OpenAI

text_chunks = [
    "The sky is blue.",
    "The sun is shining.",
    "I love chocolate.",
    "Ice cream is delicious.",
    "Roses are red.",
    "Violets are blue.",
]

# Generate embeddings
client = OpenAI()
model = "text-embedding-3-small"
def get_embedding(text):
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

embeddings_model = client  # For compatibility with later code
embeddings_list = [get_embedding(text) for text in text_chunks]

# Add data frame to collection
collection.add(
    embeddings=embeddings_list,
    documents=text_chunks,
    ids=[
        str(i) for i in range(len(text_chunks))
    ],  # Create a list of strings as index
)

现在可以通过查询 collection 测试设置。在这个例子中，你会询问天空的颜色。该问题会被 embedding，并用于查询 collection：

# Query collection
query = "What is the color of the sky?"
query_embedding = get_embedding(query)
results = collection.query(query_embeddings=[query_embedding], n_results=3)

查询返回三个最接近的 text chunks，以及它们的 cosine distance。图 6-6 展示了结果字典。它不仅包含最佳匹配的 text chunks 和 IDs，也包含 Chroma 中存储的任何 metadata。

图 6-6：在 Chroma 中处理 embeddings

Discussion

Chroma 的工作方式，是在一个简单 embedded database 中存储 vectors 及其 metadata。它会自动将数据持久化到磁盘，因此你可以重启应用而不必重建 index。这使 Chroma 适合 prototypes 和小型项目，这些项目需要比 FAISS 更多能力，但又不需要完整生产数据库。

当你需要 persistence 和 basic metadata filtering，但又希望最小设置成本时，使用 Chroma。Chroma 非常适合 prototypes、proof-of-concept projects，以及数据集少于几十万 vectors 的小型生产系统。当你想快速测试 RAG application，而不想配置 PostgreSQL 或注册 managed service 时，它很合适。

当数据集增长到超过几十万 vectors、需要 role-based access control、需要与 application data 做 SQL joins，或需要在规模化时 query latency 低于 10ms 时，不要使用 Chroma。Chroma 缺乏高级安全功能、复杂 query optimization 和 horizontal scaling 能力。

Chroma 的主要取舍是 ease of use versus scalability。你可以用最少 dependencies 立即开始写代码，但随着数据集增长，性能会下降。Chroma 位于 FAISS 和 PostgreSQL 之间：FAISS 可以保存到磁盘但缺少数据库功能，PostgreSQL 则拥有完整数据库功能。

Chroma 与 FAISS 的不同在于，它增加了 persistence 和 metadata filtering，这意味着你可以重启 app 而不用重建 index。Chroma 与 PostgreSQL with pgvector 的不同在于，它设置更简单，但缺少 SQL capabilities、transaction support 和 production-grade access control。当数据集超过几十万 vectors，或需要 role-based access 和 SQL joins 时，应迁移到 PostgreSQL with pgvector，或 Pinecone 这样的 managed vector database。

See Also

Chroma documentation 提供全面的设置和使用指南。

Chroma repository 包含代码示例和实现细节。

Chroma cookbook 提供常见 use cases 的实践示例。

6.4 使用 pgvector Extension 在 PostgreSQL 中存储 Embeddings

Problem

你想使用 PostgreSQL with pgvector extension，将 application data 和 vector data 存储在同一个数据库中。

Solution

PostgreSQL with pgvector extension 允许你在同一个数据库中存储 application data 和 high-dimensional vectors。Vector search 会成为另一个 SQL operator，因此你可以用 similarity search capabilities 扩展已有 tables 和 queries。

如果还没有安装 PostgreSQL，你可以使用 Docker 启动一个安装了 PostgreSQL 和 pgvector extension 的 container，也可以直接在机器上安装 PostgreSQL。

NOTE

Docker 是一个将 applications 及其 dependencies 打包进 containers 的平台。Containers 是轻量、可移植的单元，可以在不同环境中一致运行。Docker container 包含运行应用所需的一切，因此便于分享和部署，无需担心系统特定配置。

要直接安装 PostgreSQL，请访问 PostgreSQL 网站并下载适用于你操作系统的最新 installer。此外，你还需要安装 pgvector extension。安装步骤取决于操作系统：

Windows

Clone pgvector repository。README 包含正确安装说明。

macOS

使用 Homebrew，通过一条命令安装 pgvector。

Docker 提供更简单且更可复现的方法。使用 PostgreSQL 作为 vector store 的 RAG application，可以在带固定配置的 container 中运行，使同一套设置部署到 AWS、Google Cloud 或 Azure 时也很直接。

要在本地运行 Docker containers，请从 Docker 网站下载并安装 Docker Desktop。Docker 安装完成后，创建一个 Docker Compose 文件，定义你想运行的 container 配置，并让 Docker 启动 container。

现在我们走一遍该流程。在项目目录中创建新的 docker-compose.yml 文件，加入下面内容并保存：

version: '3.8'
services:
  db:
    image: ankane/pgvector
    container_name: postgres_with_pgvector
    restart: always
    environment:
      POSTGRES_USER: rag_cookbook_user
      POSTGRES_PASSWORD: rag_cookbook_user_pw
      POSTGRES_DB: rag_cookbook
    ports:
      - "5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

要启动 container，请进入 docker-compose.yml 文件所在目录并运行：

docker-compose up -d

这个命令会在当前目录查找 Docker Compose 文件，并下载 container image ankane/pgvector，其中包含 PostgreSQL 和 pgvector。Docker Compose 文件设置了 PostgreSQL user、password 和 database name，并暴露用于连接 PostgreSQL database 的端口。

要验证 container 正常运行，可以打开 Docker Dashboard 并检查 container 状态。图 6-7 展示了 PostgreSQL container 运行时的 Docker Dashboard。

图 6-7：检查 Docker container 是否正在运行

要从 Python 连接新创建的 PostgreSQL database，这个 recipe 使用 Psycopg 2 package，按如下方式安装：

pip install psycopg2-binary openai

在这个示例中，你会使用定义好的 username、password 和 database name 连接数据库，然后创建一个新 table，其中包含：

• text chunk 的 ID
• text chunk
• embedding vector

为了启用 vector search 并存储 embeddings，embedding column 使用具有正确维度数量的 vector data type。在这个例子中，vector 有 1,536 维，因为该 recipe 使用 OpenAI 的 text-embedding-3-small 模型，该模型生成该维度的 embeddings。请记住，只有安装了 pgvector extension，vector data type 才可用。运行以下代码：

import psycopg2

conn = psycopg2.connect(
    host="localhost",
    port=5432,
    dbname="rag_cookbook",
    user="rag_cookbook_user",
    password="rag_cookbook_user_pw",
)

cur = conn.cursor()

cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")

cur.execute("""
    CREATE TABLE IF NOT EXISTS embeddings (
        id SERIAL PRIMARY KEY,
        chunk TEXT NOT NULL,
        embedding vector(1536) NOT NULL
    );
""")

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

下一步，用示例数据填充 table。你会遍历一组 sample text chunks，通过 OpenAI embedding model 生成 embeddings，然后将 text chunks 和对应 embeddings 插入 table：

from openai import OpenAI

# Define text chunks
text_chunks = [
    "The sky is blue.",
    "The sun is shining.",
    "I love chocolate.",
    "Ice cream is delicious.",
    "Roses are red.",
    "Violets are blue.",
]

client = OpenAI()
model = "text-embedding-3-small"

def get_embedding(text):
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

index = 0
cur = conn.cursor()

# Insert the embeddings into the table
for text_chunk in text_chunks:
    embedding = get_embedding(text_chunk)
    cur.execute(
        """INSERT INTO embeddings
        (id, chunk, embedding)
        VALUES (%s, %s, %s)""",
        (index, text_chunk, embedding),
    )
    index += 1

该函数会用 text chunks 和对应 embeddings 填充 PostgreSQL table。图 6-8 展示了 table 填充结果：六个 text chunks 及其 embeddings。

图 6-8：PostgreSQL 中的 embeddings table

现在你可以像使用其他 SQL table 一样使用这个 table。你可以执行 SELECT、CREATE 和 UPDATE 操作。由于 table 包含 vector column，你也可以使用 cosine similarity、Euclidean distance 或 inner product 作为 distance metric 执行 vector searches。

Discussion

PostgreSQL with pgvector 的工作方式，是将 embeddings 作为原生 vector columns，与 application data 一起存储。pgvector extension 会直接向 SQL 添加 vector data types 和 similarity search operators，因此你可以使用熟悉的 PostgreSQL 语法查询 embeddings。这使你可以在单个查询中 join user tables、document metadata 和 embeddings，而不需要管理多个系统。

当你的应用已经使用 PostgreSQL 存储结构化数据、需要在 similarity search 之前按 metadata 过滤 embeddings，例如 user permissions、dates、categories，或希望避免管理单独 vector database 时，可以使用 PostgreSQL with pgvector。PostgreSQL 很适合拥有数千万 vectors 的 production RAG systems，尤其是在需要 ACID（atomicity、consistency、isolation、durability）transactions 和 role-based access control 时。

当你需要在数亿或数十亿 vectors 上执行个位数毫秒延迟搜索时，不要使用 PostgreSQL with pgvector。在这种规模下，Pinecone 或 Weaviate 等 purpose-built vector databases 提供更好性能。如果你的 workload 是纯 vector search，并且没有 metadata filtering，也应避免 PostgreSQL，因为 FAISS 或专用 vector databases 会更快。

PostgreSQL 的主要取舍是 integration versus raw performance。你获得 unified data management、SQL joins 和 transactional guarantees，但 specialized vector databases 在超大规模纯 vector-only workloads 上可能更快。对于大多数 RAG 应用，集成收益通常超过性能差异。

PostgreSQL with pgvector 与 FAISS 的不同在于，它提供 persistence、ACID transactions 和 SQL capabilities。它与 Chroma 的不同在于，PostgreSQL 提供 production-grade access control、query optimization，以及将 embeddings 与 structured data join 的能力。它与 specialized vector databases 的不同在于，它与你现有 PostgreSQL infrastructure 集成，但在极限规模，例如 billions of vectors，可能性能更低。

See Also

pgvector repository 包含安装说明和配置选项。

PostgreSQL documentation 提供全面的数据库设置和管理指南。

pgvector installation guide 提供平台特定设置说明。

6.5 在 PostgreSQL 中执行 Similarity Search

Problem

你想使用 PostgreSQL 查找与用户 query 最相似的 text chunks。

Solution

这个 recipe 假设你已有一个带 vector column 的 PostgreSQL table。示例使用名为 embeddings 的 table，其中包含 id（integer）、chunk（text）和 embedding（vector）列。

如果你单独跟随这个 recipe，首先建立 database connection，并确保 table 存在。下面代码会创建 connection 和 table structure：

import psycopg2

conn = psycopg2.connect(
    host="localhost",
    port=5432,
    dbname="rag_cookbook",
    user="rag_cookbook_user",
    password="rag_cookbook_user_pw",
)

cur = conn.cursor()
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
cur.execute("""
    CREATE TABLE IF NOT EXISTS embeddings (
        id SERIAL PRIMARY KEY,
        chunk TEXT NOT NULL,
        embedding vector(1536) NOT NULL
    );
""")
conn.commit()
cur.close()

完整设置细节，包括 Docker 配置，请见 Recipe 6.4。

pgvector extension 支持最常见的 similarity metrics，例如 cosine similarity、Euclidean distance 和 inner product。

这个示例使用 Psycopg 2 library 连接 PostgreSQL 并执行 SQL queries。Embeddings 由 OpenAI embedding models 创建。安装这些 dependencies：

pip install psycopg2-binary openai

在这个 recipe 中，你会使用 <=> operator 执行 cosine similarity search。

也可以使用以下方式：

• 使用 <-> operator 计算 Euclidean distance
• 使用 <#> operator 计算 Inner product
• 使用 <+> operator 计算 Taxicab（Manhattan）distance

首先，搜索刚创建的 table，其中包含关于天气、花和食物的随机 text chunks。示例 query 是 “sweet”，因此假设与食物相关的 text chunks 会排名更高：

from openai import OpenAI

# Initialize the OpenAI embeddings model
client = OpenAI()
model = "text-embedding-3-small"

# Example text chunk
text_chunk = "Sweets are delicious."

# Get the embedding
response = client.embeddings.create(input=text_chunk, model=model)
embedded_query = response.data[0].embedding

cur = conn.cursor()
cur.execute(
    f"""
    SELECT 1 - (embedding <=> '{embedded_query}') AS cosine_similarity, *
    FROM embeddings
    ORDER BY 1 - (embedding <=> '{embedded_query}') DESC
    LIMIT 20;
"""
)

results = cur.fetchall()

图 6-9 展示了排序结果，包括 similarity score、text chunk 和 embedding vector 本身。排名前两位的 text chunks 与冰淇淋和巧克力有关。

图 6-9：在 PostgreSQL 中计算 cosine similarity

Discussion

PostgreSQL 中的 similarity search 通过 SQL operators 工作，用于计算 query embedding 与 stored embeddings 之间的 distance metrics，例如 cosine similarity、Euclidean distance、inner product。<=> operator 计算 cosine distance，<-> 计算 Euclidean distance，<#> 计算 inner product。PostgreSQL 执行这些计算时，要么扫描所有 rows，要么使用 HNSW（Hierarchical Navigable Small Worlds）或 IVF（Inverted File Indexing）等 vector indexes 来缩小搜索空间。

PostgreSQL similarity search 的强大之处，在于可以与 SQL filtering 结合。你可以使用 WHERE clauses 在计算 similarity 之前按 metadata 过滤，从而减少搜索空间并提升相关性。例如，可以过滤为当前用户可访问的 documents、某日期之后发布的 documents，或带有特定 categories 的 documents，然后在过滤后的集合中寻找最相似 vectors。

当 RAG 应用需要 user-specific results、document permissions 很重要、你希望限制搜索到 recent content，或 categories 和 tags 可以提升相关性时，可以将 SQL filtering 与 similarity search 一起使用。例如，法律 RAG 系统应该只搜索用户有权限查看的合同。新闻 RAG 系统可能只搜索最近 30 天的文章。

如果所有 documents 对所有用户都同等可访问，且不需要 metadata filtering，不要使用 SQL filtering。在这种情况下，不带 filtering 的纯 vector search 更简单也更快。同时也要避免在 similarity search 的同一个 query 中使用复杂 multitable joins，因为这会拖慢性能。应先在 subquery 或 common table expression（CTE）中过滤，再运行 similarity search。

主要取舍是 flexibility versus complexity。SQL filtering 支持复杂 query logic，但需要谨慎构建查询以保持性能。Filter columns 索引不当或 joins 过度复杂，会让 queries 变慢。部署生产前，应使用真实数据规模测试 query performance。

这种方法与 pure vector search，例如 FAISS、standalone Pinecone，不同之处在于，它可以在数据库内部启用 metadata filtering。它与 hybrid search，即下一 recipe 使用的方法，也不同，因为它关注 metadata filtering，而不是 keyword matching。当你需要按 metadata 限制搜索空间时，使用简单 SQL filtering。当 exact keyword matches 也很重要时，使用 hybrid search。

See Also

pgvector repository 记录了 similarity metrics 和 query operators。

PostgreSQL indexing documentation 解释了 index types 和 optimization strategies。

pgvector query guide 提供了不同 distance functions 的示例。

6.6 使用 Indexing Techniques 加速 PostgreSQL 中的 Vector Searches

Problem

当你处理包含数百万 vectors 的数据集时，需要加速 vector searches。

Solution

使用 indexing techniques 组织 vectors，以加快 similarity searches。PostgreSQL with pgvector extension 支持两种 indexing methods：IVFFlat 和 HNSW。这个 recipe 使用一个约 200,000 text chunks 的数据集演示两种方法。

NOTE

这个 recipe 使用较大数据集，以便进行有意义的性能比较。你可以使用 Kaggle 上的示例数据集，或替换为体量相近的自有数据。关键要求是需要足够 rows，100,000+，才能看到 indexing benefits。

实践代码使用 Psycopg 2 连接 PostgreSQL、OpenAI 生成 embeddings，并使用 pandas 做数据预处理。安装三者如下：

pip install psycopg2 openai pandas

首先，创建一个 table 来保存 job descriptions 及其 embeddings。下面代码会创建必要 table structure 并建立 database connection：

import psycopg2
from psycopg2 import Error

conn = psycopg2.connect(
    dbname="rag_cookbook",
    user="rag_cookbook_user",
    password="rag_cookbook_user_pw",
    host="localhost",
    port="5432",
)

cur = conn.cursor()
cur.execute("""CREATE EXTENSION IF NOT EXISTS vector""")

cur.execute(
    """
    CREATE TABLE IF NOT EXISTS job_description_table(
        id integer PRIMARY KEY,
        text_chunk TEXT,
        embedding vector(1536)
    )
    """
)
conn.commit()

这个 recipe 会比较 full-search performance 与使用 HNSW index 和 IVFFlat index 的 indexed vector search。

下面示例从使用 cosine similarity 的 full search 开始。该方法不使用 index，会 sequentially scan 所有 entries，因此在大型数据集上最慢。对于小型数据集，full search 有时可能更快。运行如下代码：

query = "I am looking for a job as a data scientist in Berlin."
client = OpenAI()
model = "text-embedding-3-small"
response = client.embeddings.create(input=query, model=model)
query_embedding = response.data[0].embedding

full_search_sql = f"""
    EXPLAIN ANALYZE SELECT 1 - (embedding <=> '{str(query_embedding)}')
    AS cosine_similarity,
    * FROM job_description_table
    ORDER BY 1 - (embedding <=> '{str(query_embedding)}') DESC
    LIMIT 20;
    """

cur.execute(full_search_sql)

results_full_search = cur.fetchall()
# [..., ('Planning Time: 13.669 ms',), ('Execution Time: 85.582 ms',)]

这个 baseline 将用于与 indexed searches 做性能对比。

NOTE

为了确保 index 实际被使用，你可能需要强制 PostgreSQL 禁用 sequential scans。当 PostgreSQL 无法确定 index search 更优时，有时会回退到 sequential scan。

接下来，创建 IVFFlat index 并在其上执行 similarity search。首先复制已有 job description table，并在新 table 上创建 index。为了分析 query plan 和 execution time，以 EXPLAIN ANALYZE 模式运行 query：

import psycopg2
from psycopg2 import Error

ivfflat_sql = f"""
    DROP TABLE IF EXISTS test_embedding_table;
    CREATE TABLE test_embedding_table AS
        SELECT * FROM job_description_table;
    CREATE INDEX ON test_embedding_table
        USING ivfflat (embedding vector_cosine_ops)
        WITH (lists = 30);
    -- Reduce the number of probes for faster search
    SET ivfflat.probes = 3;
    EXPLAIN ANALYZE SELECT 1 - (embedding <=> '{str(query_embedding)}')
    AS cosine_similarity, *
    FROM test_embedding_table
    ORDER BY 1 - (embedding <=> '{str(query_embedding)}') DESC
    LIMIT 20;
    """

cur.execute(ivfflat_sql)
ivfflat_search = cur.fetchall()

在这个包含 200,000 条 job descriptions 的数据集测试中，index search 大约比 full search 快两倍。对于数百万或数十亿 data points 的超大数据集，这种差异会更加明显。

接下来，使用 HNSW indexing method。该代码使用 vector_cosine_ops 选项，它会创建一个针对 cosine similarity search 优化的 index。或者，也可以使用 vector_ip_ops，它针对 inner product similarity search 优化；或使用 vector_l2_ops，它针对 Euclidean（L2）distance vector search 优化。

创建 index 时，必须设置以下 hyperparameters：

m 设置 HNSW graph 中每层最大连接数。它定义每个 data point 与邻居之间有多少连接。更多连接会让 graph 更密，通常会加速 search queries，但会增加 index build time 和 memory usage。典型值范围为 16 到 64。

ef_construction 控制 index construction 期间使用的 dynamic candidate list 大小。更高值通常会提升 index quality。

ef_search 控制 search 期间考虑的 candidates 数量，使你可以在 search accuracy 和 speed 之间取舍。更高值返回更准确结果，但 queries 更慢。

下面是示例 query：

import psycopg2
from psycopg2 import Error

hnsw_test_sql = f"""
    DROP TABLE IF EXISTS test_embedding_table;
    CREATE TABLE test_embedding_table AS
        SELECT * FROM job_description_table;
    CREATE INDEX ON test_embedding_table
        USING hnsw (embedding vector_cosine_ops)
        WITH (m = 16, ef_construction = 200);
    SET hnsw.ef_search = 50;
    EXPLAIN ANALYZE SELECT 1 - (embedding <=> '{str(query_embedding)}')
    AS cosine_similarity, *
    FROM test_embedding_table
    ORDER BY 1 - (embedding <=> '{str(query_embedding)}') DESC
    LIMIT 20;
"""

cur.execute(hnsw_test_sql)
hnsw_search = cur.fetchall()

# [..., ('Planning Time: 1.109 ms',), ('Execution Time: 13.927 ms',)]

Discussion

对于少于 100,000 rows 的 tables，full scans 往往已经足够快。这种情况下可以跳过 indexing。超过 100 万 rows 后，indexing 就变得必要。

当 query speed 至关重要且你有足够 memory 时，选择 HNSW，因为 HNSW 会在 RAM 中存储更多数据。当 memory 受限，或你需要更快 index building，并能接受稍慢 query 时，选择 IVF。对于大多数拥有数百万 vectors 的生产 RAG 系统，HNSW 提供最佳 query performance。

创建 indexes 时，可以根据 similarity metric 选择不同 operators。vector_cosine_ops 选项会创建针对 cosine similarity search 优化的 index。或者，使用 vector_ip_ops 做 inner product similarity search，或使用 vector_l2_ops 做 Euclidean（L2）distance vector search。

HNSW indexes 需要调优多个 hyperparameters。m 参数设置 HNSW graph 中每层最大连接数。更多连接让 graph 更密，通常会加速 search queries，但增加 index build time 和 memory usage。典型值范围为 16 到 64。ef_construction 参数控制 index construction 期间 dynamic candidate list 的大小。更高值通常会改善 index quality。ef_search 参数控制 search 期间考虑的 candidates 数量，让你在 search accuracy 与 speed 之间权衡。更高值返回更准确结果，但 queries 更慢。

理解这些 indexing methods 如何工作，有助于你为自己的数据集调优参数。IVF 代表 inverted file。它会将 index 中的内容连接到 table 或 document 中的位置。构建 IVF index 时，系统会将 vector space 以及其中所有 data points 划分为 partitions。每个 partition 都有一个 centroid，每个 item 同一时间只属于一个 partition。当开始新 search 时，index 会先通过比较 centroids 位置与用户 query 的关系，寻找正确 partition，如图 6-10 所示。

图 6-10：通过寻找最近 centroid 查询 IVF index

一旦识别出最近 partition，搜索会继续在该 partition 内查找 nearest neighbor，如图 6-11 所示。Query vector，即蓝点，会与选定 partition 中的 data points 比较，以识别最接近匹配。请注意，虽然该 partition 的 centroid 是所有 centroids 中最近的，但如果实际 nearest neighbor 位于相邻 partition 且靠近 partition boundary，它仍可能在相邻 partition 中。

图 6-11：通过识别 partition 内最近 data point 查询 IVF index

对于所有 indexing methods，你通常会用一定准确率换取速度。对于 IVF index，挑战在于，当你只在特定 partition 中搜索时，可能会错过靠近邻近 partition 边缘的相关 data points。在实践中，这通常可以接受，因为准确率损失相对较小，而性能收益很明显。

PostgreSQL 中的另一种选择是 HNSW index，它会构建一个 multilayer graph。较高层包含较少连接，但覆盖 data space 中更长距离，像 shortcuts 一样支持高效搜索。较低层包含更密连接，用于细粒度导航。

可以把 HNSW 想象成一张多层地图。最高层像全国高速公路地图，路很少，但每条路覆盖长距离。从最高层开始，你会向下进入更低层，那里有更多道路和更多细节。最底层显示每一条本地小路。图 6-12 展示了这种结构。

图 6-12：HNSW layer approach

不同 vector databases 使用不同 indexing methods。FAISS 支持 product quantization（PQ）、IVF 和 HNSW。Annoy 使用 tree-based approach，而 HNSWlib 只专注 HNSW。Milvus 支持 HNSW、IVF 和 PQ。

See Also

pgvector indexing documentation 描述了 HNSW 和 IVFFlat 配置选项及性能特征。

HNSWlib repository 提供 HNSW indexing 的 reference implementation 和技术细节。

PostgreSQL CREATE INDEX documentation 解释了创建 indexes 的语法和选项。

6.7 在 PostgreSQL 中结合 Keyword Search 和 Similarity Search 提升检索准确率

Problem

你想通过结合 keyword search 和 semantic search，提升 retrieval step 的稳健性。

Solution

Hybrid search 会将 keyword search 和 semantic search 组合成一个结果集。这个 recipe 直接在 PostgreSQL 中实现 hybrid search：在一个 SQL query 中执行两种 search，并用 weighted scoring 合并结果。

这个 recipe 假设你已经有一个 PostgreSQL table，其中 embeddings 已经创建。示例复用 Recipe 6.6 中的 job_description_table。如果你单独跟随这个 recipe，需要具备以下条件：（1）启用 pgvector extension 的 PostgreSQL database；（2）一个包含 id、text_chunk 和 embedding vector(1536) 列的 table；（3）该 table 至少填充了几十个 text chunks 及其 embeddings。设置说明请见 Recipe 6.4 或 Recipe 6.6。

这个 recipe 使用 Psycopg 2 library 连接 PostgreSQL 并执行 SQL queries。安装命令如下：

pip install psycopg2-binary openai

接下来，复制 job description table，并添加一个新的 tsvector data type 列。tsvector 专为 full-text search 设计，会在后续步骤中用于执行 keyword search：

import psycopg2
from psycopg2 import Error

conn = psycopg2.connect(
    dbname="rag_cookbook",
    user="rag_cookbook_user",
    password="rag_cookbook_user_pw",
    host="localhost",
    port="5432",
)

cur = conn.cursor()

hybrid_search_table = f"""
    DROP TABLE IF EXISTS test_embedding_table;
    CREATE TABLE test_embedding_table AS
    SELECT *, to_tsvector(text_chunk) AS tsv
    FROM job_description_table;
    """

cur.execute(hybrid_search_table)

要运行 hybrid search query，可以使用标准 SQL SELECT statement。除了返回 rows，该 query 还会计算以下列：

• text_score，存储 full-text search ranking score
• vector_score，存储 cosine similarity score
• hybrid_score，将两种 rankings 合并为单一 score

在这个示例中，combined score 通过对两种 rankings 等权加权计算：

• Keyword ranking weight：0.5
• Vector similarity ranking weight：0.5

你可以调整这些权重。例如，为了优先 semantic search，可以使用 0.7 的 semantic similarity 和 0.3 的 keyword ranking：

import psycopg2
from psycopg2 import Error
from openai import OpenAI

query = "I am looking for a job as a data scientist in Berlin."
client = OpenAI()
model = "text-embedding-3-small"
response = client.embeddings.create(input=query, model=model)
query_embedding = response.data[0].embedding

hybrid_search_query = f"""
    WITH ranked_docs AS (
        SELECT
            id,
            text_chunk AS text,
            ts_rank(tsv, plainto_tsquery('PostgreSQL')) AS text_score,
            1 - (embedding <=> '{str(query_embedding)}') AS vector_score,
            ts_rank(tsv, plainto_tsquery('PostgreSQL'))
            * 0.5 + (1 - (embedding <=> '{str(query_embedding)}')) * 0.5
            AS hybrid_score
        FROM test_embedding_table
        -- Filter for relevant text first
        WHERE tsv @@ plainto_tsquery('PostgreSQL')
        ORDER BY hybrid_score DESC
        LIMIT 10
    )
    SELECT * FROM ranked_docs;
    """

cur.execute(hybrid_search_query)
results = cur.fetchall()

Discussion

Keyword search，也就是 full-text search，使用 PostgreSQL 的 tsvector type 和 ts_rank function 来识别包含 query terms 的 documents。Semantic search 使用 cosine similarity 来查找接近 query embedding 的 vectors。通过计算两种 scores 并用权重合并，可以检索那些包含精确词语或表达相似概念的 documents。

当 exact keyword matches 对你的领域很重要时，使用 hybrid search。这包括法律文档，例如搜索 “force majeure” 或 “indemnification”；技术文档，例如搜索特定 error codes 或 API names；医疗记录，例如搜索药品名或 diagnosis codes；或任何用户期望结果包含其 exact query terms 的场景。当 queries 包含 proper nouns、acronyms 或 domain-specific terminology，而这些内容在 embedding spaces 中可能表示不好时，hybrid search 也有帮助。

当 semantic meaning 已足够，exact terms 并不关键时，不要使用 hybrid search。例如，搜索 conversational customer support tickets 或 general knowledge articles，通常使用 pure semantic search 效果更好。如果 keyword search index 没有调好，也要避免 hybrid search，因为糟糕的 keyword results 会拉低 combined score。

Hybrid search 的主要取舍是 complexity versus robustness。你必须同时维护 text search index 和 vector index，调优两套参数，并确定合适权重。作为回报，你会获得更稳健的 retrieval，可以同时处理 semantic queries 和 exact term matching。

从等权开始，即 0.5 text、0.5 vector。如果用户抱怨相关结果因为缺少 exact term 而被漏掉，可以将 vector weight 提高到 0.7。如果结果偏离字面 query 太远，可以将 text weight 提高到 0.7。不同领域可能需要不同默认权重：法律和技术应用通常需要更高 text weights，而对话式应用通常需要更高 vector weights。

Hybrid search 与 pure semantic search 的不同在于，它还会考虑 exact term matches。它与 pure keyword search 的不同在于，也能找到不包含 query terms 但语义相似的内容。它与 SQL filtering，也就是前一个 recipe，不同之处在于，它组合两种搜索方法，而不是先过滤再搜索。当字面词语和含义都重要时，使用 hybrid search。当需要按 metadata 限制搜索范围时，使用 SQL filtering。当仅含义足够时，使用 pure semantic search。

See Also

PostgreSQL full-text search documentation 解释了 keyword search capabilities 和配置。

pgvector query documentation 提供了组合 vector 和 keyword searches 的示例。

PostgreSQL UNION documentation 描述了合并 query results 的技术。