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

vectordb_bench/frontend/components/concurrent/charts.py

@@ -0,0 +1,79 @@
+from vectordb_bench.frontend.components.check_results.expanderStyle import (
+    initMainExpanderStyle,
+)
+import plotly.express as px
+
+from vectordb_bench.frontend.config.styles import COLOR_MAP
+
+
+def drawChartsByCase(allData, showCaseNames: list[str], st):
+    initMainExpanderStyle(st)
+    for caseName in showCaseNames:
+        chartContainer = st.expander(caseName, True)
+        caseDataList = [data for data in allData if data["case_name"] == caseName]
+        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,
+        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}")
+
+    st.plotly_chart(
+        fig,
+        use_container_width=True,
+    )

vectordb_bench/frontend/components/custom/displayCustomCase.py

@@ -0,0 +1,31 @@
+
+from vectordb_bench.frontend.components.custom.getCustomConfig import CustomCaseConfig
+
+
+def displayCustomCase(customCase: CustomCaseConfig, st, key):
+
+    columns = st.columns([1, 2])
+    customCase.dataset_config.name = columns[0].text_input(
+        "Name", key=f"{key}_name", value=customCase.dataset_config.name)
+    customCase.name = f"{customCase.dataset_config.name} (Performace Case)"
+    customCase.dataset_config.dir = columns[1].text_input(
+        "Folder Path", key=f"{key}_dir", value=customCase.dataset_config.dir)
+
+    columns = st.columns(4)
+    customCase.dataset_config.dim = columns[0].number_input(
+        "dim", key=f"{key}_dim", value=customCase.dataset_config.dim)
+    customCase.dataset_config.size = columns[1].number_input(
+        "size", key=f"{key}_size", value=customCase.dataset_config.size)
+    customCase.dataset_config.metric_type = columns[2].selectbox(
+        "metric type", key=f"{key}_metric_type", options=["L2", "Cosine", "IP"])
+    customCase.dataset_config.file_count = columns[3].number_input(
+        "train file count", key=f"{key}_file_count", value=customCase.dataset_config.file_count)
+
+    columns = st.columns(4)
+    customCase.dataset_config.use_shuffled = columns[0].checkbox(
+        "use shuffled data",  key=f"{key}_use_shuffled", value=customCase.dataset_config.use_shuffled)
+    customCase.dataset_config.with_gt = columns[1].checkbox(
+        "with groundtruth",  key=f"{key}_with_gt", value=customCase.dataset_config.with_gt)
+
+    customCase.description = st.text_area(
+        "description", key=f"{key}_description", value=customCase.description)

vectordb_bench/frontend/components/custom/displaypPrams.py

@@ -0,0 +1,11 @@
+def displayParams(st):
+    st.markdown("""
+- `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.
+  - Vectors data files: The file must be named `train.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Query test vectors: The file must be named `test.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Ground truth file: The file must be named `neighbors.parquet` and should have two columns: `id` corresponding to query vectors and `neighbors_id` as an array of `int`.
+
+- `Train File Count` - If the vector file is too large, you can consider splitting it into multiple files. The naming format for the split files should be `train-[index]-of-[file_count].parquet`. For example, `train-01-of-10.parquet` represents the second file (0-indexed) among 10 split files.
+
+- `Use Shuffled Data` - If you check this option, the vector data files need to be modified. VectorDBBench will load the data labeled with `shuffle`. For example, use `shuffle_train.parquet` instead of `train.parquet` and `shuffle_train-04-of-10.parquet` instead of `train-04-of-10.parquet`. The `id` column in the shuffled data can be in any order.
+""")

vectordb_bench/frontend/components/custom/getCustomConfig.py

@@ -0,0 +1,40 @@
+import json
+
+from pydantic import BaseModel
+
+from vectordb_bench import config
+
+
+class CustomDatasetConfig(BaseModel):
+    name: str = "custom_dataset"
+    dir: str = ""
+    size: int = 0
+    dim: int = 0
+    metric_type: str = "L2"
+    file_count: int = 1
+    use_shuffled: bool = False
+    with_gt: bool = True
+
+
+class CustomCaseConfig(BaseModel):
+    name: str = "custom_dataset (Performace Case)"
+    description: str = ""
+    load_timeout: int = 36000
+    optimize_timeout: int = 36000
+    dataset_config: CustomDatasetConfig = CustomDatasetConfig()
+
+
+def get_custom_configs():
+    with open(config.CUSTOM_CONFIG_DIR, "r") as f:
+        custom_configs = json.load(f)
+        return [CustomCaseConfig(**custom_config) for custom_config in custom_configs]
+
+
+def save_custom_configs(custom_configs: list[CustomDatasetConfig]):
+    with open(config.CUSTOM_CONFIG_DIR, "w") as f:
+        json.dump([custom_config.dict()
+                   for custom_config in custom_configs], f, indent=4)
+
+
+def generate_custom_case():
+    return CustomCaseConfig()

vectordb_bench/frontend/components/custom/initStyle.py

@@ -0,0 +1,15 @@
+def initStyle(st):
+    st.markdown(
+        """<style>
+            /* expander - header */
+            .main div[data-testid='stExpander'] summary p {font-size: 20px; font-weight: 600;}
+            /* 
+            button {
+                height: auto;
+                padding-left: 8px !important;
+                padding-right: 6px !important;
+            }
+            */
+        </style>""",
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/run_test/autoRefresh.py

@@ -1,5 +1,5 @@
 from streamlit_autorefresh import st_autorefresh
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 
 
 def autoRefresh():

vectordb_bench/frontend/components/run_test/caseSelector.py

@@ -1,9 +1,13 @@
-from vectordb_bench.frontend.const.styles import *
+
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.backend.cases import CaseType
-from vectordb_bench.frontend.const.dbCaseConfigs import *
+from vectordb_bench.frontend.config.dbCaseConfigs import *
+from collections import defaultdict
+
+from vectordb_bench.frontend.utils import addHorizontalLine
 
 
-def caseSelector(st, activedDbList):
+def caseSelector(st, activedDbList: list[DB]):
     st.markdown(
         "<div style='height: 24px;'></div>",
         unsafe_allow_html=True,
@@ -14,41 +18,49 @@ def caseSelector(st, activedDbList):
         unsafe_allow_html=True,
     )
 
-    caseIsActived = {case: False for case in CASE_LIST}
-    allCaseConfigs = {db: {case: {} for case in CASE_LIST} for db in DB_LIST}
-    for caseOrDivider in CASE_LIST_WITH_DIVIDER:
-        if caseOrDivider == DIVIDER:
-            caseItemContainer.markdown(
-                "<div style='border: 1px solid #cccccc60; margin-bottom: 24px;'></div>",
-                unsafe_allow_html=True,
-            )
+    activedCaseList: list[CaseConfig] = []
+    dbToCaseClusterConfigs = defaultdict(lambda: defaultdict(dict))
+    dbToCaseConfigs = defaultdict(lambda: defaultdict(dict))
+    caseClusters = UI_CASE_CLUSTERS + [get_custom_case_cluter()]
+    for caseCluster in caseClusters:
+        activedCaseList += caseClusterExpander(
+            st, caseCluster, dbToCaseClusterConfigs, activedDbList)
+    for db in dbToCaseClusterConfigs:
+        for uiCaseItem in dbToCaseClusterConfigs[db]:
+            for case in uiCaseItem.cases:
+                dbToCaseConfigs[db][case] = dbToCaseClusterConfigs[db][uiCaseItem]
+
+    return activedCaseList, dbToCaseConfigs
+
+
+def caseClusterExpander(st, caseCluster: UICaseItemCluster, dbToCaseClusterConfigs, activedDbList: list[DB]):
+    expander = st.expander(caseCluster.label, False)
+    activedCases: list[CaseConfig] = []
+    for uiCaseItem in caseCluster.uiCaseItems:
+        if uiCaseItem.isLine:
+            addHorizontalLine(expander)
         else:
-            case = caseOrDivider
-            caseItemContainer = st.container()
-            caseIsActived[case] = caseItem(
-                caseItemContainer, allCaseConfigs, case, activedDbList
-            )
-    activedCaseList = [case for case in CASE_LIST if caseIsActived[case]]
-    return activedCaseList, allCaseConfigs
+            activedCases += caseItemCheckbox(expander,
+                                             dbToCaseClusterConfigs, uiCaseItem, activedDbList)
+    return activedCases
 
 
-def caseItem(st, allCaseConfigs, case: CaseType, activedDbList):
-    selected = st.checkbox(case.case_name)
+def caseItemCheckbox(st, dbToCaseClusterConfigs, uiCaseItem: UICaseItem, activedDbList: list[DB]):
+    selected = st.checkbox(uiCaseItem.label)
     st.markdown(
-        f"<div style='color: #1D2939; margin: -8px 0 20px {CHECKBOX_INDENT}px; font-size: 14px;'>{case.case_description}</div>",
+        f"<div style='color: #1D2939; margin: -8px 0 20px {CHECKBOX_INDENT}px; font-size: 14px;'>{uiCaseItem.description}</div>",
         unsafe_allow_html=True,
     )
 
     if selected:
-        caseConfigSettingContainer = st.container()
         caseConfigSetting(
-            caseConfigSettingContainer, allCaseConfigs, case, activedDbList
+            st.container(), dbToCaseClusterConfigs, uiCaseItem, activedDbList
         )
 
-    return selected
+    return uiCaseItem.cases if selected else []
 
 
-def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
+def caseConfigSetting(st, dbToCaseClusterConfigs, uiCaseItem: UICaseItem, activedDbList: list[DB]):
     for db in activedDbList:
         columns = st.columns(1 + CASE_CONFIG_SETTING_COLUMNS)
         # column 0 - title
@@ -57,12 +69,12 @@ def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
             f"<div style='margin: 0 0 24px {CHECKBOX_INDENT}px; font-size: 18px; font-weight: 600;'>{db.name}</div>",
             unsafe_allow_html=True,
         )
-        caseConfig = allCaseConfigs[db][case]
         k = 0
-        for config in CASE_CONFIG_MAP.get(db, {}).get(case.case_cls().label, []):
+        caseConfig = dbToCaseClusterConfigs[db][uiCaseItem]
+        for config in CASE_CONFIG_MAP.get(db, {}).get(uiCaseItem.caseLabel, []):
             if config.isDisplayed(caseConfig):
                 column = columns[1 + k % CASE_CONFIG_SETTING_COLUMNS]
-                key = "%s-%s-%s" % (db, case, config.label.value)
+                key = "%s-%s-%s" % (db, uiCaseItem.label, config.label.value)
                 if config.inputType == InputType.Text:
                     caseConfig[config.label] = column.text_input(
                         config.displayLabel if config.displayLabel else config.label.value,

vectordb_bench/frontend/components/run_test/dbConfigSetting.py

@@ -1,13 +1,9 @@
 from pydantic import ValidationError
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.frontend.utils import inputIsPassword
 
 
 def dbConfigSettings(st, activedDbList):
-    st.markdown(
-        "<style> .streamlit-expanderHeader p {font-size: 20px; font-weight: 600;}</style>",
-        unsafe_allow_html=True,
-    )
     expander = st.expander("Configurations for the selected databases", True)
 
     dbConfigs = {}

vectordb_bench/frontend/components/run_test/dbSelector.py

@@ -1,5 +1,6 @@
-from vectordb_bench.frontend.const.styles import *
-from vectordb_bench.frontend.const.dbCaseConfigs import DB_LIST
+from streamlit.runtime.media_file_storage import MediaFileStorageError
+from vectordb_bench.frontend.config.styles import DB_SELECTOR_COLUMNS, DB_TO_ICON
+from vectordb_bench.frontend.config.dbCaseConfigs import DB_LIST
 
 
 def dbSelector(st):
@@ -16,21 +17,14 @@ def dbSelector(st):
     dbContainerColumns = st.columns(DB_SELECTOR_COLUMNS, gap="small")
     dbIsActived = {db: False for db in DB_LIST}
 
-    # style - image; column gap; checkbox font;
-    st.markdown(
-        """
-        <style>
-            div[data-testid='stImage'] {margin: auto;}
-            div[data-testid='stHorizontalBlock'] {gap: 8px;}
-            .stCheckbox p { color: #000; font-size: 18px; font-weight: 600; }
-        </style>
-        """,
-        unsafe_allow_html=True,
-    )
     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/generateTasks.py

@@ -1,17 +1,15 @@
+from vectordb_bench.backend.clients import DB
 from vectordb_bench.models import CaseConfig, CaseConfigParamType, TaskConfig
 
 
-def generate_tasks(activedDbList, dbConfigs, activedCaseList, allCaseConfigs):
+def generate_tasks(activedDbList: list[DB], dbConfigs, activedCaseList: list[CaseConfig], allCaseConfigs):
     tasks = []
     for db in activedDbList:
         for case in activedCaseList:
             task = TaskConfig(
                     db=db.value,
                     db_config=dbConfigs[db],
-                    case_config=CaseConfig(
-                        case_id=case.value,
-                        custom_case={},
-                    ),
+                    case_config=case,
                     db_case_config=db.case_config_cls(
                         allCaseConfigs[db][case].get(CaseConfigParamType.IndexType, None)
                     )(**{key.value: value for key, value in allCaseConfigs[db][case].items()}),

vectordb_bench/frontend/components/run_test/initStyle.py

@@ -0,0 +1,14 @@
+def initStyle(st):
+    st.markdown(
+        """<style>
+            /* expander - header */
+            .main div[data-testid='stExpander'] p {font-size: 18px; font-weight: 600;}
+            /* db icon */
+            div[data-testid='stImage'] {margin: auto;}
+            /* db column gap */
+            div[data-testid='stHorizontalBlock'] {gap: 8px;}
+            /* check box */
+            .stCheckbox p { color: #000; font-size: 18px; font-weight: 600; }
+        </style>""",
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/run_test/submitTask.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.interface import benchMarkRunner
 
 
@@ -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