Posted on

使PostgreSQL支持向量相似度搜尋

安裝相關套件pgvector

套件Github頁面

https://github.com/pgvector/pgvector

pgvector可做到Postgres 的開源向量相似度搜索,將向量與其他數據一起存儲。

支援:

  • 精確和近似最近鄰搜索
  • 單精度(single-precision)、半精度(half-precision)、二進位(binary)和稀疏向量(sparse vectors)
  • L2 距離(L2 distance)、內積(inner product)、餘弦距離(cosine distance)、L1 距離(L1 distance)、漢明距離(Hamming distance)和 Jaccard 距離(Jaccard distance)
  • 具有 Postgres 用戶端的任何語言
  • 支持 ACID 合規性、時間點恢復、JOIN 和 Postgres 的所有其他強大功能

Linux安裝

cd /tmp
git clone --branch v0.7.2 https://github.com/pgvector/pgvector.git
cd pgvector
make
make install # may need sudo

Windows安裝

首先安裝VisualStudioSetup,此為下載連結: https://visualstudio.microsoft.com/downloads/

接著至少勾選這些套件

執行C的編譯環境

call "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"

使用nmake編譯檔案

set "PGROOT=C:\Program Files\PostgreSQL\16"
cd %TEMP%
git clone --branch v0.7.2 https://github.com/pgvector/pgvector.git
cd pgvector
nmake /F Makefile.win
nmake /F Makefile.win install

建立可儲存的向量的資料欄位

pgvector的資料格式為vector(1536),後面的(1536)代表長度,此為OpenAI的embedding長度

以下為一個建立表格的SQL範例

CREATE TABLE IF NOT EXISTS public.test_example
(
    id character varying(255) COLLATE pg_catalog."default" NOT NULL,
    embeddings vector(1536),
    CONSTRAINT test_example_pkey PRIMARY KEY (id)
)

在Object Explorer按右鍵選擇,Query Tool可開啟SQL輸入介面,要瀏覽現有資料庫表格列表可至Schemas->Tables瀏覽

新增資料至資料表

以下為一個範例

import psycopg2
df = pd.read_csv('./data.csv')
def convert_embedding(embedding_str):
    return list(map(float, embedding_str.split(',')))

df["embeddings"] = df["embeddings"].apply(literal_eval)

row = df.iloc[0]
sql = """
    INSERT INTO fare_info (
        id, embeddings 
    ) VALUES (
        %s, %s
    )
"""

conn = psycopg2.connect(dbname="myDB", user="myUser", password="xxxxxx")
cursor = conn.cursor()
cursor.execute(sql, (row['id'],row['embeddings']))
conn.commit()