diff --git a/engine/clients/pgvector/config.py b/engine/clients/pgvector/config.py index dc3b8365..5507745c 100644 --- a/engine/clients/pgvector/config.py +++ b/engine/clients/pgvector/config.py @@ -1,6 +1,6 @@ import os -PGVECTOR_PORT = int(os.getenv("PGVECTOR_PORT", 9200)) +PGVECTOR_PORT = int(os.getenv("PGVECTOR_PORT", 5432)) PGVECTOR_DB = os.getenv("PGVECTOR_DB", "postgres") PGVECTOR_USER = os.getenv("PGVECTOR_USER", "postgres") PGVECTOR_PASSWORD = os.getenv("PGVECTOR_PASSWORD", "passwd") @@ -9,6 +9,7 @@ def get_db_config(host, connection_params): return { "host": host or "localhost", + "port": PGVECTOR_PORT, "dbname": PGVECTOR_DB, "user": PGVECTOR_USER, "password": PGVECTOR_PASSWORD, diff --git a/engine/clients/pgvector/configure.py b/engine/clients/pgvector/configure.py index d5587431..0da692b2 100644 --- a/engine/clients/pgvector/configure.py +++ b/engine/clients/pgvector/configure.py @@ -9,11 +9,6 @@ class PgVectorConfigurator(BaseConfigurator): - DISTANCE_MAPPING = { - Distance.L2: "vector_l2_ops", - Distance.COSINE: "vector_cosine_ops", - } - def __init__(self, host, collection_params: dict, connection_params: dict): super().__init__(host, collection_params, connection_params) self.conn = psycopg.connect(**get_db_config(host, connection_params)) @@ -38,17 +33,6 @@ def recreate(self, dataset: Dataset, collection_params): ) self.conn.execute("ALTER TABLE items ALTER COLUMN embedding SET STORAGE PLAIN") - try: - hnsw_distance_type = self.DISTANCE_MAPPING[dataset.config.distance] - except KeyError: - raise IncompatibilityError( - f"Unsupported distance metric: {dataset.config.distance}" - ) - - self.conn.execute( - 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']})" - ) - self.conn.close() def delete_client(self): diff --git a/engine/clients/pgvector/search.py b/engine/clients/pgvector/search.py index fa8bde5a..62f53035 100644 --- a/engine/clients/pgvector/search.py +++ b/engine/clients/pgvector/search.py @@ -23,24 +23,19 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic cls.conn = psycopg.connect(**get_db_config(host, connection_params)) register_vector(cls.conn) cls.cur = cls.conn.cursor() - cls.distance = distance - cls.search_params = search_params["search_params"] - - @classmethod - def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]: - cls.cur.execute(f"SET hnsw.ef_search = {cls.search_params['hnsw_ef']}") - - if cls.distance == Distance.COSINE: - query = f"SELECT id, embedding <=> %s AS _score FROM items ORDER BY _score LIMIT {top};" - elif cls.distance == Distance.L2: - query = f"SELECT id, embedding <-> %s AS _score FROM items ORDER BY _score LIMIT {top};" + cls.cur.execute( + f"SET hnsw.ef_search = {search_params['search_params']['hnsw_ef']}" + ) + if distance == Distance.COSINE: + cls.query = f"SELECT id, embedding <=> %s AS _score FROM items ORDER BY _score LIMIT %s" + elif distance == Distance.L2: + cls.query = f"SELECT id, embedding <-> %s AS _score FROM items ORDER BY _score LIMIT %s" else: raise NotImplementedError(f"Unsupported distance metric {cls.distance}") - cls.cur.execute( - query, - (np.array(vector),), - ) + @classmethod + def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]: + cls.cur.execute(cls.query, (np.array(vector), top), binary=True, prepare=True) return cls.cur.fetchall() @classmethod diff --git a/engine/clients/pgvector/upload.py b/engine/clients/pgvector/upload.py index 8d59ee7f..c3921e95 100644 --- a/engine/clients/pgvector/upload.py +++ b/engine/clients/pgvector/upload.py @@ -4,11 +4,16 @@ import psycopg from pgvector.psycopg import register_vector +from engine.base_client.distances import Distance from engine.base_client.upload import BaseUploader from engine.clients.pgvector.config import get_db_config class PgVectorUploader(BaseUploader): + DISTANCE_MAPPING = { + Distance.L2: "vector_l2_ops", + Distance.COSINE: "vector_cosine_ops", + } conn = None cur = None upload_params = {} @@ -27,10 +32,28 @@ def upload_batch( vectors = np.array(vectors) # Copy is faster than insert - with cls.cur.copy("COPY items (id, embedding) FROM STDIN") as copy: + with cls.cur.copy( + "COPY items (id, embedding) FROM STDIN WITH (FORMAT BINARY)" + ) as copy: + copy.set_types(["integer", "vector"]) for i, embedding in zip(ids, vectors): copy.write_row((i, embedding)) + @classmethod + def post_upload(cls, distance): + try: + hnsw_distance_type = cls.DISTANCE_MAPPING[distance] + except KeyError: + raise IncompatibilityError(f"Unsupported distance metric: {distance}") + + cls.conn.execute("SET max_parallel_workers = 128") + cls.conn.execute("SET max_parallel_maintenance_workers = 128") + cls.conn.execute( + 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']})" + ) + + return {} + @classmethod def delete_client(cls): if cls.cur: diff --git a/engine/servers/pgvector-single-node/docker-compose.yaml b/engine/servers/pgvector-single-node/docker-compose.yaml index ea554d92..13e8b8e6 100644 --- a/engine/servers/pgvector-single-node/docker-compose.yaml +++ b/engine/servers/pgvector-single-node/docker-compose.yaml @@ -3,13 +3,17 @@ version: '3.7' services: pgvector: container_name: pgvector - image: ankane/pgvector:v0.5.1 + image: pgvector/pgvector:0.6.2-pg16 environment: - POSTGRES_DB=postgres - POSTGRES_USER=postgres - POSTGRES_PASSWORD=passwd - POSTGRES_HOST_AUTH_METHOD=trust - - POSTGRES_MAX_CONNECTIONS=200 + # shared_buffers should be 25% of memory + # maintenance_work_mem should be ~65% + command: postgres -c shared_buffers=6GB -c maintenance_work_mem=16GB -c max_connections=200 + # shm_size should be shared_buffers + maintenance_work_mem + shm_size: 22g ports: - 5432:5432 logging: diff --git a/experiments/configurations/pgvector-single-node.json b/experiments/configurations/pgvector-single-node.json index e1c8e33a..22ced04d 100644 --- a/experiments/configurations/pgvector-single-node.json +++ b/experiments/configurations/pgvector-single-node.json @@ -3,104 +3,88 @@ "name": "pgvector-default", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 16, "ef_construct": 128 } - }, + "collection_params": {}, "search_params": [ - { "parallel": 1, "search_params": { "hnsw_ef": 128 } } + { "parallel": 8, "search_params": { "hnsw_ef": 128 } } ], - "upload_params": { "parallel": 1, "batch_size": 1024 } + "upload_params": { "parallel": 16, "batch_size": 1024, "hnsw_config": { "m": 16, "ef_construct": 128 } } }, { "name": "pgvector-parallel", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 16, "ef_construct": 128 } - }, + "collection_params": {}, "search_params": [ { "parallel": 8, "search_params": { "hnsw_ef": 128 } }, { "parallel": 16, "search_params": { "hnsw_ef": 128 } }, { "parallel": 100, "search_params": { "hnsw_ef": 128 } } ], - "upload_params": { "parallel": 1, "batch_size": 1024 } + "upload_params": { "parallel": 1, "batch_size": 1024, "hnsw_config": { "m": 16, "ef_construct": 128 } } }, { "name": "pgvector-m-16-ef-128", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 16, "ef_construct": 128 } - }, + "collection_params": {}, "search_params": [ { "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 } }, { "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 } } ], - "upload_params": { "parallel": 16 } + "upload_params": { "parallel": 16, "hnsw_config": { "m": 16, "ef_construct": 128 } } }, { "name": "pgvector-m-32-ef-128", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 32, "ef_construct": 128 } - }, + "collection_params": {}, "search_params": [ { "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 } }, { "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 } } ], - "upload_params": { "parallel": 16 } + "upload_params": { "parallel": 16, "hnsw_config": { "m": 32, "ef_construct": 128 } } }, { "name": "pgvector-m-32-ef-256", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 32, "ef_construct": 256 } - }, + "collection_params": {}, "search_params": [ { "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 } }, { "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 } } ], - "upload_params": { "parallel": 16 } + "upload_params": { "parallel": 16, "hnsw_config": { "m": 32, "ef_construct": 256 } } }, { "name": "pgvector-m-32-ef-512", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 32, "ef_construct": 512 } - }, + "collection_params": {}, "search_params": [ { "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 } }, { "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 } } ], - "upload_params": { "parallel": 16 } + "upload_params": { "parallel": 16, "hnsw_config": { "m": 32, "ef_construct": 512 } } }, { "name": "pgvector-m-64-ef-256", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 64, "ef_construct": 256 } - }, + "collection_params": {}, "search_params": [ { "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 } }, { "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 } } ], - "upload_params": { "parallel": 16 } + "upload_params": { "parallel": 16, "hnsw_config": { "m": 64, "ef_construct": 256 } } }, { "name": "pgvector-m-64-ef-512", "engine": "pgvector", "connection_params": {}, - "collection_params": { - "hnsw_config": { "m": 64, "ef_construct": 512 } - }, + "collection_params": {}, "search_params": [ { "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 } }, { "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 } } ], - "upload_params": { "parallel": 16 } + "upload_params": { "parallel": 16, "hnsw_config": { "m": 64, "ef_construct": 512 } } } ]