vectordb-bench 0.0.11__tar.gz → 0.0.13__tar.gz

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/.gitignore

@@ -9,3 +9,4 @@ __MACOSX
 build/
 venv/
 .idea/
+results/

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.11
+Version: 0.0.13
 Summary: VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -36,12 +36,14 @@ Requires-Dist: pinecone-client; extra == "all"
 Requires-Dist: weaviate-client; extra == "all"
 Requires-Dist: elasticsearch; extra == "all"
 Requires-Dist: pgvector; extra == "all"
+Requires-Dist: pgvecto_rs[psycopg3]>=0.2.1; extra == "all"
 Requires-Dist: sqlalchemy; extra == "all"
 Requires-Dist: redis; extra == "all"
 Requires-Dist: chromadb; extra == "all"
-Requires-Dist: psycopg2; extra == "all"
 Requires-Dist: psycopg; extra == "all"
 Requires-Dist: psycopg-binary; extra == "all"
+Requires-Dist: opensearch-dsl==2.1.0; extra == "all"
+Requires-Dist: opensearch-py==2.6.0; extra == "all"
 Provides-Extra: qdrant
 Requires-Dist: qdrant-client; extra == "qdrant"
 Provides-Extra: pinecone
@@ -54,12 +56,20 @@ Provides-Extra: pgvector
 Requires-Dist: psycopg; extra == "pgvector"
 Requires-Dist: psycopg-binary; extra == "pgvector"
 Requires-Dist: pgvector; extra == "pgvector"
+Provides-Extra: pgvectorscale
+Requires-Dist: psycopg; extra == "pgvectorscale"
+Requires-Dist: psycopg-binary; extra == "pgvectorscale"
+Requires-Dist: pgvector; extra == "pgvectorscale"
 Provides-Extra: pgvecto-rs
-Requires-Dist: psycopg2; extra == "pgvecto-rs"
+Requires-Dist: pgvecto_rs[psycopg3]>=0.2.1; extra == "pgvecto-rs"
 Provides-Extra: redis
 Requires-Dist: redis; extra == "redis"
+Provides-Extra: memorydb
+Requires-Dist: memorydb; extra == "memorydb"
 Provides-Extra: chromadb
 Requires-Dist: chromadb; extra == "chromadb"
+Provides-Extra: awsopensearch
+Requires-Dist: awsopensearch; extra == "awsopensearch"
 Provides-Extra: zilliz-cloud
 
 # VectorDBBench: A Benchmark Tool for VectorDB
@@ -91,18 +101,21 @@ pip install vectordb-bench[pinecone]
 ```
 All the database client supported
 
-|Optional database client|install command|
-|---------------|---------------|
-|pymilvus(*default*)|`pip install vectordb-bench`|
-|all|`pip install vectordb-bench[all]`|
-|qdrant|`pip install vectordb-bench[qdrant]`|
-|pinecone|`pip install vectordb-bench[pinecone]`|
-|weaviate|`pip install vectordb-bench[weaviate]`|
-|elastic|`pip install vectordb-bench[elastic]`|
-|pgvector|`pip install vectordb-bench[pgvector]`|
-|pgvecto.rs|`pip install vectordb-bench[pgvecto_rs]`|
-|redis|`pip install vectordb-bench[redis]`|
-|chromadb|`pip install vectordb-bench[chromadb]`|
+| Optional database client | install command                             |
+|--------------------------|---------------------------------------------|
+| pymilvus(*default*)      | `pip install vectordb-bench`                |
+| all                      | `pip install vectordb-bench[all]`           |
+| qdrant                   | `pip install vectordb-bench[qdrant]`        |
+| pinecone                 | `pip install vectordb-bench[pinecone]`      |
+| weaviate                 | `pip install vectordb-bench[weaviate]`      |
+| elastic                  | `pip install vectordb-bench[elastic]`       |
+| pgvector                 | `pip install vectordb-bench[pgvector]`      |
+| pgvecto.rs               | `pip install vectordb-bench[pgvecto_rs]`    |
+| pgvectorscale            | `pip install vectordb-bench[pgvectorscale]` |
+| redis                    | `pip install vectordb-bench[redis]`         |
+| memorydb                 | `pip install vectordb-bench[memorydb]`      |
+| chromadb                 | `pip install vectordb-bench[chromadb]`      |
+| awsopensearch            | `pip install vectordb-bench[awsopensearch]` |
 
 ### Run
 
@@ -345,6 +358,24 @@ Case No. | Case Type | Dataset Size  | Filtering Rate | Results |
 
 Each case provides an in-depth examination of a vector database's abilities, providing you a comprehensive view of the database's performance.
 
+#### Custom Dataset for Performance case
+
+Through the `/custom` page, users can customize their own performance case using local datasets. After saving, the corresponding case can be selected from the `/run_test` page to perform the test.
+
+![image](fig/custom_dataset.png)
+![image](fig/custom_case_run_test.png)
+
+We have strict requirements for the data set format, please follow them.
+- `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.
+  - Vectors data files: The file must be named `train.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Query test vectors: The file must be named `test.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Ground truth file: The file must be named `neighbors.parquet` and should have two columns: `id` corresponding to query vectors and `neighbors_id` as an array of `int`.
+
+- `Train File Count` - If the vector file is too large, you can consider splitting it into multiple files. The naming format for the split files should be `train-[index]-of-[file_count].parquet`. For example, `train-01-of-10.parquet` represents the second file (0-indexed) among 10 split files.
+
+- `Use Shuffled Data` - If you check this option, the vector data files need to be modified. VectorDBBench will load the data labeled with `shuffle`. For example, use `shuffle_train.parquet` instead of `train.parquet` and `shuffle_train-04-of-10.parquet` instead of `train-04-of-10.parquet`. The `id` column in the shuffled data can be in any order.
+
+
 ## Goals
 Our goals of this benchmark are:
 ### Reproducibility & Usability

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/README.md

@@ -27,18 +27,21 @@ pip install vectordb-bench[pinecone]
 ```
 All the database client supported
 
-|Optional database client|install command|
-|---------------|---------------|
-|pymilvus(*default*)|`pip install vectordb-bench`|
-|all|`pip install vectordb-bench[all]`|
-|qdrant|`pip install vectordb-bench[qdrant]`|
-|pinecone|`pip install vectordb-bench[pinecone]`|
-|weaviate|`pip install vectordb-bench[weaviate]`|
-|elastic|`pip install vectordb-bench[elastic]`|
-|pgvector|`pip install vectordb-bench[pgvector]`|
-|pgvecto.rs|`pip install vectordb-bench[pgvecto_rs]`|
-|redis|`pip install vectordb-bench[redis]`|
-|chromadb|`pip install vectordb-bench[chromadb]`|
+| Optional database client | install command                             |
+|--------------------------|---------------------------------------------|
+| pymilvus(*default*)      | `pip install vectordb-bench`                |
+| all                      | `pip install vectordb-bench[all]`           |
+| qdrant                   | `pip install vectordb-bench[qdrant]`        |
+| pinecone                 | `pip install vectordb-bench[pinecone]`      |
+| weaviate                 | `pip install vectordb-bench[weaviate]`      |
+| elastic                  | `pip install vectordb-bench[elastic]`       |
+| pgvector                 | `pip install vectordb-bench[pgvector]`      |
+| pgvecto.rs               | `pip install vectordb-bench[pgvecto_rs]`    |
+| pgvectorscale            | `pip install vectordb-bench[pgvectorscale]` |
+| redis                    | `pip install vectordb-bench[redis]`         |
+| memorydb                 | `pip install vectordb-bench[memorydb]`      |
+| chromadb                 | `pip install vectordb-bench[chromadb]`      |
+| awsopensearch            | `pip install vectordb-bench[awsopensearch]` |
 
 ### Run
 
@@ -281,6 +284,24 @@ Case No. | Case Type | Dataset Size  | Filtering Rate | Results |
 
 Each case provides an in-depth examination of a vector database's abilities, providing you a comprehensive view of the database's performance.
 
+#### Custom Dataset for Performance case
+
+Through the `/custom` page, users can customize their own performance case using local datasets. After saving, the corresponding case can be selected from the `/run_test` page to perform the test.
+
+![image](fig/custom_dataset.png)
+![image](fig/custom_case_run_test.png)
+
+We have strict requirements for the data set format, please follow them.
+- `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.
+  - Vectors data files: The file must be named `train.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Query test vectors: The file must be named `test.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Ground truth file: The file must be named `neighbors.parquet` and should have two columns: `id` corresponding to query vectors and `neighbors_id` as an array of `int`.
+
+- `Train File Count` - If the vector file is too large, you can consider splitting it into multiple files. The naming format for the split files should be `train-[index]-of-[file_count].parquet`. For example, `train-01-of-10.parquet` represents the second file (0-indexed) among 10 split files.
+
+- `Use Shuffled Data` - If you check this option, the vector data files need to be modified. VectorDBBench will load the data labeled with `shuffle`. For example, use `shuffle_train.parquet` instead of `train.parquet` and `shuffle_train-04-of-10.parquet` instead of `train-04-of-10.parquet`. The `id` column in the shuffled data can be in any order.
+
+
 ## Goals
 Our goals of this benchmark are:
 ### Reproducibility & Usability

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/install/requirements_py3.11.txt

@@ -5,6 +5,7 @@ pinecone-client
 weaviate-client
 elasticsearch
 pgvector
+pgvecto_rs[psycopg3]>=0.2.1
 sqlalchemy
 redis
 chromadb

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/pyproject.toml

@@ -56,12 +56,14 @@ all = [
     "weaviate-client",
     "elasticsearch",
     "pgvector",
+    "pgvecto_rs[psycopg3]>=0.2.1",
     "sqlalchemy",
     "redis",
     "chromadb",
-    "psycopg2",
     "psycopg",
     "psycopg-binary",
+    "opensearch-dsl==2.1.0",
+    "opensearch-py==2.6.0",
 ]
 
 qdrant = [ "qdrant-client" ]
@@ -69,9 +71,12 @@ pinecone = [ "pinecone-client" ]
 weaviate = [ "weaviate-client" ]
 elastic = [ "elasticsearch" ]
 pgvector = [ "psycopg", "psycopg-binary", "pgvector" ]
-pgvecto_rs = [ "psycopg2" ]
+pgvectorscale = [ "psycopg", "psycopg-binary", "pgvector" ]
+pgvecto_rs = [ "pgvecto_rs[psycopg3]>=0.2.1" ]
 redis = [ "redis" ]
+memorydb = [ "memorydb" ]
 chromadb = [ "chromadb" ]
+awsopensearch = [ "awsopensearch" ]
 zilliz_cloud = []
 
 [project.urls]
@@ -80,4 +85,5 @@ zilliz_cloud = []
 [project.scripts]
 init_bench = "vectordb_bench.__main__:main"
 vectordbbench = "vectordb_bench.cli.vectordbbench:cli"
+
 [tool.setuptools_scm]

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/vectordb_bench/__init__.py

@@ -35,6 +35,7 @@ class config:
 
 
     K_DEFAULT = 100  # default return top k nearest neighbors during search
+    CUSTOM_CONFIG_DIR = pathlib.Path(__file__).parent.joinpath("custom/custom_case.json")
 
     CAPACITY_TIMEOUT_IN_SECONDS = 24 * 3600 # 24h
     LOAD_TIMEOUT_DEFAULT        = 2.5 * 3600 # 2.5h

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/vectordb_bench/backend/assembler.py

@@ -14,7 +14,7 @@ class Assembler:
     def assemble(cls, run_id , task: TaskConfig, source: DatasetSource) -> CaseRunner:
         c_cls = task.case_config.case_id.case_cls
 
-        c = c_cls()
+        c = c_cls(task.case_config.custom_case)
         if type(task.db_case_config) != EmptyDBCaseConfig:
             task.db_case_config.metric_type = c.dataset.data.metric_type
 

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/vectordb_bench/backend/cases.py

@@ -4,9 +4,13 @@ from enum import Enum, auto
 from typing import Type
 
 from vectordb_bench import config
+from vectordb_bench.backend.clients.api import MetricType
 from vectordb_bench.base import BaseModel
+from vectordb_bench.frontend.components.custom.getCustomConfig import (
+    CustomDatasetConfig,
+)
 
-from .dataset import Dataset, DatasetManager
+from .dataset import CustomDataset, Dataset, DatasetManager
 
 
 log = logging.getLogger(__name__)
@@ -44,25 +48,24 @@ class CaseType(Enum):
     Performance1536D50K = 50
 
     Custom = 100
+    PerformanceCustomDataset = 101
 
-    @property
     def case_cls(self, custom_configs: dict | None = None) -> Type["Case"]:
-        if self not in type2case:
-            raise NotImplementedError(f"Case {self} has not implemented. You can add it manually to vectordb_bench.backend.cases.type2case or define a custom_configs['custom_cls']")
-        return type2case[self]
+        if custom_configs is None:
+            return type2case.get(self)()
+        else:
+            return type2case.get(self)(**custom_configs)
 
-    @property
-    def case_name(self) -> str:
-        c = self.case_cls
+    def case_name(self, custom_configs: dict | None = None) -> str:
+        c = self.case_cls(custom_configs)
         if c is not None:
-            return c().name
+            return c.name
         raise ValueError("Case unsupported")
 
-    @property
-    def case_description(self) -> str:
-        c = self.case_cls
+    def case_description(self, custom_configs: dict | None = None) -> str:
+        c = self.case_cls(custom_configs)
         if c is not None:
-            return c().description
+            return c.description
         raise ValueError("Case unsupported")
 
 
@@ -289,26 +292,69 @@ Results will show index building time, recall, and maximum QPS."""
     optimize_timeout: float | int | None = 15 * 60
 
 
+def metric_type_map(s: str) -> MetricType:
+    if s.lower() == "cosine":
+        return MetricType.COSINE
+    if s.lower() == "l2" or s.lower() == "euclidean":
+        return MetricType.L2
+    if s.lower() == "ip":
+        return MetricType.IP
+    err_msg = f"Not support metric_type: {s}"
+    log.error(err_msg)
+    raise RuntimeError(err_msg)
+
+
+class PerformanceCustomDataset(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceCustomDataset
+    name: str = "Performance With Custom Dataset"
+    description: str = ""
+    dataset: DatasetManager
+
+    def __init__(
+        self,
+        name,
+        description,
+        load_timeout,
+        optimize_timeout,
+        dataset_config,
+        **kwargs,
+    ):
+        dataset_config = CustomDatasetConfig(**dataset_config)
+        dataset = CustomDataset(
+            name=dataset_config.name,
+            size=dataset_config.size,
+            dim=dataset_config.dim,
+            metric_type=metric_type_map(dataset_config.metric_type),
+            use_shuffled=dataset_config.use_shuffled,
+            with_gt=dataset_config.with_gt,
+            dir=dataset_config.dir,
+            file_num=dataset_config.file_count,
+        )
+        super().__init__(
+            name=name,
+            description=description,
+            load_timeout=load_timeout,
+            optimize_timeout=optimize_timeout,
+            dataset=DatasetManager(data=dataset),
+        )
+
+
 type2case = {
     CaseType.CapacityDim960: CapacityDim960,
     CaseType.CapacityDim128: CapacityDim128,
-
     CaseType.Performance768D100M: Performance768D100M,
     CaseType.Performance768D10M: Performance768D10M,
     CaseType.Performance768D1M: Performance768D1M,
-
     CaseType.Performance768D10M1P: Performance768D10M1P,
     CaseType.Performance768D1M1P: Performance768D1M1P,
     CaseType.Performance768D10M99P: Performance768D10M99P,
     CaseType.Performance768D1M99P: Performance768D1M99P,
-
     CaseType.Performance1536D500K: Performance1536D500K,
     CaseType.Performance1536D5M: Performance1536D5M,
-
     CaseType.Performance1536D500K1P: Performance1536D500K1P,
     CaseType.Performance1536D5M1P: Performance1536D5M1P,
-
     CaseType.Performance1536D500K99P: Performance1536D500K99P,
     CaseType.Performance1536D5M99P: Performance1536D5M99P,
     CaseType.Performance1536D50K: Performance1536D50K,
+    CaseType.PerformanceCustomDataset: PerformanceCustomDataset,
 }

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/vectordb_bench/backend/clients/__init__.py

@@ -30,8 +30,11 @@ class DB(Enum):
     WeaviateCloud = "WeaviateCloud"
     PgVector = "PgVector"
     PgVectoRS = "PgVectoRS"
+    PgVectorScale = "PgVectorScale"
     Redis = "Redis"
+    MemoryDB = "MemoryDB"
     Chroma = "Chroma"
+    AWSOpenSearch = "OpenSearch"
     Test = "test"
 
 
@@ -69,15 +72,27 @@ class DB(Enum):
         if self == DB.PgVectoRS:
             from .pgvecto_rs.pgvecto_rs import PgVectoRS
             return PgVectoRS
+        
+        if self == DB.PgVectorScale:
+            from .pgvectorscale.pgvectorscale import PgVectorScale
+            return PgVectorScale
 
         if self == DB.Redis:
             from .redis.redis import Redis
             return Redis
+        
+        if self == DB.MemoryDB:
+            from .memorydb.memorydb import MemoryDB
+            return MemoryDB
 
         if self == DB.Chroma:
             from .chroma.chroma import ChromaClient
             return ChromaClient
 
+        if self == DB.AWSOpenSearch:
+            from .aws_opensearch.aws_opensearch import AWSOpenSearch
+            return AWSOpenSearch
+
     @property
     def config_cls(self) -> Type[DBConfig]:
         """Import while in use"""
@@ -113,14 +128,26 @@ class DB(Enum):
             from .pgvecto_rs.config import PgVectoRSConfig
             return PgVectoRSConfig
 
+        if self == DB.PgVectorScale:
+            from .pgvectorscale.config import PgVectorScaleConfig
+            return PgVectorScaleConfig
+
         if self == DB.Redis:
             from .redis.config import RedisConfig
             return RedisConfig
+        
+        if self == DB.MemoryDB:
+            from .memorydb.config import MemoryDBConfig
+            return MemoryDBConfig
 
         if self == DB.Chroma:
             from .chroma.config import ChromaConfig
             return ChromaConfig
 
+        if self == DB.AWSOpenSearch:
+            from .aws_opensearch.config import AWSOpenSearchConfig
+            return AWSOpenSearchConfig
+
     def case_config_cls(self, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
         if self == DB.Milvus:
             from .milvus.config import _milvus_case_config
@@ -150,6 +177,14 @@ class DB(Enum):
             from .pgvecto_rs.config import _pgvecto_rs_case_config
             return _pgvecto_rs_case_config.get(index_type)
 
+        if self == DB.AWSOpenSearch:
+            from .aws_opensearch.config import AWSOpenSearchIndexConfig
+            return AWSOpenSearchIndexConfig
+
+        if self == DB.PgVectorScale:
+            from .pgvectorscale.config import _pgvectorscale_case_config
+            return _pgvectorscale_case_config.get(index_type)
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

{vectordb_bench-0.0.11 → vectordb_bench-0.0.13}/vectordb_bench/backend/clients/api.py

@@ -15,6 +15,7 @@ class MetricType(str, Enum):
 class IndexType(str, Enum):
     HNSW = "HNSW"
     DISKANN = "DISKANN"
+    STREAMING_DISKANN = "DISKANN"
     IVFFlat = "IVF_FLAT"
     IVFSQ8 = "IVF_SQ8"
     Flat = "FLAT"
@@ -38,6 +39,22 @@ class DBConfig(ABC, BaseModel):
     """
 
     db_label: str = ""
+    version: str = ""
+    note: str = ""
+
+    @staticmethod
+    def common_short_configs() -> list[str]:
+        """
+        short input, such as `db_label`, `version`
+        """
+        return ["version", "db_label"]
+
+    @staticmethod
+    def common_long_configs() -> list[str]:
+        """
+        long input, such as `note`
+        """
+        return ["note"]
 
     @abstractmethod
     def to_dict(self) -> dict:
@@ -45,7 +62,10 @@ class DBConfig(ABC, BaseModel):
 
     @validator("*")
     def not_empty_field(cls, v, field):
-        if field.name == "db_label":
+        if (
+            field.name in cls.common_short_configs()
+            or field.name in cls.common_long_configs()
+        ):
             return v
         if not v and isinstance(v, (str, SecretStr)):
             raise ValueError("Empty string!")

vectordb_bench-0.0.13/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py

@@ -0,0 +1,159 @@
+import logging
+from contextlib import contextmanager
+import time
+from typing import Iterable, Type
+from ..api import VectorDB, DBCaseConfig, DBConfig, IndexType
+from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig
+from opensearchpy import OpenSearch
+from opensearchpy.helpers import bulk
+
+log = logging.getLogger(__name__)
+
+
+class AWSOpenSearch(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: AWSOpenSearchIndexConfig,
+        index_name: str = "vdb_bench_index",  # must be lowercase
+        id_col_name: str = "id",
+        vector_col_name: str = "embedding",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.dim = dim
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.index_name = index_name
+        self.id_col_name = id_col_name
+        self.category_col_names = [
+            f"scalar-{categoryCount}" for categoryCount in [2, 5, 10, 100, 1000]
+        ]
+        self.vector_col_name = vector_col_name
+
+        log.info(f"AWS_OpenSearch client config: {self.db_config}")
+        client = OpenSearch(**self.db_config)
+        if drop_old:
+            log.info(f"AWS_OpenSearch client drop old index: {self.index_name}")
+            is_existed = client.indices.exists(index=self.index_name)
+            if is_existed:
+                client.indices.delete(index=self.index_name)
+            self._create_index(client)
+
+    @classmethod
+    def config_cls(cls) -> AWSOpenSearchConfig:
+        return AWSOpenSearchConfig
+
+    @classmethod
+    def case_config_cls(
+        cls, index_type: IndexType | None = None
+    ) -> AWSOpenSearchIndexConfig:
+        return AWSOpenSearchIndexConfig
+
+    def _create_index(self, client: OpenSearch):
+        settings = {
+            "index": {
+                "knn": True,
+                # "number_of_shards": 5,
+                # "refresh_interval": "600s",
+            }
+        }
+        mappings = {
+            "properties": {
+                self.id_col_name: {"type": "integer"},
+                **{
+                    categoryCol: {"type": "keyword"}
+                    for categoryCol in self.category_col_names
+                },
+                self.vector_col_name: {
+                    "type": "knn_vector",
+                    "dimension": self.dim,
+                    "method": self.case_config.index_param(),
+                },
+            }
+        }
+        try:
+            client.indices.create(
+                index=self.index_name, body=dict(settings=settings, mappings=mappings)
+            )
+        except Exception as e:
+            log.warning(f"Failed to create index: {self.index_name} error: {str(e)}")
+            raise e from None
+
+    @contextmanager
+    def init(self) -> None:
+        """connect to elasticsearch"""
+        self.client = OpenSearch(**self.db_config)
+
+        yield
+        # self.client.transport.close()
+        self.client = None
+        del self.client
+
+    def insert_embeddings(
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        **kwargs,
+    ) -> tuple[int, Exception]:
+        """Insert the embeddings to the elasticsearch."""
+        assert self.client is not None, "should self.init() first"
+
+        insert_data = []
+        for i in range(len(embeddings)):
+            insert_data.append({"index": {"_index": self.index_name, "_id": metadata[i]}})
+            insert_data.append({self.vector_col_name: embeddings[i]})
+        try:
+            resp = self.client.bulk(insert_data)
+            log.info(f"AWS_OpenSearch adding documents: {len(resp['items'])}")
+            resp = self.client.indices.stats(self.index_name)
+            log.info(f"Total document count in index: {resp['_all']['primaries']['indexing']['index_total']}")
+            return (len(embeddings), None)
+        except Exception as e:
+            log.warning(f"Failed to insert data: {self.index_name} error: {str(e)}")
+            time.sleep(10)
+            return self.insert_embeddings(embeddings, metadata)
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+    ) -> list[int]:
+        """Get k most similar embeddings to query vector.
+
+        Args:
+            query(list[float]): query embedding to look up documents similar to.
+            k(int): Number of most similar embeddings to return. Defaults to 100.
+            filters(dict, optional): filtering expression to filter the data while searching.
+
+        Returns:
+            list[tuple[int, float]]: list of k most similar embeddings in (id, score) tuple to the query embedding.
+        """
+        assert self.client is not None, "should self.init() first"
+
+        body = {
+            "size": k,
+            "query": {"knn": {self.vector_col_name: {"vector": query, "k": k}}},
+        }
+        try:
+            resp = self.client.search(index=self.index_name, body=body)
+            log.info(f'Search took: {resp["took"]}')
+            log.info(f'Search shards: {resp["_shards"]}')
+            log.info(f'Search hits total: {resp["hits"]["total"]}')
+            result = [int(d["_id"]) for d in resp["hits"]["hits"]]
+            # log.info(f'success! length={len(res)}')
+
+            return result
+        except Exception as e:
+            log.warning(f"Failed to search: {self.index_name} error: {str(e)}")
+            raise e from None
+
+    def optimize(self):
+        """optimize will be called between insertion and search in performance cases."""
+        pass
+
+    def ready_to_load(self):
+        """ready_to_load will be called before load in load cases."""
+        pass