HNSW インデックスに基づく pgvector 拡張機能を使用したパフォーマンステスト - ApsaraDB RDS

このトピックでは、pgvector 拡張機能を使用して、階層型ナビゲーション可能な小規模ワールド (HNSW) インデックスに基づく ApsaraDB RDS for PostgreSQL インスタンスのパフォーマンスをテストする方法について説明します。 pgvector 拡張機能は、ANN-Benchmarks ツールを使用して、再現率、1 秒あたりのクエリ数 (QPS)、インデックス作成時間などの主要なメトリックを評価します。

テスト環境

RDS インスタンスと Elastic Compute Service (ECS) インスタンスは、ネットワークの変動によるエラーを防ぐために、同じ仮想プライベートクラウド (VPC) 内に存在し、同じ vSwitch に属している必要があります。

テストインスタンスとテストツール	説明
RDS インスタンスの仕様	RDS インスタンスは PostgreSQL 16 を実行しています。 RDS インスタンスは、RDS High-availability Edition を実行し、pg.x8.2xlarge.2c 専用インスタンスタイプを使用する標準 RDS インスタンスです。インスタンスタイプは 16 コアと 128 GB のメモリを提供します。 pgvector 拡張機能はバージョン 0.8.0 を実行しています。
ECS インスタンスの仕様	ECS インスタンスは、4 コアと 8 GiB のメモリを提供する ecs.c6.xlarge インスタンスタイプを使用します。 ECS インスタンスは Alibaba Cloud Linux 3 を実行しています。
テストツール	ANN-Benchmarks 重要デフォルトでは、ANN-Benchmarks ツールはシングルスレッドのパフォーマンスをテストします。同時実行パフォーマンスをテストする場合は、「IVF インデックスに基づく pgvector 拡張機能を使用したパフォーマンステスト」で提供されている手順に従うことができます。

準備

RDS インスタンス

ann_testuser という名前の特権アカウントと ann_testdb という名前のデータベースを作成します。詳細については、「データベースとアカウントを作成する」をご参照ください。
ann_testdb データベースに pgvector 拡張機能をインストールします。 pgvector 拡張機能は、システムでは vector という名前です。詳細については、「拡張機能を管理する」をご参照ください。

ECS インスタンス

Docker をインストールします。詳細については、「Docker をインストールする」をご参照ください。
次のコマンドを実行して、ANN-Benchmarks ツールをダウンロードします。
```
cd ~
git clone https://github.com/erikbern/ann-benchmarks.git
```
次のコマンドを実行して、Conda を使用して ann_test という名前の仮想環境を作成してアクティブ化し、Python 3.10.6 をインストールします。
```
yum install git
yum install conda
conda create -n ann_test python=3.10.6
conda init bash
source /usr/etc/profile.d/conda.sh 
conda activate ann_test
```
次のコマンドを実行して、ANN-Benchmarks ツールの依存関係をインストールします。
```
cd ~/ann-benchmarks/
pip install -r requirements.txt
```

テスト手順

重要

テスト手順のすべての手順は、ann_test 仮想環境で実行されます。タイムアウトなどのエラーのためにログアウトを強制された場合は、conda activate ann_test コマンドを実行して、環境に再度ログインできます。

ステップ 1: テストツールの接続情報を構成する

次の接続設定をテストツールの ~/ann-benchmarks/ann_benchmarks/algorithms/pgvector/module.py ファイルに追加し、ビジネス要件に基づいて設定を構成します。

# RDS インスタンスに接続するためのパラメータを構成します。
os.environ['ANN_BENCHMARKS_PG_USER'] = 'ann_testuser'     # RDS インスタンスへの接続に使用するアカウントのユーザー名を指定します。
os.environ['ANN_BENCHMARKS_PG_PASSWORD'] = 'testPawword'  # RDS インスタンスへの接続に使用するアカウントのパスワードを指定します。
os.environ['ANN_BENCHMARKS_PG_DBNAME'] = 'ann_testdb'     # RDS インスタンスで必要なデータベースの名前を指定します。
os.environ['ANN_BENCHMARKS_PG_HOST'] = 'pgm-****.pg.rds.aliyuncs.com'  # RDS インスタンスの内部エンドポイントを指定します。
os.environ['ANN_BENCHMARKS_PG_PORT'] = '5432'             # RDS インスタンスの内部ポートを指定します。
os.environ['ANN_BENCHMARKS_PG_START_SERVICE'] = 'false'   # 自動起動を無効にします。

ステップ 2: ANN-Benchmarks ツールのテストパラメータを構成する

ビジネス要件に基づいて、テストツールの ~/ann-benchmarks/ann_benchmarks/algorithms/pgvector/config.yml ファイルを変更します。例:

float:
  any:
  - base_args: ['@metric']
    constructor: PGVector
    disabled: false
    docker_tag: ann-benchmarks-pgvector
    module: ann_benchmarks.algorithms.pgvector
    name: pgvector
    run_groups:
      M-16(100):
        arg_groups: [{M: 16, efConstruction: 100}]
        args: {}
        query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
      M-16(200):
        arg_groups: [{M: 16, efConstruction: 200}]
        args: {}
        query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
      M-24(200):
        arg_groups: [{M: 24, efConstruction: 200}]
        args: {}
        query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]

このテストは、M-16(100)、M-16(200)、および M-24(200) のグループで実行されます。各グループのテストでは、arg_groups を使用して HNSW インデックスの関連パラメータを作成し、query_args を使用して関連パラメータを取得します。

パラメータ		説明
arg_groups	M	HNSW インデックスのパラメータ M に対応し、HNSW インデックスを作成するときに各レイヤーの各ノードの隣接ノードの最大数を指定します。値が大きいほど、グラフ密度が高くなります。ほとんどの場合、グラフ密度が高いほど再現率は向上しますが、インデックスの作成とクエリに必要な時間も長くなります。
arg_groups	efConstruction	HNSW インデックスの ef_construction に対応し、HNSW インデックスを作成するときに候補セットのサイズを指定します。候補セットのサイズは、最適な接続を実現するために保持される候補ノードの数を指定します。ほとんどの場合、値が大きいほど再現率は向上しますが、インデックスの作成とクエリに必要な時間も長くなります。
query_args	ef_search	クエリ中に維持される候補セットのサイズを指定します。ほとんどの場合、値が大きいほど再現率は向上しますが、インデックスのクエリに必要な時間も長くなります。

ステップ 3: テスト Docker イメージを作成する

オプション。次の情報に基づいてテストツールの ~/ann-benchmarks/ann_benchmarks/algorithms/pgvector/Dockerfile ファイルを変更し、ツール用に自動的に構成される PostgreSQL コミュニティのテストをスキップします。
```
FROM ann-benchmarks
USER root
RUN pip install psycopg[binary] pgvector
```
--algorithm pgvector を指定し、install.py スクリプトを実行してテスト Docker イメージを作成します。
```
cd ~/ann-benchmarks/
python install.py --algorithm pgvector
```
説明
python install.py --help コマンドを実行して、サポートされている構成パラメータを表示できます。

ステップ 4: データセットを取得する

テストスクリプトを実行すると、指定されたパブリックテストデータセットが自動的にダウンロードされます。

このトピックでは、Angular 距離式を使用する nytimes-256-angular データセットが類似性検索に使用されます。詳細については、「ann-benchmarks」をご参照ください。

データセット	ディメンション	データ行数	テストベクターの数	上位 N 個の最近傍	距離式
NYTimes	256	290,000	10,000	100	Angular

説明

パブリックデータセットがテスト要件を満たしていない場合は、ベクトルデータを Hierarchical Data Format version 5 (HDF5) などの標準フォーマットに変換して、検索パフォーマンスをテストすることをお勧めします。詳細については、「付録 2: カスタムテストデータセット」をご参照ください。

ステップ 5: テストを実行して結果を取得する

次のコマンドを実行して、テストスクリプトを実行します。

cd ~/ann-benchmarks
nohup python run.py --dataset nytimes-256-angular -k 10  --algorithm pgvector --runs 1 > ann_benchmark_test.log 2>&1 &
tail -f ann_benchmark_test.log

パラメータ	説明
--dataset	テストするデータセット。
--k	SQL 文の LIMIT 値。返すレコードの数を指定します。
--algorithm	テスト対象のベクトルデータベースのアルゴリズム。この例では、パラメータは pgvector に設定されています。
--runs	最適な結果セットが選択されるテストの実行回数。
--parallelism	並列度。デフォルト値: 1。

次のコマンドを実行して、テスト結果を取得します。

cd ~/ann-benchmarks
python plot.py --dataset nytimes-256-angular --recompute

説明

次の表にテスト結果を示します。

m	ef_construction	ef_search	再現率	QPS
16	100	10	0.630	1423.985
			0.741	1131.941
			0.820	836.017
			0.871	574.733
			0.894	440.076
			0.918	297.267
			0.945	162.759
			0.969	84.268
16	200	10	0.683	1299.667
			0.781	1094.968
			0.849	790.838
			0.895	533.826
			0.914	405.975
			0.933	272.591
			0.956	148.688
			0.977	76.555
24	200	10	0.767	1182.747
			0.840	922.770
			0.887	639.899
			0.920	411.140
			0.936	303.323
			0.953	199.752
			0.973	105.506
			0.988	53.904

オプション。次のコマンドを実行して、再現率、QPS、応答時間 (RT)、インデックス作成時間、インデックスサイズなどのテスト結果の詳細を取得します。
```
cd ~/ann-benchmarks
python data_export.py --out res.csv
```
たとえば、m、ef_construction、およびインデックス作成時間の間の関係をクエリできます。
m
ef_construction
build (インデックス作成時間 (秒))
16
100
33.35161
16
200
57.66014
24
200
87.22608

テストの結論

HNSW インデックスを作成すると、次の結論が得られます。

m、ef_construction、および ef_search パラメータの値を大きくすると、再現率は向上しますが、QPS は低下します。
m および ef_construction パラメータの値を大きくすると、再現率は向上し、QPS は低下し、インデックス作成時間は長くなります。
再現率に高い要件がある場合は、インデックスパラメータのデフォルト値を使用しないことをお勧めします。 m、ef_construction、および ef_search パラメータのデフォルト値は、それぞれ 16、64、および 40 です。

付録 1: RDS インスタンスとベクターのパラメータ構成がインデックス作成に与える影響

ANN-Benchmarks ツールのさまざまなパラメータ値を構成し、テストを実行してから、テスト結果を分析して、RDS インスタンスとベクターのパラメータ構成がインデックス作成に与える影響を判断できます。

RDS インスタンスのパラメータ構成がインデックス作成に与える影響

ANN-Benchmarks ツールの m パラメータを 16 に、efConstruction パラメータを 64 に設定した場合、RDS インスタンスのパラメータ構成はインデックス作成に次のような影響を与えます。

maintenance_work_mem
このパラメータは、VACUUM や CREATE INDEX などのメンテナンス操作に使用できるメモリの最大量を指定します。単位: KB。 maintenance_work_mem パラメータの値がテストデータのサイズよりも小さい場合は、値を大きくしてインデックスの作成に必要な時間を短縮できます。 maintenance_work_mem パラメータの値がテストデータのサイズを超えても、インデックス作成時間は短縮されません。
たとえば、RDS インスタンスが 16 コアと 128 GB のメモリを提供する pg.x8.2xlarge.2c インスタンスタイプを使用し、max_parallel_maintenance_workers パラメータをデフォルト値 8 に設定し、nytimes-256-angular データセットに約 324 MB のデータが含まれているとします。この場合、maintenance_work_mem パラメータの構成は、インデックス作成に次のような影響を与えます。
maintenance_work_mem
インデックス作成時間 (単位: 秒)
64 MB (65,536 KB)
52.82
128 MB (131,072 KB)
46.79
256 MB (262,144 KB)
36.40
512 MB (524,288 KB)
18.90
1 GB (1,048,576 KB)
19.06
max_parallel_maintenance_workers
このパラメータは、単一の CREATE INDEX 操作で起動できる並列ワーカーの最大数を指定します。 max_parallel_maintenance_workers パラメータの値が大きいほど、インデックス作成時間は短縮されます。
たとえば、RDS インスタンスが 16 コアと 128 GB のメモリを提供する pg.x8.2xlarge.2c インスタンスタイプを使用し、maintenance_work_mem パラメータをデフォルト値 2048 (2 GB に相当) に設定し、nytimes-256-angular データセットに約 324 MB のデータが含まれているとします。この場合、max_parallel_maintenance_workers パラメータの構成は、インデックス作成に次のような影響を与えます。
max_parallel_maintenance_workers
インデックス作成時間 (単位: 秒)
1
76.00
2
51.34
4
32.49
8
19.66
12
14.44
16
13.07
24
13.15

ベクターのパラメータ構成がインデックス作成に与える影響

RDS インスタンスの maintenance_work_mem パラメータを 8 GB (8,388,608 KB に相当) に、max_parallel_maintenance_workers パラメータを 16 に設定した場合、RDS インスタンスのパラメータ構成はインデックス作成に次のような影響を与えます。

ベクターディメンション
1,183,514 行のデータを含む GloVe データセットが使用されます。 ANN-Benchmarks ツールの m パラメータを 16 に、efConstruction パラメータを 64 に、ef_search パラメータを 40 に設定すると、次の結果が返されます。テスト結果から、ベクターディメンションが大きくなるにつれて、インデックス作成時間が長くなり、再現率と QPS が低下し、クエリレイテンシが長くなることがわかります。
ディメンション
インデックス作成時間 (単位: 秒)
再現率
QPS
P99 (ミリ秒)
25
195.10
0.99985
192.94
7.84
50
236.92
0.99647
152.36
9.69
100
319.36
0.97231
126.89
11.14
200
529.33
0.93186
95.05
15.11
説明
P99: すべてのクエリリクエストの RT を昇順にソートした後のレイテンシの 99 パーセンタイル。これは、すべてのクエリリクエストの 99% の RT がこの値よりも低いことを示しています。

ベクターの行数

dbpedia-openai-{n}k-angular データセットが使用され、n で指定されるベクターの行数は 100 から 1,000 の範囲です。 ANN-Benchmarks ツールの m パラメータを 48 に、efConstruction パラメータを 256 に、ef_search パラメータを 200 に設定すると、次の結果が返されます。テスト結果から、ベクターの行数が増えるにつれて、インデックス作成時間が非線形的に増加し、再現率と QPS が低下し、クエリレイテンシが長くなることがわかります。

ベクターの行数	行数 (10,000)	インデックス作成時間 (単位: 秒)	再現率	QPS	P99 (ミリ秒)
100	10	54.05 秒	0.9993	171.74	8.93
200	20	137.23	0.99901	146.78	10.81
500	50	436.68	0.999	118.55	13.94
1000	100	957.26	0.99879	101.60	16.35

付録 2: カスタムテストデータセット

次のデータセットは、nytimes-256-angular データセットのフォーマットに基づいて生成され、参照用にのみ提供されています。

次のコマンドを実行して、カスタムデータセットを作成します。

説明

この例では、rds_ai がインストールされています。詳細については、「rds_ai 拡張機能で提供される AI 機能を使用する」をご参照ください。

import h5py
import numpy as np
import psycopg2
import pgvector.psycopg2

# 仮定
conn_info = {
    'host': 'pgm-****.rds.aliyuncs.com',
    'user': 'ann_testuser',
    'password': '****',
    'port': '5432',
    'dbname': 'ann_testdb'
}

embedding_len = 1024
distance_top_n = 100
query_batch_size = 100

try:
    # RDS インスタンスに接続します。
    with psycopg2.connect(**conn_info) as connection:
        pgvector.psycopg2.register_vector(connection)
        with connection.cursor() as cur:
            # ベクトルデータを取得します。
            cur.execute("select count(1) from test_rag")
            count = cur.fetchone()[0]

            train_embeddings = []
            for start in range(0, count, query_batch_size):
                query = f"SELECT embedding FROM test_rag ORDER BY id OFFSET {start} LIMIT {query_batch_size}"
                cur.execute(query)
                res = [embedding[0] for embedding in cur.fetchall()]
                train_embeddings.extend(res)
            train = np.array(train_embeddings)

            # クエリデータを取得し、埋め込みを計算します。
            with open('query.txt', 'r', encoding='utf-8') as file:
                queries = [query.strip() for query in file]
            test = []
            # rds_ai 拡張機能をインストールするか、Alibaba Cloud Model Studio SDK を使用します。
            for query in queries:
                cur.execute(f"SELECT rds_ai.embed('{query.strip()}')::vector(1024)")
                test.extend([cur.fetchone()[0]])
            test = np.array(test)

    # 上位 N 個の最近傍と必要なノードの間の距離を計算します。
    dot_product = np.dot(test, train.T)
    norm_test = np.linalg.norm(test, axis=1, keepdims=True)
    norm_train = np.linalg.norm(train, axis=1, keepdims=True)
    similarity = dot_product / (norm_test * norm_train.T)
    distance_matrix = 1 - similarity

    neighbors = np.argsort(distance_matrix, axis=1)[:, :distance_top_n]
    distances = np.take_along_axis(distance_matrix, neighbors, axis=1)

    with h5py.File('custom_dataset.hdf5', 'w') as f:
        f.create_dataset('distances', data=distances)
        f.create_dataset('neighbors', data=neighbors)
        f.create_dataset('test', data=test)
        f.create_dataset('train', data=train)
        f.attrs.update({
            "type": "dense",
            "distance": "angular",
            "dimension": embedding_len,
            "point_type": "float"
        })

    print("HDF5 ファイルが作成され、データセットが追加されました。")

except (Exception, psycopg2.DatabaseError) as error:
    print(f"エラー: {error}")

~/ann-benchmarks/ann_benchmarks/datasets.py ファイルの DATASETS セクションに、カスタムデータセットを追加します。
```
DATASETS: Dict[str, Callable[[str], None]] = {
  ......,
  "<custom_dataset>": None,
}
```
カスタムデータセットを ~/ann-benchmarks ディレクトリにアップロードし、データセットを使用して run.py テストスクリプトを実行します。

m	ef_construction	build (インデックス作成時間 (秒))
16	100	33.35161
16	200	57.66014
24	200	87.22608

maintenance_work_mem	インデックス作成時間 (単位: 秒)
64 MB (65,536 KB)	52.82
128 MB (131,072 KB)	46.79
256 MB (262,144 KB)	36.40
512 MB (524,288 KB)	18.90
1 GB (1,048,576 KB)	19.06

max_parallel_maintenance_workers	インデックス作成時間 (単位: 秒)
1	76.00
2	51.34
4	32.49
8	19.66
12	14.44
16	13.07
24	13.15

ディメンション	インデックス作成時間 (単位: 秒)	再現率	QPS	P99 (ミリ秒)
25	195.10	0.99985	192.94	7.84
50	236.92	0.99647	152.36	9.69
100	319.36	0.97231	126.89	11.14
200	529.33	0.93186	95.05	15.11