Skip to content

Commit fa5b6b6

Browse files
ankanepre-commit-ci[bot]
authored and
KShivendu
committed
pgvector improvements (#98)
* pgvector improvements * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Updated Postgres parameters * Use versioned Docker image * Updated pgvector to 0.6.2 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 0a2c99c commit fa5b6b6

File tree

6 files changed

+59
-68
lines changed

6 files changed

+59
-68
lines changed

‎engine/clients/pgvector/config.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import os
22

3-
PGVECTOR_PORT = int(os.getenv("PGVECTOR_PORT", 9200))
3+
PGVECTOR_PORT = int(os.getenv("PGVECTOR_PORT", 5432))
44
PGVECTOR_DB = os.getenv("PGVECTOR_DB", "postgres")
55
PGVECTOR_USER = os.getenv("PGVECTOR_USER", "postgres")
66
PGVECTOR_PASSWORD = os.getenv("PGVECTOR_PASSWORD", "passwd")
@@ -9,6 +9,7 @@
99
def get_db_config(host, connection_params):
1010
return {
1111
"host": host or "localhost",
12+
"port": PGVECTOR_PORT,
1213
"dbname": PGVECTOR_DB,
1314
"user": PGVECTOR_USER,
1415
"password": PGVECTOR_PASSWORD,

‎engine/clients/pgvector/configure.py

-16
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,6 @@
99

1010

1111
class PgVectorConfigurator(BaseConfigurator):
12-
DISTANCE_MAPPING = {
13-
Distance.L2: "vector_l2_ops",
14-
Distance.COSINE: "vector_cosine_ops",
15-
}
16-
1712
def __init__(self, host, collection_params: dict, connection_params: dict):
1813
super().__init__(host, collection_params, connection_params)
1914
self.conn = psycopg.connect(**get_db_config(host, connection_params))
@@ -38,17 +33,6 @@ def recreate(self, dataset: Dataset, collection_params):
3833
)
3934
self.conn.execute("ALTER TABLE items ALTER COLUMN embedding SET STORAGE PLAIN")
4035

41-
try:
42-
hnsw_distance_type = self.DISTANCE_MAPPING[dataset.config.distance]
43-
except KeyError:
44-
raise IncompatibilityError(
45-
f"Unsupported distance metric: {dataset.config.distance}"
46-
)
47-
48-
self.conn.execute(
49-
f"CREATE INDEX on items USING hnsw(embedding {hnsw_distance_type}) WITH (m = {collection_params['hnsw_config']['m']}, ef_construction = {collection_params['hnsw_config']['ef_construct']})"
50-
)
51-
5236
self.conn.close()
5337

5438
def delete_client(self):

‎engine/clients/pgvector/search.py

+10-15
Original file line numberDiff line numberDiff line change
@@ -23,24 +23,19 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic
2323
cls.conn = psycopg.connect(**get_db_config(host, connection_params))
2424
register_vector(cls.conn)
2525
cls.cur = cls.conn.cursor()
26-
cls.distance = distance
27-
cls.search_params = search_params["search_params"]
28-
29-
@classmethod
30-
def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
31-
cls.cur.execute(f"SET hnsw.ef_search = {cls.search_params['hnsw_ef']}")
32-
33-
if cls.distance == Distance.COSINE:
34-
query = f"SELECT id, embedding <=> %s AS _score FROM items ORDER BY _score LIMIT {top};"
35-
elif cls.distance == Distance.L2:
36-
query = f"SELECT id, embedding <-> %s AS _score FROM items ORDER BY _score LIMIT {top};"
26+
cls.cur.execute(
27+
f"SET hnsw.ef_search = {search_params['search_params']['hnsw_ef']}"
28+
)
29+
if distance == Distance.COSINE:
30+
cls.query = f"SELECT id, embedding <=> %s AS _score FROM items ORDER BY _score LIMIT %s"
31+
elif distance == Distance.L2:
32+
cls.query = f"SELECT id, embedding <-> %s AS _score FROM items ORDER BY _score LIMIT %s"
3733
else:
3834
raise NotImplementedError(f"Unsupported distance metric {cls.distance}")
3935

40-
cls.cur.execute(
41-
query,
42-
(np.array(vector),),
43-
)
36+
@classmethod
37+
def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
38+
cls.cur.execute(cls.query, (np.array(vector), top), binary=True, prepare=True)
4439
return cls.cur.fetchall()
4540

4641
@classmethod

‎engine/clients/pgvector/upload.py

+24-1
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,16 @@
44
import psycopg
55
from pgvector.psycopg import register_vector
66

7+
from engine.base_client.distances import Distance
78
from engine.base_client.upload import BaseUploader
89
from engine.clients.pgvector.config import get_db_config
910

1011

1112
class PgVectorUploader(BaseUploader):
13+
DISTANCE_MAPPING = {
14+
Distance.L2: "vector_l2_ops",
15+
Distance.COSINE: "vector_cosine_ops",
16+
}
1217
conn = None
1318
cur = None
1419
upload_params = {}
@@ -27,10 +32,28 @@ def upload_batch(
2732
vectors = np.array(vectors)
2833

2934
# Copy is faster than insert
30-
with cls.cur.copy("COPY items (id, embedding) FROM STDIN") as copy:
35+
with cls.cur.copy(
36+
"COPY items (id, embedding) FROM STDIN WITH (FORMAT BINARY)"
37+
) as copy:
38+
copy.set_types(["integer", "vector"])
3139
for i, embedding in zip(ids, vectors):
3240
copy.write_row((i, embedding))
3341

42+
@classmethod
43+
def post_upload(cls, distance):
44+
try:
45+
hnsw_distance_type = cls.DISTANCE_MAPPING[distance]
46+
except KeyError:
47+
raise IncompatibilityError(f"Unsupported distance metric: {distance}")
48+
49+
cls.conn.execute("SET max_parallel_workers = 128")
50+
cls.conn.execute("SET max_parallel_maintenance_workers = 128")
51+
cls.conn.execute(
52+
f"CREATE INDEX ON items USING hnsw (embedding {hnsw_distance_type}) WITH (m = {cls.upload_params['hnsw_config']['m']}, ef_construction = {cls.upload_params['hnsw_config']['ef_construct']})"
53+
)
54+
55+
return {}
56+
3457
@classmethod
3558
def delete_client(cls):
3659
if cls.cur:

‎engine/servers/pgvector-single-node/docker-compose.yaml

+6-2
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,17 @@ version: '3.7'
33
services:
44
pgvector:
55
container_name: pgvector
6-
image: ankane/pgvector:v0.5.1
6+
image: pgvector/pgvector:0.6.2-pg16
77
environment:
88
- POSTGRES_DB=postgres
99
- POSTGRES_USER=postgres
1010
- POSTGRES_PASSWORD=passwd
1111
- POSTGRES_HOST_AUTH_METHOD=trust
12-
- POSTGRES_MAX_CONNECTIONS=200
12+
# shared_buffers should be 25% of memory
13+
# maintenance_work_mem should be ~65%
14+
command: postgres -c shared_buffers=6GB -c maintenance_work_mem=16GB -c max_connections=200
15+
# shm_size should be shared_buffers + maintenance_work_mem
16+
shm_size: 22g
1317
ports:
1418
- 5432:5432
1519
logging:

‎experiments/configurations/pgvector-single-node.json

+17-33
Original file line numberDiff line numberDiff line change
@@ -3,104 +3,88 @@
33
"name": "pgvector-default",
44
"engine": "pgvector",
55
"connection_params": {},
6-
"collection_params": {
7-
"hnsw_config": { "m": 16, "ef_construct": 128 }
8-
},
6+
"collection_params": {},
97
"search_params": [
10-
{ "parallel": 1, "search_params": { "hnsw_ef": 128 } }
8+
{ "parallel": 8, "search_params": { "hnsw_ef": 128 } }
119
],
12-
"upload_params": { "parallel": 1, "batch_size": 1024 }
10+
"upload_params": { "parallel": 16, "batch_size": 1024, "hnsw_config": { "m": 16, "ef_construct": 128 } }
1311
},
1412
{
1513
"name": "pgvector-parallel",
1614
"engine": "pgvector",
1715
"connection_params": {},
18-
"collection_params": {
19-
"hnsw_config": { "m": 16, "ef_construct": 128 }
20-
},
16+
"collection_params": {},
2117
"search_params": [
2218
{ "parallel": 8, "search_params": { "hnsw_ef": 128 } },
2319
{ "parallel": 16, "search_params": { "hnsw_ef": 128 } },
2420
{ "parallel": 100, "search_params": { "hnsw_ef": 128 } }
2521
],
26-
"upload_params": { "parallel": 1, "batch_size": 1024 }
22+
"upload_params": { "parallel": 1, "batch_size": 1024, "hnsw_config": { "m": 16, "ef_construct": 128 } }
2723
},
2824
{
2925
"name": "pgvector-m-16-ef-128",
3026
"engine": "pgvector",
3127
"connection_params": {},
32-
"collection_params": {
33-
"hnsw_config": { "m": 16, "ef_construct": 128 }
34-
},
28+
"collection_params": {},
3529
"search_params": [
3630
{ "parallel": 1, "search_params": { "hnsw_ef": 64 } }, { "parallel": 1, "search_params": { "hnsw_ef": 128 } }, { "parallel": 1, "search_params": { "hnsw_ef": 256 } }, { "parallel": 1, "search_params": { "hnsw_ef": 512 } },
3731
{ "parallel": 100, "search_params": { "hnsw_ef": 64 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 256 } }, { "parallel": 100, "search_params": { "hnsw_ef": 512 } }
3832
],
39-
"upload_params": { "parallel": 16 }
33+
"upload_params": { "parallel": 16, "hnsw_config": { "m": 16, "ef_construct": 128 } }
4034
},
4135
{
4236
"name": "pgvector-m-32-ef-128",
4337
"engine": "pgvector",
4438
"connection_params": {},
45-
"collection_params": {
46-
"hnsw_config": { "m": 32, "ef_construct": 128 }
47-
},
39+
"collection_params": {},
4840
"search_params": [
4941
{ "parallel": 1, "search_params": { "hnsw_ef": 64 } }, { "parallel": 1, "search_params": { "hnsw_ef": 128 } }, { "parallel": 1, "search_params": { "hnsw_ef": 256 } }, { "parallel": 1, "search_params": { "hnsw_ef": 512 } },
5042
{ "parallel": 100, "search_params": { "hnsw_ef": 64 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 256 } }, { "parallel": 100, "search_params": { "hnsw_ef": 512 } }
5143
],
52-
"upload_params": { "parallel": 16 }
44+
"upload_params": { "parallel": 16, "hnsw_config": { "m": 32, "ef_construct": 128 } }
5345
},
5446
{
5547
"name": "pgvector-m-32-ef-256",
5648
"engine": "pgvector",
5749
"connection_params": {},
58-
"collection_params": {
59-
"hnsw_config": { "m": 32, "ef_construct": 256 }
60-
},
50+
"collection_params": {},
6151
"search_params": [
6252
{ "parallel": 1, "search_params": { "hnsw_ef": 64 } }, { "parallel": 1, "search_params": { "hnsw_ef": 128 } }, { "parallel": 1, "search_params": { "hnsw_ef": 256 } }, { "parallel": 1, "search_params": { "hnsw_ef": 512 } },
6353
{ "parallel": 100, "search_params": { "hnsw_ef": 64 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 256 } }, { "parallel": 100, "search_params": { "hnsw_ef": 512 } }
6454
],
65-
"upload_params": { "parallel": 16 }
55+
"upload_params": { "parallel": 16, "hnsw_config": { "m": 32, "ef_construct": 256 } }
6656
},
6757
{
6858
"name": "pgvector-m-32-ef-512",
6959
"engine": "pgvector",
7060
"connection_params": {},
71-
"collection_params": {
72-
"hnsw_config": { "m": 32, "ef_construct": 512 }
73-
},
61+
"collection_params": {},
7462
"search_params": [
7563
{ "parallel": 1, "search_params": { "hnsw_ef": 64 } }, { "parallel": 1, "search_params": { "hnsw_ef": 128 } }, { "parallel": 1, "search_params": { "hnsw_ef": 256 } }, { "parallel": 1, "search_params": { "hnsw_ef": 512 } },
7664
{ "parallel": 100, "search_params": { "hnsw_ef": 64 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 256 } }, { "parallel": 100, "search_params": { "hnsw_ef": 512 } }
7765
],
78-
"upload_params": { "parallel": 16 }
66+
"upload_params": { "parallel": 16, "hnsw_config": { "m": 32, "ef_construct": 512 } }
7967
},
8068
{
8169
"name": "pgvector-m-64-ef-256",
8270
"engine": "pgvector",
8371
"connection_params": {},
84-
"collection_params": {
85-
"hnsw_config": { "m": 64, "ef_construct": 256 }
86-
},
72+
"collection_params": {},
8773
"search_params": [
8874
{ "parallel": 1, "search_params": { "hnsw_ef": 64 } }, { "parallel": 1, "search_params": { "hnsw_ef": 128 } }, { "parallel": 1, "search_params": { "hnsw_ef": 256 } }, { "parallel": 1, "search_params": { "hnsw_ef": 512 } },
8975
{ "parallel": 100, "search_params": { "hnsw_ef": 64 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 256 } }, { "parallel": 100, "search_params": { "hnsw_ef": 512 } }
9076
],
91-
"upload_params": { "parallel": 16 }
77+
"upload_params": { "parallel": 16, "hnsw_config": { "m": 64, "ef_construct": 256 } }
9278
},
9379
{
9480
"name": "pgvector-m-64-ef-512",
9581
"engine": "pgvector",
9682
"connection_params": {},
97-
"collection_params": {
98-
"hnsw_config": { "m": 64, "ef_construct": 512 }
99-
},
83+
"collection_params": {},
10084
"search_params": [
10185
{ "parallel": 1, "search_params": { "hnsw_ef": 64 } }, { "parallel": 1, "search_params": { "hnsw_ef": 128 } }, { "parallel": 1, "search_params": { "hnsw_ef": 256 } }, { "parallel": 1, "search_params": { "hnsw_ef": 512 } },
10286
{ "parallel": 100, "search_params": { "hnsw_ef": 64 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 256 } }, { "parallel": 100, "search_params": { "hnsw_ef": 512 } }
10387
],
104-
"upload_params": { "parallel": 16 }
88+
"upload_params": { "parallel": 16, "hnsw_config": { "m": 64, "ef_construct": 512 } }
10589
}
10690
]

0 commit comments

Comments
 (0)