vectordb-bench 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

vectordb_bench/__init__.py

@@ -18,12 +18,23 @@ class config:
     USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)
 
     RESULTS_LOCAL_DIR = pathlib.Path(__file__).parent.joinpath("results")
-    CASE_TIMEOUT_IN_SECOND = 24 * 60 * 60
+
+    CAPACITY_TIMEOUT_IN_SECONDS =  24 * 3600 # 24h
+    LOAD_TIMEOUT_1M             = 2.5 * 3600 # 2.5h
+    LOAD_TIMEOUT_10M            =  25 * 3600 # 25h
+    LOAD_TIMEOUT_100M           = 250 * 3600 # 10.41d
+
+    OPTIMIZE_TIMEOUT_1M         =  15 * 60   # 15min
+    OPTIMIZE_TIMEOUT_10M        = 2.5 * 3600 # 2.5h
+    OPTIMIZE_TIMEOUT_100M       =  25 * 3600 # 1.04d
 
 
     def display(self) -> str:
-        tmp = [i for i in inspect.getmembers(self)
-            if not inspect.ismethod(i[1]) and not i[0].startswith('_') \
+        tmp = [
+            i for i in inspect.getmembers(self)
+            if not inspect.ismethod(i[1])
+            and not i[0].startswith('_')
+            and "TIMEOUT" not in i[0]
         ]
         return tmp
 

vectordb_bench/backend/assembler.py

@@ -1,4 +1,4 @@
-from .cases import type2case, CaseLabel
+from .cases import CaseLabel
 from .task_runner import CaseRunner, RunningStatus, TaskRunner
 from ..models import TaskConfig
 from ..backend.clients import EmptyDBCaseConfig
@@ -11,7 +11,7 @@ log = logging.getLogger(__name__)
 class Assembler:
     @classmethod
     def assemble(cls, run_id , task: TaskConfig) -> CaseRunner:
-        c_cls = type2case.get(task.case_config.case_id)
+        c_cls = task.case_config.case_id.case_cls
 
         c = c_cls()
         if type(task.db_case_config) != EmptyDBCaseConfig:

vectordb_bench/backend/cases.py

@@ -1,13 +1,59 @@
+import typing
 import logging
 from enum import Enum, auto
 
-from . import dataset as ds
-from ..base import BaseModel
-from ..models import CaseType
+from vectordb_bench import config
+from vectordb_bench.base import BaseModel
+
+from .dataset import Dataset, DatasetManager
 
 
 log = logging.getLogger(__name__)
 
+Case = typing.TypeVar("Case")
+
+
+class CaseType(Enum):
+    """
+    Example:
+        >>> case_cls = CaseType.CapacityDim128.case_cls
+        >>> assert c is not None
+        >>> CaseType.CapacityDim128.case_name
+        "Capacity Test (128 Dim Repeated)"
+    """
+
+    CapacityDim128 = 1
+    CapacityDim960 = 2
+
+    Performance100M = 3
+    Performance10M = 4
+    Performance1M = 5
+
+    Performance10M1P = 6
+    Performance1M1P = 7
+    Performance10M99P = 8
+    Performance1M99P = 9
+
+    Custom = 100
+
+    @property
+    def case_cls(self, custom_configs: dict | None = None) -> Case:
+        return type2case.get(self)
+
+    @property
+    def case_name(self) -> str:
+        c = self.case_cls
+        if c is not None:
+            return c().name
+        raise ValueError("Case unsupported")
+
+    @property
+    def case_description(self) -> str:
+        c = self.case_cls
+        if c is not None:
+            return c().description
+        raise ValueError("Case unsupported")
+
 
 class CaseLabel(Enum):
     Load = auto()
@@ -15,10 +61,10 @@ class CaseLabel(Enum):
 
 
 class Case(BaseModel):
-    """ Undifined case
+    """Undifined case
 
     Fields:
-        case_id(CaseType): default 11 case type plus one custom cases.
+        case_id(CaseType): default 9 case type plus one custom cases.
         label(CaseLabel): performance or load.
         dataset(DataSet): dataset for this case runner.
         filter_rate(float | None): one of 99% | 1% | None
@@ -27,7 +73,12 @@ class Case(BaseModel):
 
     case_id: CaseType
     label: CaseLabel
-    dataset: ds.DataSet
+    name: str
+    description: str
+    dataset: DatasetManager
+
+    load_timeout: float | int
+    optimize_timeout: float | int | None
 
     filter_rate: float | None
 
@@ -45,80 +96,118 @@ class Case(BaseModel):
 
 class CapacityCase(Case, BaseModel):
     label: CaseLabel = CaseLabel.Load
-    filter_rate: float | int | None = None
+    filter_rate: float | None = None
+    load_timeout: float | int = config.CAPACITY_TIMEOUT_IN_SECONDS
+    optimize_timeout: float | int | None = None
+
 
 class PerformanceCase(Case, BaseModel):
     label: CaseLabel = CaseLabel.Performance
-    filter_rate: float | int | None = None
+    filter_rate: float | None = None
 
-class CapacityLDimCase(CapacityCase):
-    case_id: CaseType = CaseType.CapacityLDim
-    dataset: ds.DataSet = ds.get(ds.Name.GIST, ds.Label.SMALL)
 
-class CapacitySDimCase(CapacityCase):
-    case_id: CaseType = CaseType.CapacitySDim
-    dataset: ds.DataSet = ds.get(ds.Name.SIFT, ds.Label.SMALL)
+class CapacityDim960(CapacityCase):
+    case_id: CaseType = CaseType.CapacityDim960
+    dataset: DatasetManager = Dataset.GIST.manager(100_000)
+    name: str = "Capacity Test (960 Dim Repeated)"
+    description: str = """This case tests the vector database's loading capacity by repeatedly inserting large-dimension vectors (GIST 100K vectors, <b>960 dimensions</b>) until it is fully loaded.
+Number of inserted vectors will be reported."""
 
-class PerformanceLZero(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceLZero
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.LARGE)
 
-class PerformanceMZero(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceMZero
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.MEDIUM)
+class CapacityDim128(CapacityCase):
+    case_id: CaseType = CaseType.CapacityDim128
+    dataset: DatasetManager = Dataset.SIFT.manager(500_000)
+    name: str = "Capacity Test (128 Dim Repeated)"
+    description: str = """This case tests the vector database's loading capacity by repeatedly inserting small-dimension vectors (SIFT 100K vectors, <b>128 dimensions</b>) until it is fully loaded.
+Number of inserted vectors will be reported."""
 
-class PerformanceSZero(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceSZero
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.SMALL)
 
-class PerformanceLLow(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceLLow
-    filter_rate: float | int | None = 0.01
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.LARGE)
+class Performance10M(PerformanceCase):
+    case_id: CaseType = CaseType.Performance10M
+    dataset: DatasetManager = Dataset.COHERE.manager(10_000_000)
+    name: str = "Search Performance Test (10M Dataset, 768 Dim)"
+    description: str = """This case tests the search performance of a vector database with a large dataset (<b>Cohere 10M vectors</b>, 768 dimensions) at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_10M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_10M
+
+
+class Performance1M(PerformanceCase):
+    case_id: CaseType = CaseType.Performance1M
+    dataset: DatasetManager = Dataset.COHERE.manager(1_000_000)
+    name: str = "Search Performance Test (1M Dataset, 768 Dim)"
+    description: str = """This case tests the search performance of a vector database with a medium dataset (<b>Cohere 1M vectors</b>, 768 dimensions) at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_1M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1M
+
 
-class PerformanceMLow(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceMLow
+class Performance10M1P(PerformanceCase):
+    case_id: CaseType = CaseType.Performance10M1P
     filter_rate: float | int | None = 0.01
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.MEDIUM)
+    dataset: DatasetManager = Dataset.COHERE.manager(10_000_000)
+    name: str = "Filtering Search Performance Test (10M Dataset, 768 Dim, Filter 1%)"
+    description: str = """This case tests the search performance of a vector database with a large dataset (<b>Cohere 10M vectors</b>, 768 dimensions) under a low filtering rate (<b>1% vectors</b>), at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_10M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_10M
 
-class PerformanceSLow(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceSLow
+
+class Performance1M1P(PerformanceCase):
+    case_id: CaseType = CaseType.Performance1M1P
     filter_rate: float | int | None = 0.01
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.SMALL)
+    dataset: DatasetManager = Dataset.COHERE.manager(1_000_000)
+    name: str = "Filtering Search Performance Test (1M Dataset, 768 Dim, Filter 1%)"
+    description: str = """This case tests the search performance of a vector database with a medium dataset (<b>Cohere 1M vectors</b>, 768 dimensions) under a low filtering rate (<b>1% vectors</b>), at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_1M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1M
 
-class PerformanceLHigh(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceLHigh
-    filter_rate: float | int | None = 0.99
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.LARGE)
 
-class PerformanceMHigh(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceMHigh
+class Performance10M99P(PerformanceCase):
+    case_id: CaseType = CaseType.Performance10M99P
     filter_rate: float | int | None = 0.99
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.MEDIUM)
+    dataset: DatasetManager = Dataset.COHERE.manager(10_000_000)
+    name: str = "Filtering Search Performance Test (10M Dataset, 768 Dim, Filter 99%)"
+    description: str = """This case tests the search performance of a vector database with a large dataset (<b>Cohere 10M vectors</b>, 768 dimensions) under a high filtering rate (<b>99% vectors</b>), at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_10M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_10M
 
-class PerformanceSHigh(PerformanceCase):
-    case_id: CaseType = CaseType.PerformanceSLow
+
+class Performance1M99P(PerformanceCase):
+    case_id: CaseType = CaseType.Performance1M99P
     filter_rate: float | int | None = 0.99
-    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.SMALL)
+    dataset: DatasetManager = Dataset.COHERE.manager(1_000_000)
+    name: str = "Filtering Search Performance Test (1M Dataset, 768 Dim, Filter 99%)"
+    description: str = """This case tests the search performance of a vector database with a medium dataset (<b>Cohere 1M vectors</b>, 768 dimensions) under a high filtering rate (<b>99% vectors</b>), at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_1M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1M
+
+
 
 class Performance100M(PerformanceCase):
     case_id: CaseType = CaseType.Performance100M
     filter_rate: float | int | None = None
-    dataset: ds.DataSet = ds.get(ds.Name.LAION, ds.Label.LARGE)
+    dataset: DatasetManager = Dataset.LAION.manager(100_000_000)
+    name: str = "Search Performance Test (100M Dataset, 768 Dim)"
+    description: str = """This case tests the search performance of a vector database with a large 100M dataset (<b>LAION 100M vectors</b>, 768 dimensions), at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = config.LOAD_TIMEOUT_100M
+    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_100M
+
 
 type2case = {
-    CaseType.CapacityLDim: CapacityLDimCase,
-    CaseType.CapacitySDim: CapacitySDimCase,
-
-    CaseType.PerformanceLZero: PerformanceLZero,
-    CaseType.PerformanceMZero: PerformanceMZero,
-    CaseType.PerformanceSZero: PerformanceSZero,
-
-    CaseType.PerformanceLLow: PerformanceLLow,
-    CaseType.PerformanceMLow: PerformanceMLow,
-    CaseType.PerformanceSLow: PerformanceSLow,
-    CaseType.PerformanceLHigh: PerformanceLHigh,
-    CaseType.PerformanceMHigh: PerformanceMHigh,
-    CaseType.PerformanceSHigh: PerformanceSHigh,
+    CaseType.CapacityDim960: CapacityDim960,
+    CaseType.CapacityDim128: CapacityDim128,
+
     CaseType.Performance100M: Performance100M,
+    CaseType.Performance10M: Performance10M,
+    CaseType.Performance1M: Performance1M,
+
+    CaseType.Performance10M1P: Performance10M1P,
+    CaseType.Performance1M1P: Performance1M1P,
+    CaseType.Performance10M99P: Performance10M99P,
+    CaseType.Performance1M99P: Performance1M99P,
 }

vectordb_bench/backend/clients/__init__.py

@@ -15,7 +15,7 @@ from .pinecone.pinecone import Pinecone
 from .weaviate_cloud.weaviate_cloud import WeaviateCloud
 from .qdrant_cloud.qdrant_cloud import QdrantCloud
 from .zilliz_cloud.zilliz_cloud import ZillizCloud
-
+from .pgvector.pgvector import PgVector
 
 class DB(Enum):
     """Database types
@@ -35,6 +35,7 @@ class DB(Enum):
     ElasticCloud = "ElasticCloud"
     QdrantCloud = "QdrantCloud"
     WeaviateCloud = "WeaviateCloud"
+    PgVector = "PgVector"
 
 
     @property
@@ -49,8 +50,12 @@ db2client = {
     DB.ElasticCloud: ElasticCloud,
     DB.QdrantCloud: QdrantCloud,
     DB.Pinecone: Pinecone,
+    DB.PgVector: PgVector
 }
 
+for db in DB:
+    assert issubclass(db.init_cls, VectorDB)
+
 
 __all__ = [
     "DB", "VectorDB", "DBConfig", "DBCaseConfig", "IndexType", "MetricType", "EmptyDBCaseConfig",

vectordb_bench/backend/clients/api.py

@@ -3,7 +3,7 @@ from enum import Enum
 from typing import Any, Type
 from contextlib import contextmanager
 
-from pydantic import BaseModel
+from pydantic import BaseModel, validator, SecretStr
 
 
 class MetricType(str, Enum):
@@ -32,12 +32,20 @@ class DBConfig(ABC, BaseModel):
             ZillizCloudConfig.db_label = 1cu-perf
     """
 
-    db_label: str | None = None
+    db_label: str = ""
 
     @abstractmethod
     def to_dict(self) -> dict:
         raise NotImplementedError
 
+    @validator("*")
+    def not_empty_field(cls, v, field):
+        if field.name == "db_label":
+            return v
+        if isinstance(v, (str, SecretStr)) and len(v) == 0:
+            raise ValueError("Empty string!")
+        return v
+
 
 class DBCaseConfig(ABC):
     """Case specific vector database configs, usually uesed for index params like HNSW"""
@@ -65,7 +73,7 @@ class VectorDB(ABC):
 
     In each process, the benchmark cases ensure VectorDB.init() calls before any other methods operations
 
-    insert_embeddings, search_embedding, and, ready_to_search will be timed for each call.
+    insert_embeddings, search_embedding, and, optimize will be timed for each call.
 
     Examples:
         >>> milvus = Milvus()
@@ -82,9 +90,12 @@ class VectorDB(ABC):
         db_case_config: DBCaseConfig | None,
         collection_name: str,
         drop_old: bool = False,
-        **kwargs
+        **kwargs,
     ) -> None:
-        """Initialize wrapper around the vector database client
+        """Initialize wrapper around the vector database client.
+
+        Please drop the existing collection if drop_old is True. And create collection
+        if collection not in the Vector Database
 
         Args:
             dim(int): the dimension of the dataset
@@ -122,15 +133,15 @@ class VectorDB(ABC):
         self,
         embeddings: list[list[float]],
         metadata: list[int],
-        kwargs: Any,
-    ) -> int:
+        **kwargs,
+    ) -> (int, Exception):
         """Insert the embeddings to the vector database. The default number of embeddings for
         each insert_embeddings is 5000.
 
         Args:
             embeddings(list[list[float]]): list of embedding to add to the vector database.
             metadatas(list[int]): metadata associated with the embeddings, for filtering.
-            kwargs(Any): vector database specific parameters.
+            **kwargs(Any): vector database specific parameters.
 
         Returns:
             int: inserted data count
@@ -158,13 +169,14 @@ class VectorDB(ABC):
 
     # TODO: remove
     @abstractmethod
-    def ready_to_search(self):
-        """ready_to_search will be called between insertion and search in performance cases.
+    def optimize(self):
+        """optimize will be called between insertion and search in performance cases.
 
         Should be blocked until the vectorDB is ready to be tested on
         heavy performance cases.
 
-        Time(insert the dataset) + Time(ready_to_search) will be recorded as "load_duration" metric
+        Time(insert the dataset) + Time(optimize) will be recorded as "load_duration" metric
+        Optimize's execution time is limited, the limited time is based on cases.
         """
         raise NotImplementedError
 

vectordb_bench/backend/clients/elastic_cloud/config.py

@@ -4,9 +4,9 @@ from pydantic import SecretStr, BaseModel
 from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
 
 
-class ElasticsearchConfig(DBConfig, BaseModel):
+class ElasticCloudConfig(DBConfig, BaseModel):
     cloud_id: SecretStr
-    password: SecretStr | None = None
+    password: SecretStr
 
     def to_dict(self) -> dict:
         return {
@@ -20,7 +20,7 @@ class ESElementType(str, Enum):
     byte = "byte"  # 1 byte, -128 to 127
 
 
-class ElasticsearchIndexConfig(BaseModel, DBCaseConfig):
+class ElasticCloudIndexConfig(BaseModel, DBCaseConfig):
     element_type: ESElementType = ESElementType.float
     index: IndexType = IndexType.ES_HNSW  # ES only support 'hnsw'
 
@@ -45,8 +45,8 @@ class ElasticsearchIndexConfig(BaseModel, DBCaseConfig):
             "index_options": {
                 "type": self.index.value,
                 "m": self.M,
-                "ef_construction": self.efConstruction
-            }
+                "ef_construction": self.efConstruction,
+            },
         }
         return params
 

vectordb_bench/backend/clients/elastic_cloud/elastic_cloud.py

@@ -2,7 +2,7 @@ import logging
 from contextlib import contextmanager
 from typing import Iterable, Type
 from ..api import VectorDB, DBCaseConfig, DBConfig, IndexType
-from .config import ElasticsearchIndexConfig, ElasticsearchConfig
+from .config import ElasticCloudIndexConfig, ElasticCloudConfig
 from elasticsearch.helpers import bulk
 
 
@@ -16,11 +16,12 @@ class ElasticCloud(VectorDB):
         self,
         dim: int,
         db_config: dict,
-        db_case_config: ElasticsearchIndexConfig,
+        db_case_config: ElasticCloudIndexConfig,
         indice: str = "vdb_bench_indice",  # must be lowercase
         id_col_name: str = "id",
         vector_col_name: str = "vector",
         drop_old: bool = False,
+        **kwargs,
     ):
         self.dim = dim
         self.db_config = db_config
@@ -43,12 +44,12 @@ class ElasticCloud(VectorDB):
 
     @classmethod
     def config_cls(cls) -> Type[DBConfig]:
-        return ElasticsearchConfig
+        return ElasticCloudConfig
 
 
     @classmethod
     def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
-        return ElasticsearchIndexConfig
+        return ElasticCloudIndexConfig
 
 
     @contextmanager
@@ -83,7 +84,8 @@ class ElasticCloud(VectorDB):
         self,
         embeddings: Iterable[list[float]],
         metadata: list[int],
-    ) -> int:
+        **kwargs,
+    ) -> (int, Exception):
         """Insert the embeddings to the elasticsearch."""
         assert self.client is not None, "should self.init() first"
 
@@ -99,10 +101,10 @@ class ElasticCloud(VectorDB):
         ]
         try:
             bulk_insert_res = bulk(self.client, insert_data)
-            return bulk_insert_res[0]
+            return (bulk_insert_res[0], None)
         except Exception as e:
             log.warning(f"Failed to insert data: {self.indice} error: {str(e)}")
-            raise e from None
+            return (0, e)
 
     def search_embedding(
         self,
@@ -143,8 +145,8 @@ class ElasticCloud(VectorDB):
             log.warning(f"Failed to search: {self.indice} error: {str(e)}")
             raise e from None
 
-    def ready_to_search(self):
-        """ready_to_search will be called between insertion and search in performance cases."""
+    def optimize(self):
+        """optimize will be called between insertion and search in performance cases."""
         pass
 
     def ready_to_load(self):

vectordb_bench/backend/clients/milvus/config.py

@@ -2,14 +2,13 @@ from pydantic import BaseModel, SecretStr
 from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
 
 
-class MilvusConfig(DBConfig, BaseModel):
-    uri: SecretStr | None = "http://localhost:19530"
+class MilvusConfig(DBConfig):
+    uri: SecretStr = "http://localhost:19530"
 
     def to_dict(self) -> dict:
         return {"uri": self.uri.get_secret_value()}
 
 
-
 class MilvusIndexConfig(BaseModel):
     """Base config for milvus"""
 

vectordb_bench/backend/clients/milvus/milvus.py

@@ -2,7 +2,7 @@
 
 import logging
 from contextlib import contextmanager
-from typing import Any, Iterable, Type
+from typing import Iterable, Type
 
 from pymilvus import Collection, utility
 from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusException
@@ -13,6 +13,7 @@ from .config import MilvusConfig, _milvus_case_config
 
 log = logging.getLogger(__name__)
 
+MILVUS_LOAD_REQS_SIZE = 1.5 * 1024 *1024
 
 class Milvus(VectorDB):
     def __init__(
@@ -23,12 +24,14 @@ class Milvus(VectorDB):
         collection_name: str = "VectorDBBenchCollection",
         drop_old: bool = False,
         name: str = "Milvus",
+        **kwargs,
     ):
         """Initialize wrapper around the milvus vector database."""
         self.name = name
         self.db_config = db_config
         self.case_config = db_case_config
         self.collection_name = collection_name
+        self.batch_size = int(MILVUS_LOAD_REQS_SIZE / (dim *4))
 
         self._primary_field = "pk"
         self._scalar_field = "id"
@@ -51,7 +54,7 @@ class Milvus(VectorDB):
             log.info(f"{self.name} create collection: {self.collection_name}")
 
             # Create the collection
-            coll = Collection(
+            Collection(
                 name=self.collection_name,
                 schema=CollectionSchema(fields),
                 consistency_level="Session",
@@ -105,6 +108,14 @@ class Milvus(VectorDB):
 
     def _optimize(self):
         log.info(f"{self.name} optimizing before search")
+        try:
+            self.col.load()
+        except Exception as e:
+            log.warning(f"{self.name} optimize error: {e}")
+            raise e from None
+
+    def _post_insert(self):
+        log.info(f"{self.name} post insert before optimize")
         try:
             self.col.flush()
             self.col.compact()
@@ -117,10 +128,6 @@ class Milvus(VectorDB):
                 index_name=self._index_name,
             )
             utility.wait_for_index_building_complete(self.collection_name)
-            self.col.load()
-            #  self.col.load(_refresh=True)
-            #  utility.wait_for_loading_complete(self.collection_name)
-            #  import time; time.sleep(10)
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
             raise e from None
@@ -130,7 +137,7 @@ class Milvus(VectorDB):
         self._pre_load(self.col)
         pass
 
-    def ready_to_search(self):
+    def optimize(self):
         assert self.col, "Please call self.init() before"
         self._optimize()
 
@@ -138,23 +145,29 @@ class Milvus(VectorDB):
         self,
         embeddings: Iterable[list[float]],
         metadata: list[int],
-        **kwargs: Any,
-    ) -> int:
+        **kwargs,
+    ) -> (int, Exception):
         """Insert embeddings into Milvus. should call self.init() first"""
         # use the first insert_embeddings to init collection
         assert self.col is not None
-        insert_data = [
-                metadata,
-                metadata,
-                embeddings,
-        ]
-
+        assert len(embeddings) == len(metadata)
+        insert_count = 0
         try:
-            res = self.col.insert(insert_data, **kwargs)
-            return len(res.primary_keys)
+            for batch_start_offset in range(0, len(embeddings), self.batch_size):
+                batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
+                insert_data = [
+                        metadata[batch_start_offset : batch_end_offset],
+                        metadata[batch_start_offset : batch_end_offset],
+                        embeddings[batch_start_offset : batch_end_offset],
+                ]
+                res = self.col.insert(insert_data)
+                insert_count += len(res.primary_keys)
+            if kwargs.get("last_batch"):
+                self._post_insert()
         except MilvusException as e:
-            log.warning("Failed to insert data")
-            raise e from None
+            log.info(f"Failed to insert data: {e}")
+            return (insert_count, e)
+        return (insert_count, None)
 
     def search_embedding(
         self,