vectordb-bench 0.0.9__py3-none-any.whl → 0.0.11__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,20 @@
+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 .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)
+
+
+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/frontend/components/check_results/data.py

@@ -87,15 +87,18 @@ def mergeMetrics(metrics_1: dict, metrics_2: dict) -> dict:
 
 
 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/concurrent/charts.py

@@ -0,0 +1,82 @@
+
+
+from vectordb_bench.backend.cases import Case
+from vectordb_bench.frontend.components.check_results.expanderStyle import initMainExpanderStyle
+import plotly.express as px
+
+from vectordb_bench.frontend.const.styles import COLOR_MAP
+
+
+def drawChartsByCase(allData, cases: list[Case], st):
+    initMainExpanderStyle(st)
+    for case in cases:
+        chartContainer = st.expander(case.name, True)
+        caseDataList = [
+            data for data in allData if data["case_name"] == case.name]
+        data = [{
+            "conc_num": caseData["conc_num_list"][i],
+            "qps": caseData["conc_qps_list"][i],
+            "latency_p99": caseData["conc_latency_p99_list"][i] * 1000,
+            "db_name": caseData["db_name"],
+            "db": caseData["db"]
+
+        } for caseData in caseDataList for i in range(len(caseData["conc_num_list"]))]
+        drawChart(data, chartContainer)
+
+
+def getRange(metric, data, padding_multipliers):
+    minV = min([d.get(metric, 0) for d in data])
+    maxV = max([d.get(metric, 0) for d in data])
+    padding = maxV - minV
+    rangeV = [
+        minV - padding * padding_multipliers[0],
+        maxV + padding * padding_multipliers[1],
+    ]
+    return rangeV
+
+
+def drawChart(data, st):
+    if len(data) == 0:
+        return
+    
+    x = "latency_p99"
+    xrange = getRange(x, data, [0.05, 0.1])
+
+    y = "qps"
+    yrange = getRange(y, data, [0.2, 0.1])
+
+    color = "db"
+    color_discrete_map = COLOR_MAP
+    color = "db_name"
+    color_discrete_map = None
+    line_group = "db_name"
+    text = "conc_num"
+
+    data.sort(key=lambda a: a["conc_num"])
+
+    fig = px.line(
+        data,
+        x=x,
+        y=y,
+        color=color,
+        color_discrete_map=color_discrete_map,
+        line_group=line_group,
+        text=text,
+        markers=True,
+        # color_discrete_map=color_discrete_map,
+        hover_data={
+            "conc_num": True,
+        },
+        height=720,
+    )
+    fig.update_xaxes(range=xrange, title_text="Latency P99 (ms)")
+    fig.update_yaxes(range=yrange, title_text="QPS")
+    fig.update_traces(textposition="bottom right",
+                      texttemplate="conc-%{text:,.4~r}")
+    # fig.update_layout(
+    #     margin=dict(l=0, r=0, t=40, b=0, pad=8),
+    #     legend=dict(
+    #         orientation="h", yanchor="bottom", y=1, xanchor="right", x=1, title=""
+    #     ),
+    # )
+    st.plotly_chart(fig, use_container_width=True,)

vectordb_bench/frontend/components/run_test/dbSelector.py

@@ -1,3 +1,5 @@
+from streamlit.runtime.media_file_storage import MediaFileStorageError
+
 from vectordb_bench.frontend.const.styles import *
 from vectordb_bench.frontend.const.dbCaseConfigs import DB_LIST
 
@@ -30,7 +32,11 @@ def dbSelector(st):
     for i, db in enumerate(DB_LIST):
         column = dbContainerColumns[i % DB_SELECTOR_COLUMNS]
         dbIsActived[db] = column.checkbox(db.name)
-        column.image(DB_TO_ICON.get(db, ""))
+        try:
+            column.image(DB_TO_ICON.get(db, ""))
+        except MediaFileStorageError as e:
+            column.warning(f"{db.name} image not available")
+            pass
     activedDbList = [db for db in DB_LIST if dbIsActived[db]]
 
     return activedDbList

vectordb_bench/frontend/components/run_test/submitTask.py

@@ -37,22 +37,30 @@ def taskLabelInput(st):
 def advancedSettings(st):
     container = st.columns([1, 2])
     index_already_exists = container[0].checkbox("Index already exists", value=False)
-    container[1].caption("if actived, inserting and building will be skipped.")
+    container[1].caption("if selected, inserting and building will be skipped.")
 
     container = st.columns([1, 2])
     use_aliyun = container[0].checkbox("Dataset from Aliyun (Shanghai)", value=False)
     container[1].caption(
-        "if actived, the dataset will be downloaded from Aliyun OSS shanghai, default AWS S3 aws-us-west."
+        "if selected, the dataset will be downloaded from Aliyun OSS shanghai, default AWS S3 aws-us-west."
     )
 
-    return index_already_exists, use_aliyun
+    container = st.columns([1, 2])
+    k = container[0].number_input("k",min_value=1, value=100, label_visibility="collapsed")
+    container[1].caption(
+        "K value for number of nearest neighbors to search"
+    )
+
+    return index_already_exists, use_aliyun, k
 
 
 def controlPanel(st, tasks, taskLabel, isAllValid):
-    index_already_exists, use_aliyun = advancedSettings(st)
+    index_already_exists, use_aliyun, k = advancedSettings(st)
 
     def runHandler():
         benchMarkRunner.set_drop_old(not index_already_exists)
+        for task in tasks:
+            task.case_config.k = k
         benchMarkRunner.set_download_address(use_aliyun)
         benchMarkRunner.run(tasks, taskLabel)
 

vectordb_bench/frontend/components/tables/data.py

@@ -0,0 +1,44 @@
+from dataclasses import asdict
+from vectordb_bench.backend.cases import CaseType
+from vectordb_bench.interface import benchMarkRunner
+from vectordb_bench.models import CaseResult, ResultLabel
+import pandas as pd
+
+
+def getNewResults():
+    allResults = benchMarkRunner.get_results()
+    newResults: list[CaseResult] = []
+
+    for res in allResults:
+        results = res.results
+        for result in results:
+            if result.label == ResultLabel.NORMAL:
+                newResults.append(result)
+                    
+
+    df = pd.DataFrame(formatData(newResults))
+    return df
+
+
+def formatData(caseResults: list[CaseResult]):
+    data = []
+    for caseResult in caseResults:
+        db = caseResult.task_config.db.value
+        db_label = caseResult.task_config.db_config.db_label
+        case_config = caseResult.task_config.case_config
+        db_case_config = caseResult.task_config.db_case_config
+        case = case_config.case_id.case_cls()
+        filter_rate = case.filter_rate
+        dataset = case.dataset.data.name
+        metrics = asdict(caseResult.metrics)
+        data.append(
+            {
+                "db": db,
+                "db_label": db_label,
+                "case_name": case.name,
+                "dataset": dataset,
+                "filter_rate": filter_rate,
+                **metrics,
+            }
+        )
+    return data

vectordb_bench/frontend/const/dbCaseConfigs.py

@@ -9,7 +9,7 @@ from vectordb_bench.models import CaseConfigParamType
 
 MAX_STREAMLIT_INT = (1 << 53) - 1
 
-DB_LIST = [d for d in DB]
+DB_LIST = [d for d in DB if d != DB.Test]
 
 DIVIDER = "DIVIDER"
 CASE_LIST_WITH_DIVIDER = [
@@ -19,6 +19,7 @@ CASE_LIST_WITH_DIVIDER = [
     DIVIDER,
     CaseType.Performance1536D5M,
     CaseType.Performance1536D500K,
+    CaseType.Performance1536D50K,
     DIVIDER,
     CaseType.Performance768D10M1P,
     CaseType.Performance768D1M1P,

vectordb_bench/frontend/pages/concurrent.py

@@ -0,0 +1,72 @@
+
+
+
+import streamlit as st
+from vectordb_bench.backend.cases import CaseType
+from vectordb_bench.frontend.components.check_results.footer import footer
+from vectordb_bench.frontend.components.check_results.expanderStyle import initMainExpanderStyle
+from vectordb_bench.frontend.components.check_results.priceTable import priceTable
+from vectordb_bench.frontend.components.check_results.headerIcon import drawHeaderIcon
+from vectordb_bench.frontend.components.check_results.nav import NavToResults, NavToRunTest
+from vectordb_bench.frontend.components.check_results.charts import drawMetricChart
+from vectordb_bench.frontend.components.check_results.filters import getshownData
+from vectordb_bench.frontend.components.concurrent.charts import drawChartsByCase
+from vectordb_bench.frontend.components.get_results.saveAsImage import getResults
+from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.interface import benchMarkRunner
+from vectordb_bench.models import TestResult
+
+
+def main():
+    # set page config
+    st.set_page_config(
+        page_title="VDBBench Conc Perf",
+        page_icon=FAVICON,
+        layout="wide",
+        # initial_sidebar_state="collapsed",
+    )
+
+    # header
+    drawHeaderIcon(st)
+
+    allResults = benchMarkRunner.get_results()
+    
+    def check_conc_data(res: TestResult):
+        case_results = res.results
+        count = 0
+        for case_result in case_results:
+            if len(case_result.metrics.conc_num_list) > 0:
+                count += 1
+        
+        return count > 0
+    
+    checkedResults = [res for res in allResults if check_conc_data(res)]
+        
+
+    st.title("VectorDB Benchmark (Concurrent Performance)")
+
+    # results selector
+    resultSelectorContainer = st.sidebar.container()
+    shownData, _, showCases = getshownData(
+        checkedResults, resultSelectorContainer)
+    
+
+    resultSelectorContainer.divider()
+
+    # nav
+    navContainer = st.sidebar.container()
+    NavToRunTest(navContainer)
+    NavToResults(navContainer)
+
+    # save or share
+    resultesContainer = st.sidebar.container()
+    getResults(resultesContainer, "vectordb_bench_concurrent")
+    
+    drawChartsByCase(shownData, showCases, st.container())
+
+    # footer
+    footer(st.container())
+
+
+if __name__ == "__main__":
+    main()