高維向量相似度搜索(pgvector)
RDS PostgreSQL支持pgvector插件,提供了一個新的數(shù)據(jù)類型,能夠方便快捷地對高維向量進行檢索,是一款功能強大的向量相似度匹配搜索插件。
您可以加入RDS PostgreSQL插件交流釘釘群(103525002795),進行咨詢、交流和反饋,獲取更多關于插件的信息。
背景
RDS PostgreSQL支持pgvector插件,能夠存儲向量類型數(shù)據(jù),并實現(xiàn)向量相似度匹配,為AI產(chǎn)品提供底層數(shù)據(jù)支持。
pgvector主要提供如下能力:
支持數(shù)據(jù)類型vector,能夠對向量數(shù)據(jù)存儲以及查詢。
支持精確和近似最近鄰搜索(ANN,Approximate Nearest Neighbor),其距離或相似度度量方法包括歐氏距離(L2)、余弦相似度(Cosine)以及內積運算(Inner Product)。索引構建支持HNSW索引、并行索引IVFFlat、向量的逐元素乘法、L1距離函數(shù)以及求和聚合。
最大支持創(chuàng)建16000維度的向量,最大支持對2000維度的向量建立索引。
相關概念及實現(xiàn)原理
嵌入
嵌入(embedding)是指將高維數(shù)據(jù)映射為低維表示的過程。在機器學習和自然語言處理中,嵌入通常用于將離散的符號或對象表示為連續(xù)的向量空間中的點。
在自然語言處理中,詞嵌入(word embedding)是一種常見的技術,它將單詞映射到實數(shù)向量,以便計算機可以更好地理解和處理文本。通過詞嵌入,單詞之間的語義和語法關系可以在向量空間中得到反映。
實現(xiàn)原理
嵌入可以將文本、圖像、音視頻等信息在多個維度上抽象,轉化為向量數(shù)據(jù)。
pgvector提供vector數(shù)據(jù)類型,使RDS PostgreSQL數(shù)據(jù)庫具備了存儲向量數(shù)據(jù)的能力。
pgvector可以對存儲的向量數(shù)據(jù)進行精確搜索以及近似最近鄰搜索。
假設需要將蘋果、香蕉、貓三個對象存儲到數(shù)據(jù)庫中,并使用pgvector計算相似度,實現(xiàn)步驟如下:
先使用嵌入,將蘋果、香蕉、貓三個對象轉化為向量,假設以二維嵌入為例,結果如下:
蘋果:embedding[1,1] 香蕉:embedding[1.2,0.8] 貓:embedding[6,0.4]
將嵌入轉化的向量數(shù)據(jù)存儲到數(shù)據(jù)庫中。如何將二維向量數(shù)據(jù)存儲到數(shù)據(jù)庫中,具體請參見使用示例。
在二維平面中,三個對象分布如下:
對于蘋果和香蕉,都屬于水果,因此在二維坐標視圖中二者的距離更接近,而香蕉與貓屬于兩個完全不同的物種,因此距離較遠。
可以對水果的屬性進一步細化,比如水果的顏色,產(chǎn)地,味道等,每一個屬性都是一個維度,也就代表了維度越高,對于該信息的分類就更細,也就越有可能搜索出更精確的結果。
應用場景
存儲向量類型數(shù)據(jù)。
向量相似度匹配搜索。
前提條件
RDS PostgreSQL實例需滿足以下要求:
實例大版本為PostgreSQL 14或以上。
實例內核小版本為20230430或以上。
如需升級實例大版本或內核小版本,請參見升級數(shù)據(jù)庫大版本或升級內核小版本。
插件管理
創(chuàng)建插件
CREATE EXTENSION IF NOT EXISTS vector;
刪除插件
DROP EXTENSION vector;
更新插件
ALTER EXTENSION vector UPDATE [ TO new_version ]
說明new_version配置為pgvector的版本,pgvector的最新版本號及相關特性,請參見pgvector官方文檔。
使用示例
如下僅是對pgvector的簡單使用示例,更多使用方法,請參見pgvector官方文檔。
創(chuàng)建一個存儲vector類型的表(items),用于存儲embeddings。
CREATE TABLE items ( id bigserial PRIMARY KEY, item text, embedding vector(2) );
說明上述示例中,以二維為例,pgvector最大支持創(chuàng)建16000維度的向量。
將向量數(shù)據(jù)插入表中。
INSERT INTO items (item, embedding) VALUES ('蘋果', '[1, 1]'), ('香蕉', '[1.2, 0.8]'), ('貓', '[6, 0.4]');
使用余弦相似度操作符
<=>
計算香蕉與蘋果、貓之間的相似度。SELECT item, 1 - (embedding <=> '[1.2, 0.8]') AS cosine_similarity FROM items ORDER BY cosine_similarity DESC;
說明在上述示例中,使用公式
cosine_similarity = 1 - cosine_distance
進行計算,距離越近,相似度越高。您也可以使用歐氏距離操作符
<->
或內積運算操作符<#>
計算相似度。
結果示例:
item | cosine_similarity ------+-------------------- 香蕉 | 1 蘋果 | 0.980580680748848 貓 | 0.867105556566985
在上述結果中:
香蕉結果為1,表示完全匹配。
蘋果的結果為0.98,表示蘋果與香蕉高度相似。
貓的結果為0.86,表示貓與香蕉相似度較低。
說明您可以在實際業(yè)務中設置一個合適的相似度閾值,將相似度較低的結果直接排除。
為了提高相似度的查詢效率,pgvector支持為向量數(shù)據(jù)建立索引,執(zhí)行如下語句,為embedding字段建立索引。
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
各參數(shù)說明如下:
參數(shù)/取值
說明
items
添加索引的表名。
embedding
添加索引的列名。
vector_cosine_ops
向量索引方法中指定的訪問方法。
余弦相似性搜索,使用
vector_cosine_ops
。歐氏距離,使用
vector_l2_ops
。內積相似性,使用
vector_ip_ops
。
lists = 100
lists參數(shù)表示將數(shù)據(jù)集分成的列表數(shù),該值越大,表示數(shù)據(jù)集被分割得越多,每個子集的大小相對較小,索引查詢速度越快。但隨著lists值的增加,查詢的召回率可能會下降。
說明召回率是指在信息檢索或分類任務中,正確檢索或分類的樣本數(shù)量與所有相關樣本數(shù)量之比。召回率衡量了系統(tǒng)能夠找到所有相關樣本的能力,它是一個重要的評估指標。
構建索引需要的內存較多,當lists參數(shù)值超過2000時,會直接報錯
ERROR: memory required is xxx MB, maintenance_work_mem is xxx MB
,您需要設置更大的maintenance_work_mem才能為向量數(shù)據(jù)建立索引,該值設置過大實例會有很高的OOM風險。設置方法,請參見設置實例參數(shù)。您需要通過調整lists參數(shù)的值,在查詢速度和召回率之間進行權衡,以滿足具體應用場景的需求。
您可以使用如下兩種方式之一來設置ivfflat.probes參數(shù),指定在索引中搜索的列表數(shù)量,通過增加ivfflat.probes的值,將搜索更多的列表,可以提高查詢結果的召回率,即找到更多相關的結果。
會話級別
SET ivfflat.probes = 10;
事務級別
BEGIN; SET LOCAL ivfflat.probes = 10; SELECT ... COMMIT;
ivfflat.probes的值越大,查詢結果的召回率越高,但是查詢的速度會降低,根據(jù)具體的應用需求和數(shù)據(jù)集的特性,lists和ivfflat.probes的值可能需要進行調整以獲得最佳的查詢性能和召回率。
說明如果ivfflat.probes的值與創(chuàng)建索引時指定的lists值相等時,查詢將會忽略向量索引并進行全表掃描。在這種情況下,索引不會被使用,而是直接對整個表進行搜索,可能會降低查詢性能。
性能數(shù)據(jù)
為向量數(shù)據(jù)設置索引時,需要根據(jù)實際業(yè)務數(shù)據(jù)量及應用場景,在查詢速度和召回率之間進行權衡,您可以參考如下測試結果進行性能調優(yōu)。
以下基于RDS PostgreSQL實例,分別展示向量數(shù)據(jù)以及索引在不同數(shù)據(jù)量下占用的存儲空間情況,以及在設置不同的lists值以及probes值對查詢效率以及召回率的影響。
測試數(shù)據(jù)準備
創(chuàng)建測試數(shù)據(jù)庫。
CREATE DATABASE testdb;
安裝插件。
CREATE EXTENSION IF NOT EXISTS vector;
生成固定長度的隨機向量作為測試數(shù)據(jù)。
CREATE OR REPLACE FUNCTION random_array(dim integer) RETURNS DOUBLE PRECISION[] AS $$ SELECT array_agg(random()) FROM generate_series(1, dim); $$ LANGUAGE SQL VOLATILE COST 1;
創(chuàng)建一個存儲1536維向量的表。
CREATE TABLE vtest(id BIGINT, v VECTOR(1536));
向表中插入數(shù)據(jù)。
INSERT INTO vtest SELECT i, random_array(1536)::VECTOR(1536) FROM generate_series(1, 100000) AS i;
建立索引。
CREATE INDEX ON vtest USING ivfflat(v vector_cosine_ops) WITH(lists = 100);
測試步驟
為避免網(wǎng)絡延遲等因素對測試數(shù)據(jù)的影響,推薦使用內網(wǎng)連接地址,本示例是在與RDS PostgreSQL同地域、同VPC下的ECS中進行測試。
使用一個隨機向量,與vtest表中的數(shù)據(jù)進行相似度比對,獲取比對結果中最相似的50條記錄。
您需要創(chuàng)建一個sql文件,然后寫入如下內容,用于后續(xù)壓測時使用。
WITH tmp AS ( SELECT random_array(1536)::VECTOR(1536) AS vec ) SELECT id FROM vtest ORDER BY v <=> (SELECT vec FROM tmp) LIMIT FLOOR(RANDOM() * 50);
使用pgbench進行壓測。
如下命令需要在命令行窗口執(zhí)行,請確保已安裝PostgreSQL客戶端(本示例以15.1為例),pgbench是在PostgreSQL上運行基準測試的簡單程序。該命令的更多用法,請參見PostgreSQL官方文檔。
pgbench -f ./test.sql -c6 -T60 -P5 -U testuser -h pgm-bp****.pg.rds.aliyuncs.com -p 5432 -d testdb
各參數(shù)及說明如下:
參數(shù)/取值
說明
-f ./test.sql
指定測試腳本文件的路徑和文件名。
./test.sql
僅為示例,您需要根據(jù)實際情況修改路徑及文件名。-c6
設置并發(fā)客戶端數(shù)。-c表示指定并發(fā)客戶端數(shù),6表示本示例指定了6個并發(fā)客戶端來執(zhí)行測試。
-T60
設置測試時間。-T表示指定測試的運行時間,60表示本示例指定測試將運行60秒。
-P5
設置腳本參數(shù)。表示本示例中每5秒顯示一次進程報告。
-U testuser
指定數(shù)據(jù)庫用戶。testuser需要替換為您的數(shù)據(jù)庫用戶名。
-h pgm-bp****.pg.rds.aliyuncs.com
指定RDS PostgreSQL實例的內網(wǎng)連接地址。
-p 5432
指定RDS PostgreSQL實例的內網(wǎng)端口。
-d testdb
指定連接的數(shù)據(jù)庫,本示例以testdb為例。
測試結果
數(shù)據(jù)量(單位:萬行) | table size(單位:MB) | index size(單位:MB) | Latency(單位:ms) | TPS(單位:個) |
10 | 796 | 782 | 15.7 | 380 |
30 | 2388 | 2345 | 63 | 94 |
50 | 3979 | 3907 | 74 | 80 |
80 | 6367 | 6251 | 90 | 66 |
100 | 7958 | 7813 | 105 | 56 |
當lists固定為2000,表中數(shù)據(jù)量為100萬行時,probes越大召回率越高,TPS越低。
當probes固定為20,表中數(shù)據(jù)量為100萬行時,lists越大,召回率越低,TPS越高
測試結論
lists的值對索引占用的存儲空間影響微乎其微,和表中的數(shù)據(jù)量有直接的關系。
lists和probes對查詢效率以及召回率起著相反的作用,因此合理地設置這兩個值可以在查詢效率以及召回率上達到一個平衡。
根據(jù)表中行數(shù)(rows)的不同,建議設置的lists和probes值如下:
小于等于100萬行:
lists = rows / 1000
、probes = lists / 10
大于100萬行:
lists = sqrt(rows)
、probes = sqrt(lists)
說明sqrt表示開方運算。