為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. 召回率）；注意後設資料大小限制和空值處理。整個流程可重複使用，適用於不同規模、模型和硬體配置。