为Pinecone生成测试数据 | Pinecone
本文介绍了如何为Pinecone生成测试数据，包括使用Parquet文件生成器、向量嵌入生成器等工具，以及如何将数据导入Pinecone。文章强调了模块化工作流的重要性，并讨论了数据规模对测试的影响。
作为一名Pinecone解决方案工程师，John Ward在过去一年中开发了多个内部工具和示例，用于测试平台的各个部分。本文是该系列的第一篇，重点介绍如何生成测试数据——一个看似简单但在大规模下至关重要的环节。
生成测试数据的需求源于不同测试类型对数据集的不同要求。对于向量搜索测试，随机向量仅适用于测试写入吞吐量或查询延迟，而召回率和准确性测试则需要基于真实文本的嵌入。为此，Ward设计了一套工具集，包括Parquet文件生成器、向量嵌入生成器、使用本地LLM进行Parquet分类的工具，以及按类别命名空间重新分区Parquet的工具。
工作流的第一步是生成源数据并准备嵌入和导入。Ward选择了Hugging Face上的stanford-oval/ccnews数据集，因为它提供广泛的新闻文章和丰富的元数据（标题、作者、来源URL、发布日期、类别等）。数据以Parquet格式存储，并符合Pinecone批量导入的结构要求：id字段使用UUID，values字段存储稠密向量，metadata字段为JSON字符串。
在嵌入模型的选择上，Ward使用了BAAI/bge-large-en-v1.5，生成1024维稠密向量。虽然Pinecone也提供托管嵌入模型，但本地生成可以模拟客户常见的自生成向量工作流。
第一个版本的脚本直接流式处理数据集、分块文本、本地生成嵌入并直接upsert到Pinecone索引。这种方法适合小规模测试，但存在元数据大小限制（40KB）和空值处理问题。脚本中包含了迭代截断长文本和过滤空值的逻辑。
然而，对于数百万或数亿向量，直接upsert的效率极低。因此，Ward将工作流拆分为两个阶段：第一阶段仅将文本记录写入Parquet文件（values列为空列表），第二阶段读取这些文件、生成嵌入并写入新的Parquet文件。这种分离使得嵌入生成可以在多台机器上并行处理，支持Apple M系列、NVIDIA GPU等异构环境。
第二阶段使用PyArrow定义明确的Parquet模式，确保values列为LIST类型。每个记录包含id、空values和元数据JSON字符串。分区大小通常为10,000到20,000条记录，便于分发和错误重试。
第三阶段是嵌入生成脚本，它自动检测设备（使用mlx_embedding_models在Apple Silicon上运行，sentence-transformers在CUDA上，或CPU回退），读取Parquet文件中的文本，生成嵌入并填充values列。输出文件可直接用于Pinecone批量导入。
批量导入时，将Parquet文件复制到对象存储（如Amazon S3），然后运行Pinecone的批量导入命令。对于小规模测试，可以直接upsert；对于大规模数据集，批量导入更高效。
Ward总结了几点经验：保持工作流模块化；明确测试目标（延迟/吞吐量 vs. 召回率）；注意元数据大小限制和空值处理。整个流程可重复使用，适用于不同规模、模型和硬件配置。