PineconeのソリューションエンジニアであるJohn Wardは、過去1年間にプラットフォームのさまざまな部分をテストするための内部ツールやサンプルを開発してきました。本記事はそのシリーズの第一弾であり、テストデータの生成に焦点を当てています。これは一見単純ですが、規模が大きくなると重要になる作業です。
テストデータの生成は、さまざまなテストタイプの要件に応じて異なるデータセットが必要になることから始まりました。例えば、ベクトル検索のテストでは、ランダムベクトルは書き込みスループットやクエリレイテンシのテストにしか使えませんが、再現率や精度のテストには実際のテキストに基づく埋め込みが必要です。そこでWardは、Parquetファイル生成器、ベクトル埋め込み生成器、ローカルLLMを使用したParquet分類器、およびカテゴリを名前空間としてPineconeの一括インポート用に再パーティションするツールを含む一連のユーティリティを設計しました。
ワークフローの最初のステップは、ソースデータを生成し、埋め込みとインポートに備えることです。WardはHugging Faceのstanford-oval/ccnewsデータセットを選択しました。これは広範なニュース記事と、タイトル、著者、ソースURL、公開日、カテゴリなどの有用なメタデータを提供します。データはParquet形式で保存され、Pineconeの一括インポート形式に準拠しています。IDフィールドにはUUID、valuesフィールドには密ベクトル、metadataフィールドにはJSON文字列が含まれます。
埋め込みモデルとしては、BAAI/bge-large-en-v1.5を使用し、1024次元の密ベクトルを生成します。Pineconeホスト型の埋め込みモデルも利用可能ですが、このプロジェクトではローカルでの埋め込み生成を試し、顧客が自分で埋め込みを生成してPineconeにロードする一般的なワークフローをシミュレートしました。
最初のバージョンのスクリプトは、CC Newsデータセットをストリーム処理し、記事テキストをチャンク分割し、ローカルで埋め込みを生成し、直接Pineconeインデックスにアップサートします。これは小規模なテストには便利ですが、メタデータサイズの制限（40KB）やnull値の処理など、いくつかの注意点があります。スクリプトは長いテキストを反復的に切り詰め、null値をフィルタリングするロジックを含んでいます。
しかし、数百万または数十億のベクトルになると、直接アップサートは非効率的です。そこでWardはワークフローを2段階に分割しました。最初の段階ではテキストレコードをParquetファイルに書き込むだけ（values列は空のリスト）、次の段階でこれらのファイルを読み込み、埋め込みを生成し、新しいParquetファイルにvalues列を追加します。これにより、埋め込み生成を複数のマシンに分散でき、Apple Mシリーズ、NVIDIA GPUなど異なるハードウェア環境を活用できます。
第2段階ではPyArrowを使用してParquetスキーマを明示的に定義し、values列がLIST型であることを保証します。各レコードにはid、空のvalues、およびJSON文字列のメタデータが含まれます。パーティションサイズは通常10,000〜20,000レコードで、分散処理とエラーリトライが容易になります。
第3段階は埋め込み生成スクリプトで、デバイスを自動検出し（Apple Siliconではmlx_embedding_models、CUDAではsentence-transformers、それ以外はCPU）、Parquetファイルのテキストから埋め込みを生成してvalues列を埋めます。出力ファイルはPineconeの一括インポートに直接使用できます。
一括インポートでは、ParquetファイルをAmazon S3などのオブジェクトストレージにコピーし、Pineconeの一括インポートを実行します。小規模テストには直接アップサート、大規模データセットには一括インポートが効率的です。
Wardはいくつかの教訓を挙げています：ワークフローをモジュール化すること、テスト目的を明確にすること（レイテンシ/スループット vs. 再現率）、メタデータサイズ制限とnull値に注意すること。このプロセスは反復可能であり、さまざまな規模、モデル、ハードウェア構成に対応できます。