Support GPU_BRUTE_FORCE index for Milvus (zilliztech#476)

Signed-off-by: Rachit Chaudhary <rachit.chaudhary@outlook.com>
Co-authored-by: Signed-off-by: Rachit Chaudhary - r0c0axe <Rachit.Chaudhary@walmart.com>

Author: Rachit-Chaudhary11

vectordb_bench/backend/clients/api.py

@@ -25,6 +25,7 @@ class IndexType(str, Enum):
     ES_HNSW = "hnsw"
     ES_IVFFlat = "ivfflat"
     GPU_IVF_FLAT = "GPU_IVF_FLAT"
+    GPU_BRUTE_FORCE = "GPU_BRUTE_FORCE"
     GPU_IVF_PQ = "GPU_IVF_PQ"
     GPU_CAGRA = "GPU_CAGRA"
     SCANN = "scann"

vectordb_bench/backend/clients/milvus/cli.py

@@ -194,6 +194,26 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
         **parameters,
     )
 
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUBruteForceTypedDict)
+def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):
+    from .config import GPUBruteForceConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
+        ),
+        db_case_config=GPUBruteForceConfig(
+            metric_type=parameters["metric_type"],
+            limit=parameters["limit"],  # top-k for search
+        ),
+        **parameters,
+    )
+
 
 class MilvusGPUIVFPQTypedDict(
     CommonTypedDict,

vectordb_bench/backend/clients/milvus/config.py

@@ -40,6 +40,7 @@ def is_gpu_index(self) -> bool:
             IndexType.GPU_CAGRA,
             IndexType.GPU_IVF_FLAT,
             IndexType.GPU_IVF_PQ,
+            IndexType.GPU_BRUTE_FORCE,
         ]
 
     def parse_metric(self) -> str:
@@ -184,6 +185,37 @@ def search_param(self) -> dict:
         }
 
 
+class GPUBruteForceConfig(MilvusIndexConfig, DBCaseConfig):
+    limit: int = 10  # Default top-k for search
+    metric_type: str  # Metric type (e.g., 'L2', 'IP', etc.)
+    index: IndexType = IndexType.GPU_BRUTE_FORCE  # Index type set to GPU_BRUTE_FORCE
+
+    def index_param(self) -> dict:
+        """
+        Returns the parameters for creating the GPU_BRUTE_FORCE index.
+        No additional parameters required for index building.
+        """
+        return {
+            "metric_type": self.parse_metric(),  # Metric type for distance calculation (L2, IP, etc.)
+            "index_type": self.index.value,  # GPU_BRUTE_FORCE index type
+            "params": {},  # No additional parameters for GPU_BRUTE_FORCE
+        }
+
+    def search_param(self) -> dict:
+        """
+        Returns the parameters for performing a search on the GPU_BRUTE_FORCE index.
+        Only metric_type and top-k (limit) are needed for search.
+        """
+        return {
+            "metric_type": self.parse_metric(),  # Metric type for search
+            "params": {
+                "nprobe": 1,  # For GPU_BRUTE_FORCE, set nprobe to 1 (brute force search)
+                "limit": self.limit,  # Top-k for search
+            },
+        }
+
+
+
 class GPUIVFPQConfig(MilvusIndexConfig, DBCaseConfig):
     nlist: int = 1024
     m: int = 0
@@ -261,4 +293,5 @@ def search_param(self) -> dict:
     IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
     IndexType.GPU_IVF_PQ: GPUIVFPQConfig,
     IndexType.GPU_CAGRA: GPUCAGRAConfig,
+    IndexType.GPU_BRUTE_FORCE: GPUBruteForceConfig,
 }

vectordb_bench/frontend/config/dbCaseConfigs.py

@@ -173,6 +173,7 @@ class CaseConfigInput(BaseModel):
             IndexType.GPU_IVF_FLAT.value,
             IndexType.GPU_IVF_PQ.value,
             IndexType.GPU_CAGRA.value,
+            IndexType.GPU_BRUTE_FORCE.value,
         ],
     },
 )
@@ -562,6 +563,7 @@ class CaseConfigInput(BaseModel):
         IndexType.IVFSQ8.value,
         IndexType.GPU_IVF_FLAT.value,
         IndexType.GPU_IVF_PQ.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -579,6 +581,7 @@ class CaseConfigInput(BaseModel):
         IndexType.IVFSQ8.value,
         IndexType.GPU_IVF_FLAT.value,
         IndexType.GPU_IVF_PQ.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -703,6 +706,7 @@ class CaseConfigInput(BaseModel):
         IndexType.GPU_CAGRA.value,
         IndexType.GPU_IVF_PQ.value,
         IndexType.GPU_IVF_FLAT.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )
 
@@ -720,6 +724,7 @@ class CaseConfigInput(BaseModel):
         IndexType.GPU_CAGRA.value,
         IndexType.GPU_IVF_PQ.value,
         IndexType.GPU_IVF_FLAT.value,
+        IndexType.GPU_BRUTE_FORCE.value,
     ],
 )