博客
POC避坑指南｜VectorDBBench自定义数据集测试实战教程

POC避坑指南｜VectorDBBench自定义数据集测试实战教程

2025-11-07

By 蔡一凡

本文是 “向量数据库 POC 指南” 的系列文章之一。

大模型落地，如何做配套的技术设施选型，已经成为了困扰很多企业的一大难题。但哪怕只是一个向量数据库，纸面参数就已经膨胀到需要重新统一度量衡。这些被吹的天花乱坠的成绩，往往会干扰企业的判断，甚至导致企业浪费了大量企业的人力、物力之后，后期产品性能掉链子、运维一堆麻烦、成本越花越高…… 全部推倒重来。

在这个过程中，一个更科学、更符合实际生产的POC 可以帮我们规避大部分问题。

那如何做好POC？接下来，我们将产出“向量数据库 POC 指南” 系列文章，对 POC 环节的重点问题做逐一拆解，本文为该系列的第一篇文章。

企业大模型落地，做POC的时候最怕什么？

答案是，供应商给的参数都是真的，但成绩都是基于 SIFT、GloVe、LAION 这些老旧数据集刷出来的，但生产上用的都是最新的模型，性能效果完全不一样。

比如，SIFT是向量数据库领域的常见测试数据集，其向量数据维度是128 维，而 OpenAI、Cohere 的主流模型的embedding模型维度是 768~3072 维。也就是说，主流benchmark的测试结果，对真实业务场景已经失去参考意义。

0602-1.webp

而这些差异可能导致：

QPS（查询吞吐量）和延迟在业务数据上与测试数据差别很大
Recall（召回率） 受向量分布影响，可能和真实效果不一致
索引构建、内存占用等成本评估不准确

更多数据集老旧导致的测试问题，详见历史文章：开源|VDBBench 1.0正式官宣，完全复刻业务场景，支持用户自定义数据集

而解决问题的办法也很简单，那就是使用业务里真实产生的向量数据测试，也就是用自定义数据集。

但如何实操呢？

Zilliz 团队开源的VectorDBBench 向量数据库性能评测工具，支持多种数据库（Milvus、Elasticsearch、pgvector 等）和真实模拟生产中的流量特征（比如边插入边查询），并允许用户接入自定义数据集，这就很符合我们的需求了。

接下来，我们将手把手教你如何用VectorDBBench 基于真实数据，完成一场高质量的POC。

01 准备工作

安装与环境要求

Python >= 3.11

安装

可选仅安装 Milvus/Zilliz Cloud 客户端：

pip install vectordb-bench

也可安装所有支持的数据库客户端（如需对比多种数据库）：

pip install vectordb-bench[all]

或者安装指定数据库客户端（如 Elasticsearch）：

pip install vectordb-bench[elastic]

支持的数据库客户端及安装命令如下：

0602-2.webp

启动

安装完成后，直接运行：

init_bench

默认会启动本地 Web 服务：

0602-3.webp

界面示例：

0602-4.webp

02 准备自定义数据集

（1） VectorDBBench 支持的数据格式

VectorDBBench 官方要求的数据集必须是特定结构的 Parquet 文件，通常包括：

0602-5.webp

向量数据要求：

训练向量文件（train.parquet）
- 必须包含：id（递增整数）、向量列（如 emb，类型为 float32 数组）
- 列名可以自定义
测试向量文件（test.parquet）
- 必须包含：id（递增整数）、向量列（如 emb）
- 注意：id 列名必须是 id，向量列名可自定义
Ground Truth 文件（neighbors.parquet）
- 必须包含： id（对应测试向量 ID）、近邻 ID 数组列（如 neighbors_id）
- id 列名必须是 id

（可选）标签文件（scalar_labels.parquet）

必须包含：id（对应训练向量 ID）、 labels (标量字符串)

（2）普通用户面临的困难

虽然 Parquet 格式和结构有明确要求，但大多数业务数据是：

存在 CSV 文件（向量以字符串形式存储）
存在 NPY 文件（纯向量矩阵，没有 ID、标签）
存在 数据库表（MySQL、MongoDB 中存储的 embedding）

这些格式和 VectorDBBench 直接使用的 Parquet 文件差异较大，普通用户要自己生成符合规范的 train/test/ground truth 三个文件会遇到：

格式转换难（CSV/NPY → Parquet，需要代码）
Ground Truth 计算难（需要用 FAISS 或其他库做精确 KNN 检索）
数据切分麻烦（区分 train/test，并保证匹配关系正确）

（3）解决方案：一键生成工具脚本

为了简化流程，笔者准备了一个 Python 脚本，可以将用户已有的 CSV 或 NPY 数据一键转换成 VectorDBBench 能直接识别的完整数据集目录，包括：

train.parquet
test.parquet
neighbors.parquet
（可选）scalar_labels.parquet

优势：

支持 CSV 和 NPY 两种输入
自动切分 train/test 数据
自动生成 Ground Truth（精确搜索计算）
可自定义 top_k、相似度度量方式（L2 / Inner Product）

（4）脚本输入要求

脚本接受两种输入形式：

格式一：CSV 格式

第一列为 id（唯一标识符）

第二列为 vector（字符串形式的浮点数组，如 [0.1, 0.2, 0.3, ...]）

其他列可选（metadata、标签）

示例：

id,emb,label
1,"[0.12,0.56,0.89,...]",A
2,"[0.33,0.48,0.90,...]",B

格式二：NPY 格式

一个二维数组，shape = (num_vectors, dim)

向量顺序默认从 0 开始分配 id

标签可单独提供一个 CSV（id,label）

示例：

import numpy as np
vectors = np.random.rand(10000, 768).astype('float32')
np.save("vectors.npy", vectors)

（5）运行脚本

安装依赖：

pip install numpy pandas faiss-cpu

启动命令：

python convert_to_vdb_format.py   --train data/train.csv   --test data/test.csv   --out datasets/custom   --topk 10

参数说明：

0602-6.webp

（6）输出目录结构

datasets/custom/
├── train.parquet          # 训练向量
├── test.parquet           # 查询向量
├── neighbors.parquet      # Ground Truth
└── scalar_labels.parquet  # 可选标量标签

（7）脚本代码

import os
import argparse
import numpy as np
import pandas as pd
import faiss
from ast import literal_eval

def load_csv(path: str):
    df = pd.read_csv(path)
    df['emb'] = df['emb'].apply(literal_eval)
    if 'id' not in df.columns:
        df.insert(0, 'id', range(len(df)))
    return df

def load_npy(path: str):
    arr = np.load(path)
    return pd.DataFrame({'id': range(arr.shape[0]), 'emb': arr.tolist()})

def compute_ground_truth(train_vectors, test_vectors, top_k=10):
    index = faiss.IndexFlatL2(train_vectors.shape[1])
    index.add(train_vectors)
    _, indices = index.search(test_vectors, top_k)
    return indices

def save_ground_truth(df_path, indices):
    pd.DataFrame({
        "id": np.arange(indices.shape[0]),
        "neighbors_id": indices.tolist()
    }).to_parquet(df_path, index=False)
    print(f"✅ Ground truth 保存成功: {df_path}")

def main(train_path, test_path, output_dir, top_k=10):
    os.makedirs(output_dir, exist_ok=True)
    train_df = load_csv(train_path) if train_path.endswith('.csv') else load_npy(train_path)
    test_df = load_csv(test_path) if test_path.endswith('.csv') else load_npy(test_path)
    train_vectors = np.array(train_df['emb'].to_list(), dtype='float32')
    test_vectors = np.array(test_df['emb'].to_list(), dtype='float32')
    train_df.to_parquet(os.path.join(output_dir, 'train.parquet'), index=False)
    test_df.to_parquet(os.path.join(output_dir, 'test.parquet'), index=False)
    indices = compute_ground_truth(train_vectors, test_vectors, top_k)
    save_ground_truth(os.path.join(output_dir, 'neighbors.parquet'), indices)
    print("✅ 所有文件生成完毕")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--train", required=True)
    parser.add_argument("--test", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--topk", type=int, default=10)
    args = parser.parse_args()
    main(args.train, args.test, args.out, args.topk)

（8）使用示例

运行脚本后输出：

0602-7.webp

文件输出： 640 (1).webp

03 加载自定义数据集

在第二章中，我们已经用脚本生成了符合 VectorDBBench 要求的 Custom 数据集目录（train.parquet, test.parquet, neighbors.parquet 等）。接下来，我们将在 VectorDBBench Web UI 中加载并运行测试。

进入 Web UI 首页，选择主页中的 Custom Dataset：

0602-8.webp

选择之后，我们可以看到 Custom Dataset 相关的解释以及需要填写的内容： 0602-9.webp

参数详解 0602-10.webp

点击 Save 保存。

04 配置测试方案并运行测试

在 Web UI 中进入 Run Test 页面：

06-02-11.webp

勾选并填写要测试的向量数据库，本文以 milvus 为例：

06-02-12.webp

选择我们创建的 Custom 数据集：

0602-13.webp

设置任务标签

0602-14.webp

开始测试

06-02-15.webp

05 查看结果

在 Web UI 的 Results 页面，可以看到测试结果：

640.webp

测试说明：

并发度：由于测试机器性能有限，本次仅测试了 1、5、10 三种并发情况
测试数据集：
- 维度：768
- 数量：train 和 test 各 3000 条向量
- 标量标签：未使用

06 踩坑与经验

维度不一致会直接报错 → 确保 train/test 文件维度相同
Ground Truth 必须匹配 test 文件 → 否则 Recall 计算错误
数据量太小 → QPS 测试结果波动大，建议至少几万条向量
Docker 资源限制 → 测试大规模数据时，增加内存和 CPU 配额
要注意 VectorDBBench 输出的日志，可能在任务中报错了，Web UI 中没有提示

本文也适用于不同数据库之间的横向对比

07 总结

本文通过对 VectorDBBench 自定义（Custom）数据集测试全流程的详细讲解，展示了如何将业务真实向量数据高效转换为符合 VectorDBBench 标准的数据格式，并在其 Web UI 中完成全流程性能测试。

相比于使用官方默认数据集，基于业务场景的自定义数据集测试能够更准确地反映实际生产环境中的性能表现，包括查询吞吐量（QPS）、延迟、召回率以及索引构建等关键指标，帮助我们更科学地评估和优化向量数据库的部署方案。

此外，借助本文提供的一键转换脚本，用户能够轻松解决格式转换、数据切分和 Ground Truth 计算等常见难题，大幅降低自定义测试的门槛，实现高效、可重复的性能验证流程。

未来，随着向量数据库应用场景的不断丰富，定制化性能测试的重要性将愈发凸显。希望本实践经验能够为广大开发者和运维人员在向量数据库落地过程中提供实用参考和技术支持，推动行业整体性能测试水平的提升。

08 彩蛋

如果只是想生成指定条数和维度的人造数据集，那么可以使用以下代码来生成：

import pandas as pd
import numpy as np

def generate_csv(num_records: int, dim: int, filename: str):
    ids = range(num_records)
    vectors = np.random.rand(num_records, dim).round(6)  # 保留6位小数
    emb_str = [str(list(vec)) for vec in vectors]
    df = pd.DataFrame({
        'id': ids,
        'emb': emb_str
    })
    df.to_csv(filename, index=False)
    print(f"生成文件 {filename} ，共 {num_records} 条数据，向量维度 {dim}")

if __name__ == "__main__":
    num_records = 3000  # 生成数据的数量
    dim = 768           # 向量维度
    generate_csv(num_records, dim, "train.csv")
    generate_csv(num_records, dim, "test.csv")