vectordb-bench 0.0.11__py3-none-any.whl → 0.0.12__py3-none-any.whl

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/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/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/backend/clients/__init__.py

@@ -32,6 +32,7 @@ class DB(Enum):
     PgVectoRS = "PgVectoRS"
     Redis = "Redis"
     Chroma = "Chroma"
+    AWSOpenSearch = "OpenSearch"
     Test = "test"
 
 
@@ -78,6 +79,10 @@ class DB(Enum):
             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"""
@@ -121,6 +126,10 @@ class DB(Enum):
             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 +159,10 @@ 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
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

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

vectordb_bench/backend/clients/aws_opensearch/cli.py

@@ -0,0 +1,44 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor2,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class AWSOpenSearchTypedDict(TypedDict):
+    host: Annotated[
+        str, click.option("--host", type=str, help="Db host", required=True)
+    ]
+    port: Annotated[int, click.option("--port", type=int, default=443, help="Db Port")]
+    user: Annotated[str, click.option("--user", type=str, default="admin", help="Db User")]
+    password: Annotated[str, click.option("--password", type=str, help="Db password")]
+
+
+class AWSOpenSearchHNSWTypedDict(CommonTypedDict, AWSOpenSearchTypedDict, HNSWFlavor2):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(AWSOpenSearchHNSWTypedDict)
+def AWSOpenSearch(**parameters: Unpack[AWSOpenSearchHNSWTypedDict]):
+    from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig
+    run(
+        db=DB.AWSOpenSearch,
+        db_config=AWSOpenSearchConfig(
+            host=parameters["host"],
+            port=parameters["port"],
+            user=parameters["user"],
+            password=SecretStr(parameters["password"]),
+        ),
+        db_case_config=AWSOpenSearchIndexConfig(
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/aws_opensearch/config.py

@@ -0,0 +1,58 @@
+from enum import Enum
+from pydantic import SecretStr, BaseModel
+
+from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+
+
+class AWSOpenSearchConfig(DBConfig, BaseModel):
+    host: str = ""
+    port: int = 443
+    user: str = ""
+    password: SecretStr = ""
+
+    def to_dict(self) -> dict:
+        return {
+            "hosts": [{'host': self.host, 'port': self.port}],
+            "http_auth": (self.user, self.password.get_secret_value()),
+            "use_ssl": True,
+            "http_compress": True,
+            "verify_certs": True,
+            "ssl_assert_hostname": False,
+            "ssl_show_warn": False,
+            "timeout": 600,
+        }
+
+
+class AWSOS_Engine(Enum):
+    nmslib = "nmslib"
+    faiss = "faiss"
+    lucene = "Lucene"
+
+
+class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType = MetricType.L2
+    engine: AWSOS_Engine = AWSOS_Engine.nmslib
+    efConstruction: int = 360
+    M: int = 30
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.IP:
+            return "innerproduct"  # only support faiss / nmslib, not for Lucene.
+        elif self.metric_type == MetricType.COSINE:
+            return "cosinesimil"
+        return "l2"
+
+    def index_param(self) -> dict:
+        params = {
+            "name": "hnsw",
+            "space_type": self.parse_metric(),
+            "engine": self.engine.value,
+            "parameters": {
+                "ef_construction": self.efConstruction,
+                "m": self.M
+            }
+        }
+        return params
+
+    def search_param(self) -> dict:
+        return {}

vectordb_bench/backend/clients/aws_opensearch/run.py

@@ -0,0 +1,125 @@
+import time, random
+from opensearchpy import OpenSearch
+from opensearch_dsl import Search, Document, Text, Keyword
+
+_HOST = 'xxxxxx.us-west-2.es.amazonaws.com'
+_PORT = 443
+_AUTH = ('admin', 'xxxxxx') # For testing only. Don't store credentials in code.
+
+_INDEX_NAME = 'my-dsl-index'
+_BATCH = 100
+_ROWS = 100
+_DIM = 128
+_TOPK = 10
+
+
+def create_client():
+    client = OpenSearch(
+        hosts=[{'host': _HOST, 'port': _PORT}],
+        http_compress=True, # enables gzip compression for request bodies
+        http_auth=_AUTH,
+        use_ssl=True,
+        verify_certs=True,
+        ssl_assert_hostname=False,
+        ssl_show_warn=False,
+    )
+    return client
+
+
+def create_index(client, index_name):
+    settings = {
+        "index": {
+            "knn": True,
+            "number_of_shards": 1,
+            "refresh_interval": "5s",
+        }
+    }
+    mappings = {
+        "properties": {
+            "embedding": {
+                "type": "knn_vector",
+                "dimension": _DIM,
+                "method": {
+                    "engine": "nmslib",
+                    "name": "hnsw",
+                    "space_type": "l2",
+                    "parameters": {
+                        "ef_construction": 128,
+                        "m": 24,
+                    }
+                }
+            }
+        }
+    }
+
+    response = client.indices.create(index=index_name, body=dict(settings=settings, mappings=mappings))
+    print('\nCreating index:')
+    print(response)
+
+
+def delete_index(client, index_name):
+    response = client.indices.delete(index=index_name)
+    print('\nDeleting index:')
+    print(response)
+
+
+def bulk_insert(client, index_name):
+    # Perform bulk operations
+    ids = [i for i in range(_ROWS)]
+    vec = [[random.random() for _ in range(_DIM)] for _ in range(_ROWS)]
+
+    docs = []
+    for i in range(0, _ROWS, _BATCH):
+        docs.clear()
+        for j in range(0, _BATCH):
+            docs.append({"index": {"_index": index_name, "_id": ids[i+j]}})
+            docs.append({"embedding": vec[i+j]})
+        response = client.bulk(docs)
+        print('\nAdding documents:', len(response['items']), response['errors'])
+        response = client.indices.stats(index_name)
+        print('\nTotal document count in index:', response['_all']['primaries']['indexing']['index_total'])
+
+
+def search(client, index_name):
+    # Search for the document.
+    search_body = {
+        "size": _TOPK,
+        "query": {
+            "knn": {
+                "embedding": {
+                    "vector": [random.random() for _ in range(_DIM)],
+                    "k": _TOPK,
+                }
+            }
+        }
+    }
+    while True:
+        response = client.search(index=index_name, body=search_body)
+        print(f'\nSearch took: {response["took"]}')
+        print(f'\nSearch shards: {response["_shards"]}')
+        print(f'\nSearch hits total: {response["hits"]["total"]}')
+        result = response["hits"]["hits"]
+        if len(result) != 0:
+            print('\nSearch results:')
+            for hit in response["hits"]["hits"]:
+                print(hit["_id"], hit["_score"])
+            break
+        else:
+            print('\nSearch not ready, sleep 1s')
+            time.sleep(1)
+
+
+def main():
+    client = create_client()
+    try:
+        create_index(client, _INDEX_NAME)
+        bulk_insert(client, _INDEX_NAME)
+        search(client, _INDEX_NAME)
+        delete_index(client, _INDEX_NAME)
+    except Exception as e:
+        print(e)
+        delete_index(client, _INDEX_NAME)
+
+
+if __name__ == '__main__':
+    main()

vectordb_bench/backend/dataset.py

@@ -33,6 +33,7 @@ class BaseDataset(BaseModel):
     use_shuffled: bool
     with_gt: bool = False
     _size_label: dict[int, SizeLabel] = PrivateAttr()
+    isCustom: bool = False
 
     @validator("size")
     def verify_size(cls, v):
@@ -52,7 +53,27 @@ class BaseDataset(BaseModel):
     def file_count(self) -> int:
         return self._size_label.get(self.size).file_count
 
+class CustomDataset(BaseDataset):
+    dir: str
+    file_num: int
+    isCustom: bool = True
+    
+    @validator("size")
+    def verify_size(cls, v):
+        return v
+    
+    @property
+    def label(self) -> str:
+        return "Custom"
 
+    @property
+    def dir_name(self) -> str:
+        return self.dir
+
+    @property
+    def file_count(self) -> int:
+        return self.file_num
+    
 class LAION(BaseDataset):
     name: str = "LAION"
     dim: int = 768
@@ -186,11 +207,12 @@ class DatasetManager(BaseModel):
             gt_file, test_file = utils.compose_gt_file(filters), "test.parquet"
             all_files.extend([gt_file, test_file])
 
-        source.reader().read(
-            dataset=self.data.dir_name.lower(),
-            files=all_files,
-            local_ds_root=self.data_dir,
-        )
+        if not self.data.isCustom:
+            source.reader().read(
+                dataset=self.data.dir_name.lower(),
+                files=all_files,
+                local_ds_root=self.data_dir,
+            )
 
         if gt_file is not None and test_file is not None:
             self.test_data = self._read_file(test_file)

vectordb_bench/cli/vectordbbench.py

@@ -4,6 +4,7 @@ from ..backend.clients.test.cli import Test
 from ..backend.clients.weaviate_cloud.cli import Weaviate
 from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex
 from ..backend.clients.milvus.cli import MilvusAutoIndex
+from ..backend.clients.aws_opensearch.cli import AWSOpenSearch
 
 
 from .cli import cli
@@ -14,6 +15,7 @@ cli.add_command(Weaviate)
 cli.add_command(Test)
 cli.add_command(ZillizAutoIndex)
 cli.add_command(MilvusAutoIndex)
+cli.add_command(AWSOpenSearch)
 
 
 if __name__ == "__main__":

vectordb_bench/custom/custom_case.json

@@ -0,0 +1,18 @@
+[
+    {
+        "name": "My Dataset (Performace Case)",
+        "description": "this is a customized dataset.",
+        "load_timeout": 36000,
+        "optimize_timeout": 36000,
+        "dataset_config": {
+            "name": "My Dataset",
+            "dir": "/my_dataset_path",
+            "size": 1000000,
+            "dim": 1024,
+            "metric_type": "L2",
+            "file_count": 1,
+            "use_shuffled": false,
+            "with_gt": true
+        }
+    }
+]

vectordb_bench/frontend/components/check_results/charts.py

@@ -1,19 +1,19 @@
 from vectordb_bench.backend.cases import Case
 from vectordb_bench.frontend.components.check_results.expanderStyle import initMainExpanderStyle
 from vectordb_bench.metric import metricOrder, isLowerIsBetterMetric, metricUnitMap
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.models import ResultLabel
 import plotly.express as px
 
 
-def drawCharts(st, allData, failedTasks, cases: list[Case]):
+def drawCharts(st, allData, failedTasks, caseNames: list[str]):
     initMainExpanderStyle(st)
-    for case in cases:
-        chartContainer = st.expander(case.name, True)
-        data = [data for data in allData if data["case_name"] == case.name]
+    for caseName in caseNames:
+        chartContainer = st.expander(caseName, True)
+        data = [data for data in allData if data["case_name"] == caseName]
         drawChart(data, chartContainer)
 
-        errorDBs = failedTasks[case.name]
+        errorDBs = failedTasks[caseName]
         showFailedDBs(chartContainer, errorDBs)