PQ support for pgdiskann

config-bq.json

@@ -0,0 +1,341 @@
+{
+  "database": {
+    "host": "localhost",
+    "username": "postgres",
+    "password": "admin123",
+    "db-name": "postgres"
+  },
+  "benchmark-info": {
+    "name": "hnsw-bq-param-sweep",
+    "instance-size": "Standard_D8ds_v5",
+    "instance-service": "azure-flex",
+    "provider": "azure",
+    "description": "Running param sweep for HNSW binary quantization with reranking index alogrithms"
+  },
+  "cases": [
+    {
+      "db-label": "index-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": true,
+      "load": true,
+      "search-serial": false,
+      "search-concurrent": false,
+      "index-params": {
+        "m": 16,
+        "ef-construction": 64,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64],
+        "quantized-fetch-limit": [30]
+      },
+      "reranking": true,
+      "num-concurrency": "1",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 1
+    },
+    {
+      "db-label": "run-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": false,
+      "load": false,
+      "search-serial": true,
+      "search-concurrent": true,
+      "index-params": {
+        "m": 16,
+        "ef-construction": 64,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64, 128, 256],
+        "quantized-fetch-limit": [30, 40, 50]
+      },
+      "reranking": true,
+      "num-concurrency": "1,10,20,30,40,50,60,70,80,90,100",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 3
+    },
+    {
+      "db-label": "index-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": true,
+      "load": true,
+      "search-serial": false,
+      "search-concurrent": false,
+      "index-params": {
+        "m": 16,
+        "ef-construction": 128,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64],
+        "quantized-fetch-limit": [30]
+      },
+      "reranking": true,
+      "num-concurrency": "1",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 1
+    },
+    {
+      "db-label": "run-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": false,
+      "load": false,
+      "search-serial": true,
+      "search-concurrent": true,
+      "index-params": {
+        "m": 16,
+        "ef-construction": 128,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64, 128, 256],
+        "quantized-fetch-limit": [30, 40, 50]
+      },
+      "reranking": true,
+      "num-concurrency": "1,10,20,30,40,50,60,70,80,90,100",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 3
+    },
+    {
+      "db-label": "index-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": true,
+      "load": true,
+      "search-serial": false,
+      "search-concurrent": false,
+      "index-params": {
+        "m": 16,
+        "ef-construction": 256,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64],
+        "quantized-fetch-limit": [30]
+      },
+      "reranking": true,
+      "num-concurrency": "1",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 1
+    },
+    {
+      "db-label": "run-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": false,
+      "load": false,
+      "search-serial": true,
+      "search-concurrent": true,
+      "index-params": {
+        "m": 16,
+        "ef-construction": 256,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64, 128, 256],
+        "quantized-fetch-limit": [30, 40, 50]
+      },
+      "reranking": true,
+      "num-concurrency": "1,10,20,30,40,50,60,70,80,90,100",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 3
+    },
+    {
+      "db-label": "index-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": true,
+      "load": true,
+      "search-serial": false,
+      "search-concurrent": false,
+      "index-params": {
+        "m": 32,
+        "ef-construction": 64,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64],
+        "quantized-fetch-limit": [30]
+      },
+      "reranking": true,
+      "num-concurrency": "1",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 1
+    },
+    {
+      "db-label": "run-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": false,
+      "load": false,
+      "search-serial": true,
+      "search-concurrent": true,
+      "index-params": {
+        "m": 32,
+        "ef-construction": 64,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64, 128, 256],
+        "quantized-fetch-limit": [30, 40, 50]
+      },
+      "reranking": true,
+      "num-concurrency": "1,10,20,30,40,50,60,70,80,90,100",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 3
+    },
+    {
+      "db-label": "index-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": true,
+      "load": true,
+      "search-serial": false,
+      "search-concurrent": false,
+      "index-params": {
+        "m": 32,
+        "ef-construction": 128,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64],
+        "quantized-fetch-limit": [30]
+      },
+      "reranking": true,
+      "num-concurrency": "1",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 1
+    },
+    {
+      "db-label": "run-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": false,
+      "load": false,
+      "search-serial": true,
+      "search-concurrent": true,
+      "index-params": {
+        "m": 32,
+        "ef-construction": 128,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64, 128, 256],
+        "quantized-fetch-limit": [30, 40, 50]
+      },
+      "reranking": true,
+      "num-concurrency": "1,10,20,30,40,50,60,70,80,90,100",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 3
+    },
+    {
+      "db-label": "index-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": true,
+      "load": true,
+      "search-serial": false,
+      "search-concurrent": false,
+      "index-params": {
+        "m": 32,
+        "ef-construction": 256,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64],
+        "quantized-fetch-limit": [30]
+      },
+      "reranking": true,
+      "num-concurrency": "1",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 1
+    },
+    {
+      "db-label": "run-hnsw-bq-param-sweep",
+      "vdb-command": "pgvectorhnsw",
+      "vector-ext": "vector",
+      "index-type": "hnsw-bq",
+      "case-type": "Performance1536D500K",
+      "drop-old": false,
+      "load": false,
+      "search-serial": true,
+      "search-concurrent": true,
+      "index-params": {
+        "m": 32,
+        "ef-construction": 256,
+        "maintenance-work-mem": "8GB",
+        "max-parallel-workers": 7,
+        "quantization-type": "bit"
+      },
+      "search-params": {
+        "ef-search": [64, 128, 256],
+        "quantized-fetch-limit": [30, 40, 50]
+      },
+      "reranking": true,
+      "num-concurrency": "1,10,20,30,40,50,60,70,80,90,100",
+      "concurrency-duration": 30,
+      "k": 10,
+      "run-count": 3
+    }
+  ]
+}

run.py

@@ -26,7 +26,7 @@ def main():
     start_time = time.time()
     start_timeh = time.strftime('%Y-%m-%d %H:%M:%S')
     logger.info(f"Benchmark run start time: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-    
+
     for case in config['cases']:
         print(f"Running case: {case['db-label']}")
         setup_database(config)
@@ -50,11 +50,6 @@ def run_benchmark(case, db_config, benchmark_info, dry_run=False):
         print(f"Starting run {run + 1} of {run_count} for case: {case['db-label']}")
         for i, search_params in enumerate(generate_combinations(case["search-params"])):
             command = base_command + search_params
-            if case["index-type"] == "hnsw-bq" and "reranking" in case:
-                if case.get("half-quantized-fetch-limit", False):
-                    command += ["--quantized-fetch-limit", str(int(int(search_params[1]) / 2))]
-                else:
-                    command += ["--quantized-fetch-limit", search_params[1]]
 
             if i > 0 or run > 0:
                 command = handle_drop_old_load_flags(command)

vectordb_bench/__init__.py

@@ -14,11 +14,12 @@ class config:
     ALIYUN_OSS_URL = "assets.zilliz.com.cn/benchmark/"
     AWS_S3_URL = "assets.zilliz.com/benchmark/"
 
-    LOG_LEVEL = env.str("LOG_LEVEL", "INFO")
+    LOG_LEVEL = env.str("LOG_LEVEL", "DEBUG")
 
     DEFAULT_DATASET_URL = env.str("DEFAULT_DATASET_URL", AWS_S3_URL)
     DATASET_LOCAL_DIR = env.path("DATASET_LOCAL_DIR", f"/home/{os.getenv('USER')}/vectordb_bench/dataset")
-    NUM_PER_BATCH = env.int("NUM_PER_BATCH", 100)
+
+    NUM_PER_BATCH = env.int("NUM_PER_BATCH", 1000)
 
     DROP_OLD = env.bool("DROP_OLD", True)
     USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)