MilvusとPythonによる大規模セマンティック検索の構築:数百万件のベクトルに対応するデータパーティショニングとクエリ最適化

AI tutorial - IT technology blog
AI tutorial - IT technology blog

現実世界のボトルネック:10万ベクトルを超えるセマンティック検索のスケーリング

ローカルのサンドボックス環境ではすべてが順調に動いていました。しかし、ステージングクラスタに500万件のエンタープライズドキュメントを投入した瞬間、検索パイプラインは崩壊しました。かつては15ミリ秒で返っていた単純な類似度検索が突如として停滞し、p99レイテンシは1,200ミリ秒を超えて急上昇。ワーカーノードは瞬く間にメモリ枯渇(OOM)に陥りました。

多くのチュートリアルでは、数千件程度のトイデータセットを使ってセマンティック検索を解説しています。その規模であれば、全探索(Flatなブルートフォーススキャン)や最適化されていないインメモリインデックスでも問題なく動作します。しかし、本番システムでは話が別です。数百万件規模のレコードへのスケールには、主に3つの根深い障壁が存在します。

  • メモリの肥大化: 1536次元のfloat32埋め込みベクトル(OpenAIのtext-embedding-3-smallなど)は、1ドキュメントあたり約6.14 KBを消費します。1,000万件のベクトルでは、グラフのリンク構造やインデックスのオーバーヘッドを考慮する前の生データだけで60 GB以上のRAMが必要になります。
  • 無駄なコレクションスキャン: マルチテナントデータを扱うアプリケーションにおいて、単一の巨大なベクトル空間全体を検索すると、ユーザーが閲覧権限を持たない無関係な数百万件の埋め込みベクトルに対しても無駄にCPUサイクルを消費してしまいます。
  • インデックス構築によるリソース枯渇: 高スループットなリアルタイム更新を取り込みながら同時に密な近傍グラフを構築すると、CPU使用率が100%に跳ね上がり、アクティブユーザーのクエリタイムアウトを引き起こします。

こうした規模の壁を乗り越えるには、専用設計されたベクトルデータベース、自動化された物理パーティショニング、そして現実のメモリとクエリの制約に合わせてチューニングされたインデックス構成が不可欠です。

主要な概念:なぜMilvus、Partition Keys、HNSWが重要なのか

Milvusは、ストレージとコンピュートの分離アーキテクチャを中心に構築されたオープンソースの分散ベクトルデータベースです。リレーショナルエンジンに無理やりベクトル演算を追加するのではなく、Milvusはコレクションを個別の物理セグメントに分割し、ハードウェアレベルのSIMDやGPUアクセラレーションを活用して類似度を計算します。

1. 物理パーティション vs. Partition Keys

Milvusで手動パーティションを作成する方法はシンプルな構成には有効ですが、アプリケーションコード内で何千もの個別パーティションタグを管理するのは運用上の大きな負担となります。Milvus 2.2以降では、Partition Keysによってこの課題が解決されました。

tenant_idやworkspace_idなどのスカラーフィールドにis_partition_key=Trueを設定します。これにより、Milvusはそのスカラー値を自動的にハッシュ化し、基盤となる固定プール(通常64〜128個)の物理パーティションにレコードを割り振ります。クエリ実行時には、テナントフィルターを渡すだけでベクトル検索が該当セグメントに直接ルーティングされ、不要なベクトル比較を90%以上削減できます。

2. インデックスのトレードオフ:IVF_FLAT vs. HNSW

近似最近傍探索(ANN)アルゴリズムの選定は、レイテンシの目標値とハードウェアコストのバランスをどう取るかに帰着します。

  • IVF_FLAT(Inverted File Flat): ベクトル空間をボロノイ領域にクラスタリングします。生の埋め込みベクトルに加える追加のRAMオーバーヘッドはほぼゼロです。しかし、高い再現率(Recall)を得るには複数のクラスタをスキャンする必要があるため(高いnprobe値)、スループットが増加するにつれてクエリレイテンシが急激に悪化します。
  • HNSW(Hierarchical Navigable Small World): 上位層で高速なルーティングを行い、下位層で局所的な近傍関係を捉える多層グラフを構築します。98%以上の再現率を維持しながら、クエリ時間を安定して10ミリ秒未満に抑えられます。そのトレードオフとして、ノードの接続リストを保持するために生ベクトルサイズの1.2〜1.5倍の追加RAMが必要となります。

高いクエリスループットと20ミリ秒未満のp99 SLAが求められる場合、HNSWとPartition Keysの組み合わせが本番環境において最も信頼性の高い構成となります。

実践演習:パイプラインの構築とチューニング

以下に完全なセットアップ手順を示します。Milvus Standaloneのデプロイ、Pythonでのパーティションキー対応スキーマの定義、HNSWインデックスの構築、そしてフィルター付きベクトル検索の実行までを行います。

ステップ 1: Docker ComposeによるMilvus Standaloneの起動

この設定をdocker-compose.ymlとして保存します:

version: '3.5'

services:
  etcd:
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
    volumes:
      - ./volumes/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls=http://0.0.0.0:2379 --data-dir=/etcd

  minio:
    container_name: milvus-minio
    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
    environment:
      MINIO_ACCESS_KEY: minioadmin
      MINIO_SECRET_KEY: minioadmin
    volumes:
      - ./volumes/minio:/minio_data
    command: minio server /minio_data

  standalone:
    container_name: milvus-standalone
    image: milvusdb/milvus:v2.3.4
    command: ["milvus", "run", "standalone"]
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ./volumes/milvus:/var/lib/milvus
    ports:
      - "19530:19530"
      - "9091:9091"
    depends_on:
      - "etcd"
      - "minio"

スタックを起動します:

docker compose up -d
docker compose ps

ステップ 2: Partition Keysを使用したスキーマの定義

必要なPythonの依存パッケージをインストールします:

pip install pymilvus sentence-transformers numpy

接続の初期化とスキーマ定義を行うmilvus_semantic_search.pyを作成します:

import numpy as np
from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    utility
)

# Milvusサーバーへ接続
connections.connect(alias="default", host="localhost", port="19530")

COLLECTION_NAME = "enterprise_docs"
DIMENSION = 384  # all-MiniLM-L6-v2の標準出力次元数

if utility.has_collection(COLLECTION_NAME):
    utility.drop_collection(COLLECTION_NAME)

# tenant_idをパーティションキーとしてフィールドを設定
fields = [
    FieldSchema(name="doc_id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="tenant_id", dtype=DataType.VARCHAR, max_length=64, is_partition_key=True),
    FieldSchema(name="doc_category", dtype=DataType.VARCHAR, max_length=64),
    FieldSchema(name="raw_text", dtype=DataType.VARCHAR, max_length=2048),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]

schema = CollectionSchema(fields=fields, description="テナントごとにパーティション分割されたセマンティック知識ベース")
collection = Collection(name=COLLECTION_NAME, schema=schema, num_partitions=64)
print(f"コレクション '{COLLECTION_NAME}' を64個のハッシュパーティションで作成しました。")

ステップ 3: HNSWインデックスの設定

HNSWの構築パフォーマンスと再現率の精度は、主に2つのパラメータによって制御されます:

  • M(接続数の上限、通常は16〜64): 各ノードからの最大出リンク数。値を大きくするとノイズの多い埋め込みでも精度が向上しますが、メモリ消費量が増加します。通常はM=16が最適なスイートスポットです。
  • efConstruction(構築時の探索深度、64〜512): ベクトル挿入時の探索深度を制御します。efConstruction=200に設定することで、構築時間とインデックス品質のバランスを良好に保てます。
index_params = {
    "metric_type": "COSINE",
    "index_type": "HNSW",
    "params": {
        "M": 16,
        "efConstruction": 200
    }
}

collection.create_index(field_name="vector", index_params=index_params)
print("HNSWインデックスの構築が完了しました。")

ステップ 4: データの取り込みとパーティション分離検索

サンプルドキュメントをベクトル化し、複数のテナントIDにまたがって挿入した上で、対象を絞り込んだクエリを実行してみましょう:

from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("all-MiniLM-L6-v2")

sample_docs = [
    {"tenant": "tenant_a", "cat": "network", "text": "エンタープライズエッジゲートウェイでのBGPルーティングの設定。"},
    {"tenant": "tenant_a", "cat": "security", "text": "リモートエンジニア向けMFA認証ポリシーの設定。"},
    {"tenant": "tenant_b", "cat": "database", "text": "PostgreSQLのコネクションプール最適化とクエリキャッシング戦略。"},
    {"tenant": "tenant_b", "cat": "network", "text": "マルチリージョンクラスタにおけるVPCピアリングのレイテンシデバッグ。"}
]

# データバッチの準備
tenant_ids = [d["tenant"] for d in sample_docs]
categories = [d["cat"] for d in sample_docs]
texts = [d["text"] for d in sample_docs]
vectors = encoder.encode(texts).tolist()

# データの取り込みとディスクへのフラッシュ
collection.insert([
    tenant_ids,
    categories,
    texts,
    vectors
])
collection.flush()

# 検索用にコレクションをRAMへロード
collection.load()

# tenant_bにスコープを絞って検索を実行
query_text = "マルチリージョンのVPCネットワーク遅延を解決するにはどうすればよいですか?"
query_vector = encoder.encode([query_text]).tolist()

search_params = {
    "metric_type": "COSINE",
    "params": {"ef": 64}  # efは実行時の探索深度とレイテンシのバランスを制御
}

results = collection.search(
    data=query_vector,
    anns_field="vector",
    param=search_params,
    limit=2,
    expr="tenant_id == 'tenant_b'",
    output_fields=["tenant_id", "doc_category", "raw_text"]
)

for hits in results:
    for hit in hits:
        print(f"スコア: {hit.distance:.4f} | テナント: {hit.entity.get('tenant_id')} | カテゴリ: {hit.entity.get('doc_category')} | テキスト: {hit.entity.get('raw_text')}")

運用上のベストプラクティス

本番環境で数百万件規模のベクトルを運用するには、綿密なリソース計画とクエリの規律が求められます。以下の原則を念頭に置いてください:

  • 検索フィルターには必ずパーティションキーを含める: キーを省略すると、すべてのセグメントに対するグローバルスキャンが発生します。1,000万ベクトルのベンチマークでは、テナントのパーティションキーでクエリを絞り込むことにより、p99レイテンシが185ミリ秒から14ミリ秒に短縮されました。
  • 実行時のefを動的に調整する: SLAに応じてefを適合させます。通常、ef=32またはef=64に設定すると、検索レイテンシを15ミリ秒未満に抑えつつ95%以上の再現率が得られます。速度よりも再現率の精度が優先されるタスクではef=128を使用してください。
  • 事前にメモリ容量を見積もる: RAMの必要容量は次の計算式で見積もります:総RAM容量 ≈ ベクトル数 × (次元数 × 4バイト + M × 8バイト) × 1.35(OSバッファやクエリ実行用として35%の安全マージンを含む)。
  • 単一レコード挿入ごとのcollection.flush()呼び出しを避ける: 小さなペイロードごとにフラッシュを実行すると、何百もの極小セグメントファイルが生成され、深刻なセグメントの断片化とクエリパフォーマンスの低下を招きます。代わりに、挿入データを2,000〜10,000件のバッチにバッファリングしてください。

Milvusのパーティションキーと適切にパラメータ調整されたHNSWインデックスを組み合わせることで、データが数千件から数千万件へと拡大しても、検索レイテンシを予測可能な範囲に維持できます。

Share: