Add nbits parameter to IVF_PQ index and adapt new filter logic (#576)

* support oceanbase

* fix bug

* fix bug

* support IP

* support multiple zone

* support filter case

* support ivf

* support hnsw_bq and extra info

* optimized formatting

* optimized formatting

* Revert dbCaseConfigs.py to match remote main branch

* resolve review comments

* Update __init__.py

* Addressed review comments

* Fixed code style issues

* Add nbits parameter to IVF_PQ index and adapt new filter logic

Author: wyfanxiao

vectordb_bench/backend/clients/oceanbase/cli.py

@@ -93,6 +93,7 @@ def OceanBaseIVF(**parameters: Unpack[OceanBaseIVFTypedDict]):
             m=input_m,
             nlist=parameters["nlist"],
             sample_per_nlist=parameters["sample_per_nlist"],
+            nbits=parameters["nbits"],
             index=input_index_type,
             ivf_nprobes=parameters["ivf_nprobes"],
         ),

vectordb_bench/backend/clients/oceanbase/config.py

@@ -85,6 +85,7 @@ def search_param(self) -> dict:
 class OceanBaseIVFConfig(OceanBaseIndexConfig, DBCaseConfig):
     m: int
     sample_per_nlist: int
+    nbits: int | None = None
     nlist: int
     index: IndexType
     ivf_nprobes: int | None = None
@@ -96,8 +97,9 @@ def index_param(self) -> dict:
                 "metric_type": self.parse_metric(),
                 "index_type": self.index.value,
                 "params": {
-                    "m": self.M,
+                    "m": self.m,
                     "sample_per_nlist": self.sample_per_nlist,
+                    "nbits": self.nbits,
                     "nlist": self.nlist,
                 },
             }

vectordb_bench/backend/clients/oceanbase/oceanbase.py

@@ -7,6 +7,8 @@
 
 import mysql.connector as mysql
 
+from vectordb_bench.backend.filter import Filter, FilterOp
+
 from ..api import IndexType, VectorDB
 from .config import OceanBaseConfigDict, OceanBaseHNSWConfig
 
@@ -16,6 +18,12 @@
 
 
 class OceanBase(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
     def __init__(
         self,
         dim: int,
@@ -187,22 +195,30 @@ def insert_embeddings(
 
         return insert_count, None
 
+    def prepare_filter(self, filters: Filter):
+        if filters.type == FilterOp.NonFilter:
+            self.expr = ""
+        elif filters.type == FilterOp.NumGE:
+            self.expr = f"WHERE id >= {filters.int_value}"
+        elif filters.type == FilterOp.StrEqual:
+            self.expr = f"WHERE id == '{filters.label_value}'"
+        else:
+            msg = f"Not support Filter for Oceanbase - {filters}"
+            raise ValueError(msg)
+
     def search_embedding(
         self,
         query: list[float],
         k: int = 100,
-        filters: dict[str, Any] | None = None,
-        timeout: int | None = None,
     ) -> list[int]:
         if not self._cursor:
             raise ValueError("Cursor is not initialized")
 
         packed = struct.pack(f"<{len(query)}f", *query)
         hex_vec = packed.hex()
-        filter_clause = f"WHERE id >= {filters['id']}" if filters else ""
         query_str = (
             f"SELECT id FROM {self.table_name} "  # noqa: S608
-            f"{filter_clause} ORDER BY "
+            f"{self.expr} ORDER BY "
             f"{self.db_case_config.parse_metric_func_str()}(embedding, X'{hex_vec}') "
             f"APPROXIMATE LIMIT {k}"
         )

vectordb_bench/cli/cli.py

@@ -501,6 +501,15 @@ class OceanBaseIVFTypedDict(TypedDict):
         int | None,
         click.option("--nlist", "nlist", type=int, help="Number of cluster centers", required=True),
     ]
+    nbits: Annotated[
+        int | None,
+        click.option(
+            "--nbits",
+            "nbits",
+            type=int,
+            help="Number of bits used to encode the index of a sub-vector's centroid in the compressed representation",
+        ),
+    ]
     sample_per_nlist: Annotated[
         int | None,
         click.option(