构建你的向量数据库（从零手写实现 Mini Vector DB）

在前五章中，我们系统学习了向量数据库的核心概念、底层算法和技术原理。本章将聚焦“实践”，带领大家用Python手动实现一个简化版向量数据库，理解其核心工作流程。通过本次实践，你将掌握向量存储、索引构建、相似度检索等核心功能的实现逻辑，为后续使用工业级向量数据库打下坚实基础。

学习目标：

本章目标：掌握向量数据库的基本架构，并通过 Python 手写实现一个可执行的迷你版本，包括：向量存储、相似度搜索、IVF 与 HNSW 简易模拟结构。

1.核心目标与技术选型

1.1 实现目标

我们将构建一个具备以下核心功能的向量数据库原型：

• 向量数据的增删改查（基础数据操作）
• 基于余弦相似度的精确检索
• 基于IVF（倒排文件）的近似检索（提升检索效率）
• 向量数据的持久化存储（避免程序退出后数据丢失）

1.2 技术选型

为简化实现并聚焦核心逻辑，我们选用以下轻量级工具库：

• numpy：核心数值计算库，用于向量的存储、运算和相似度计算
• scikit-learn：提供KMeans聚类算法，用于IVF索引的构建（对应Chapter 5的IVF算法原理）
• pickle：Python内置序列化工具，用于向量数据和索引的持久化存储
• uuid：生成唯一标识符，用于向量数据的主键标识
  安装命令：

pip install numpy scikit-learn sentence-transformers modelscope
​

向量数据库不是一个“算法集合”，而是一整套系统：

向量数据库核心模块

2.分步实现向量数据库

2.1初始化数据库与数据存储层

数据存储层的核心作用是管理原始向量数据及对应的元数据（如唯一ID、描述信息等）。我们将用numpy数组存储向量集合，用字典建立ID与向量下标的映射（便于快速定位）。
以下代码需整合到同一个类中

import numpy as np
import pickle
import uuid
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity

class SimpleVectorDB:
    def __init__(self, vector_dim, db_path="vector_db.pkl"):
        """
        初始化向量数据库
        :param vector_dim: 向量维度（所有入库向量必须保持一致）
        :param db_path: 数据持久化存储路径
        """
        self.vector_dim = vector_dim  # 向量维度
        self.db_path = db_path        # 持久化路径
        
        # 数据存储核心结构
        self.vectors = np.array([])               # 存储所有向量（numpy数组，形状为[N, vector_dim]）
        self.vector_ids = []                      # 存储向量唯一ID，与vectors下标一一对应
        self.id_to_index = dict()                 # 映射：向量ID → 向量在vectors中的下标
        self.metadata = dict()                    # 存储向量元数据（key: 向量ID, value: 元数据字典）
        
        # 索引相关（后续初始化）
        self.ivf_index = None                     # IVF索引结构
        self.ivf_kmeans = None
        
    def _check_vector_dim(self, vector):
        """校验向量维度是否符合要求"""
        if len(vector) != self.vector_dim:
            raise ValueError(f"向量维度错误，需为{self.vector_dim}维，当前为{len(vector)}维")
    
​

2.2 实现增删改查 API

这部分将实现向量的插入、查询、更新和删除功能，是数据库最基础的交互能力。

    def insert(self, vector, metadata=None):
        """
        插入向量数据
        :param vector: 向量数据（list或numpy.ndarray）
        :param metadata: 元数据（可选，如文本描述、来源等）
        :return: 向量唯一ID（用于后续查询/更新）
        """
        # 类型转换与维度校验
        vector = np.array(vector, dtype=np.float32).flatten()
        self._check_vector_dim(vector)
        
        # 生成唯一ID
        vector_id = str(uuid.uuid4())
        
        # 插入数据
        if len(self.vectors) == 0:
            self.vectors = np.expand_dims(vector, axis=0)
        else:
            self.vectors = np.vstack([self.vectors, vector])
        self.vector_ids.append(vector_id)
        self.id_to_index[vector_id] = len(self.vector_ids) - 1
        if metadata:
            self.metadata[vector_id] = metadata
        
        return vector_id

    def get_by_id(self, vector_id):
        """通过ID查询向量及元数据"""
        if vector_id not in self.id_to_index:
            raise KeyError(f"未找到ID为{vector_id}的向量")
        index = self.id_to_index[vector_id]
        return {
            "vector_id": vector_id,
            "vector": self.vectors[index].tolist(),
            "metadata": self.metadata.get(vector_id, {})
        }

    def update(self, vector_id, new_vector=None, new_metadata=None):
        """更新向量数据或元数据"""
        if vector_id not in self.id_to_index:
            raise KeyError(f"未找到ID为{vector_id}的向量")
        index = self.id_to_index[vector_id]
        
        # 更新向量（若提供新向量）
        if new_vector is not None:
            new_vector = np.array(new_vector, dtype=np.float32).flatten()
            self._check_vector_dim(new_vector)
            self.vectors[index] = new_vector
        
        # 更新元数据（若提供新元数据）
        if new_metadata is not None:
            self.metadata[vector_id] = new_metadata

    def delete(self, vector_id):
        """删除向量数据"""
        if vector_id not in self.id_to_index:
            raise KeyError(f"未找到ID为{vector_id}的向量")
        index = self.id_to_index[vector_id]
        
        # 删除核心数据
        self.vectors = np.delete(self.vectors, index, axis=0)
        self.vector_ids.pop(index)
        del self.id_to_index[vector_id]
        if vector_id in self.metadata:
            del self.metadata[vector_id]
        
        # 重新构建ID与下标的映射（因删除后下标发生变化）
        self.id_to_index = {vid: idx for idx, vid in enumerate(self.vector_ids)}
    
​

2.3 实现相似度检索（暴力 + IVF）

检索是向量数据库的核心价值，我们将先实现精确检索（暴力检索），再基于IVF算法实现近似检索，对比两者的效率差异。

def brute_force_search(self, query_vector, top_k=5):
        """
        暴力检索（精确匹配）：计算查询向量与所有向量的余弦相似度
        :param query_vector: 查询向量
        :param top_k: 返回相似度最高的前k个结果
        :return: 检索结果（按相似度降序排列）
        """
        if len(self.vectors) == 0:
            return []
        
        # 预处理查询向量
        query_vector = np.array(query_vector, dtype=np.float32).flatten()
        self._check_vector_dim(query_vector)
        
        # 计算余弦相似度（利用sklearn简化实现，也可手动实现：(a·b)/(||a||·||b||)）
        similarities = cosine_similarity([query_vector], self.vectors)[0]
        
        # 按相似度降序排序，取前top_k
        top_indices = np.argsort(similarities)[::-1][:top_k]
        
        # 组装结果
        results = []
        for idx in top_indices:
            vector_id = self.vector_ids[idx]
            results.append({
                "vector_id": vector_id,
                "similarity": float(similarities[idx]),
                "vector": self.vectors[idx].tolist(),
                "metadata": self.metadata.get(vector_id, {})
            })
        return results

    def build_ivf_index(self, n_clusters=8):
        """
        构建IVF索引（基于KMeans聚类）
        核心逻辑：将向量聚类到n_clusters个桶中，检索时先找查询向量所属的桶，再在桶内暴力检索
        """
        if len(self.vectors) == 0:
            raise ValueError("数据库中无向量数据，无法构建索引")
        
        # 转换为 float64
        vectors_for_kmeans = self.vectors.astype(np.float64)
        
        # KMeans 聚类
        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        cluster_labels = kmeans.fit_predict(vectors_for_kmeans)
        
        # 构建IVF索引
        self.ivf_index = {i: [] for i in range(n_clusters)}
        for idx, label in enumerate(cluster_labels):
            self.ivf_index[label].append(idx)
        
        self.ivf_kmeans = kmeans

    def ivf_search(self, query_vector, top_k=5):
        """基于IVF索引的近似检索"""
        if self.ivf_index is None:
            raise ValueError("请先调用build_ivf_index()构建IVF索引")
        if len(self.vectors) == 0:
            return []
        
        # 1. 预处理查询向量，确定其所属的聚类（桶）
        query_vector = np.array(query_vector, dtype=np.float64).flatten()
        self._check_vector_dim(query_vector)
        cluster_id = self.ivf_kmeans.predict([query_vector])[0]
        
        # 2. 获取该聚类下的所有向量下标，提取对应向量
        cluster_indices = self.ivf_index[cluster_id]
        if not cluster_indices:
            return []
        cluster_vectors = self.vectors[cluster_indices]
        
        # 3. 在聚类内计算相似度并排序
        similarities = cosine_similarity([query_vector], cluster_vectors)[0]
        top_cluster_indices = np.argsort(similarities)[::-1][:top_k]
        
        # 4. 组装结果（映射回原数据库的向量ID）
        results = []
        for idx in top_cluster_indices:
            original_idx = cluster_indices[idx]
            vector_id = self.vector_ids[original_idx]
            results.append({
                "vector_id": vector_id,
                "similarity": float(similarities[idx]),
                "vector": self.vectors[original_idx].tolist(),
                "metadata": self.metadata.get(vector_id, {}),
                "cluster_id": int(cluster_id)  # 标注所属聚类，便于理解
            })
        return results
    
​

2.4 数据持久化

为避免程序退出后数据丢失，我们需要将核心数据结构序列化到本地文件，下次启动时再加载。

def save(self):
    """将数据库数据与索引持久化到本地文件"""
    data = {
        "vector_dim": self.vector_dim,
        "vectors": self.vectors,
        "vector_ids": self.vector_ids,
        "id_to_index": self.id_to_index,
        "metadata": self.metadata,
        "ivf_index": self.ivf_index,
        "ivf_kmeans": self.ivf_kmeans  # 存储聚类模型
    }
    with open(self.db_path, "wb") as f:
        pickle.dump(data, f)
    print(f"数据库已保存至{self.db_path}")

@classmethod
def load(cls, db_path="vector_db.pkl"):
    """从本地文件加载数据库"""
    with open(db_path, "rb") as f:
        data = pickle.load(f)
    
    # 重建数据库实例
    db = cls(vector_dim=data["vector_dim"], db_path=db_path)
    db.vectors = data["vectors"]
    db.vector_ids = data["vector_ids"]
    db.id_to_index = data["id_to_index"]
    db.metadata = data["metadata"]
    db.ivf_index = data["ivf_index"]
    db.ivf_kmeans = data.get("ivf_kmeans")
    
    print(f"已从{db_path}加载数据库，共包含{len(db.vectors)}个向量")
    return db
    
​

3.实战：使用自定义向量数据库实习文本向量检索

我们通过一个“文本向量检索”的实例，演示如何使用上面实现的向量数据库。

文本向量使用GTE文本向量-中文-通用领域模型，将文本转换为向量。
需要安装的第三方包

pip install sentence-transformers modelscope
​

3.1 下载模型

首先是使用modelscope 下载对应的模型

#模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('iic/nlp_gte_sentence-embedding_chinese-base',cache_dir='./model')
​

3.2 加载模型

#加载模型
import os
import jieba
from sentence_transformers import SentenceTransformer

def read_txt_and_split(file_path):
    """读取txt文件内容并切分为句子列表"""
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"指定的txt文件不存在：{file_path}")
    # 读取文件（默认UTF-8编码，若有乱码可尝试gbk）
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    # 句子切分
    sentences = content.split("。")
    print(f"成功读取txt文件，共切分出 {len(sentences)} 个句子")
    return sentences

def text_to_vector(text, model):
    """将文本转换为向量（基于SentenceTransformer模型）"""
    # encode方法直接返回向量，convert_to_numpy=True确保输出为numpy数组
    return model.encode(text, convert_to_numpy=True)

if __name__ == "__main__":
    # 1. 配置参数（请根据实际情况修改txt文件路径和模型路径）
    txt_file = "./data/Datawhale社区介绍大模型改写版.txt"  # 你的txt文件路径
    model_dir = "./model/iic/nlp_gte_sentence-embedding_chinese-base"  # 预训练模型本地存储路径
    VECTOR_DIM = 768  # 主流中文语义模型输出维度多为768，可根据实际模型调整
    
    # 2. 加载中文语义向量模型（两种加载方式可选）
    print("正在加载中文文本嵌入模型...")
    model = SentenceTransformer(model_dir)
    print(f"成功加载本地模型：{model_dir}")

    # 3. 初始化向量数据库
    db = SimpleVectorDB(vector_dim=VECTOR_DIM)
    
    # 4. 读取txt文件并切分句子
    print(f"\n正在读取并处理txt文件：{txt_file}")
    sentences = read_txt_and_split(txt_file)
    if not sentences:
        raise ValueError("未从txt文件中提取到有效句子")
    
    # 5. 句子转向量并插入数据库（附带元数据：原始句子）
    print("\n正在将句子转向量并插入数据库...")
    embeddings = model.encode(sentences)  # 批量生成向量，效率更高
    for idx, (sentence, embedding) in enumerate(zip(sentences, embeddings), 1):
        # 元数据包含句子内容和序号，便于后续查看
        metadata = {"sentence": sentence, "sequence": idx}
        db.insert(embedding, metadata=metadata)
    
    # 6. 持久化数据库
    db.save()
    
    # 7. 测试暴力检索（查询与向量数据库相关的内容）
    print("\n=== 暴力检索结果（查询：'Datawhale有多个学习者参与活动'）===")
    query_text = "Datawhale有多个学习者参与活动"
    query_vector = text_to_vector(query_text, model)
    brute_results = db.brute_force_search(query_vector, top_k=3)
    for res in brute_results:
        print(f"相似度：{res['similarity']:.4f} | 句子：{res['metadata']['sentence']}")
    
    # 8. 构建IVF索引并测试近似检索
    print("\n=== IVF近似检索结果（查询：'Datawhale有多个学习者参与活动'）===")
    # 根据句子数量调整聚类数（一般为数据量的平方根左右）
    n_clusters = max(2, int(len(sentences)**0.5))
    db.build_ivf_index(n_clusters=n_clusters)
    ivf_results = db.ivf_search(query_vector, top_k=3)
    for res in ivf_results:
        print(f"相似度：{res['similarity']:.4f} | 聚类ID：{res['cluster_id']} | 句子：{res['metadata']['sentence']}")
    
    # 9. 测试数据更新与查询
    print("\n=== 数据更新与查询测试 ===")
    # 获取第一个向量的ID（即第一个句子对应的向量）
    first_vector_id = db.vector_ids[0]
    first_sentence = db.get_by_id(first_vector_id)['metadata']['sentence']
    print(f"待更新的原始句子：{first_sentence}")
    # 更新其元数据（模拟句子修正）
    new_metadata = {"sentence": f"【修正】{first_sentence}", "sequence": 1, "updated": True}
    db.update(first_vector_id, new_metadata=new_metadata)
    # 按ID查询更新结果
    updated_res = db.get_by_id(first_vector_id)
    print(f"更新后的数据：{updated_res['metadata']['sentence']}")
    
    # 10. 测试数据删除
    db.delete(first_vector_id)
    print(f"\n删除后数据库向量总数：{len(db.vectors)}")
    
    # 11. 从本地加载数据库验证持久化功能
    print("\n=== 从本地加载数据库 ===")
    loaded_db = SimpleVectorDB.load()
    print(f"加载的数据库向量总数：{len(loaded_db.vectors)}")
    # 验证加载的数据
    if loaded_db.vector_ids:
        sample_id = loaded_db.vector_ids[0]
        sample_data = loaded_db.get_by_id(sample_id)
        print(f"加载数据示例：{sample_data['metadata']['sentence']}")
​

运行结果：

正在加载中文文本嵌入模型...
No sentence-transformers model found with name ./model/iic/nlp_gte_sentence-embedding_chinese-base. Creating a new one with mean pooling.
成功加载本地模型：./model/iic/nlp_gte_sentence-embedding_chinese-base

正在读取并处理txt文件：./data/Datawhale社区介绍大模型改写版.txt
成功读取txt文件，共切分出 68 个句子

正在将句子转向量并插入数据库...
数据库已保存至vector_db.pkl

=== 暴力检索结果（查询：'Datawhale有多个学习者参与活动'）===
相似度：0.9302 | 句子：
Datawhale 目前有超过六万名核心学习者和志愿者参与社区活动
相似度：0.9168 | 句子：
Datawhale 将学习者组织成不同兴趣方向的小组
相似度：0.9065 | 句子：
Datawhale 通过协作机制吸引大量志愿者参与学习内容的共建

=== IVF近似检索结果（查询：'Datawhale有多个学习者参与活动'）===
相似度：0.9302 | 聚类ID：7 | 句子：
Datawhale 目前有超过六万名核心学习者和志愿者参与社区活动
相似度：0.9017 | 聚类ID：7 | 句子：
Datawhale 每年举行超过二十场技术主题分享
相似度：0.8981 | 聚类ID：7 | 句子：
Datawhale 每年举办十次以上 AI 社区活动

=== 数据更新与查询测试 ===
待更新的原始句子：Datawhale 是一个开源、开放、公益的学习型组织
更新后的数据：【修正】Datawhale 是一个开源、开放、公益的学习型组织

删除后数据库向量总数：67

=== 从本地加载数据库 ===
已从vector_db.pkl加载数据库，共包含68个向量
加载的数据库向量总数：68
加载数据示例：Datawhale 是一个开源、开放、公益的学习型组织
​

4.结果分析与扩展方向

4.1 核心功能验证

运行上述代码后，你将观察到：

1. 暴力检索能精准返回与查询文本语义最相关的结果（相似度最高）；
2. IVF检索结果与暴力检索接近，但速度更快（尤其当向量数量庞大时，优势更明显）；
3. 增删改查及持久化功能正常工作，数据可跨程序会话保留。

4.2 与工业级向量数据库的差距

我们的简化版实现仅包含核心逻辑，工业级向量数据库（如Milvus、Pinecone）还具备以下高级特性：

• 分布式存储与并行计算（支持亿级向量）；
• 更多索引算法（如HNSW、PQ的优化实现）；
• 向量与结构化数据的混合查询；
• 高可用与容错机制（副本、备份）；
• RESTful API或SDK封装（便于多语言调用）。

4.3 扩展方向（课后练习）

基于本实现，你可以尝试扩展以下功能，深化理解：

1. 实现HNSW索引（参考Chapter 5的HNSW算法原理，用图结构存储向量邻居关系）；
2. 添加向量维度压缩功能（基于PQ算法，减少存储开销）；
3. 实现批量插入/批量检索接口（提升数据操作效率）；
4. 增加向量过滤功能（如按元数据中的关键词筛选后再检索）。

5.总结

本章通过Python手动实现向量数据库，将前序章节的理论知识（如向量相似度、IVF索引）落地为可运行的代码。核心收获包括：

• 理解向量数据库“存储-索引-检索”的核心流程；
• 掌握余弦相似度计算、IVF索引构建的实际代码实现；
• 明确简化版与工业级实现的差异，为后续学习奠定基础。

恭喜！你已经成功完成本章的实践，并成功实现了一个向量数据库。