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

vectordb_bench/backend/clients/oss_opensearch/run.py

@@ -0,0 +1,166 @@
+import logging
+import random
+import time
+
+from opensearchpy import OpenSearch
+
+log = logging.getLogger(__name__)
+
+_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():
+    return 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,
+    )
+
+
+def create_index(client: OpenSearch, index_name: str):
+    settings = {
+        "index": {
+            "knn": True,
+            "number_of_shards": 1,
+            "refresh_interval": "5s",
+        },
+    }
+    mappings = {
+        "properties": {
+            "embedding": {
+                "type": "knn_vector",
+                "dimension": _DIM,
+                "method": {
+                    "engine": "faiss",
+                    "name": "hnsw",
+                    "space_type": "l2",
+                    "parameters": {
+                        "ef_construction": 256,
+                        "m": 16,
+                    },
+                },
+            },
+        },
+    }
+
+    response = client.indices.create(
+        index=index_name,
+        body={"settings": settings, "mappings": mappings},
+    )
+    log.info("\nCreating index:")
+    log.info(response)
+
+
+def delete_index(client: OpenSearch, index_name: str):
+    response = client.indices.delete(index=index_name)
+    log.info("\nDeleting index:")
+    log.info(response)
+
+
+def bulk_insert(client: OpenSearch, index_name: str):
+    # Perform bulk operations
+    ids = list(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(_BATCH):
+            docs.append({"index": {"_index": index_name, "_id": ids[i + j]}})
+            docs.append({"embedding": vec[i + j]})
+        response = client.bulk(docs)
+        log.info(f"Adding documents: {len(response['items'])}, {response['errors']}")
+        response = client.indices.stats(index_name)
+        log.info(
+            f'Total document count in index: { response["_all"]["primaries"]["indexing"]["index_total"] }',
+        )
+
+
+def search(client: OpenSearch, index_name: str):
+    # 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)
+        log.info(f'\nSearch took: {response["took"]}')
+        log.info(f'\nSearch shards: {response["_shards"]}')
+        log.info(f'\nSearch hits total: {response["hits"]["total"]}')
+        result = response["hits"]["hits"]
+        if len(result) != 0:
+            log.info("\nSearch results:")
+            for hit in response["hits"]["hits"]:
+                log.info(hit["_id"], hit["_score"])
+            break
+        log.info("\nSearch not ready, sleep 1s")
+        time.sleep(1)
+
+
+SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30
+WAITINT_FOR_REFRESH_SEC = 30
+
+
+def optimize_index(client: OpenSearch, index_name: str):
+    log.info(f"Starting force merge for index {index_name}")
+    force_merge_endpoint = f"/{index_name}/_forcemerge?max_num_segments=1&wait_for_completion=false"
+    force_merge_task_id = client.transport.perform_request("POST", force_merge_endpoint)["task"]
+    while True:
+        time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC)
+        task_status = client.tasks.get(task_id=force_merge_task_id)
+        if task_status["completed"]:
+            break
+    log.info(f"Completed force merge for index {index_name}")
+
+
+def refresh_index(client: OpenSearch, index_name: str):
+    log.info(f"Starting refresh for index {index_name}")
+    while True:
+        try:
+            log.info("Starting the Refresh Index..")
+            client.indices.refresh(index=index_name)
+            break
+        except Exception as e:
+            log.info(
+                f"Refresh errored out. Sleeping for {WAITINT_FOR_REFRESH_SEC} sec and then Retrying : {e}",
+            )
+            time.sleep(WAITINT_FOR_REFRESH_SEC)
+            continue
+    log.info(f"Completed refresh for index {index_name}")
+
+
+def main():
+    client = create_client()
+    try:
+        create_index(client, _INDEX_NAME)
+        bulk_insert(client, _INDEX_NAME)
+        optimize_index(client, _INDEX_NAME)
+        refresh_index(client, _INDEX_NAME)
+        search(client, _INDEX_NAME)
+        delete_index(client, _INDEX_NAME)
+    except Exception as e:
+        log.info(e)
+        delete_index(client, _INDEX_NAME)
+
+
+if __name__ == "__main__":
+    main()

vectordb_bench/backend/clients/s3_vectors/config.py

@@ -0,0 +1,41 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+
+class S3VectorsConfig(DBConfig):
+    region_name: str = "us-west-2"
+    access_key_id: SecretStr
+    secret_access_key: SecretStr
+    bucket_name: str
+    index_name: str = "vdbbench-index"
+
+    def to_dict(self) -> dict:
+        return {
+            "region_name": self.region_name,
+            "access_key_id": self.access_key_id.get_secret_value() if self.access_key_id else "",
+            "secret_access_key": self.secret_access_key.get_secret_value() if self.secret_access_key else "",
+            "bucket_name": self.bucket_name,
+            "index_name": self.index_name,
+        }
+
+
+class S3VectorsIndexConfig(DBCaseConfig, BaseModel):
+    """Base config for s3-vectors"""
+
+    metric_type: MetricType | None = None
+    data_type: str = "float32"
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.COSINE:
+            return "cosine"
+        if self.metric_type == MetricType.L2:
+            return "euclidean"
+        msg = f"Unsupported metric type: {self.metric_type}"
+        raise ValueError(msg)
+
+    def index_param(self) -> dict:
+        return {}
+
+    def search_param(self) -> dict:
+        return {}

vectordb_bench/backend/clients/s3_vectors/s3_vectors.py

@@ -0,0 +1,171 @@
+"""Wrapper around the Milvus vector database over VectorDB"""
+
+import logging
+from collections.abc import Iterable
+from contextlib import contextmanager
+
+import boto3
+
+from vectordb_bench.backend.filter import Filter, FilterOp
+
+from ..api import VectorDB
+from .config import S3VectorsIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class S3Vectors(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: S3VectorsIndexConfig,
+        drop_old: bool = False,
+        with_scalar_labels: bool = False,
+        **kwargs,
+    ):
+        """Initialize wrapper around the s3-vectors client."""
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.with_scalar_labels = with_scalar_labels
+
+        self.batch_size = 500
+
+        self._scalar_id_field = "id"
+        self._scalar_label_field = "label"
+        self._vector_field = "vector"
+
+        self.region_name = self.db_config.get("region_name")
+        self.access_key_id = self.db_config.get("access_key_id")
+        self.secret_access_key = self.db_config.get("secret_access_key")
+        self.bucket_name = self.db_config.get("bucket_name")
+        self.index_name = self.db_config.get("index_name")
+
+        client = boto3.client(
+            service_name="s3vectors",
+            region_name=self.region_name,
+            aws_access_key_id=self.access_key_id,
+            aws_secret_access_key=self.secret_access_key,
+        )
+
+        if drop_old:
+            # delete old index if exists
+            response = client.list_indexes(vectorBucketName=self.bucket_name)
+            index_names = [index["indexName"] for index in response["indexes"]]
+            if self.index_name in index_names:
+                log.info(f"drop old index: {self.index_name}")
+                client.delete_index(vectorBucketName=self.bucket_name, indexName=self.index_name)
+
+            # create the index
+            client.create_index(
+                vectorBucketName=self.bucket_name,
+                indexName=self.index_name,
+                dataType=self.case_config.data_type,
+                dimension=dim,
+                distanceMetric=self.case_config.parse_metric(),
+            )
+
+        client.close()
+
+    @contextmanager
+    def init(self):
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+        self.client = boto3.client(
+            service_name="s3vectors",
+            region_name=self.region_name,
+            aws_access_key_id=self.access_key_id,
+            aws_secret_access_key=self.secret_access_key,
+        )
+
+        yield
+        self.client.close()
+
+    def optimize(self, **kwargs):
+        return
+
+    def need_normalize_cosine(self) -> bool:
+        """Wheather this database need to normalize dataset to support COSINE"""
+        return False
+
+    def insert_embeddings(
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        labels_data: list[str] | None = None,
+        **kwargs,
+    ) -> tuple[int, Exception]:
+        """Insert embeddings into s3-vectors. should call self.init() first"""
+        # use the first insert_embeddings to init collection
+        assert self.client is not None
+        assert len(embeddings) == len(metadata)
+        insert_count = 0
+        try:
+            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 = [
+                    {
+                        "key": str(metadata[i]),
+                        "data": {self.case_config.data_type: embeddings[i]},
+                        "metadata": (
+                            {self._scalar_label_field: labels_data[i], self._scalar_id_field: metadata[i]}
+                            if self.with_scalar_labels
+                            else {self._scalar_id_field: metadata[i]}
+                        ),
+                    }
+                    for i in range(batch_start_offset, batch_end_offset)
+                ]
+                self.client.put_vectors(
+                    vectorBucketName=self.bucket_name,
+                    indexName=self.index_name,
+                    vectors=insert_data,
+                )
+                insert_count += len(insert_data)
+        except Exception as e:
+            log.info(f"Failed to insert data: {e}")
+            return insert_count, e
+        return insert_count, None
+
+    def prepare_filter(self, filters: Filter):
+        if filters.type == FilterOp.NonFilter:
+            self.filter = None
+        elif filters.type == FilterOp.NumGE:
+            self.filter = {self._scalar_id_field: {"$gte": filters.int_value}}
+        elif filters.type == FilterOp.StrEqual:
+            self.filter = {self._scalar_label_field: filters.label_value}
+        else:
+            msg = f"Not support Filter for S3Vectors - {filters}"
+            raise ValueError(msg)
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        timeout: int | None = None,
+    ) -> list[int]:
+        """Perform a search on a query embedding and return results."""
+        assert self.client is not None
+
+        # Perform the search.
+        res = self.client.query_vectors(
+            vectorBucketName=self.bucket_name,
+            indexName=self.index_name,
+            queryVector={"float32": query},
+            topK=k,
+            filter=self.filter,
+            returnDistance=False,
+            returnMetadata=False,
+        )
+
+        # Organize results.
+        return [int(result["key"]) for result in res["vectors"]]

vectordb_bench/backend/clients/tidb/cli.py

@@ -17,7 +17,6 @@ class TiDBTypedDict(CommonTypedDict):
             help="Username",
             default="root",
             show_default=True,
-            required=True,
         ),
     ]
     password: Annotated[
@@ -37,7 +36,6 @@ class TiDBTypedDict(CommonTypedDict):
             type=str,
             default="127.0.0.1",
             show_default=True,
-            required=True,
             help="Db host",
         ),
     ]
@@ -48,7 +46,6 @@ class TiDBTypedDict(CommonTypedDict):
             type=int,
             default=4000,
             show_default=True,
-            required=True,
             help="Db Port",
         ),
     ]
@@ -59,7 +56,6 @@ class TiDBTypedDict(CommonTypedDict):
             type=str,
             default="test",
             show_default=True,
-            required=True,
             help="Db name",
         ),
     ]

vectordb_bench/backend/clients/tidb/config.py

@@ -1,8 +1,20 @@
-from pydantic import BaseModel, SecretStr
+from typing import TypedDict
+
+from pydantic import BaseModel, SecretStr, validator
 
 from ..api import DBCaseConfig, DBConfig, MetricType
 
 
+class TiDBConfigDict(TypedDict):
+    host: str
+    port: int
+    user: str
+    password: str
+    database: str
+    ssl_verify_cert: bool
+    ssl_verify_identity: bool
+
+
 class TiDBConfig(DBConfig):
     user_name: str = "root"
     password: SecretStr
@@ -11,7 +23,7 @@ class TiDBConfig(DBConfig):
     db_name: str = "test"
     ssl: bool = False
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> TiDBConfigDict:
         pwd_str = self.password.get_secret_value()
         return {
             "host": self.host,
@@ -23,6 +35,14 @@ class TiDBConfig(DBConfig):
             "ssl_verify_identity": self.ssl,
         }
 
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        if field.name in ["password", "db_label"]:
+            return v
+        if isinstance(v, str | SecretStr) and len(v) == 0:
+            raise ValueError("Empty string!")
+        return v
+
 
 class TiDBIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None

vectordb_bench/backend/dataset.py

@@ -48,6 +48,7 @@ class BaseDataset(BaseModel):
     scalar_labels_file_separated: bool = True
     scalar_labels_file: str = "scalar_labels.parquet"
     scalar_label_percentages: list[float] = []
+    scalar_int_rates: list[float] = []
     train_id_field: str = "id"
     train_vector_field: str = "emb"
     test_file: str = "test.parquet"
@@ -164,6 +165,29 @@ class Cohere(BaseDataset):
     }
     with_scalar_labels: bool = True
     scalar_label_percentages: list[float] = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
+    scalar_int_rates: list[float] = [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.6,
+        0.7,
+        0.8,
+        0.9,
+        0.95,
+        0.98,
+        0.99,
+        0.995,
+        0.998,
+        0.999,
+    ]
 
 
 class Bioasq(BaseDataset):
@@ -178,6 +202,29 @@ class Bioasq(BaseDataset):
     }
     with_scalar_labels: bool = True
     scalar_label_percentages: list[float] = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
+    scalar_int_rates: list[float] = [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.6,
+        0.7,
+        0.8,
+        0.9,
+        0.95,
+        0.98,
+        0.99,
+        0.995,
+        0.998,
+        0.999,
+    ]
 
 
 class Glove(BaseDataset):
@@ -217,6 +264,29 @@ class OpenAI(BaseDataset):
     }
     with_scalar_labels: bool = True
     scalar_label_percentages: list[float] = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
+    scalar_int_rates: list[float] = [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.6,
+        0.7,
+        0.8,
+        0.9,
+        0.95,
+        0.98,
+        0.99,
+        0.995,
+        0.998,
+        0.999,
+    ]
 
 
 class DatasetManager(BaseModel):

vectordb_bench/backend/filter.py

@@ -51,6 +51,23 @@ class IntFilter(Filter):
         raise RuntimeError(msg)
 
 
+class NewIntFilter(Filter):
+    type: FilterOp = FilterOp.NumGE
+    int_field: str = "id"
+    int_value: int
+
+    @property
+    def int_rate(self) -> str:
+        r = self.filter_rate * 100
+        if 1 <= r <= 99:
+            return f"int_{int(r)}p"
+        return f"int_{r:.1f}p"
+
+    @property
+    def groundtruth_file(self) -> str:
+        return f"neighbors_{self.int_rate}.parquet"
+
+
 class LabelFilter(Filter):
     """
     filter expr: label_field == label_value, like `color == "red"`

vectordb_bench/backend/runner/mp_runner.py

@@ -103,6 +103,7 @@ class MultiProcessingSearchRunner:
         conc_num_list = []
         conc_qps_list = []
         conc_latency_p99_list = []
+        conc_latency_p95_list = []
         conc_latency_avg_list = []
         try:
             for conc in self.concurrencies:
@@ -125,6 +126,7 @@ class MultiProcessingSearchRunner:
                         all_count = sum([r.result()[0] for r in future_iter])
                         latencies = sum([r.result()[2] for r in future_iter], start=[])
                         latency_p99 = np.percentile(latencies, 99)
+                        latency_p95 = np.percentile(latencies, 95)
                         latency_avg = np.mean(latencies)
                         cost = time.perf_counter() - start
 
@@ -132,6 +134,7 @@ class MultiProcessingSearchRunner:
                         conc_num_list.append(conc)
                         conc_qps_list.append(qps)
                         conc_latency_p99_list.append(latency_p99)
+                        conc_latency_p95_list.append(latency_p95)
                         conc_latency_avg_list.append(latency_avg)
                         log.info(f"End search in concurrency {conc}: dur={cost}s, total_count={all_count}, qps={qps}")
 
@@ -156,6 +159,7 @@ class MultiProcessingSearchRunner:
             conc_num_list,
             conc_qps_list,
             conc_latency_p99_list,
+            conc_latency_p95_list,
             conc_latency_avg_list,
         )
 

vectordb_bench/backend/runner/read_write_runner.py

@@ -98,10 +98,10 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         log.info("Search after write - Serial search start")
         test_time = round(time.perf_counter(), 4)
         res, ssearch_dur = self.serial_search_runner.run()
-        recall, ndcg, p99_latency = res
+        recall, ndcg, p99_latency, p95_latency = res
         log.info(
             f"Search after write - Serial search - recall={recall}, ndcg={ndcg}, "
-            f"p99={p99_latency}, dur={ssearch_dur:.4f}",
+            f"p99={p99_latency}, p95={p95_latency}, dur={ssearch_dur:.4f}",
         )
         log.info(
             f"Search after wirte - Conc search start, dur for each conc={self.read_dur_after_write}",
@@ -109,7 +109,7 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         max_qps, conc_failed_rate = self.run_by_dur(self.read_dur_after_write)
         log.info(f"Search after wirte - Conc search finished, max_qps={max_qps}")
 
-        return [(perc, test_time, max_qps, recall, ndcg, p99_latency, conc_failed_rate)]
+        return [(perc, test_time, max_qps, recall, ndcg, p99_latency, p95_latency, conc_failed_rate)]
 
     def run_read_write(self) -> Metric:
         """
@@ -157,7 +157,8 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
                     m.st_recall_list = [d[3] for d in r]
                     m.st_ndcg_list = [d[4] for d in r]
                     m.st_serial_latency_p99_list = [d[5] for d in r]
-                    m.st_conc_failed_rate_list = [d[6] for d in r]
+                    m.st_serial_latency_p95_list = [d[6] for d in r]
+                    m.st_conc_failed_rate_list = [d[7] for d in r]
 
                 except Exception as e:
                     log.warning(f"Read and write error: {e}")
@@ -201,7 +202,7 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         """
         result, start_batch = [], 0
         total_batch = math.ceil(self.data_volume / self.insert_rate)
-        recall, ndcg, p99_latency = None, None, None
+        recall, ndcg, p99_latency, p95_latency = None, None, None, None
 
         def wait_next_target(start: int, target_batch: int) -> bool:
             """Return False when receive True or None"""
@@ -224,15 +225,15 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
 
             log.info(f"Insert {perc}% done, total batch={total_batch}")
             test_time = round(time.perf_counter(), 4)
-            max_qps, recall, ndcg, p99_latency, conc_failed_rate = 0, 0, 0, 0, 0
+            max_qps, recall, ndcg, p99_latency, p95_latency, conc_failed_rate = 0, 0, 0, 0, 0, 0
             try:
                 log.info(f"[{target_batch}/{total_batch}] Serial search - {perc}% start")
                 res, ssearch_dur = self.serial_search_runner.run()
                 ssearch_dur = round(ssearch_dur, 4)
-                recall, ndcg, p99_latency = res
+                recall, ndcg, p99_latency, p95_latency = res
                 log.info(
                     f"[{target_batch}/{total_batch}] Serial search - {perc}% done, "
-                    f"recall={recall}, ndcg={ndcg}, p99={p99_latency}, dur={ssearch_dur}"
+                    f"recall={recall}, ndcg={ndcg}, p99={p99_latency}, p95={p95_latency}, dur={ssearch_dur}"
                 )
 
                 each_conc_search_dur = self.get_each_conc_search_dur(
@@ -250,7 +251,7 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
                     log.warning(f"Skip concurrent tests, each_conc_search_dur={each_conc_search_dur} less than 10s.")
             except Exception as e:
                 log.warning(f"Streaming Search Failed at stage={stage}. Exception: {e}")
-            result.append((perc, test_time, max_qps, recall, ndcg, p99_latency, conc_failed_rate))
+            result.append((perc, test_time, max_qps, recall, ndcg, p99_latency, p95_latency, conc_failed_rate))
             start_batch = target_batch
 
         # Drain the queue