pgvector效能測試（基於HNSW索引） - ApsaraDB RDS

本文介紹RDS PostgreSQL執行個體pgvector外掛程式基於HNSW索引的效能測試，使用ann-benchmarks工具進行評估，涵蓋了召回率、QPS、索引構建時間等關鍵計量。

測試環境

RDS PostgreSQL執行個體與用戶端ECS應位於同一VPC和交換器，以避免網路波動帶來的誤差。

測試執行個體及工具	說明
RDS PostgreSQL執行個體規格	大版本：RDS PostgreSQL 16 規格：標準版高可用系列獨享型規格，pg.x8.2xlarge.2c（16核 128GB） pgvector外掛程式版本：0.8.0
用戶端ECS規格	規格：ecs.c6.xlarge（4核 8GiB）作業系統：Alibaba Cloud Linux 3
測試載入器	ann-benchmarks 重要 ann-benchmarks預設測試單線程效能。如需測試執行個體的並發效能，請參見pgvector效能測試（基於IVF索引）。

準備工作

RDS PostgreSQL執行個體側

建立高許可權帳號ann_testuser和測試資料庫ann_testdb，詳情請參見建立帳號和資料庫。
在測試資料庫ann_testdb中安裝外掛程式pgvector（即vector），詳情請參見管理外掛程式。

用戶端ECS側

安裝Docker，詳情請參見安裝Docker。

執行如下命令，下載ann-benchmark測試載入器。

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-benchmark測試載入器的依賴。
```
cd ~/ann-benchmarks/
pip install -r requirements.txt
```

測試流程

重要

測試流程中的所有步驟均在虛擬環境ann_test中進行。如果因逾時或其他原因退出，請使用命令conda activate ann_test重新進入該環境。

步驟一：配置測試載入器串連資訊

編輯測試載入器中的~/ann-benchmarks/ann_benchmarks/algorithms/pgvector/module.py檔案，增加如下的串連配置，並根據實際情況填寫配置資訊。

# 設定 PostgreSQL 的串連參數
os.environ['ANN_BENCHMARKS_PG_USER'] = 'ann_testuser'     # RDS PostgreSQL執行個體的帳號
os.environ['ANN_BENCHMARKS_PG_PASSWORD'] = 'testPawword'  # RDS PostgreSQL執行個體帳號的密碼
os.environ['ANN_BENCHMARKS_PG_DBNAME'] = 'ann_testdb'     # RDS PostgreSQL執行個體的資料庫名稱
os.environ['ANN_BENCHMARKS_PG_HOST'] = 'pgm-****.pg.rds.aliyuncs.com'  # RDS PostgreSQL執行個體的內網地址
os.environ['ANN_BENCHMARKS_PG_PORT'] = '5432'             # RDS PostgreSQL執行個體的內網連接埠號碼
os.environ['ANN_BENCHMARKS_PG_START_SERVICE'] = 'false'   # 禁用自動啟動服務

步驟二：配置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	查詢階段設定，表示搜尋過程中維護候選集的大小。該值越大，通常召回率越高，但查詢時間也會相應增加。

步驟三：構建測試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查看支援配置參數。

步驟四：擷取資料集

執行測試指令碼時，將自動下載指定的公開測試資料集。

本文以文本相似性檢索為例，採用Angular距離類型的資料集nytimes-256-angular進行測試。更多公開的資料集請參見ann-benchmarks。

資料集	維度	資料行數	測試向量數	最近鄰Top N	距離類型
NYTimes	256	290,000	10,000	100	Angular

說明

如果公開資料集無法滿足測試需求，建議將實際的向量資料轉換為標準格式HDF5或其他標準格式，以用於測試檢索效能。詳情請參見附錄二：自訂測試資料集。

步驟五：進行測試並擷取測試結果

執行如下命令，運行測試指令碼進行測試。

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
		20	0.741	1131.941
		40	0.820	836.017
		80	0.871	574.733
		120	0.894	440.076
		200	0.918	297.267
		400	0.945	162.759
		800	0.969	84.268
16	200	10	0.683	1299.667
		20	0.781	1094.968
		40	0.849	790.838
		80	0.895	533.826
		120	0.914	405.975
		200	0.933	272.591
		400	0.956	148.688
		800	0.977	76.555
24	200	10	0.767	1182.747
		20	0.840	922.770
		40	0.887	639.899
		80	0.920	411.140
		120	0.936	303.323
		200	0.953	199.752
		400	0.973	105.506
		800	0.988	53.904

（可選）執行如下命令，擷取詳細的測試結果，包括召回率、QPS、回應時間、索引構建時間和索引大小等指標的結果。
```
cd ~/ann-benchmarks
python data_export.py --out res.csv
```
例如，可以擷取參數m、ef_construction和索引構建時間之間的關係：
m
ef_construction
build（索引構建時間，單位：s）
16
100
33.35161
16
200
57.66014
24
200
87.22608

測試結論

在構建HNSW索引時：

增加m、ef_construction和ef_search可以提高召回率，但會導致QPS降低。
增加m和ef_construction同樣會提高召回率，同時也會導致QPS降低，並延長索引構建時間。
如果對召回率的要求較高，不建議使用預設的索引參數（m=16、ef_construction=64和ef_search=40）。

附錄一：RDS PostgreSQL參數、向量參數對索引構建的影響

在不同ann-benchmarks的測試參數下，通過對測試結果的分析，可以確定RDS PostgreSQL參數及向量參數對索引構建的影響。

RDS PostgreSQL參數對索引構建的影響

例如，當ann-benchmarks的測試參數設定為m=16和efConstruction=64時，RDS PostgreSQL參數對索引構建的影響如下。

maintenance_work_mem
用於維護操作的最大記憶體，包括VACUUM、CREATE INDEX等操作，單位為KB。當maintenance_work_mem值小於測試資料大小時，增大該值會縮短索引構建時間；然而，當maintenance_work_mem值超過資料大小後，索引構建時間將不再減少。
以pg.x8.2xlarge.2c（16核 128GB）規格的RDS PostgreSQL為例，執行個體參數max_parallel_maintenance_workers設定為預設值8，nytimes-256-angular資料集資料量約324MB，maintenance_work_mem參數對索引構建的影響：
maintenance_work_mem
索引構建時間（s）
64 MB（65536 KB）
52.82
128 MB（131072 KB）
46.79
256 MB（262144 KB）
36.40
512 MB（524288 KB）
18.90
1 GB（1048576 KB）
19.06
max_parallel_maintenance_workers
設定CREATE INDEX並行工作的最大數量。索引構建時間隨max_parallel_maintenance_workers增加而減少。
以pg.x8.2xlarge.2c（16核 128GB）規格的RDS PostgreSQL為例，執行個體參數maintenance_work_mem設定為預設值2048（即2 GB），nytimes-256-angular資料集資料量約324 MB，max_parallel_maintenance_workers參數對索引構建的影響：
max_parallel_maintenance_workers
索引構建時間（s）
1
76.00
2
51.34
4
32.49
8
19.66
12
14.44
16
13.07
24
13.15

向量參數對索引構建的影響

例如，當RDS PostgreSQL參數設定為maintenance_work_mem=8 GB（8388608 KB）和max_parallel_maintenance_workers=16時，RDS PostgreSQL參數對索引構建的影響如下。

向量維度
使用GloVe資料集，行數固定為1183514。ann-benchmarks的測試參數設定為m=16、efConstruction=64和ef_search=40時，測試結果如下。根據測試結果可以得出，索引構建時間隨著向量維度增加而相應增加；同時，召回率和QPS均呈下降趨勢，查詢延遲則有所增加。
維度
構建時間（s）
召回率
QPS
p99（ms）
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：將所有查詢請求的回應時間按升序排序後，位於第99百分位位置的延遲值。即表示在所有查詢請求中，99%的請求回應時間低於此值，也就是說，只有1%的請求回應時間超過這個值。
向量行數
使用dbpedia-openai-{n}k-angular資料集，向量行數（n）支援範圍為100至1000。ann-benchmarks的測試參數設定為m=48、efConstruction=256和ef_search=200時，測試結果如下。根據測試結果可以得出，索引構建時間隨向量行數的增加呈現出非線性增長趨勢；同時，召回率和QPS均有所下降，查詢延遲也隨之增加。
向量行數
行數（萬）
構建時間（s）
召回率
QPS
p99（ms）
100
10
54.05s
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

附錄二：自訂測試資料集

如下資料集製作參考nytimes-256-angular資料格式，僅供參考。

執行如下命令，製作自訂資料集。

說明

本樣本中需要安裝rds_ai外掛程式，詳情請參見AI（rds_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 pg資料庫
    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外掛程式或者使用百鍊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)

    # 計算最近Top 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: {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（索引構建時間，單位：s）
16	100	33.35161
16	200	57.66014
24	200	87.22608

maintenance_work_mem	索引構建時間（s）
64 MB（65536 KB）	52.82
128 MB（131072 KB）	46.79
256 MB（262144 KB）	36.40
512 MB（524288 KB）	18.90
1 GB（1048576 KB）	19.06

max_parallel_maintenance_workers	索引構建時間（s）
1	76.00
2	51.34
4	32.49
8	19.66
12	14.44
16	13.07
24	13.15

維度	構建時間（s）	召回率	QPS	p99（ms）
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

向量行數	行數（萬）	構建時間（s）	召回率	QPS	p99（ms）
100	10	54.05s	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