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

vectordb_bench/cli/cli.py

@@ -0,0 +1,362 @@
+import logging
+import time
+from concurrent.futures import wait
+from datetime import datetime
+from pprint import pformat
+from typing import (
+    Annotated,
+    Callable,
+    List,
+    Optional,
+    Type,
+    TypedDict,
+    Unpack,
+    get_origin,
+    get_type_hints,
+    Dict,
+    Any,
+)
+import click
+from .. import config
+from ..backend.clients import DB
+from ..interface import benchMarkRunner, global_result_future
+from ..models import (
+    CaseConfig,
+    CaseType,
+    ConcurrencySearchConfig,
+    DBCaseConfig,
+    DBConfig,
+    TaskConfig,
+    TaskStage,
+)
+import os
+from yaml import load
+try:
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader
+
+
+def click_get_defaults_from_file(ctx, param, value):
+    if value:
+        if os.path.exists(value):
+            input_file = value
+        else:
+            input_file = os.path.join(config.CONFIG_LOCAL_DIR, value)
+        try:
+            with open(input_file, 'r') as f:
+                _config: Dict[str, Dict[str, Any]] = load(f.read(), Loader=Loader)
+                ctx.default_map = _config.get(ctx.command.name, {})
+        except Exception as e:
+            raise click.BadParameter(f"Failed to load config file: {e}")
+    return value
+
+
+def click_parameter_decorators_from_typed_dict(
+    typed_dict: Type,
+) -> Callable[[click.decorators.FC], click.decorators.FC]:
+    """A convenience method decorator that will read in a TypedDict with parameters defined by Annotated types.
+    from .models import CaseConfig, CaseType, DBCaseConfig, DBConfig, TaskConfig, TaskStage
+    The click.options will be collected and re-composed as a single decorator to apply to the click.command.
+
+    Args:
+        typed_dict (TypedDict) with Annotated[..., click.option()] keys
+
+    Returns:
+        a fully decorated method
+
+
+    For clarity, the key names of the TypedDict will be used to determine the type hints for the input parameters.
+    The actual function parameters are controlled by the click.option definitions. You must manually ensure these are aligned in a sensible way!
+
+    Example:
+    ```
+    class CommonTypedDict(TypedDict):
+        z: Annotated[int, click.option("--z/--no-z", is_flag=True, type=bool, help="help z", default=True, show_default=True)]
+        name: Annotated[str, click.argument("name", required=False, default="Jeff")]
+
+    class FooTypedDict(CommonTypedDict):
+        x: Annotated[int, click.option("--x", type=int, help="help x", default=1, show_default=True)]
+        y: Annotated[str, click.option("--y", type=str, help="help y", default="foo", show_default=True)]
+
+    @cli.command()
+    @click_parameter_decorators_from_typed_dict(FooTypedDict)
+    def foo(**parameters: Unpack[FooTypedDict]):
+        "Foo docstring"
+        print(f"input parameters: {parameters["x"]}")
+    ```
+    """
+    decorators = []
+    for _, t in get_type_hints(typed_dict, include_extras=True).items():
+        assert get_origin(t) is Annotated
+        if (
+            len(t.__metadata__) == 1
+            and t.__metadata__[0].__module__ == "click.decorators"
+        ):
+            # happy path -- only accept Annotated[..., Union[click.option,click.argument,...]] with no additional metadata defined (len=1)
+            decorators.append(t.__metadata__[0])
+        else:
+            raise RuntimeError(
+                "Click-TypedDict decorator parsing must only contain root type and a click decorator like click.option. See docstring"
+            )
+
+    def deco(f):
+        for dec in reversed(decorators):
+            f = dec(f)
+        return f
+
+    return deco
+
+
+def click_arg_split(ctx: click.Context, param: click.core.Option, value):
+    """Will split a comma-separated list input into an actual list.
+
+    Args:
+        ctx (...): unused click arg
+        param (...): unused click arg
+        value (str): input comma-separated list
+
+    Returns:
+        value (List[str]): list of original
+    """
+    # split columns by ',' and remove whitespace
+    if value is None:
+        return []
+    return [c.strip() for c in value.split(",") if c.strip()]
+
+
+def parse_task_stages(
+    drop_old: bool,
+    load: bool,
+    search_serial: bool,
+    search_concurrent: bool,
+) -> List[TaskStage]:
+    stages = []
+    if load and not drop_old:
+        raise RuntimeError("Dropping old data cannot be skipped if loading data")
+    elif drop_old and not load:
+        raise RuntimeError("Load cannot be skipped if dropping old data")
+    if drop_old:
+        stages.append(TaskStage.DROP_OLD)
+    if load:
+        stages.append(TaskStage.LOAD)
+    if search_serial:
+        stages.append(TaskStage.SEARCH_SERIAL)
+    if search_concurrent:
+        stages.append(TaskStage.SEARCH_CONCURRENT)
+    return stages
+
+
+log = logging.getLogger(__name__)
+
+
+class CommonTypedDict(TypedDict):
+    config_file: Annotated[
+        bool,
+        click.option('--config-file',
+                     type=click.Path(),
+                     callback=click_get_defaults_from_file,
+                     is_eager=True,
+                     expose_value=False,
+                     help='Read configuration from yaml file'),
+    ]
+    drop_old: Annotated[
+        bool,
+        click.option(
+            "--drop-old/--skip-drop-old",
+            type=bool,
+            default=True,
+            help="Drop old or skip",
+            show_default=True,
+        ),
+    ]
+    load: Annotated[
+        bool,
+        click.option(
+            "--load/--skip-load",
+            type=bool,
+            default=True,
+            help="Load or skip",
+            show_default=True,
+        ),
+    ]
+    search_serial: Annotated[
+        bool,
+        click.option(
+            "--search-serial/--skip-search-serial",
+            type=bool,
+            default=True,
+            help="Search serial or skip",
+            show_default=True,
+        ),
+    ]
+    search_concurrent: Annotated[
+        bool,
+        click.option(
+            "--search-concurrent/--skip-search-concurrent",
+            type=bool,
+            default=True,
+            help="Search concurrent or skip",
+            show_default=True,
+        ),
+    ]
+    case_type: Annotated[
+        str,
+        click.option(
+            "--case-type",
+            type=click.Choice([ct.name for ct in CaseType if ct.name != "Custom"]),
+            default="Performance1536D50K",
+            help="Case type",
+        ),
+    ]
+    db_label: Annotated[
+        str,
+        click.option(
+            "--db-label", type=str, help="Db label, default: date in ISO format",
+            show_default=True,
+            default=datetime.now().isoformat()
+        ),
+    ]
+    dry_run: Annotated[
+        bool,
+        click.option(
+            "--dry-run",
+            type=bool,
+            default=False,
+            is_flag=True,
+            help="Print just the configuration and exit without running the tasks",
+        ),
+    ]
+    k: Annotated[
+        int,
+        click.option(
+            "--k",
+            type=int,
+            default=config.K_DEFAULT,
+            show_default=True,
+            help="K value for number of nearest neighbors to search",
+        ),
+    ]
+    concurrency_duration: Annotated[
+        int,
+        click.option(
+            "--concurrency-duration",
+            type=int,
+            default=config.CONCURRENCY_DURATION,
+            show_default=True,
+            help="Adjusts the duration in seconds of each concurrency search",
+        ),
+    ]
+    num_concurrency: Annotated[
+        List[str],
+        click.option(
+            "--num-concurrency",
+            type=str,
+            help="Comma-separated list of concurrency values to test during concurrent search",
+            show_default=True,
+            default=",".join(map(str, config.NUM_CONCURRENCY)),
+            callback=lambda *args: list(map(int, click_arg_split(*args))),
+        ),
+    ]
+
+
+class HNSWBaseTypedDict(TypedDict):
+    m: Annotated[Optional[int], click.option("--m", type=int, help="hnsw m")]
+    ef_construction: Annotated[
+        Optional[int],
+        click.option("--ef-construction", type=int, help="hnsw ef-construction"),
+    ]
+
+
+class HNSWBaseRequiredTypedDict(TypedDict):
+    m: Annotated[Optional[int], click.option("--m", type=int, help="hnsw m", required=True)]
+    ef_construction: Annotated[
+        Optional[int],
+        click.option("--ef-construction", type=int, help="hnsw ef-construction", required=True),
+    ]
+
+
+class HNSWFlavor1(HNSWBaseTypedDict):
+    ef_search: Annotated[
+        Optional[int], click.option("--ef-search", type=int, help="hnsw ef-search")
+    ]
+
+
+class HNSWFlavor2(HNSWBaseTypedDict):
+    ef_runtime: Annotated[
+        Optional[int], click.option("--ef-runtime", type=int, help="hnsw ef-runtime")
+    ]
+
+
+class HNSWFlavor3(HNSWBaseRequiredTypedDict):
+    ef_search: Annotated[
+        Optional[int], click.option("--ef-search", type=int, help="hnsw ef-search", required=True)
+    ]
+
+
+class IVFFlatTypedDict(TypedDict):
+    lists: Annotated[
+        Optional[int], click.option("--lists", type=int, help="ivfflat lists")
+    ]
+    probes: Annotated[
+        Optional[int], click.option("--probes", type=int, help="ivfflat probes")
+    ]
+
+
+class IVFFlatTypedDictN(TypedDict):
+    nlist: Annotated[
+        Optional[int], click.option("--lists", "nlist", type=int, help="ivfflat lists", required=True)
+    ]
+    nprobe: Annotated[
+        Optional[int], click.option("--probes", "nprobe", type=int, help="ivfflat probes", required=True)
+    ]
+
+
+@click.group()
+def cli():
+    ...
+
+
+def run(
+    db: DB,
+    db_config: DBConfig,
+    db_case_config: DBCaseConfig,
+    **parameters: Unpack[CommonTypedDict],
+):
+    """Builds a single VectorDBBench Task and runs it, awaiting the task until finished.
+
+    Args:
+        db (DB)
+        db_config (DBConfig)
+        db_case_config (DBCaseConfig)
+        **parameters: expects keys from CommonTypedDict
+    """
+
+    task = TaskConfig(
+        db=db,
+        db_config=db_config,
+        db_case_config=db_case_config,
+        case_config=CaseConfig(
+            case_id=CaseType[parameters["case_type"]],
+            k=parameters["k"],
+            concurrency_search_config=ConcurrencySearchConfig(
+                concurrency_duration=parameters["concurrency_duration"],
+                num_concurrency=[int(s) for s in parameters["num_concurrency"]],
+            ),
+        ),
+        stages=parse_task_stages(
+            (
+                False if not parameters["load"] else parameters["drop_old"]
+            ),  # only drop old data if loading new data
+            parameters["load"],
+            parameters["search_serial"],
+            parameters["search_concurrent"],
+        ),
+    )
+
+    log.info(f"Task:\n{pformat(task)}\n")
+    if not parameters["dry_run"]:
+        benchMarkRunner.run([task])
+        time.sleep(5)
+        if global_result_future:
+            wait([global_result_future])

vectordb_bench/cli/vectordbbench.py

@@ -0,0 +1,22 @@
+from ..backend.clients.pgvector.cli import PgVectorHNSW
+from ..backend.clients.redis.cli import Redis
+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(Redis)
+cli.add_command(Weaviate)
+cli.add_command(Test)
+cli.add_command(ZillizAutoIndex)
+cli.add_command(MilvusAutoIndex)
+cli.add_command(AWSOpenSearch)
+
+
+if __name__ == "__main__":
+    cli()

vectordb_bench/config-files/sample_config.yml

@@ -0,0 +1,17 @@
+pgvectorhnsw:
+  db_label: pgConfigTest
+  user_name: vectordbbench
+  db_name:  vectordbbench
+  host: localhost
+  m: 16
+  ef_construction: 128
+  ef_search: 128
+milvushnsw:
+  skip_search_serial: True
+  case_type: Performance1536D50K
+  uri: http://localhost:19530
+  m: 16
+  ef_construction: 128
+  ef_search: 128
+  drop_old: False
+  load: False

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)
 
 

vectordb_bench/frontend/components/check_results/data.py

@@ -8,9 +8,9 @@ from vectordb_bench.models import CaseResult, ResultLabel
 def getChartData(
     tasks: list[CaseResult],
     dbNames: list[str],
-    cases: list[Case],
+    caseNames: list[str],
 ):
-    filterTasks = getFilterTasks(tasks, dbNames, cases)
+    filterTasks = getFilterTasks(tasks, dbNames, caseNames)
     mergedTasks, failedTasks = mergeTasks(filterTasks)
     return mergedTasks, failedTasks
 
@@ -18,14 +18,13 @@ def getChartData(
 def getFilterTasks(
     tasks: list[CaseResult],
     dbNames: list[str],
-    cases: list[Case],
+    caseNames: list[str],
 ) -> list[CaseResult]:
-    case_ids = [case.case_id for case in cases]
     filterTasks = [
         task
         for task in tasks
         if task.task_config.db_name in dbNames
-        and task.task_config.case_config.case_id in case_ids
+        and task.task_config.case_config.case_id.case_cls(task.task_config.case_config.custom_case).name in caseNames
     ]
     return filterTasks
 
@@ -36,16 +35,17 @@ 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_id = task.task_config.case_config.case_id
-        dbCaseMetricsMap[db_name][case_id] = {
+        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,
             "metrics": mergeMetrics(
-                dbCaseMetricsMap[db_name][case_id].get("metrics", {}),
+                dbCaseMetricsMap[db_name][case.name].get("metrics", {}),
                 asdict(task.metrics),
             ),
             "label": getBetterLabel(
-                dbCaseMetricsMap[db_name][case_id].get("label", ResultLabel.FAILED),
+                dbCaseMetricsMap[db_name][case.name].get(
+                    "label", ResultLabel.FAILED),
                 task.label,
             ),
         }
@@ -53,12 +53,11 @@ def mergeTasks(tasks: list[CaseResult]):
     mergedTasks = []
     failedTasks = defaultdict(lambda: defaultdict(str))
     for db_name, caseMetricsMap in dbCaseMetricsMap.items():
-        for case_id, metricInfo in caseMetricsMap.items():
+        for case_name, metricInfo in caseMetricsMap.items():
             metrics = metricInfo["metrics"]
             db = metricInfo["db"]
             db_label = metricInfo["db_label"]
             label = metricInfo["label"]
-            case_name = case_id.case_name
             if label == ResultLabel.NORMAL:
                 mergedTasks.append(
                     {
@@ -80,22 +79,26 @@ 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
 
 
 def getBetterMetric(metric, value_1, value_2):
-    if value_1 < 1e-7:
-        return value_2
-    if value_2 < 1e-7:
+    try:
+        if value_1 < 1e-7:
+            return value_2
+        if value_2 < 1e-7:
+            return value_1
+        return (
+            min(value_1, value_2)
+            if isLowerIsBetterMetric(metric)
+            else max(value_1, value_2)
+        )
+    except Exception:
         return value_1
-    return (
-        min(value_1, value_2)
-        if isLowerIsBetterMetric(metric)
-        else max(value_1, value_2)
-    )
 
 
 def getBetterLabel(label_1: ResultLabel, label_2: ResultLabel):

vectordb_bench/frontend/components/check_results/expanderStyle.py

@@ -1,7 +1,7 @@
 def initMainExpanderStyle(st):
     st.markdown(
         """<style>
-            .main .streamlit-expanderHeader p {font-size: 20px; font-weight: 600;}
+            .main div[data-testid='stExpander'] p {font-size: 18px; font-weight: 600;}
             .main div[data-testid='stExpander'] {
                 background-color: #F6F8FA;
                 border: 1px solid #A9BDD140;

vectordb_bench/frontend/components/check_results/filters.py

@@ -1,8 +1,8 @@
 from vectordb_bench.backend.cases import Case
 from vectordb_bench.frontend.components.check_results.data import getChartData
 from vectordb_bench.frontend.components.check_results.expanderStyle import initSidebarExanderStyle
-from vectordb_bench.frontend.const.dbCaseConfigs import CASE_LIST
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.dbCaseConfigs import CASE_NAME_ORDER
+from vectordb_bench.frontend.config.styles import *
 import streamlit as st
 
 from vectordb_bench.models import CaseResult, TestResult
@@ -18,11 +18,12 @@ def getshownData(results: list[TestResult], st):
     st.header("Filters")
 
     shownResults = getshownResults(results, st)
-    showDBNames, showCases = getShowDbsAndCases(shownResults, st)
+    showDBNames, showCaseNames = getShowDbsAndCases(shownResults, st)
 
-    shownData, failedTasks = getChartData(shownResults, showDBNames, showCases)
+    shownData, failedTasks = getChartData(
+        shownResults, showDBNames, showCaseNames)
 
-    return shownData, failedTasks, showCases
+    return shownData, failedTasks, showCaseNames
 
 
 def getshownResults(results: list[TestResult], st) -> list[CaseResult]:
@@ -52,12 +53,18 @@ def getshownResults(results: list[TestResult], st) -> list[CaseResult]:
     return selectedResult
 
 
-def getShowDbsAndCases(result: list[CaseResult], st) -> tuple[list[str], list[Case]]:
+def getShowDbsAndCases(result: list[CaseResult], st) -> tuple[list[str], list[str]]:
     initSidebarExanderStyle(st)
     allDbNames = list(set({res.task_config.db_name for res in result}))
     allDbNames.sort()
-    allCasesSet = set({res.task_config.case_config.case_id for res in result})
-    allCases: list[Case] = [case.case_cls() for case in CASE_LIST if case in allCasesSet]
+    allCases: list[Case] = [
+        res.task_config.case_config.case_id.case_cls(
+            res.task_config.case_config.custom_case)
+        for res in result
+    ]
+    allCaseNameSet = set({case.name for case in allCases})
+    allCaseNames = [case_name for case_name in CASE_NAME_ORDER if case_name in allCaseNameSet] + \
+        [case_name for case_name in allCaseNameSet if case_name not in CASE_NAME_ORDER]
 
     # DB Filter
     dbFilterContainer = st.container()
@@ -70,15 +77,14 @@ def getShowDbsAndCases(result: list[CaseResult], st) -> tuple[list[str], list[Ca
 
     # Case Filter
     caseFilterContainer = st.container()
-    showCases = filterView(
+    showCaseNames = filterView(
         caseFilterContainer,
         "Case Filter",
-        [case for case in allCases],
+        [caseName for caseName in allCaseNames],
         col=1,
-        optionLables=[case.name for case in allCases],
     )
 
-    return showDBNames, showCases
+    return showDBNames, showCaseNames
 
 
 def filterView(container, header, options, col, optionLables=None):
@@ -114,7 +120,8 @@ def filterView(container, header, options, col, optionLables=None):
     )
     if optionLables is None:
         optionLables = options
-    isActive = {option: st.session_state[selectAllState] for option in optionLables}
+    isActive = {option: st.session_state[selectAllState]
+                for option in optionLables}
     for i, option in enumerate(optionLables):
         isActive[option] = columns[i % col].checkbox(
             optionLables[i],

vectordb_bench/frontend/components/check_results/headerIcon.py

@@ -1,4 +1,4 @@
-from vectordb_bench.frontend.const.styles import HEADER_ICON
+from vectordb_bench.frontend.config.styles import HEADER_ICON
 
 
 def drawHeaderIcon(st):

vectordb_bench/frontend/components/check_results/priceTable.py

@@ -3,7 +3,7 @@ import pandas as pd
 from collections import defaultdict
 import streamlit as st
 
-from vectordb_bench.frontend.const.dbPrices import DB_DBLABEL_TO_PRICE
+from vectordb_bench.frontend.config.dbPrices import DB_DBLABEL_TO_PRICE
 
 
 def priceTable(container, data):

vectordb_bench/frontend/components/check_results/stPageConfig.py

@@ -1,4 +1,4 @@
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 
 
 def initResultsPageConfig(st):