fix: pgvector-hnsw-support

Signed-off-by: min.tian <min.tian.cn@gmail.com>
zilliztech · Apr 26, 2024 · 0d75990 · 0d75990
1 parent a07a8f0
commit 0d75990
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 12 deletions.
diff --git a/vectordb_bench/backend/clients/pgvector/config.py b/vectordb_bench/backend/clients/pgvector/config.py
@@ -54,17 +54,10 @@ class HNSWConfig(PgVectorIndexConfig):
     ef: int | None = None
     index: IndexType = IndexType.HNSW
 
-    def index_param(self) -> dict:
-        return {
-            "metric_type": self.parse_metric(),
-            "index_type": self.index.value,
-            "params": {"M": self.M, "efConstruction": self.efConstruction},
-        }
-
     def index_param(self) -> dict:
         return {
             "m" : self.M,
-            "efConstruction" : self.efConstruction,
+            "ef_construction" : self.efConstruction,
             "metric" : self.parse_metric()
         }
 

diff --git a/vectordb_bench/frontend/const/dbCaseConfigs.py b/vectordb_bench/frontend/const/dbCaseConfigs.py
@@ -371,14 +371,27 @@ class CaseConfigInput(BaseModel):
     ],
 )
 
+CaseConfigParamInput_IndexType_PG = CaseConfigInput(
+    label=CaseConfigParamType.IndexType,
+    inputType=InputType.Option,
+    inputConfig={
+        "options": [
+            IndexType.HNSW.value,
+            IndexType.IVFFlat.value,
+        ],
+    },
+)
+
 CaseConfigParamInput_Lists = CaseConfigInput(
     label=CaseConfigParamType.lists,
     inputType=InputType.Number,
     inputConfig={
         "min": 1,
         "max": 65536,
-        "value": 10,
+        "value": 1000,
     },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    in [IndexType.IVFFlat.value],
 )
 
 CaseConfigParamInput_Probes = CaseConfigInput(
@@ -387,8 +400,34 @@ class CaseConfigInput(BaseModel):
     inputConfig={
         "min": 1,
         "max": 65536,
-        "value": 1,
+        "value": 10,
+    },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    in [IndexType.IVFFlat.value],
+)
+
+CaseConfigParamInput_EF_PG = CaseConfigInput(
+    label=CaseConfigParamType.EF,
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 1,
+        "max": 65536,
+        "value": 128,
     },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    in [IndexType.HNSW.value],
+)
+
+CaseConfigParamInput_EFC_PG = CaseConfigInput(
+    label=CaseConfigParamType.EFConstruction,
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 1,
+        "max": 65536,
+        "value": 300,
+    },
+    isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
+    in [IndexType.HNSW.value],
 )
 
 CaseConfigParamInput_QuantizationType_PgVectoRS = CaseConfigInput(
@@ -479,8 +518,20 @@ class CaseConfigInput(BaseModel):
     CaseConfigParamInput_NumCandidates_ES,
 ]
 
-PgVectorLoadingConfig = [CaseConfigParamInput_Lists]
-PgVectorPerformanceConfig = [CaseConfigParamInput_Lists, CaseConfigParamInput_Probes]
+PgVectorLoadingConfig = [
+    CaseConfigParamInput_IndexType_PG,
+    CaseConfigParamInput_Lists,
+    CaseConfigParamInput_M,
+    CaseConfigParamInput_EFC_PG
+]
+PgVectorPerformanceConfig = [
+    CaseConfigParamInput_IndexType_PG,
+    CaseConfigParamInput_Lists,
+    CaseConfigParamInput_Probes,
+    CaseConfigParamInput_M,
+    CaseConfigParamInput_EF_PG,
+    CaseConfigParamInput_EFC_PG
+]
 
 PgVectoRSLoadingConfig = [
     CaseConfigParamInput_IndexType,