Azure Cosmos DB for PostgreSQL で pgvector を使用する場合にパフォーマンスを最適化する方法

[アーティクル]
08/15/2024

適用対象: Azure Cosmos DB for PostgreSQL (PostgreSQL の Citus データベース拡張機能を利用)

pgvector 拡張機能により、オープンソースベクトルの類似性検索が PostgreSQL に追加されます。

この記事では、pgvector の制限事項とトレードオフについて説明し、パーティション分割、インデックス作成、検索の設定を使用してパフォーマンスを向上させる方法について説明します。

拡張機能自体の詳細については、「pgvector の基本」を参照してください。また、プロジェクトの公式 README を参照することもできます。

パフォーマンス

常にクエリプランを調査することから始める必要があります。クエリが適切な速度で終了する場合は、EXPLAIN (ANALYZE,VERBOSE, BUFFERS) を実行します。

EXPLAIN (ANALYZE, VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

実行に時間がかかりすぎるクエリの場合は、ANALYZE キーワードを削除することを検討してください。結果に含まれる詳細は少なくなりますが、すぐに提供されます。

EXPLAIN (VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

explain.depesz.com などのサードパーティのサイトは、クエリプランを理解するのに役立ちます。回答したほうがよいいくつかの質問は次のとおりです。

クエリは並列化されましたか?
インデックスは使用されましたか?
WHERE 句で部分インデックス定義と同じ条件を使用しましたか?
パーティション分割を使用する場合、不要なパーティションは取り除かれましたか?

ベクトルが長さ 1 に正規化されている場合 (OpenAI 埋め込みなど)。最適なパフォーマンスを得るには、内積 (<#>) の使用を検討する必要があります。

並列実行

Explain プランの出力で、Workers Planned と Workers Launched を探します (後者は、ANALYZE キーワードが使用された場合のみ)。 max_parallel_workers_per_gather PostgreSQL パラメータは、データベースが Gather および Gather Merge プランノードごとに起動できるバックグラウンドワーカーの数を定義します。この値を増やすと、インデックスを作成しなくても、正確な検索クエリが高速化される可能性があります。ただし、この値が大きい場合でも、データベースがプランを並列で実行することを決定しない場合があることに注意してください。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 3;

                                        QUERY PLAN
------------------------------------------------------------------------------------------
 Limit  (cost=4214.82..4215.16 rows=3 width=33)
   ->  Gather Merge  (cost=4214.82..13961.30 rows=84752 width=33)
         Workers Planned: 1
         ->  Sort  (cost=3214.81..3426.69 rows=84752 width=33)
               Sort Key: ((embedding <-> '[1,2,3]'::vector))
               ->  Parallel Seq Scan on t_test  (cost=0.00..2119.40 rows=84752 width=33)
(6 rows)

インデックス作成

インデックスが存在しない場合、拡張機能は完全な検索を実行するため、パフォーマンスは犠牲になりますが、完全な再現が可能になります。

近似ニアレストネイバー検索を実行するには、データにインデックスを作成します。これにより、実行パフォーマンスの呼び戻しが行われます。

可能な場合は、インデックスを作成する前に常にデータを読み込みます。このようにインデックスを作成する方が高速であり、結果として得られるレイアウトもより最適になります。

サポートされているインデックスの種類は 2 つあります。

フラット圧縮を使用した反転ファイル (IVVFlat)
階層ナビゲーション可能な小さい世界 (HNSW)

HNSW と比較して、IVFFlat インデックスのビルド時間は短縮され、使用するメモリは少ないですが、クエリパフォーマンスは低くなります (速度/リコールのトレードオフの点で)。

制限

列にインデックスを付けるには、列にディメンションを定義する必要があります。 col vector として定義された列にインデックスを付けようとすると、エラー ERROR: column does not have dimensions が発生します。
インデックスを作成できるのは、最大 2,000 個のディメンションを持つ列のみです。より多くのディメンションを持つ列にインデックスを付けようとすると、エラー ERROR: column cannot have more than 2000 dimensions for INDEX_TYPE index が発生します。ここで INDEX_TYPE は ivfflat または hnsw です。

2,000 を超える次元のベクトルを格納できますが、インデックスを付けることはできません。次元削減を使用して、制限内に収めることができます。または、インデックス作成なしで許容できるパフォーマンスを実現するために、Azure Cosmos DB for PostgreSQL でのパーティション分割やシャーディングに依存します。

フラット圧縮を使用した反転ファイル (IVVFlat)

ivfflat は、近似ニアレストネイバー (ANN) 検索のインデックスです。このメソッドでは、反転ファイルインデックスを使用して、データセットを複数のリストにパーティション分割します。 probes パラメーターは、検索されるリストの数を制御します。これにより、検索速度は遅くなりますが、検索結果の精度が向上します。

probes パラメーターがインデックス内のリストの数に設定されている場合、すべてのリストが検索され、検索は正確なニアレストネイバー検索になります。この場合、すべてのリストを検索することは、データセット全体に対してブルートフォース検索を実行することと同じであるため、プランナーはインデックスを使用しません。

インデックス作成方法では、k-means クラスタリングアルゴリズムを使用してデータセットを複数のリストにパーティション分割します。各リストには、特定のクラスター中心に最も近いベクトルが含まれています。検索中に、クエリベクトルがクラスターの中心と比較され、ニアレストネイバーが含まれる可能性が最も高いリストを判断します。 probes パラメーターが 1 に設定されている場合は、最も近いクラスター中心に対応するリストのみが検索されます。

インデックスオプション

実行するプローブの数に正しい値を選択することと、リストのサイズは、検索のパフォーマンスに影響する可能性があります。以下から始めることをお勧めします。

最大 100 万行のテーブルには rows / 1000 と等しい lists を使用し、より大きなデータセットには sqrt(rows) を使用します。
probes の場合、最大 100 万行のテーブルの場合は lists / 10 から開始し、より大きなデータセットの場合は sqrt(lists) から始めます。

lists の量は、インデックスの作成時に lists オプションを使用して定義されます。

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 5000);

プローブは、接続全体またはトランザクションごとに設定できます (トランザクションブロック内で SET LOCAL を使用)。

SET ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

BEGIN;

SET LOCAL ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, one probe

インデックス作成の進行状況

PostgreSQL 12 以降では、インデックス作成の進行状況を確認するために pg_stat_progress_create_index を使用できます。

SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

IVFFlat インデックスを構築するためのフェーズは次のとおりです。

initializing
performing k-means
assigning tuples
loading tuples

Note

進行状況の割合 (%) は、loading tuples フェーズ中にのみ設定されます。

階層ナビゲーション可能な小さい世界 (HNSW)

hnsw は、階層ナビゲーション可能な小さい世界アルゴリズムを使用した近似最近傍 (ANN) 検索のインデックスです。これは、ランダムに選択されたエントリポイントの周囲にグラフを作成して最も近い近傍を見つけることで機能し、グラフは複数のレイヤーで拡張され、各下位レイヤーにはより多くのポイントが含まれます。この多層グラフは、検索時に上部から開始され、クエリの最近傍を含む最下位レイヤーにヒットするまで絞り込まれます。

このインデックスの構築には IVFFlat よりも多くの時間とメモリが必要ですが、速度/リコールのトレードオフはより優れています。さらに、IVFFlat のようなトレーニング手順がないため、空のテーブルにインデックスを作成できます。

インデックスオプション

インデックスを作成するときは、次の 2 つのパラメータを調整できます。

m - レイヤーあたりの接続の最大数 (既定値は 16)
ef_construction - グラフ構築に使用される動的候補リストのサイズ (既定値は 64)

CREATE INDEX t_test_hnsw_l2_idx ON t_test USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);

クエリ中に、検索の動的候補リストを指定できます (既定値は 40)。

検索の動的候補リストは、接続全体またはトランザクションごとに設定できます (トランザクションブロック内で SET LOCAL を使用)。

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

BEGIN;

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, 40 candidates

インデックス作成の進行状況

PostgreSQL 12 以降では、インデックス作成の進行状況を確認するために pg_stat_progress_create_index を使用できます。

SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

HNSW インデックスを構築するためのフェーズは次のとおりです。

initializing
loading tuples

インデックスアクセス関数の選択

vector 型を使用すると、保存されているベクトルに対して 3 種類の検索を実行できます。クエリの実行時にデータベースでインデックスを考慮するには、インデックスの正しいアクセス関数を選択する必要があります。例ではインデックスの種類 ivfflat について示していますが、hnsw インデックスに対しても同じことができます。 lists オプションは ivfflat インデックスにのみ適用されます。

コサイン距離

コサイン類似性検索の場合は、vector_cosine_ops アクセス方法を使用します。

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

上記のインデックスを使用するには、クエリでコサイン類似性検索を実行する必要があります。これは、<=> 演算子を使用して実行されます。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5;

                                              QUERY PLAN
------------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_cosine_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <=> '[1,2,3]'::vector)
(3 rows)

L2 距離

L2 距離 (ユークリッド距離とも呼ばれる) の場合は、vector_l2_ops アクセス方法を使用します。

CREATE INDEX t_test_embedding_l2_idx ON t_test USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

上記のインデックスを使用するには、クエリで L2 距離検索を実行する必要があります。これは、<-> 演算子を使用して実行されます。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_l2_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <-> '[1,2,3]'::vector)
(3 rows)

内積

内積類似性の場合は、vector_ip_ops アクセス方法を使用します。

CREATE INDEX t_test_embedding_ip_idx ON t_test USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);

上記のインデックスを使用するには、クエリで内積類似性検索を実行する必要があります。これは、<#> 演算子を使用して実行されます。

EXPLAIN SELECT * FROM t_test ORDER BY embedding <#> '[1,2,3]' LIMIT 5;

                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_ip_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <#> '[1,2,3]'::vector)
(3 rows)

部分インデックス

一部のシナリオでは、データの一部のセットのみを対象とするインデックスを作成することをお勧めします。たとえば、Premium ユーザーに対してのみインデックスを作成することができます。

CREATE INDEX t_premium ON t_test USING ivfflat (vec vector_ip_ops) WITH (lists = 100) WHERE tier = 'premium';

これで、Premium レベルでインデックスが使用されていることがわかります。

explain select * from t_test where tier = 'premium' order by vec <#> '[2,2,2]';

                                     QUERY PLAN
------------------------------------------------------------------------------------
 Index Scan using t_premium on t_test  (cost=65.57..25638.05 rows=245478 width=39)
   Order By: (vec <#> '[2,2,2]'::vector)
(2 rows)

Free レベルのユーザーには利点がありません。

explain select * from t_test where tier = 'free' order by vec <#> '[2,2,2]';

                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44019.01..44631.37 rows=244941 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15395.25 rows=244941 width=39)
         Filter: (tier = 'free'::text)
(4 rows)

データのサブセットのみにインデックスが付けられているということは、インデックスがディスク上で占有するスペースが少なくなり、検索が高速になることを意味します。

部分インデックス定義の WHERE 句で使用されているフォームがクエリで使用されているフォームと一致しない場合、PostgreSQL はインデックスが安全に使用できるかどうかを認識できない可能性があります。このサンプルデータセットでは、レベル列の個別の値として、正確な値 'free'、'test' および 'premium' のみが含まれています。 tier LIKE 'premium' PostgreSQL を使用するクエリの場合でも、インデックスは使用されません。

explain select * from t_test where tier like 'premium' order by vec <#> '[2,2,2]';

                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44086.30..44700.00 rows=245478 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15396.59 rows=245478 width=39)
         Filter: (tier ~~ 'premium'::text)
(4 rows)

パーティション分割

パフォーマンスを向上させる 1 つの方法は、データセットを複数のパーティションに分割することです。今年またはおそらく過去 2 年間のデータのみを参照するのが当たり前のようなシステムが想像できます。このようなシステムでは、データを日付範囲でパーティション分割し、システムがクエリ対象の年で定義された関連するパーティションのみを読み取ることができる場合、パフォーマンスの向上を活用できます。

パーティションテーブルを定義してみましょう。

CREATE TABLE t_test_partitioned(vec vector(3), vec_date date default now()) partition by range (vec_date);

毎年のパーティションを手動で作成することも、Citus ユーティリティ関数 (Cosmos DB for PostgreSQL で使用可能) を使用することもできます。

    select create_time_partitions(
      table_name         := 't_test_partitioned',
      partition_interval := '1 year',
      start_from         := '2020-01-01'::timestamptz,
      end_at             := '2024-01-01'::timestamptz
    );

作成されたパーティションを確認します。

\d+ t_test_partitioned

                                Partitioned table "public.t_test_partitioned"
  Column  |   Type    | Collation | Nullable | Default | Storage  | Compression | Stats target | Description
----------+-----------+-----------+----------+---------+----------+-------------+--------------+-------------
 vec      | vector(3) |           |          |         | extended |             |              |
 vec_date | date      |           |          | now()   | plain    |             |              |
Partition key: RANGE (vec_date)
Partitions: t_test_partitioned_p2020 FOR VALUES FROM ('2020-01-01') TO ('2021-01-01'),
            t_test_partitioned_p2021 FOR VALUES FROM ('2021-01-01') TO ('2022-01-01'),
            t_test_partitioned_p2022 FOR VALUES FROM ('2022-01-01') TO ('2023-01-01'),
            t_test_partitioned_p2023 FOR VALUES FROM ('2023-01-01') TO ('2024-01-01')

パーティションを手動で作成するには:

CREATE TABLE t_test_partitioned_p2019 PARTITION OF t_test_partitioned FOR VALUES FROM ('2019-01-01') TO ('2020-01-01');

次に、クエリが、実際に使用可能なパーティションのサブセットに絞り込まれていることを確認します。たとえば、次のクエリでは、次の 2 つのパーティションへとフィルター処理されています。

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01';

                                                                  QUERY PLAN
-----------------------------------------------------------------------------------------------------------------------------------------------
 Append  (cost=0.00..58.16 rows=12 width=36) (actual time=0.014..0.018 rows=3 loops=1)
   ->  Seq Scan on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.00..29.05 rows=6 width=36) (actual time=0.013..0.014 rows=1 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
   ->  Seq Scan on t_test_partitioned_p2023 t_test_partitioned_2  (cost=0.00..29.05 rows=6 width=36) (actual time=0.002..0.003 rows=2 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.125 ms
 Execution Time: 0.036 ms

パーティションテーブルのインデックスを作成できます。

CREATE INDEX ON t_test_partitioned USING ivfflat (vec vector_cosine_ops) WITH (lists = 100);

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01' order by vec <=> '[1,2,3]' limit 5;

                                                                                         QUERY PLAN
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Limit  (cost=4.13..12.20 rows=2 width=44) (actual time=0.040..0.042 rows=1 loops=1)
   ->  Merge Append  (cost=4.13..12.20 rows=2 width=44) (actual time=0.039..0.040 rows=1 loops=1)
         Sort Key: ((t_test_partitioned.vec <=> '[1,2,3]'::vector))
         ->  Index Scan using t_test_partitioned_p2022_vec_idx on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.04..4.06 rows=1 width=44) (actual time=0.022..0.023 rows=0 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
         ->  Index Scan using t_test_partitioned_p2023_vec_idx on t_test_partitioned_p2023 t_test_partitioned_2  (cost=4.08..8.11 rows=1 width=44) (actual time=0.015..0.016 rows=1 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.167 ms
 Execution Time: 0.139 ms
(11 rows)

まとめ

おめでとうございます。pgvector で最高のパフォーマンスを実現するためのトレードオフ、制限事項、ベストプラクティスについて学習しました。

次の方法で共有

Azure Cosmos DB for PostgreSQL で pgvector を使用する場合にパフォーマンスを最適化する方法

パフォーマンス

並列実行

インデックス作成

制限

フラット圧縮を使用した反転ファイル (IVVFlat)

インデックスオプション

インデックス作成の進行状況

階層ナビゲーション可能な小さい世界 (HNSW)

インデックスオプション

インデックス作成の進行状況

インデックスアクセス関数の選択

コサイン距離

L2 距離

内積

部分インデックス

パーティション分割

まとめ

フィードバック

その他のリソース

次の方法で共有

Azure Cosmos DB for PostgreSQL で pgvector を使用する場合にパフォーマンスを最適化する方法

パフォーマンス

並列実行

インデックス作成

制限

フラット圧縮を使用した反転ファイル (IVVFlat)

インデックス オプション

インデックス作成の進行状況

階層ナビゲーション可能な小さい世界 (HNSW)

インデックス オプション

インデックス作成の進行状況

インデックス アクセス関数の選択

コサイン距離

L2 距離

内積

部分インデックス

パーティション分割

まとめ

フィードバック

その他のリソース

インデックスオプション

インデックスオプション

インデックスアクセス関数の選択