テキスト埋め込みをポータブルに扱う最良の方法はParquetとPolars
- テキスト埋め込みは大規模言語モデルによって生成されたベクトルで、単語、文、文書を数値的に表現する方法
- 2025年2月時点で、合計32,254件の「マジック：ザ・ギャザリング」カード埋め込みを生成
- これにより、カードのデザインおよびメカニクス上の特性に基づく類似性を数学的に分析できる
- 生成した埋め込みは2DのUMAP次元削減によって可視化できる
- 使用した埋め込みモデルは gte-modernbert-base で、詳細な手順は GitHubリポジトリ にまとめられている
- この埋め込みデータセットは Hugging Face で提供されている
ベクトルデータベースの必要性を見直す
- 一般的にはベクトルデータベース（faiss、qdrant、Pinecone）を使って埋め込みを保存・検索する
- しかしベクトルデータベースは設定が複雑で、クラウドサービスはコストが高くなることがある
- 小規模データ（数万件レベル）であれば、ベクトルデータベースなしでもnumpyを使って高速な類似度検索が可能
- numpyの
dot product 演算を使えばシンプルなコサイン類似度計算ができ、32,254件の埋め込みに対して平均1.08msで処理できる
def fast_dot_product(query, matrix, k=3):
dot_products = query @ matrix.T
idx = np.argpartition(dot_products, -k)[-k:]
idx = idx[np.argsort(dot_products[idx])[::-1]]
score = dot_products[idx]
return idx, score
- ベクトルデータベースを使うと、特定のライブラリやサービスに依存する可能性が高い
- GPUサーバーで埋め込みを生成してからローカルにダウンロードする場合、効率的なデータ保存・転送方式が必要になる
最悪の埋め込み保存方式
- CSVファイル
- 浮動小数点（
float32）データをテキストとして保存すると、サイズが6倍以上に増える
- OpenAIの公式チュートリアルでも、小さなデータセットにのみCSVの使用を推奨している
- numpyの
.savetxt() を使って保存すると、ファイルサイズは 631.5MB に増加する
- pickleファイル
- 高速に保存・読み込みできるが、セキュリティリスクがあり、バージョン互換性も低い
- ファイルサイズは 94.49MB で元のメモリサイズと同じだが、移植性は低い
悪くはないが最適ではない保存方式
- numpyの
.npy 形式
allow_pickle=False 設定によってpickle保存を防げる- ファイルサイズと速度はpickle方式と同じで、個別のメタデータを一緒に保存しにくい
- メタデータと分離された保存構造の問題点
- numpy配列（
.npy）で保存すると、カード情報（名前、テキストなど）と埋め込みが分離される
- データが変更（追加・削除）された場合、メタデータと埋め込みの対応付けが難しくなる
- ベクトルデータベースではメタデータとベクトルを一緒に保存し、フィルタリング機能も提供する
最適な埋め込み保存方式: Parquet + polars
Parquetファイル形式の紹介
- Apache Parquet はカラム指向のデータ保存形式で、各カラムのデータ型を明確に指定できる
- リスト形式（
float32 配列）のデータを保存できるため、埋め込みの保存に適している
- CSVより高速な保存・読み込み性能を提供し、一部のデータだけを選択的に読み込める
- 圧縮機能もあるが、埋め込みデータは重複が少ないため圧縮効果は小さい
PythonでのParquetファイル活用
- pandas を使ったParquetファイルの保存・読み込み:
df = pd.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])
df
- pandasはネストされたデータ（リスト）を効率的に処理できず、numpy
object に変換される
- numpy配列に変換する際に追加の演算（
np.vstack()）が必要となり、性能低下が起こりうる
- polars を使ったParquetファイルの保存・読み込み:
df = pl.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])
df
- polarsは
float32 配列をそのまま維持し、to_numpy() 呼び出し時に即座に2Dのnumpy配列を返せる
allow_copy=False 設定によって不要なデータコピーを防げる
embeddings = df["embedding"].to_numpy(allow_copy=False)
- 新しい埋め込みを追加する場合も、簡単にカラムを追加して保存できる
df = df.with_columns(embedding=embeddings)
df.write_parquet("mtg-embeddings.parquet")
Parquet + polarsを使った類似度検索とフィルタリング
- 特定条件を満たすデータだけをフィルタリングしたうえで、類似度検索を実行できる
- 例: 特定のカード（
query_embed）に似たカードを探す際に、タイプが 'Sorcery' で色に 'Black' を含むカードだけを検索する
df_filter = df.filter(
pl.col("type").str.contains("Sorcery"),
pl.col("manaCost").str.contains("B"),
)
embeddings_filter = df_filter["embedding"].to_numpy(allow_copy=False)
idx, _ = fast_dot_product(query_embed, embeddings_filter, k=4)
related_cards = df_filter[idx]
- 平均実行時間は 1.48ms で、全データ検索より37%遅いが、それでも十分高速
大規模ベクトルデータ処理のための代替案
- Parquetとdot product方式は 数十万件の埋め込み までは十分に処理できる
- より大きなデータセットを扱う場合は、ベクトルデータベースの利用が必要になる可能性がある
- 代替案として、SQLite ベースの sqlite-vec を使えば、追加のベクトル検索やフィルタリングも可能
結論
- ベクトルデータベースは必須ではない
- Parquet + polars の組み合わせは、埋め込みを効率的に保存・検索・フィルタリングできる強力な代替手段
- 特に小規模プロジェクトでは、Parquetファイルを活用するほうが高速でコスト効率も高い
- プロジェクトに応じて、Parquetとベクトルデータベースのどちらが適切かを選ぶことが重要
- GitHubリポジトリ でコードとデータを確認できる
1件のコメント
Hacker Newsのコメント
Parquetの問題は静的であること。継続的な書き込みや更新が必要な場合には向いていない。ただし、DuckDBとオブジェクトストレージ上のParquetファイルを使ったときは良い結果が得られた。ロード時間も速い
本当に素晴らしい記事。長い間あなたの仕事を楽しんできた。SQLite実装に飛び込む人向けに付け加えると、DuckDBはParquetを読み込み、このユースケースを完璧に扱ういくつかのベクトル類似度機能を提供し始めている
依然としてデータフレームはあまり好きではないが、Polarsはpandasよりはるかに良い
Unumのusearchを見てみてほしい。何にでも勝てて、とても使いやすい。必要なことを正確にこなしてくれる
試してみたいなら、HFから遅延ロードしてフィルタリングを適用できる
POLARS_MAX_THREADSを適用して、単一ノードの飽和度に応じて調整できる優れた発見がたくさんある
Vespaのドキュメントには、ベクトルを2進数に変換したあと16進表現を使う巧妙なテクニックがある
Polars + Parquetは可搬性と性能の両面で素晴らしい。この投稿はPythonの可搬性に重点を置いていたが、Polarsにはエンジンをさまざまな場所に組み込める使いやすいRust APIがある
Polarsの熱心なファンだが、これを使って埋め込みを保存する方法は考えたことがなかった（sqlite-vecで実験していた）。本当に興味深いアイデアだと思う
優れた性能と、全文インデックス作成や変更履歴のバージョン管理といった機能を備えた別のライブラリとして、lancedbを勧める