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>
qdrant · Apr 11, 2024 · f4436e4 · f4436e4
1 parent e28cc02
commit f4436e4
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 68 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 } }
     }
 ]