vectordb-bench 1.0.4__py3-none-any.whl → 1.0.7__py3-none-any.whl

vectordb_bench/backend/clients/hologres/config.py

@@ -0,0 +1,121 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class HologresConfig(DBConfig):
+    user_name: SecretStr = SecretStr("hologres")
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 5432
+    db_name: str
+
+    def to_dict(self) -> dict:
+        user_str = self.user_name.get_secret_value()
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "dbname": self.db_name,
+            "user": user_str,
+            "password": pwd_str,
+        }
+
+
+class HologresIndexConfig(BaseModel, DBCaseConfig):
+    index: IndexType = IndexType.Hologres_HGraph
+    metric_type: MetricType | None = None
+
+    create_index_before_load: bool = False
+    create_index_after_load: bool = True
+
+    min_flush_proxima_row_count: int = 1000
+    min_compaction_proxima_row_count: int = 1000
+    max_total_size_to_merge_mb: int = 4096
+    full_compact_max_file_size_mb: int = 4096
+
+    base_quantization_type: str = "sq8_uniform"
+    precise_quantization_type: str = "fp32"
+    use_reorder: bool = True
+    build_thread_count: int = 16
+    max_degree: int = 64
+    ef_construction: int = 400
+
+    ef_search: int = 51
+
+    def index_param(self) -> dict:
+        return {
+            "algorithm": self.algorithm(),
+            "distance_method": self.distance_method(),
+            "builder_params": self.builder_params(),
+            "full_compact_max_file_size_mb": self.full_compact_max_file_size_mb,
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "distance_function": self.distance_function(),
+            "order_direction": self.order_direction(),
+            "searcher_params": self.search_params(),
+        }
+
+    def algorithm(self) -> str:
+        return self.index.value
+
+    def is_proxima(self) -> bool:
+        return self.index == IndexType.Hologres_Graph
+
+    def distance_method(self) -> str:
+        if self.metric_type == MetricType.L2:
+            if self.index == IndexType.Hologres_Graph:
+                return "SquaredEuclidean"
+            return "Euclidean"
+        if self.metric_type == MetricType.IP:
+            return "InnerProduct"
+        if self.metric_type == MetricType.COSINE:
+            if self.index == IndexType.Hologres_Graph:
+                return "InnerProduct"
+            return "Cosine"
+        return "Euclidean"
+
+    def distance_function(self) -> str:
+        if self.metric_type == MetricType.L2:
+            if self.index == IndexType.Hologres_Graph:
+                return "approx_squared_euclidean_distance"
+            return "approx_euclidean_distance"
+        if self.metric_type == MetricType.IP:
+            return "approx_inner_product_distance"
+        if self.metric_type == MetricType.COSINE:
+            if self.index == IndexType.Hologres_Graph:
+                return "approx_inner_product_distance"
+            return "approx_cosine_distance"
+        return "approx_euclidean_distance"
+
+    def order_direction(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "ASC"
+        if self.metric_type in {MetricType.IP, MetricType.COSINE}:
+            return "DESC"
+        return "ASC"
+
+    def builder_params(self) -> dict:
+        if self.use_reorder:
+            self.base_quantization_type = "sq8_uniform"
+        else:
+            self.base_quantization_type = "fp32"
+
+        return {
+            "min_flush_proxima_row_count": self.min_flush_proxima_row_count,
+            "min_compaction_proxima_row_count": self.min_compaction_proxima_row_count,
+            "max_total_size_to_merge_mb": self.max_total_size_to_merge_mb,
+            "build_thread_count": self.build_thread_count,
+            "base_quantization_type": self.base_quantization_type,
+            "max_degree": self.max_degree,
+            "ef_construction": self.ef_construction,
+            "precise_quantization_type": self.precise_quantization_type,
+            "use_reorder": self.use_reorder,
+        }
+
+    def searcher_params(self) -> dict:
+        return {
+            "ef_search": self.ef_search,
+        }

vectordb_bench/backend/clients/hologres/hologres.py

@@ -0,0 +1,365 @@
+"""Wrapper around the Hologres vector database over VectorDB"""
+
+import json
+import logging
+from collections.abc import Generator
+from contextlib import contextmanager
+from io import StringIO
+from typing import Any
+
+import psycopg
+from psycopg import Connection, Cursor, sql
+
+from ..api import VectorDB
+from .config import HologresConfig, HologresIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class Hologres(VectorDB):
+    """Use psycopg instructions"""
+
+    conn: psycopg.Connection[Any] | None = None
+    cursor: psycopg.Cursor[Any] | None = None
+
+    _tg_name: str = "vdb_bench_tg_1"
+
+    def __init__(
+        self,
+        dim: int,
+        db_config: HologresConfig,
+        db_case_config: HologresIndexConfig,
+        collection_name: str = "vector_collection",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.name = "Alibaba Cloud Hologres"
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+
+        self._primary_field = "id"
+        self._vector_field = "embedding"
+
+        # construct basic units
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
+        # create vector extension
+        if self.case_config.is_proxima():
+            self.cursor.execute("CREATE EXTENSION proxima;")
+            self.conn.commit()
+
+        log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}")
+        if not any(
+            (
+                self.case_config.create_index_before_load,
+                self.case_config.create_index_after_load,
+            ),
+        ):
+            msg = (
+                f"{self.name} config must create an index using create_index_before_load or create_index_after_load"
+                f"{self.name} config values: {self.db_config}\n{self.case_config}"
+            )
+            log.error(msg)
+            raise RuntimeError(msg)
+
+        if drop_old:
+            self._drop_table()
+            self._create_table(dim)
+            if self.case_config.create_index_before_load:
+                self._create_index()
+
+        self.cursor.close()
+        self.conn.close()
+        self.cursor = None
+        self.conn = None
+
+    @staticmethod
+    def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
+        conn = psycopg.connect(**kwargs)
+        conn.autocommit = True
+        cursor = conn.cursor()
+
+        assert conn is not None, "Connection is not initialized"
+        assert cursor is not None, "Cursor is not initialized"
+
+        return conn, cursor
+
+    @contextmanager
+    def init(self) -> Generator[None, None, None]:
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
+        self._set_search_guc()
+
+        try:
+            yield
+        finally:
+            self.cursor.close()
+            self.conn.close()
+            self.cursor = None
+            self.conn = None
+
+    def _set_search_guc(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        sql_guc = sql.SQL(f"SET hg_vector_ef_search = {self.case_config.ef_search};")
+        log.info(f"{self.name} client set search guc: {sql_guc.as_string()}")
+        self.cursor.execute(sql_guc)
+        self.conn.commit()
+
+    def _drop_table(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        log.info(f"{self.name} client drop table : {self.table_name}")
+        self.cursor.execute(
+            sql.SQL("DROP TABLE IF EXISTS {table_name};").format(
+                table_name=sql.Identifier(self.table_name),
+            ),
+        )
+        self.conn.commit()
+
+        try:
+            log.info(f"{self.name} client drop table group : {self._tg_name}")
+            self.cursor.execute(sql.SQL(f"CALL HG_DROP_TABLE_GROUP('{self._tg_name}');"))
+        except Exception as e:
+            log.info(f"{self.name} client drop table group : {self._tg_name} failed, error: {e}, ignore.")
+        finally:
+            self.conn.commit()
+
+    def optimize(self, data_size: int | None = None):
+        if self.case_config.create_index_after_load:
+            self._create_index()
+        self._full_compact()
+        self._analyze()
+
+    def _vacuum(self):
+        log.info(f"{self.name} client vacuum table : {self.table_name}")
+        try:
+            # VACUUM cannot run inside a transaction block
+            # it's better to new a connection
+            self.conn.autocommit = True
+            with self.conn.cursor() as cursor:
+                cursor.execute(
+                    sql.SQL(
+                        """
+                        VACUUM {table_name};
+                        """
+                    ).format(
+                        table_name=sql.Identifier(self.table_name),
+                    )
+                )
+            log.info(f"{self.name} client vacuum table : {self.table_name} done")
+        except Exception as e:
+            log.warning(f"Failed to vacuum table: {self.table_name} error: {e}")
+            raise e from None
+        finally:
+            self.conn.autocommit = True
+
+    def _analyze(self):
+        log.info(f"{self.name} client analyze table : {self.table_name}")
+        self.cursor.execute(sql.SQL(f"ANALYZE {self.table_name};"))
+        log.info(f"{self.name} client analyze table : {self.table_name} done")
+
+    def _full_compact(self):
+        log.info(f"{self.name} client full compact table : {self.table_name}")
+        self.cursor.execute(
+            sql.SQL(
+                """
+                SELECT hologres.hg_full_compact_table(
+                    '{table_name}',
+                    'max_file_size_mb={full_compact_max_file_size_mb}'
+                );
+                """
+            ).format(
+                table_name=sql.SQL(self.table_name),
+                full_compact_max_file_size_mb=sql.SQL(str(self.case_config.full_compact_max_file_size_mb)),
+            )
+        )
+        log.info(f"{self.name} client full compact table : {self.table_name} done")
+
+    def _create_index(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        sql_index = sql.SQL(
+            """
+            CALL set_table_property ('{table_name}', 'vectors', '{{
+                "embedding": {{
+                    "algorithm": "{algorithm}",
+                    "distance_method": "{distance_method}",
+                    "builder_params": {builder_params}
+                }}
+            }}');
+            """
+        ).format(
+            table_name=sql.Identifier(self.table_name),
+            algorithm=sql.SQL(self.case_config.algorithm()),
+            distance_method=sql.SQL(self.case_config.distance_method()),
+            builder_params=sql.SQL(json.dumps(self.case_config.builder_params())),
+        )
+
+        log.info(f"{self.name} client create index on table : {self.table_name}, with sql: {sql_index.as_string()}")
+        try:
+            self.cursor.execute(sql_index)
+            self.conn.commit()
+        except Exception as e:
+            log.warning(f"Failed to create index on table: {self.table_name} error: {e}")
+            raise e from None
+
+    def _set_replica_count(self, replica_count: int = 2):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        try:
+            # non-warehouse mode by default
+            sql_tg_replica = sql.SQL(
+                f"CALL hg_set_table_group_property('{self._tg_name}', 'replica_count', '{replica_count}');"
+            )
+
+            # check warehouse mode
+            sql_check = sql.SQL("select count(*) from hologres.hg_warehouses;")
+            log.info(f"check warehouse mode with sql: {sql_check}")
+            self.cursor.execute(sql_check)
+            result_check = self.cursor.fetchone()[0]
+            if result_check > 0:
+                # get warehouse name
+                sql_get_warehouse_name = sql.SQL("select current_warehouse();")
+                log.info(f"get warehouse name with sql: {sql_get_warehouse_name}")
+                self.cursor.execute(sql_get_warehouse_name)
+                sql_tg_replica = sql.SQL(
+                    """
+                    CALL hg_table_group_set_warehouse_replica_count (
+                        '{dbname}.{tg_name}',
+                        {replica_count},
+                        '{warehouse_name}'
+                    );
+                    """
+                ).format(
+                    tg_name=sql.SQL(self._tg_name),
+                    warehouse_name=sql.SQL(self.cursor.fetchone()[0]),
+                    dbname=sql.SQL(self.db_config["dbname"]),
+                    replica_count=replica_count,
+                )
+            log.info(f"{self.name} client set table group replica: {self._tg_name}, with sql: {sql_tg_replica}")
+            self.cursor.execute(sql_tg_replica)
+        except Exception as e:
+            log.warning(f"Failed to set replica count, error: {e}, ignore")
+        finally:
+            self.conn.commit()
+
+    def _create_table(self, dim: int):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        sql_tg = sql.SQL(f"CALL HG_CREATE_TABLE_GROUP ('{self._tg_name}', 1);")
+        log.info(f"{self.name} client create table group : {self._tg_name}, with sql: {sql_tg}")
+        try:
+            self.cursor.execute(sql_tg)
+        except Exception as e:
+            log.warning(f"Failed to create table group : {self._tg_name} error: {e}, ignore")
+        finally:
+            self.conn.commit()
+
+        self._set_replica_count(replica_count=2)
+
+        sql_table = sql.SQL(
+            """
+            CREATE TABLE IF NOT EXISTS {table_name} (
+                id BIGINT PRIMARY KEY,
+                embedding FLOAT4[] CHECK (array_ndims(embedding) = 1 AND array_length(embedding, 1) = {dim})
+            )
+            WITH (table_group = {tg_name});
+            """
+        ).format(
+            table_name=sql.Identifier(self.table_name),
+            dim=dim,
+            tg_name=sql.SQL(self._tg_name),
+        )
+        log.info(f"{self.name} client create table : {self.table_name}, with sql: {sql_table.as_string()}")
+        try:
+            self.cursor.execute(sql_table)
+            self.conn.commit()
+        except Exception as e:
+            log.warning(f"Failed to create table : {self.table_name} error: {e}")
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> tuple[int, Exception | None]:
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        try:
+            buffer = StringIO()
+            for i in range(len(metadata)):
+                buffer.write("%d\t%s\n" % (metadata[i], "{" + ",".join("%f" % x for x in embeddings[i]) + "}"))
+            buffer.seek(0)
+
+            with self.cursor.copy(
+                sql.SQL("COPY {table_name} FROM STDIN").format(table_name=sql.Identifier(self.table_name))
+            ) as copy:
+                copy.write(buffer.getvalue())
+            self.conn.commit()
+
+            return len(metadata), None
+        except Exception as e:
+            log.warning(f"Failed to insert data into table ({self.table_name}), error: {e}")
+            return 0, e
+
+    def _compose_query_and_params(self, vec: list[float], topk: int, ge_id: int | None = None):
+        params = []
+
+        where_clause = sql.SQL("")
+        if ge_id is not None:
+            where_clause = sql.SQL(" WHERE id >= %s ")
+            params.append(ge_id)
+
+        vec_float4 = [psycopg._wrappers.Float4(i) for i in vec]
+        params.append(vec_float4)
+        params.append(topk)
+
+        query = sql.SQL(
+            """
+            SELECT id
+            FROM {table_name}
+            {where_clause}
+            ORDER BY {distance_function}(embedding, %b)
+            {order_direction}
+            LIMIT %s;
+            """
+        ).format(
+            table_name=sql.Identifier(self.table_name),
+            distance_function=sql.SQL(self.case_config.distance_function()),
+            where_clause=where_clause,
+            order_direction=sql.SQL(self.case_config.order_direction()),
+        )
+
+        return query, params
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        ge = filters.get("id") if filters else None
+        q, params = self._compose_query_and_params(query, k, ge)
+        result = self.cursor.execute(q, params, prepare=True, binary=True)
+        return [int(i[0]) for i in result.fetchall()]

vectordb_bench/backend/clients/lancedb/lancedb.py

@@ -65,6 +65,7 @@ class LanceDB(VectorDB):
         self,
         embeddings: list[list[float]],
         metadata: list[int],
+        **kwargs,
     ) -> tuple[int, Exception | None]:
         try:
             data = [{"id": meta, "vector": emb} for meta, emb in zip(metadata, embeddings, strict=False)]

vectordb_bench/backend/clients/milvus/cli.py

@@ -40,6 +40,17 @@ class MilvusTypedDict(TypedDict):
             show_default=True,
         ),
     ]
+    replica_number: Annotated[
+        int,
+        click.option(
+            "--replica-number",
+            type=int,
+            help="Number of replicas",
+            required=False,
+            default=1,
+            show_default=True,
+        ),
+    ]
 
 
 class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): ...
@@ -58,6 +69,7 @@ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=AutoIndexConfig(),
         **parameters,
@@ -77,6 +89,7 @@ def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=FLATConfig(),
         **parameters,
@@ -99,6 +112,7 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=HNSWConfig(
             M=parameters["m"],
@@ -139,19 +153,14 @@ class MilvusRefineTypedDict(TypedDict):
     ]
 
 
-class MilvusHNSWPQTypedDict(
-    CommonTypedDict,
-    MilvusTypedDict,
-    MilvusHNSWTypedDict,
-    MilvusRefineTypedDict
-):
+class MilvusHNSWPQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusHNSWTypedDict, MilvusRefineTypedDict):
     nbits: Annotated[
         int,
         click.option(
             "--nbits",
             type=int,
             required=True,
-        )
+        ),
     ]
 
 
@@ -168,6 +177,7 @@ def MilvusHNSWPQ(**parameters: Unpack[MilvusHNSWPQTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=HNSWPQConfig(
             M=parameters["m"],
@@ -194,7 +204,7 @@ class MilvusHNSWPRQTypedDict(
             type=int,
             help="The number of residual subquantizers.",
             required=True,
-        )
+        ),
     ]
 
 
@@ -211,6 +221,7 @@ def MilvusHNSWPRQ(**parameters: Unpack[MilvusHNSWPRQTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=HNSWPRQConfig(
             M=parameters["m"],
@@ -220,7 +231,7 @@ def MilvusHNSWPRQ(**parameters: Unpack[MilvusHNSWPRQTypedDict]):
             refine=parameters["refine"],
             refine_type=parameters["refine_type"],
             refine_k=parameters["refine_k"],
-            nrq=parameters["nrq"]
+            nrq=parameters["nrq"],
         ),
         **parameters,
     )
@@ -251,6 +262,7 @@ def MilvusHNSWSQ(**parameters: Unpack[MilvusHNSWSQTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=HNSWSQConfig(
             M=parameters["m"],
@@ -281,6 +293,7 @@ def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=IVFFlatConfig(
             nlist=parameters["nlist"],
@@ -303,6 +316,7 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=IVFSQ8Config(
             nlist=parameters["nlist"],
@@ -364,6 +378,7 @@ def MilvusIVFRabitQ(**parameters: Unpack[MilvusIVFRABITQTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=IVFRABITQConfig(
             nlist=parameters["nlist"],
@@ -394,6 +409,7 @@ def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=DISKANNConfig(
             search_list=parameters["search_list"],
@@ -423,6 +439,7 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=GPUIVFFlatConfig(
             nlist=parameters["nlist"],
@@ -458,6 +475,7 @@ def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=GPUBruteForceConfig(
             metric_type=parameters["metric_type"],
@@ -490,6 +508,7 @@ def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=GPUIVFPQConfig(
             nlist=parameters["nlist"],
@@ -530,6 +549,7 @@ def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
             num_shards=int(parameters["num_shards"]),
+            replica_number=int(parameters["replica_number"]),
         ),
         db_case_config=GPUCAGRAConfig(
             intermediate_graph_degree=parameters["intermediate_graph_degree"],

vectordb_bench/backend/clients/milvus/config.py

@@ -8,6 +8,7 @@ class MilvusConfig(DBConfig):
     user: str | None = None
     password: SecretStr | None = None
     num_shards: int = 1
+    replica_number: int = 1
 
     def to_dict(self) -> dict:
         return {
@@ -15,6 +16,7 @@ class MilvusConfig(DBConfig):
             "user": self.user if self.user else None,
             "password": self.password.get_secret_value() if self.password else None,
             "num_shards": self.num_shards,
+            "replica_number": self.replica_number,
         }
 
     @validator("*")

vectordb_bench/backend/clients/milvus/milvus.py

@@ -92,7 +92,7 @@ class Milvus(VectorDB):
             )
 
             self.create_index()
-            col.load()
+            col.load(replica_number=self.db_config.get("replica_number", 1))
 
         connections.disconnect("default")
 

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"],
         ),