extended PGVECTOR configurables

Author: fcostaoliveira

engine/clients/pgvector/config.py

@@ -5,6 +5,14 @@
 PGVECTOR_USER = os.getenv("PGVECTOR_USER", "postgres")
 PGVECTOR_PASSWORD = os.getenv("PGVECTOR_PASSWORD", "passwd")
 
+# Parallel execution settings
+PGVECTOR_MAX_PARALLEL_WORKERS_PER_GATHER = int(os.getenv("PGVECTOR_MAX_PARALLEL_WORKERS_PER_GATHER", 8))
+PGVECTOR_PARALLEL_TUPLE_COST = float(os.getenv("PGVECTOR_PARALLEL_TUPLE_COST", 0.01))
+PGVECTOR_PARALLEL_SETUP_COST = int(os.getenv("PGVECTOR_PARALLEL_SETUP_COST", 100))
+PGVECTOR_MIN_PARALLEL_TABLE_SCAN_SIZE = os.getenv("PGVECTOR_MIN_PARALLEL_TABLE_SCAN_SIZE", "1MB")
+PGVECTOR_FORCE_PARALLEL_MODE = os.getenv("PGVECTOR_FORCE_PARALLEL_MODE", "on")
+PGVECTOR_WORK_MEM = os.getenv("PGVECTOR_WORK_MEM", "4GB")
+
 
 def get_db_config(host, connection_params):
     return {

engine/clients/pgvector/search.py

@@ -7,7 +7,15 @@
 
 from engine.base_client.distances import Distance
 from engine.base_client.search import BaseSearcher
-from engine.clients.pgvector.config import get_db_config
+from engine.clients.pgvector.config import (
+    get_db_config,
+    PGVECTOR_MAX_PARALLEL_WORKERS_PER_GATHER,
+    PGVECTOR_PARALLEL_TUPLE_COST,
+    PGVECTOR_PARALLEL_SETUP_COST,
+    PGVECTOR_MIN_PARALLEL_TABLE_SCAN_SIZE,
+    PGVECTOR_FORCE_PARALLEL_MODE,
+    PGVECTOR_WORK_MEM,
+)
 from engine.clients.pgvector.parser import PgVectorConditionParser
 
 
@@ -26,10 +34,28 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic
         cls.distance = distance
         cls.search_params = search_params["search_params"]
 
-        # For FLAT searches, disable index usage to force full scan
+        # For FLAT searches, optimize for parallel execution
         if "force_flat" in cls.search_params and cls.search_params["force_flat"]:
-            cls.cur.execute("SET enable_indexscan = off")
-            cls.cur.execute("SET enable_bitmapscan = off")
+            # Note: We don't disable indexes globally since:
+            # 1. No vector index exists in FLAT mode (not created)
+            # 2. PostgreSQL will naturally do sequential scan for vector queries
+            # 3. Other indexes (for filtering, etc.) remain useful
+
+            # Configurable parallel execution settings for FLAT searches
+            # Priority: Environment variables (from config) > search_params > defaults
+            parallel_workers = cls.search_params.get("max_parallel_workers_per_gather", PGVECTOR_MAX_PARALLEL_WORKERS_PER_GATHER)
+            parallel_tuple_cost = cls.search_params.get("parallel_tuple_cost", PGVECTOR_PARALLEL_TUPLE_COST)
+            parallel_setup_cost = cls.search_params.get("parallel_setup_cost", PGVECTOR_PARALLEL_SETUP_COST)
+            min_parallel_size = cls.search_params.get("min_parallel_table_scan_size", PGVECTOR_MIN_PARALLEL_TABLE_SCAN_SIZE)
+            force_parallel = cls.search_params.get("force_parallel_mode", PGVECTOR_FORCE_PARALLEL_MODE)
+            work_mem = cls.search_params.get("work_mem", PGVECTOR_WORK_MEM)
+
+            cls.cur.execute(f"SET max_parallel_workers_per_gather = {parallel_workers}")
+            cls.cur.execute(f"SET parallel_tuple_cost = {parallel_tuple_cost}")
+            cls.cur.execute(f"SET parallel_setup_cost = {parallel_setup_cost}")
+            cls.cur.execute(f"SET min_parallel_table_scan_size = '{min_parallel_size}'")
+            cls.cur.execute(f"SET force_parallel_mode = {force_parallel}")
+            cls.cur.execute(f"SET work_mem = '{work_mem}'")
 
     @classmethod
     def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
@@ -73,12 +99,5 @@ def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
     @classmethod
     def delete_client(cls):
         if cls.cur:
-            # Reset index settings if they were disabled for FLAT searches
-            if "force_flat" in cls.search_params and cls.search_params["force_flat"]:
-                try:
-                    cls.cur.execute("SET enable_indexscan = on")
-                    cls.cur.execute("SET enable_bitmapscan = on")
-                except:
-                    pass  # Connection might be closed already
             cls.cur.close()
             cls.conn.close()

engine/clients/pgvector/upload.py

@@ -5,7 +5,7 @@
 from pgvector.psycopg import register_vector
 
 from engine.base_client.upload import BaseUploader
-from engine.clients.pgvector.config import get_db_config
+from engine.clients.pgvector.config import get_db_config, PGVECTOR_WORK_MEM
 
 
 class PgvectorUploader(BaseUploader):
@@ -43,7 +43,7 @@ def init_client(cls, host, distance, connection_params, upload_params):
             max_maintenance_workers = 8
 
         # Optimize memory settings for large uploads based on AWS recommendations
-        cls.conn.execute("SET maintenance_work_mem = '2GB'")
+        cls.conn.execute(f"SET maintenance_work_mem = '{PGVECTOR_WORK_MEM}'")
         cls.conn.execute(f"SET max_parallel_maintenance_workers = {max_maintenance_workers}")
 
     @classmethod

experiments/configurations/pgvector-single-node.json

@@ -120,5 +120,20 @@
           { "parallel": 100, "search_params": { "force_flat": true } }
         ],
         "upload_params": { "parallel": 16 }
+    },
+    {
+        "name": "pgvector-flat-high-perf",
+        "engine": "pgvector",
+        "connection_params": {},
+        "collection_params": {
+          "flat_config": {
+            "create_index": false,
+            "max_parallel_workers": "auto"
+          }
+        },
+        "search_params": [
+          { "parallel": 1, "search_params": { "force_flat": true } }
+        ],
+        "upload_params": { "parallel": 16 }
     }
 ]