1
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

[AWS]DynamoDBベクトル検索のTopK:100制限をシャーディングで限界突破してみる

1
Posted at

🎯 対象者

  • DynamoDBでベクトルデータを扱ってみたい人
  • でも一回で最大100件しか抽出できないから、集約もリランキングもできないと嘆いている人

🆕 DynamoDBでベクトルデータが扱えるようになった

DynamoDBでは、2026年8月にベクトルデータを扱えるように対応されました。この記事ではベクトルデータの取り扱いそのものについては触れません。詳細は公式記事をご参照ください。

🚧 TopK 100件問題

RAG検索界隈では、広く取ってアプリ側でリランキングや集約をして出典を出すという用途がポピュラーになってきています。ベクトルDBの代表格であるS3Vectorsでは2026年6月に、TopKを100件から1万件まで拡張するアップデートが入りました。これによって広く取って、アプリ側で様々な加工ができるようになったため、柔軟な検索が行えるようになっています。

fig_s3vectors_topk.png

一方DynamoDBベクトルデータではTopK(抽出できる件数)の最大は100件です:sweat:。100件抽出したところで1ドキュメントのチャンクが100件内を占有していた場合、出典にドキュメントを複数表示できなくなるという課題があります。

fig_topk100_problem.png

こちらの記事でも謳われている通り。重複する内容(例えばドキュメント単位情報)も取得したい場合は実質目減りするという事象が発生してしまいます。

📚 DynamoDBベクトルデータベースの構造のおさらい

まずはDynamoDBのベクトルデータベースの構造についておさらいしましょう。テーブルを作る際にベクトルデータの定義を行います。格納対象となる属性の指定と、ベクトルインデックスのパーティションキーを設定できます。パーティションキーを指定することでデータの格納位置が分離されパーティション内での検索が行えるという構造になっています。

fig_vector_index_structure.png

検索を行う際はパーティションキーを指定した検索を行うことでパーティションに絞った抽出が行えます。

🧩 シャーディングでTop:k 100件問題を超えられるか?

1回で100件しか取れないのであれば複数回に分けて取ればいいじゃない。ということでシャーディングしてみました。パーティションキーの値を複数に分けそれぞれを検索して結果を合わせる作戦です。

DynamoDBのベクトルインデックスもスループットの上限はパーティションキーの値ごとに決まっています。そこでハッシュ値を使って均等にシャーディングします。

fig_vector_index_sharded.png

パーティションキーの指定

パーティションキーにはDynamoDBのIDのハッシュを使います。ハッシュを使って、パーティションのグループ1から5に均等に割り振ります。再作成をする際にもハッシュを使えばまた同じグループに割り振られるため便利です。

common.py
def shard_index_of(chunk_id: str) -> int:
    digest = hashlib.sha256(chunk_id.encode("utf-8")).hexdigest()
    return int(digest, 16) % N_SHARDS

データの抽出

データの抽出は並列で行います。パーティションキーを1から5にそれぞれ設定して並列で投げます。

scatter_gather.py
    def one(shard_key: str) -> dict:
        return search_shard(
            client,
            vector,
            table_name=table_name,
            index_name=index_name,
            shard_key=shard_key,
            top_k=top_k,
            search_vector=search_vector,
            **kwargs,
        )

    t0 = time.perf_counter()
    if parallel:
        workers = max_workers or len(keys)
        with ThreadPoolExecutor(max_workers=workers) as pool:
            per_shard = list(pool.map(one, keys))
    else:
        per_shard = [one(key) for key in keys]

データ抽出後の結合

抽出したらデータを全て結合した上で、結果をスコアで並べ替えます。この時点でTopK:500で取得した状態と同じ状態になります。

scatter_gather.py
def sort_key(record: dict) -> tuple[float, str]:
    """検索結果の並び順。距離が同じときは chunk_id の辞書順で決める。
    """
    return (float(record["distance"]), record["chunk_id"])


def merge(*result_lists: Iterable[dict]) -> list[dict]:
    """シャードごとの結果を 1 本にまとめる(重複排除 → (distance, chunk_id) 昇順)。
    """
    flat: list[dict] = []
    for one in result_lists:
        flat.extend(one)
    return sorted(dedupe(flat), key=sort_key)

:horse_racing: いざ検証!

fig_verification_architecture.png

検証構成はローカルにRAG検索の画面を立て、APIキー付きでバックエンドを呼び出して回答と出典を受け取ります。同じ質問を単一検索とシャーディング検索の両方で投げ出典の数が増えるかを見ました。DynamoDBのデータはどちらも同一で、ベクトルインデックスだけシャードキーあり・なしの2本を同じテーブルに張っています。

シングルシャード検索 topk:100(従来の検索)

image.png

複数シャード検索 topk:500(今回対応した検索)

image.png

うまくいきました:open_hands:候補チャンクが500件取得できています。また、ドキュメントの出典の件数も10件に増えていることがわかります。

⏱️ 並列検索の速度は呼び出し元のCPUで決まる

ベクトル検索のレート上限はパーティションキーの値ごとに決まっていて、値を分ければ合計の上限は線形に増える、と公式に書かれています。つまり5本のシャードへ同時に投げてもキャパシティを食い合うことはありません。

Throughput scales linearly across partition key values.

ならば5並列は1回分に近い時間で終わるはずですが、同一リージョンのLambdaから実際に測ったら見事に予想が外れました。

呼び出し側 単一(ms) 5並列(ms) 並列÷単一
Lambda 1,024 MB 23.5 119.9 5.1倍
Lambda 1,769 MB 17.6 79.1 4.5倍
Lambda 3,008 MB 15.7 40.6 2.6倍

1,024MBでは5並列に単一の5倍かかっていて並列にした意味がありません。
5本同時に投げればこの計算も5本分同時に走らせる必要があります。CPUが足りなければ順番待ちになり直列と変わらなくなります。つまり並列検索は単一検索よりも呼び出し元に高いスペックを要求します。

⚠️ シャーディングを使うときの注意点

シャード数を変えると全データの入れ替えが要る

chunk_id のハッシュをNで割った余りで振り分けているのでシャード数を変更する際は再登録が必要です。シャード割り振りの定数だけ変更しても運用は可能ですが、パーディションに偏りが出て抽出が不均一になるため望ましくないです。

検索コストがシャード数分かかる

検索の回数がシャード数ぶん増えるのでDynamoDBの課金バイトも増えます。実測では単一検索1回が112 KB、5シャード合計で367 KBと3.27倍でした。ただし月100万検索でも$0.24が$0.78になる程度なので大差はありません(東京、2026年9月時点の料金)。書き込みコストとストレージは属性が1つ増えるだけなので変わりません。

本来のTopk:500とは完全には一致しない。

fig_shard_topk_gap.png

各シャード振り分けによっては各シャードのtopk:100における足切りライン付近において本来のトップ計500と異なるパターンもあります。

山ちゃん(筆者)的には問題ない許容範囲だという認識です。なぜならこの事象が影響するのはTopKで抽出した下位の部分で、つまり全体でいうとTopK:500の下の部分に対するデータ欠落なので検索精度(LLM回答の体感)への影響は極小だからです。影響するのであればシャード数を増やせば解決します。

シャード数の決め方

TopKで何件取りたいかで決めます。300件取りたいなら3シャード、500件なら5シャードです。取った結果それでも出典が偏っているならシャードを足して取る件数を増やします。

🍜 締めr

今回の全体像
fig1_architecture.png

今回はTopK:100の制限に対して、本来は検索範囲を絞るための道具であるパーティションキーをシャーディングに転用し、取得件数を広げる荒技(?)で突破してみました。

広く取ってアプリで絞るといった典型的なRAG検索の手法をDynamoDBベクトルデータでも可能となるため、OpenSearchでリランキングしてるけどコストを抑えたい。S3Vectorとの二重管理をしたくないといったユースケースにはマッチするかもしれません。本音を言えばS3VectorのようにDynamoDBでもアップデートが入って一括取得もしくはページング取得できるようになれば...と思ったりしてます:relaxed:

1
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
1
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?