初めに
RAG の回答を速くしたくて、検索処理を調べてみました。
時間がかかっていたのは、LLM ではなく Oracle AI Database に送る SQL。ベクトル検索 1 本に約 9 秒かかっていました。
HNSW 索引は作成済みです。でも、実行計画を見るとチャンクの全件走査になっていました。修正 PR #2804
原因は SQL の組み立て方だった
当時の SQL には、次の 3 つが重なっていました。
| SQL の書き方 | 確認した問題 |
|---|---|
文書の所属・有効期間などの条件を同じ WHERE に並べる |
文書側の副問い合わせが行ごとの FILTER になる |
| 距離の後に文書 ID・チャンク番号も並べる | 近似検索で索引を使うための並び方を崩す |
| 近傍検索の内側で類似度の下限を判定する | 全件走査の計画になることがある |
検索結果を安定させるための並び順や、必要な絞り込みが、検索を重くしていました。
修正は「近傍を探す処理」と「結果を整える処理」を分ける
修正後の考え方を、簡略化した SQL で示します。コメント部分の条件は省略しています。
SELECT *
FROM (
SELECT
c.chunk_id,
c.document_id,
c.chunk_index,
VECTOR_DISTANCE(c.embedding, :embedding, COSINE) AS distance
FROM rag_chunks c
WHERE c.document_id IN (
SELECT d.document_id
FROM rag_documents d
WHERE /* 所属・権限・有効期間などの文書条件 */
)
AND /* チャンク側の条件 */
ORDER BY VECTOR_DISTANCE(c.embedding, :embedding, COSINE)
FETCH APPROX FIRST 10 ROWS ONLY
)
WHERE 1 - distance >= :min_similarity
ORDER BY distance, document_id, chunk_index, chunk_id;
ポイントは 3 つです。
- 文書側の条件を分け、検索対象の文書集合を求める形にする。
- 内側の
ORDER BYはベクトル距離だけにする。 - 類似度の下限と同距離の並べ替えは、候補取得後の外側で行う。
権限や検索範囲の条件は、候補を取得する前に適用します。
なお、外側の並べ替えで安定するのは取得済み候補の順序です。上位件数の境界に同距離の候補がある場合、その候補集合まで固定するものではありません。
また、当時は複数の検索文を CTE で 1 本の SQL にまとめた場合も全件走査になったため、検索文ごとに SQL を分けました。実装差分
どれくらい変わったか
当時の検証環境では、ベクトル検索 1 本が次のように改善しました。
| 修正前 | 修正後 |
|---|---|
| 8.8〜9.3 秒 | 0.6〜0.8 秒 |
実際の質問 5 件 × 条件 2 種の計 10 パターンでは、上位 10 件が修正前後で一致しました。これは、その検証条件での結果です。検証記録
今回学んだのは、索引を作るだけでなく、実際に発行する SQL の実行計画まで見ることでした。
RAG が遅いときは、LLM の呼び出し時間と DB の検索時間を分けて測る。SQL が支配的なら、索引の利用状況、行ごとの副問い合わせ、並び順から確認すると、原因を絞り込みやすくなります。