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

vectordb_bench/backend/clients/redis/config.py

@@ -1,14 +1,45 @@
-from pydantic import SecretStr
-from ..api import DBConfig
+from pydantic import SecretStr, BaseModel
+from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
 
 class RedisConfig(DBConfig):
-    password: SecretStr
+    password: SecretStr | None = None
     host: SecretStr
-    port: int = None 
+    port: int | None = None 
 
     def to_dict(self) -> dict:
         return {
             "host": self.host.get_secret_value(),
             "port": self.port,
-            "password": self.password.get_secret_value(),
-        }
+            "password": self.password.get_secret_value() if self.password is not None else None,
+        }
+        
+
+
+class RedisIndexConfig(BaseModel):
+    """Base config for milvus"""
+
+    metric_type: MetricType | None = None
+
+    def parse_metric(self) -> str:
+        if not self.metric_type:
+            return ""
+        return self.metric_type.value
+    
+class RedisHNSWConfig(RedisIndexConfig, DBCaseConfig):
+    M: int
+    efConstruction: int
+    ef: int | None = None
+    index: IndexType = IndexType.HNSW
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"M": self.M, "efConstruction": self.efConstruction},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"ef": self.ef},
+        }

vectordb_bench/backend/runner/mp_runner.py

@@ -2,6 +2,7 @@ import time
 import traceback
 import concurrent
 import multiprocessing as mp
+import random
 import logging
 from typing import Iterable
 import numpy as np
@@ -46,7 +47,7 @@ class MultiProcessingSearchRunner:
             cond.wait()
 
         with self.db.init():
-            num, idx = len(test_data), 0
+            num, idx = len(test_data), random.randint(0, len(test_data) - 1)
 
             start_time = time.perf_counter()
             count = 0

vectordb_bench/cli/cli.py

@@ -17,6 +17,8 @@ from typing import (
     Any,
 )
 import click
+
+from vectordb_bench.backend.clients.api import MetricType
 from .. import config
 from ..backend.clients import DB
 from ..interface import benchMarkRunner, global_result_future
@@ -147,6 +149,37 @@ def parse_task_stages(
     return stages
 
 
+def check_custom_case_parameters(ctx, param, value):
+    if ctx.params.get("case_type") == "PerformanceCustomDataset":
+        if value is None:
+            raise click.BadParameter("Custom case parameters\
+                                     \n--custom-case-name\n--custom-dataset-name\n--custom-dataset-dir\n--custom-dataset-size \
+                                     \n--custom-dataset-dim\n--custom-dataset-file-count\n are required")
+    return value
+
+
+def get_custom_case_config(parameters: dict) -> dict:
+    custom_case_config = {}
+    if parameters["case_type"] == "PerformanceCustomDataset":
+        custom_case_config = {
+            "name": parameters["custom_case_name"],
+            "description": parameters["custom_case_description"],
+            "load_timeout": parameters["custom_case_load_timeout"],
+            "optimize_timeout": parameters["custom_case_optimize_timeout"],
+            "dataset_config": {
+                "name": parameters["custom_dataset_name"],
+                "dir": parameters["custom_dataset_dir"],
+                "size": parameters["custom_dataset_size"],
+                "dim": parameters["custom_dataset_dim"],
+                "metric_type": parameters["custom_dataset_metric_type"],
+                "file_count": parameters["custom_dataset_file_count"],
+                "use_shuffled": parameters["custom_dataset_use_shuffled"],
+                "with_gt": parameters["custom_dataset_with_gt"],
+            }
+        }
+    return custom_case_config
+
+
 log = logging.getLogger(__name__)
 
 
@@ -205,6 +238,7 @@ class CommonTypedDict(TypedDict):
         click.option(
             "--case-type",
             type=click.Choice([ct.name for ct in CaseType if ct.name != "Custom"]),
+            is_eager=True,
             default="Performance1536D50K",
             help="Case type",
         ),
@@ -258,6 +292,108 @@ class CommonTypedDict(TypedDict):
             callback=lambda *args: list(map(int, click_arg_split(*args))),
         ),
     ]
+    custom_case_name: Annotated[
+        str,
+        click.option(
+            "--custom-case-name",
+            help="Custom dataset case name",
+            callback=check_custom_case_parameters,
+        )
+    ]
+    custom_case_description: Annotated[
+        str,
+        click.option(
+            "--custom-case-description",
+            help="Custom dataset case description",
+            default="This is a customized dataset.",
+            show_default=True,
+        )
+    ]
+    custom_case_load_timeout: Annotated[
+        int,
+        click.option(
+            "--custom-case-load-timeout",
+            help="Custom dataset case load timeout",
+            default=36000,
+            show_default=True,
+        )
+    ]
+    custom_case_optimize_timeout: Annotated[
+        int,
+        click.option(
+            "--custom-case-optimize-timeout",
+            help="Custom dataset case optimize timeout",
+            default=36000,
+            show_default=True,
+        )
+    ]
+    custom_dataset_name: Annotated[
+        str,
+        click.option(
+            "--custom-dataset-name",
+            help="Custom dataset name",
+            callback=check_custom_case_parameters,
+        ),
+    ]
+    custom_dataset_dir: Annotated[
+        str,
+        click.option(
+            "--custom-dataset-dir",
+            help="Custom dataset directory",
+            callback=check_custom_case_parameters,
+        ),
+    ]
+    custom_dataset_size: Annotated[
+        int,
+        click.option(
+            "--custom-dataset-size",
+            help="Custom dataset size",
+            callback=check_custom_case_parameters,
+        ),
+    ]
+    custom_dataset_dim: Annotated[
+        int,
+        click.option(
+            "--custom-dataset-dim",
+            help="Custom dataset dimension",
+            callback=check_custom_case_parameters,
+        ),
+    ]
+    custom_dataset_metric_type: Annotated[
+        str,
+        click.option(
+            "--custom-dataset-metric-type",
+            help="Custom dataset metric type",
+            default=MetricType.COSINE.name,
+            show_default=True,
+        ),
+    ]
+    custom_dataset_file_count: Annotated[
+        int,
+        click.option(
+            "--custom-dataset-file-count",
+            help="Custom dataset file count",
+            callback=check_custom_case_parameters,
+        ),
+    ]
+    custom_dataset_use_shuffled: Annotated[
+        bool,
+        click.option(
+            "--custom-dataset-use-shuffled/--skip-custom-dataset-use-shuffled",
+            help="Custom dataset use shuffled",
+            default=False,
+            show_default=True,
+        ),
+    ]
+    custom_dataset_with_gt: Annotated[
+        bool,
+        click.option(
+            "--custom-dataset-with-gt/--skip-custom-dataset-with-gt",
+            help="Custom dataset with ground truth",
+            default=True,
+            show_default=True,
+        ),
+    ]
 
 
 class HNSWBaseTypedDict(TypedDict):
@@ -343,6 +479,7 @@ def run(
                 concurrency_duration=parameters["concurrency_duration"],
                 num_concurrency=[int(s) for s in parameters["num_concurrency"]],
             ),
+            custom_case=parameters.get("custom_case", {}),
         ),
         stages=parse_task_stages(
             (

vectordb_bench/cli/vectordbbench.py

@@ -1,21 +1,27 @@
 from ..backend.clients.pgvector.cli import PgVectorHNSW
+from ..backend.clients.pgvecto_rs.cli import PgVectoRSHNSW, PgVectoRSIVFFlat
+from ..backend.clients.pgvectorscale.cli import PgVectorScaleDiskAnn
 from ..backend.clients.redis.cli import Redis
+from ..backend.clients.memorydb.cli import MemoryDB
 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
 
 cli.add_command(PgVectorHNSW)
+cli.add_command(PgVectoRSHNSW)
+cli.add_command(PgVectoRSIVFFlat)
 cli.add_command(Redis)
+cli.add_command(MemoryDB)
 cli.add_command(Weaviate)
 cli.add_command(Test)
 cli.add_command(ZillizAutoIndex)
 cli.add_command(MilvusAutoIndex)
 cli.add_command(AWSOpenSearch)
+cli.add_command(PgVectorScaleDiskAnn)
 
 
 if __name__ == "__main__":

vectordb_bench/frontend/components/check_results/charts.py

@@ -1,5 +1,7 @@
 from vectordb_bench.backend.cases import Case
-from vectordb_bench.frontend.components.check_results.expanderStyle import initMainExpanderStyle
+from vectordb_bench.frontend.components.check_results.expanderStyle import (
+    initMainExpanderStyle,
+)
 from vectordb_bench.metric import metricOrder, isLowerIsBetterMetric, metricUnitMap
 from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.models import ResultLabel
@@ -11,7 +13,7 @@ def drawCharts(st, allData, failedTasks, caseNames: list[str]):
     for caseName in caseNames:
         chartContainer = st.expander(caseName, True)
         data = [data for data in allData if data["case_name"] == caseName]
-        drawChart(data, chartContainer)
+        drawChart(data, chartContainer, key_prefix=caseName)
 
         errorDBs = failedTasks[caseName]
         showFailedDBs(chartContainer, errorDBs)
@@ -35,7 +37,7 @@ def showFailedText(st, text, dbs):
         )
 
 
-def drawChart(data, st):
+def drawChart(data, st, key_prefix: str):
     metricsSet = set()
     for d in data:
         metricsSet = metricsSet.union(d["metricsSet"])
@@ -43,7 +45,8 @@ def drawChart(data, st):
 
     for i, metric in enumerate(showMetrics):
         container = st.container()
-        drawMetricChart(data, metric, container)
+        key = f"{key_prefix}-{metric}"
+        drawMetricChart(data, metric, container, key=key)
 
 
 def getLabelToShapeMap(data):
@@ -75,7 +78,7 @@ def getLabelToShapeMap(data):
     return labelToShapeMap
 
 
-def drawMetricChart(data, metric, st):
+def drawMetricChart(data, metric, st, key: str):
     dataWithMetric = [d for d in data if d.get(metric, 0) > 1e-7]
     # dataWithMetric = data
     if len(dataWithMetric) == 0:
@@ -161,4 +164,4 @@ def drawMetricChart(data, metric, st):
         ),
     )
 
-    chart.plotly_chart(fig, use_container_width=True)
+    chart.plotly_chart(fig, use_container_width=True, key=key)

vectordb_bench/frontend/components/check_results/data.py

@@ -24,7 +24,10 @@ def getFilterTasks(
         task
         for task in tasks
         if task.task_config.db_name in dbNames
-        and task.task_config.case_config.case_id.case_cls(task.task_config.case_config.custom_case).name in caseNames
+        and task.task_config.case_config.case_id.case_cls(
+            task.task_config.case_config.custom_case
+        ).name
+        in caseNames
     ]
     return filterTasks
 
@@ -35,17 +38,20 @@ def mergeTasks(tasks: list[CaseResult]):
         db_name = task.task_config.db_name
         db = task.task_config.db.value
         db_label = task.task_config.db_config.db_label or ""
-        case = task.task_config.case_config.case_id.case_cls(task.task_config.case_config.custom_case)
+        version = task.task_config.db_config.version or ""
+        case = task.task_config.case_config.case_id.case_cls(
+            task.task_config.case_config.custom_case
+        )
         dbCaseMetricsMap[db_name][case.name] = {
             "db": db,
             "db_label": db_label,
+            "version": version,
             "metrics": mergeMetrics(
                 dbCaseMetricsMap[db_name][case.name].get("metrics", {}),
                 asdict(task.metrics),
             ),
             "label": getBetterLabel(
-                dbCaseMetricsMap[db_name][case.name].get(
-                    "label", ResultLabel.FAILED),
+                dbCaseMetricsMap[db_name][case.name].get("label", ResultLabel.FAILED),
                 task.label,
             ),
         }
@@ -57,6 +63,7 @@ def mergeTasks(tasks: list[CaseResult]):
             metrics = metricInfo["metrics"]
             db = metricInfo["db"]
             db_label = metricInfo["db_label"]
+            version = metricInfo["version"]
             label = metricInfo["label"]
             if label == ResultLabel.NORMAL:
                 mergedTasks.append(
@@ -64,6 +71,7 @@ def mergeTasks(tasks: list[CaseResult]):
                         "db_name": db_name,
                         "db": db,
                         "db_label": db_label,
+                        "version": version,
                         "case_name": case_name,
                         "metricsSet": set(metrics.keys()),
                         **metrics,
@@ -79,8 +87,7 @@ def mergeMetrics(metrics_1: dict, metrics_2: dict) -> dict:
     metrics = {**metrics_1}
     for key, value in metrics_2.items():
         metrics[key] = (
-            getBetterMetric(
-                key, value, metrics[key]) if key in metrics else value
+            getBetterMetric(key, value, metrics[key]) if key in metrics else value
         )
 
     return metrics

vectordb_bench/frontend/components/concurrent/charts.py

@@ -22,7 +22,7 @@ def drawChartsByCase(allData, showCaseNames: list[str], st):
             for caseData in caseDataList
             for i in range(len(caseData["conc_num_list"]))
         ]
-        drawChart(data, chartContainer)
+        drawChart(data, chartContainer, key=f"{caseName}-qps-p99")
 
 
 def getRange(metric, data, padding_multipliers):
@@ -36,7 +36,7 @@ def getRange(metric, data, padding_multipliers):
     return rangeV
 
 
-def drawChart(data, st):
+def drawChart(data, st, key: str):
     if len(data) == 0:
         return
 
@@ -73,7 +73,4 @@ def drawChart(data, st):
     fig.update_yaxes(range=yrange, title_text="QPS")
     fig.update_traces(textposition="bottom right", texttemplate="conc-%{text:,.4~r}")
 
-    st.plotly_chart(
-        fig,
-        use_container_width=True,
-    )
+    st.plotly_chart(fig, use_container_width=True, key=key)

vectordb_bench/frontend/components/run_test/caseSelector.py

@@ -100,6 +100,16 @@ def caseConfigSetting(st, dbToCaseClusterConfigs, uiCaseItem: UICaseItem, active
                         value=config.inputConfig["value"],
                         help=config.inputHelp,
                     )
+                elif config.inputType == InputType.Float:
+                    caseConfig[config.label] = column.number_input(
+                        config.displayLabel if config.displayLabel else config.label.value,
+                        step=config.inputConfig.get("step", 0.1),
+                        min_value=config.inputConfig["min"],
+                        max_value=config.inputConfig["max"],
+                        key=key,
+                        value=config.inputConfig["value"],
+                        help=config.inputHelp,
+                    )
                 k += 1
         if k == 0:
             columns[1].write("Auto")

vectordb_bench/frontend/components/run_test/dbConfigSetting.py

@@ -1,9 +1,10 @@
 from pydantic import ValidationError
-from vectordb_bench.frontend.config.styles import *
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.frontend.config.styles import DB_CONFIG_SETTING_COLUMNS
 from vectordb_bench.frontend.utils import inputIsPassword
 
 
-def dbConfigSettings(st, activedDbList):
+def dbConfigSettings(st, activedDbList: list[DB]):
     expander = st.expander("Configurations for the selected databases", True)
 
     dbConfigs = {}
@@ -27,7 +28,7 @@ def dbConfigSettings(st, activedDbList):
     return dbConfigs, isAllValid
 
 
-def dbConfigSettingItem(st, activeDb):
+def dbConfigSettingItem(st, activeDb: DB):
     st.markdown(
         f"<div style='font-weight: 600; font-size: 20px; margin-top: 16px;'>{activeDb.value}</div>",
         unsafe_allow_html=True,
@@ -36,20 +37,41 @@ def dbConfigSettingItem(st, activeDb):
 
     dbConfigClass = activeDb.config_cls
     properties = dbConfigClass.schema().get("properties")
-    propertiesItems = list(properties.items())
-    moveDBLabelToLast(propertiesItems)
     dbConfig = {}
-    for j, property in enumerate(propertiesItems):
-        column = columns[j % DB_CONFIG_SETTING_COLUMNS]
-        key, value = property
+    idx = 0
+
+    # db config (unique)
+    for key, property in properties.items():
+        if (
+            key not in dbConfigClass.common_short_configs()
+            and key not in dbConfigClass.common_long_configs()
+        ):
+            column = columns[idx % DB_CONFIG_SETTING_COLUMNS]
+            idx += 1
+            dbConfig[key] = column.text_input(
+                key,
+                key="%s-%s" % (activeDb.name, key),
+                value=property.get("default", ""),
+                type="password" if inputIsPassword(key) else "default",
+            )
+    # db config (common short labels)
+    for key in dbConfigClass.common_short_configs():
+        column = columns[idx % DB_CONFIG_SETTING_COLUMNS]
+        idx += 1
         dbConfig[key] = column.text_input(
             key,
-            key="%s-%s" % (activeDb, key),
-            value=value.get("default", ""),
-            type="password" if inputIsPassword(key) else "default",
+            key="%s-%s" % (activeDb.name, key),
+            value="",
+            type="default",
+            placeholder="optional, for labeling results",
         )
-    return dbConfig
-
 
-def moveDBLabelToLast(propertiesItems):
-    propertiesItems.sort(key=lambda x: 1 if x[0] == "db_label" else 0)
+    # db config (common long text_input)
+    for key in dbConfigClass.common_long_configs():
+        dbConfig[key] = st.text_area(
+            key,
+            key="%s-%s" % (activeDb.name, key),
+            value="",
+            placeholder="optional",
+        )
+    return dbConfig

vectordb_bench/frontend/components/run_test/initStyle.py

@@ -9,6 +9,8 @@ def initStyle(st):
             div[data-testid='stHorizontalBlock'] {gap: 8px;}
             /* check box */
             .stCheckbox p { color: #000; font-size: 18px; font-weight: 600; }
+            /* db selector - db_name should not wrap */
+            div[data-testid="stVerticalBlockBorderWrapper"] div[data-testid="stCheckbox"] div[data-testid="stWidgetLabel"] p { white-space: nowrap; }
         </style>""",
         unsafe_allow_html=True,
-    )
+    )