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

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/clients/zilliz_cloud/cli.py

@@ -36,6 +36,17 @@ class ZillizTypedDict(CommonTypedDict):
         str,
         click.option("--level", type=str, help="Zilliz index level", required=False),
     ]
+    num_shards: Annotated[
+        int,
+        click.option(
+            "--num-shards",
+            type=int,
+            help="Number of shards",
+            required=False,
+            default=1,
+            show_default=True,
+        ),
+    ]
 
 
 @cli.command()
@@ -50,9 +61,11 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]),
+            num_shards=parameters["num_shards"],
         ),
         db_case_config=AutoIndexConfig(
-            params={parameters["level"]},
+            level=int(parameters["level"]) if parameters["level"] else 1,
+            num_shards=parameters["num_shards"],
         ),
         **parameters,
     )

vectordb_bench/backend/clients/zilliz_cloud/config.py

@@ -8,24 +8,27 @@ class ZillizCloudConfig(DBConfig):
     uri: SecretStr
     user: str
     password: SecretStr
+    num_shards: int = 1
 
     def to_dict(self) -> dict:
         return {
             "uri": self.uri.get_secret_value(),
             "user": self.user,
             "password": self.password.get_secret_value(),
+            "num_shards": self.num_shards,
         }
 
 
 class AutoIndexConfig(MilvusIndexConfig, DBCaseConfig):
     index: IndexType = IndexType.AUTOINDEX
     level: int = 1
+    num_shards: int = 1
 
     def index_param(self) -> dict:
         return {
             "metric_type": self.parse_metric(),
             "index_type": self.index.value,
-            "params": {},
+            "params": {"shardsNum": self.num_shards},
         }
 
     def search_param(self) -> dict:

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/rate_runner.py

@@ -3,9 +3,11 @@ import logging
 import multiprocessing as mp
 import time
 from concurrent.futures import ThreadPoolExecutor
+from copy import deepcopy
 
 from vectordb_bench import config
 from vectordb_bench.backend.clients import api
+from vectordb_bench.backend.clients.pgvector.pgvector import PgVector
 from vectordb_bench.backend.dataset import DataSetIterator
 from vectordb_bench.backend.utils import time_it
 
@@ -33,17 +35,27 @@ class RatedMultiThreadingInsertRunner:
         self.executing_futures = []
         self.sig_idx = 0
 
-    def send_insert_task(self, db: api.VectorDB, emb: list[list[float]], metadata: list[str], retry_idx: int = 0):
-        _, error = db.insert_embeddings(emb, metadata)
-        if error is not None:
-            log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {error}")
-            retry_idx += 1
-            if retry_idx <= config.MAX_INSERT_RETRY:
-                time.sleep(retry_idx)
-                self.send_insert_task(db, emb=emb, metadata=metadata, retry_idx=retry_idx)
-            else:
-                msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
-                raise RuntimeError(msg) from None
+    def send_insert_task(self, db: api.VectorDB, emb: list[list[float]], metadata: list[str]):
+        def _insert_embeddings(db: api.VectorDB, emb: list[list[float]], metadata: list[str], retry_idx: int = 0):
+            _, error = db.insert_embeddings(emb, metadata)
+            if error is not None:
+                log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {error}")
+                retry_idx += 1
+                if retry_idx <= config.MAX_INSERT_RETRY:
+                    time.sleep(retry_idx)
+                    _insert_embeddings(db, emb=emb, metadata=metadata, retry_idx=retry_idx)
+                else:
+                    msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
+                    raise RuntimeError(msg) from None
+
+        if isinstance(db, PgVector):
+            # pgvector is not thread-safe for concurrent insert,
+            #   so we need to copy the db object, make sure each thread has its own connection
+            db_copy = deepcopy(db)
+            with db_copy.init():
+                _insert_embeddings(db_copy, emb, metadata, retry_idx=0)
+        else:
+            _insert_embeddings(db, emb, metadata, retry_idx=0)
 
     @time_it
     def run_with_rate(self, q: mp.Queue):

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

vectordb_bench/backend/runner/serial_runner.py

@@ -241,7 +241,7 @@ class SerialSearchRunner:
 
         return results
 
-    def search(self, args: tuple[list, list[list[int]]]) -> tuple[float, float, float]:
+    def search(self, args: tuple[list, list[list[int]]]) -> tuple[float, float, float, float]:
         log.info(f"{mp.current_process().name:14} start search the entire test_data to get recall and latency")
         with self.db.init():
             self.db.prepare_filter(self.filters)
@@ -281,6 +281,7 @@ class SerialSearchRunner:
         avg_ndcg = round(np.mean(ndcgs), 4)
         cost = round(np.sum(latencies), 4)
         p99 = round(np.percentile(latencies, 99), 4)
+        p95 = round(np.percentile(latencies, 95), 4)
         log.info(
             f"{mp.current_process().name:14} search entire test_data: "
             f"cost={cost}s, "
@@ -288,20 +289,35 @@ class SerialSearchRunner:
             f"avg_recall={avg_recall}, "
             f"avg_ndcg={avg_ndcg}, "
             f"avg_latency={avg_latency}, "
-            f"p99={p99}"
+            f"p99={p99}, "
+            f"p95={p95}"
         )
-        return (avg_recall, avg_ndcg, p99)
+        return (avg_recall, avg_ndcg, p99, p95)
 
-    def _run_in_subprocess(self) -> tuple[float, float]:
+    def _run_in_subprocess(self) -> tuple[float, float, float, float]:
         with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
             future = executor.submit(self.search, (self.test_data, self.ground_truth))
             return future.result()
 
     @utils.time_it
-    def run(self) -> tuple[float, float, float]:
+    def run(self) -> tuple[float, float, float, float]:
+        log.info(f"{mp.current_process().name:14} start serial search")
+        if self.test_data is None:
+            msg = "empty test_data"
+            raise RuntimeError(msg)
+
+        return self._run_in_subprocess()
+
+    @utils.time_it
+    def run_with_cost(self) -> tuple[tuple[float, float, float, float], float]:
         """
+        Search all test data in serial.
         Returns:
-            tuple[tuple[float, float, float], float]: (avg_recall, avg_ndcg, p99_latency), cost
-
+            tuple[tuple[float, float, float, float], float]: (avg_recall, avg_ndcg, p99_latency, p95_latency), cost
         """
+        log.info(f"{mp.current_process().name:14} start serial search")
+        if self.test_data is None:
+            msg = "empty test_data"
+            raise RuntimeError(msg)
+
         return self._run_in_subprocess()

vectordb_bench/backend/task_runner.py

@@ -186,11 +186,12 @@ class CaseRunner(BaseModel):
                         m.conc_num_list,
                         m.conc_qps_list,
                         m.conc_latency_p99_list,
+                        m.conc_latency_p95_list,
                         m.conc_latency_avg_list,
                     ) = search_results
                 if TaskStage.SEARCH_SERIAL in self.config.stages:
                     search_results = self._serial_search()
-                    m.recall, m.ndcg, m.serial_latency_p99 = search_results
+                    m.recall, m.ndcg, m.serial_latency_p99, m.serial_latency_p95 = search_results
 
         except Exception as e:
             log.warning(f"Failed to run performance case, reason = {e}")
@@ -230,12 +231,12 @@ class CaseRunner(BaseModel):
         finally:
             runner = None
 
-    def _serial_search(self) -> tuple[float, float, float]:
+    def _serial_search(self) -> tuple[float, float, float, float]:
         """Performance serial tests, search the entire test data once,
-        calculate the recall, serial_latency_p99
+        calculate the recall, serial_latency_p99, serial_latency_p95
 
         Returns:
-            tuple[float, float, float]: recall, ndcg, serial_latency_p99
+            tuple[float, float, float, float]: recall, ndcg, serial_latency_p99, serial_latency_p95
         """
         try:
             results, _ = self.serial_search_runner.run()

vectordb_bench/cli/cli.py

@@ -471,6 +471,33 @@ class HNSWFlavor4(HNSWBaseRequiredTypedDict):
     ]
 
 
+class HNSWFlavor5(HNSWBaseRequiredTypedDict):
+    ef_search: Annotated[
+        int | None,
+        click.option("--ef-search", type=int, help="hnsw ef-search", required=True),
+    ]
+    index_type: Annotated[
+        str | None,
+        click.option(
+            "--index-type",
+            type=click.Choice(["HGraph"], case_sensitive=True),
+            help="Type of index to use. Supported values: HGraph",
+            required=True,
+        ),
+    ]
+    use_reorder: Annotated[
+        bool,
+        click.option(
+            "--use-reorder/--no-use-reorder",
+            is_flag=True,
+            type=bool,
+            help="use reorder index",
+            default=True,
+            show_default=True,
+        ),
+    ]
+
+
 class IVFFlatTypedDict(TypedDict):
     lists: Annotated[int | None, click.option("--lists", type=int, help="ivfflat lists")]
     probes: Annotated[int | None, click.option("--probes", type=int, help="ivfflat probes")]
@@ -501,6 +528,15 @@ class OceanBaseIVFTypedDict(TypedDict):
         int | None,
         click.option("--nlist", "nlist", type=int, help="Number of cluster centers", required=True),
     ]
+    nbits: Annotated[
+        int | None,
+        click.option(
+            "--nbits",
+            "nbits",
+            type=int,
+            help="Number of bits used to encode the index of a sub-vector's centroid in the compressed representation",
+        ),
+    ]
     sample_per_nlist: Annotated[
         int | None,
         click.option(

vectordb_bench/cli/vectordbbench.py

@@ -1,11 +1,13 @@
 from ..backend.clients.alloydb.cli import AlloyDBScaNN
 from ..backend.clients.aws_opensearch.cli import AWSOpenSearch
 from ..backend.clients.clickhouse.cli import Clickhouse
+from ..backend.clients.hologres.cli import HologresHGraph
 from ..backend.clients.lancedb.cli import LanceDB
 from ..backend.clients.mariadb.cli import MariaDBHNSW
 from ..backend.clients.memorydb.cli import MemoryDB
 from ..backend.clients.milvus.cli import MilvusAutoIndex
 from ..backend.clients.oceanbase.cli import OceanBaseHNSW, OceanBaseIVF
+from ..backend.clients.oss_opensearch.cli import OSSOpenSearch
 from ..backend.clients.pgdiskann.cli import PgDiskAnn
 from ..backend.clients.pgvecto_rs.cli import PgVectoRSHNSW, PgVectoRSIVFFlat
 from ..backend.clients.pgvector.cli import PgVectorHNSW
@@ -31,6 +33,7 @@ cli.add_command(Test)
 cli.add_command(ZillizAutoIndex)
 cli.add_command(MilvusAutoIndex)
 cli.add_command(AWSOpenSearch)
+cli.add_command(OSSOpenSearch)
 cli.add_command(PgVectorScaleDiskAnn)
 cli.add_command(PgDiskAnn)
 cli.add_command(AlloyDBScaNN)
@@ -41,6 +44,7 @@ cli.add_command(TiDB)
 cli.add_command(Clickhouse)
 cli.add_command(Vespa)
 cli.add_command(LanceDB)
+cli.add_command(HologresHGraph)
 cli.add_command(QdrantCloud)
 cli.add_command(QdrantLocal)
 cli.add_command(BatchCli)

vectordb_bench/frontend/components/check_results/filters.py

@@ -89,7 +89,7 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t
             col=1,
         )
 
-    if filter_type == FilterOp.StrEqual:
+    if filter_type == FilterOp.StrEqual or filter_type == FilterOp.NumGE:
         container = st.container()
         datasetWithSizeTypes = [dataset_with_size_type for dataset_with_size_type in DatasetWithSizeType]
         showDatasetWithSizeTypes = filterView(
@@ -102,9 +102,6 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t
         datasets = [dataset_with_size_type.get_manager() for dataset_with_size_type in showDatasetWithSizeTypes]
         showCaseNames = list(set([case.name for case in allCases if case.dataset in datasets]))
 
-    if filter_type == FilterOp.NumGE:
-        raise NotImplementedError
-
     return showDBNames, showCaseNames
 
 

vectordb_bench/frontend/components/check_results/nav.py

@@ -19,7 +19,7 @@ def NavToQuriesPerDollar(st):
 def NavToResults(st, key="nav-to-results"):
     navClick = st.button("< &nbsp;&nbsp;Back to Results", key=key)
     if navClick:
-        switch_page("vdb benchmark")
+        switch_page("results")
 
 
 def NavToPages(st):
@@ -29,6 +29,7 @@ def NavToPages(st):
         {"name": "Quries Per Dollar", "link": "quries_per_dollar"},
         {"name": "Concurrent", "link": "concurrent"},
         {"name": "Label Filter", "link": "label_filter"},
+        {"name": "Int Filter", "link": "int_filter"},
         {"name": "Streaming", "link": "streaming"},
         {"name": "Tables", "link": "tables"},
         {"name": "Custom Dataset", "link": "custom"},

vectordb_bench/frontend/components/concurrent/charts.py

@@ -20,6 +20,11 @@ def drawChartsByCase(allData, showCaseNames: list[str], st, latency_type: str):
                     if 0 <= i < len(caseData["conc_latency_p99_list"])
                     else 0
                 ),
+                "latency_p95": (
+                    caseData["conc_latency_p95_list"][i] * 1000
+                    if "conc_latency_p95_list" in caseData and 0 <= i < len(caseData["conc_latency_p95_list"])
+                    else 0
+                ),
                 "latency_avg": (
                     caseData["conc_latency_avg_list"][i] * 1000
                     if 0 <= i < len(caseData["conc_latency_avg_list"])