vectordb-bench 0.0.1__1-py3-none-any.whl

vectordb_bench/frontend/components/check_results/charts.py

@@ -0,0 +1,175 @@
+from vectordb_bench.metric import metricOrder, isLowerIsBetterMetric, metricUnitMap
+from vectordb_bench.frontend.const import *
+from vectordb_bench.models import ResultLabel
+import plotly.express as px
+
+
+def drawCharts(st, allData, failedTasks, cases):
+    st.markdown(
+        "<style> .main .streamlit-expanderHeader p {font-size: 20px; font-weight: 600;} </style>",
+        unsafe_allow_html=True,
+    )
+    st.markdown(
+        """<style>
+            .main div[data-testid='stExpander'] {
+                background-color: #F6F8FA;
+                border: 1px solid #A9BDD140;
+                border-radius: 8px;
+            }
+        </style>""",
+        unsafe_allow_html=True,
+    )
+    for case in cases:
+        chartContainer = st.expander(case, True)
+        data = [data for data in allData if data["case"] == case]
+        drawChart(data, chartContainer)
+
+        errorDBs = failedTasks[case]
+        showFailedDBs(chartContainer, errorDBs)
+
+
+def showFailedDBs(st, errorDBs):
+    failedDBs = [db for db, label in errorDBs.items() if label == ResultLabel.FAILED]
+    timeoutDBs = [
+        db for db, label in errorDBs.items() if label == ResultLabel.OUTOFRANGE
+    ]
+
+    showFailedText(st, "Failed", failedDBs)
+    showFailedText(st, "Timeout", timeoutDBs)
+
+
+def showFailedText(st, text, dbs):
+    if len(dbs) > 0:
+        st.markdown(
+            f"<div style='margin: -16px 0 12px 8px; font-size: 16px; font-weight: 600;'>{text}: &nbsp;&nbsp;{', '.join(dbs)}</div>",
+            unsafe_allow_html=True,
+        )
+
+
+def drawChart(data, st):
+    metricsSet = set()
+    for d in data:
+        metricsSet = metricsSet.union(d["metricsSet"])
+    showMetrics = [metric for metric in metricOrder if metric in metricsSet]
+
+    for i, metric in enumerate(showMetrics):
+        container = st.container()
+        drawMetricChart(data, metric, container)
+
+
+def getLabelToShapeMap(data):
+    labelIndexMap = {}
+
+    dbSet = {d["db"] for d in data}
+    for db in dbSet:
+        labelSet = {d["db_label"] for d in data if d["db"] == db}
+        labelList = list(labelSet)
+
+        usedShapes = set()
+        i = 0
+        for label in labelList:
+            if label not in labelIndexMap:
+                loopCount = 0
+                while i % len(PATTERN_SHAPES) in usedShapes:
+                    i += 1
+                    loopCount += 1
+                    if loopCount > len(PATTERN_SHAPES):
+                        break
+                labelIndexMap[label] = i
+                i += 1
+            else:
+                usedShapes.add(labelIndexMap[label] % len(PATTERN_SHAPES))
+
+    labelToShapeMap = {
+        label: getPatternShape(index) for label, index in labelIndexMap.items()
+    }
+    return labelToShapeMap
+
+
+def drawMetricChart(data, metric, st):
+    dataWithMetric = [d for d in data if d.get(metric, 0) > 1e-7]
+    # dataWithMetric = data
+    if len(dataWithMetric) == 0:
+        return
+
+    # title = st.container()
+    # title.markdown(
+    #     f"**{metric}** ({'less' if isLowerIsBetterMetric(metric) else 'more'} is better)"
+    # )
+    chart = st.container()
+
+    height = len(dataWithMetric) * 24 + 48
+    xmin = 0
+    xmax = max([d.get(metric, 0) for d in dataWithMetric])
+    xpadding = (xmax - xmin) / 16
+    xpadding_multiplier = 1.6
+    xrange = [xmin, xmax + xpadding * xpadding_multiplier]
+    unit = metricUnitMap.get(metric, "")
+    labelToShapeMap = getLabelToShapeMap(dataWithMetric)
+    categoryorder = (
+        "total descending" if isLowerIsBetterMetric(metric) else "total ascending"
+    )
+    fig = px.bar(
+        dataWithMetric,
+        x=metric,
+        y="db_name",
+        color="db",
+        height=height,
+        # pattern_shape="db_label",
+        # pattern_shape_sequence=SHAPES,
+        pattern_shape_map=labelToShapeMap,
+        orientation="h",
+        hover_data={
+            "db": False,
+            "db_label": False,
+            "db_name": True,
+        },
+        color_discrete_map=COLOR_MAP,
+        text_auto=True,
+        title=f"{metric.capitalize()} ({'less' if isLowerIsBetterMetric(metric) else 'more'} is better)",
+    )
+    fig.update_xaxes(showticklabels=False, visible=False, range=xrange)
+    fig.update_yaxes(
+        # showticklabels=False,
+        # visible=False,
+        title=dict(
+            font=dict(
+                size=1,
+            ),
+            # text="",
+        )
+    )
+    fig.update_traces(
+        textposition="outside",
+        textfont=dict(
+            color="#333",
+            size=12,
+        ),
+        marker=dict(
+            pattern=dict(fillmode="overlay", fgcolor="#fff", fgopacity=1, size=7)
+        ),
+        texttemplate="%{x:,.4~r}" + unit,
+    )
+    fig.update_layout(
+        margin=dict(l=0, r=0, t=48, b=12, pad=8),
+        bargap=0.25,
+        showlegend=False,
+        legend=dict(
+            orientation="h", yanchor="bottom", y=1, xanchor="right", x=1, title=""
+        ),
+        # legend=dict(orientation="v", title=""),
+        yaxis={"categoryorder": categoryorder},
+        title=dict(
+            font=dict(
+                size=16,
+                color="#666",
+                # family="Arial, sans-serif",
+            ),
+            pad=dict(l=16),
+            # y=0.95,
+            # yanchor="top",
+            # yref="container",
+        ),
+    )
+
+    chart.plotly_chart(fig, use_container_width=True)

vectordb_bench/frontend/components/check_results/data.py

@@ -0,0 +1,86 @@
+from collections import defaultdict
+from dataclasses import asdict
+from vectordb_bench.metric import isLowerIsBetterMetric
+from vectordb_bench.models import ResultLabel
+
+
+def getChartData(tasks, dbNames, cases):
+    filterTasks = getFilterTasks(tasks, dbNames, cases)
+    mergedTasks, failedTasks = mergeTasks(filterTasks)
+    return mergedTasks, failedTasks
+
+
+def getFilterTasks(tasks, dbNames, cases):
+    filterTasks = [
+        task
+        for task in tasks
+        if task.task_config.db_name in dbNames
+        and task.task_config.case_config.case_id.value in cases
+    ]
+    return filterTasks
+
+
+def mergeTasks(tasks):
+    dbCaseMetricsMap = defaultdict(lambda: defaultdict(dict))
+    for task in tasks:
+        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.value
+        dbCaseMetricsMap[db_name][case] = {
+            "db": db,
+            "db_label": db_label,
+            "metrics": mergeMetrics(
+                dbCaseMetricsMap[db_name][case].get("metrics", {}), asdict(task.metrics)
+            ),
+            "label": getBetterLabel(dbCaseMetricsMap[db_name][case].get("label", ResultLabel.FAILED), task.label)
+        }
+
+    mergedTasks = []
+    failedTasks = defaultdict(lambda: defaultdict(str))
+    for db_name, caseMetricsMap in dbCaseMetricsMap.items():
+        for case, metricInfo in caseMetricsMap.items():
+            metrics = metricInfo["metrics"]
+            db = metricInfo["db"]
+            db_label = metricInfo["db_label"]
+            label = metricInfo["label"]
+            if label == ResultLabel.NORMAL:
+                mergedTasks.append(
+                    {
+                        "db_name": db_name,
+                        "db": db,
+                        "db_label": db_label,
+                        "case": case,
+                        "metricsSet": set(metrics.keys()),
+                        **metrics,
+                    }
+                )
+            else: 
+                failedTasks[case][db_name] = label
+                
+    return mergedTasks, failedTasks
+
+
+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
+        )
+
+    return metrics
+
+
+def getBetterMetric(metric, value_1, value_2):
+    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)
+    )
+    
+def getBetterLabel(label_1: ResultLabel, label_2: ResultLabel):
+    return label_2 if label_1 != ResultLabel.NORMAL else label_1

vectordb_bench/frontend/components/check_results/filters.py

@@ -0,0 +1,97 @@
+from vectordb_bench.frontend.components.check_results.data import getChartData
+from vectordb_bench.frontend.const import *
+
+
+def getshownData(results, st):
+    # hide the nav
+    st.markdown(
+        "<style> div[data-testid='stSidebarNav'] {display: none;} </style>",
+        unsafe_allow_html=True,
+    )
+
+    st.header("Filters")
+
+    shownResults = getshownResults(results, st)
+    showDBNames, showCases = getShowDbsAndCases(shownResults, st)
+
+    shownData, failedTasks = getChartData(shownResults, showDBNames, showCases)
+
+    return shownData, failedTasks, showCases
+
+
+def getshownResults(results, st):
+    resultSelectOptions = [
+        result.task_label
+        if result.task_label != result.run_id
+        else f"res-{result.run_id[:4]}"
+        for result in results
+    ]
+    if len(resultSelectOptions) == 0:
+        st.write(
+            "There are no results to display. Please wait for the task to complete or run a new task."
+        )
+        return []
+
+    selectedResultSelectedOptions = st.multiselect(
+        "Select the task results you need to analyze.",
+        resultSelectOptions,
+        # label_visibility="hidden",
+        default=resultSelectOptions,
+    )
+    selectedResult = []
+    for option in selectedResultSelectedOptions:
+        result = results[resultSelectOptions.index(option)].results
+        selectedResult += result
+
+    return selectedResult
+
+
+def getShowDbsAndCases(result, st):
+    # expanderStyles
+    st.markdown("<style> section[data-testid='stSidebar'] div[data-testid='stExpander'] div[data-testid='stVerticalBlock'] { gap: 0.2rem; }  </style>", unsafe_allow_html=True,)
+    st.markdown(
+        "<style> div[data-testid='stExpander'] {background-color: #ffffff;} </style>",
+        unsafe_allow_html=True,
+    )
+    st.markdown(
+        "<style> section[data-testid='stSidebar'] .streamlit-expanderHeader p {font-size: 16px; font-weight: 600;} </style>",
+        unsafe_allow_html=True,
+    )
+
+    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 = [case["name"].value for case in CASE_LIST if case["name"] in allCasesSet]
+
+    # dbFilterContainer = st.container()
+    # dbFilterContainer.subheader("DB Filter")
+    dbFilterContainer = st.expander("DB Filter", True)
+    showDBNames = filterView(allDbNames, dbFilterContainer, col=1)
+
+    # caseFilterContainer = st.container()
+    # caseFilterContainer.subheader("Case Filter")
+    caseFilterContainer = st.expander("Case Filter", True)
+    showCases = filterView(
+        allCases,
+        caseFilterContainer,
+        col=1,
+        optionLables=[case for case in allCases],
+    )
+
+    return showDBNames, showCases
+
+
+def filterView(options, st, col, optionLables=None):
+    columns = st.columns(
+        col,
+        gap="small",
+    )
+    isActive = {option: True for option in options}
+    if optionLables is None:
+        optionLables = options
+    for i, option in enumerate(options):
+        isActive[option] = columns[i % col].checkbox(
+            optionLables[i], value=isActive[option]
+        )
+
+    return [option for option in options if isActive[option]]

vectordb_bench/frontend/components/check_results/headerIcon.py

@@ -0,0 +1,18 @@
+def drawHeaderIcon(st):
+    st.markdown("""
+<div class="headerIconContainer"></div>
+
+<style>
+.headerIconContainer {
+    position: absolute;
+    top: -50px;
+    height: 50px;
+    width: 100%;
+    border-bottom: 2px solid #E8EAEE;
+    background-image: url(https://assets.zilliz.com/vdb_benchmark_db790b5387.png);
+    background-repeat: no-repeat;
+}
+</style
+""",
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/check_results/nav.py

@@ -0,0 +1,21 @@
+from streamlit_extras.switch_page_button import switch_page
+
+
+def NavToRunTest(st):
+    st.header("Run your test")
+    st.write("You can set the configs and run your own test.")
+    navClick = st.button("Run Your Test   >")
+    if navClick:
+        switch_page("run test")
+        
+        
+def NavToQPSWithPrice(st):
+    navClick = st.button("QPS with Price   >")
+    if navClick:
+        switch_page("qps with price")
+        
+        
+def NavToResults(st):
+    navClick = st.button("< &nbsp;&nbsp;Back to Results")
+    if navClick:
+        switch_page("vdb benchmark")

vectordb_bench/frontend/components/check_results/priceTable.py

@@ -0,0 +1,48 @@
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.frontend.const import DB_DBLABEL_TO_PRICE
+import pandas as pd
+from collections import defaultdict
+import streamlit as st
+
+
+def priceTable(container, data):
+    dbAndLabelSet = {
+        (d["db"], d["db_label"]) for d in data if d["db"] != DB.Milvus.value
+    }
+
+    dbAndLabelList = list(dbAndLabelSet)
+    dbAndLabelList.sort()
+
+    table = pd.DataFrame(
+        [
+            {
+                "DB": db,
+                "Label": db_label,
+                "Price per hour": DB_DBLABEL_TO_PRICE.get(db, {}).get(db_label, 0),
+            }
+            for db, db_label in dbAndLabelList
+        ]
+    )
+    height = len(table) * 35 + 38
+
+    expander = container.expander("You can edit the price.")
+    editTable = expander.data_editor(
+        table,
+        use_container_width=True,
+        hide_index=True,
+        height=height,
+        disabled=("DB", "Label"),
+        column_config={
+            "Price per hour": st.column_config.NumberColumn(
+                min_value=0,
+                format="$ %f",
+            )
+        },
+    )
+
+    priceMap = defaultdict(lambda: defaultdict(float))
+    for _, row in editTable.iterrows():
+        db, db_label, price = row
+        priceMap[db][db_label] = price
+
+    return priceMap

vectordb_bench/frontend/components/run_test/autoRefresh.py

@@ -0,0 +1,10 @@
+from streamlit_autorefresh import st_autorefresh
+from vectordb_bench.frontend.const import *
+
+
+def autoRefresh():
+    auto_refresh_count = st_autorefresh(
+        interval=MAX_AUTO_REFRESH_INTERVAL,
+        limit=MAX_AUTO_REFRESH_COUNT,
+        key="streamlit-auto-refresh",
+    )

vectordb_bench/frontend/components/run_test/caseSelector.py

@@ -0,0 +1,87 @@
+from vectordb_bench.frontend.const import *
+
+
+def caseSelector(st, activedDbList):
+    st.markdown(
+        "<div style='height: 24px;'></div>",
+        unsafe_allow_html=True,
+    )
+    st.subheader("STEP 2: Choose the case(s)")
+    st.markdown(
+        "<div style='color: #647489; margin-bottom: 24px; margin-top: -12px;'>Choose at least one case you want to run the test for. </div>",
+        unsafe_allow_html=True,
+    )
+
+    caseIsActived = {case["name"]: False for case in CASE_LIST}
+    allCaseConfigs = {db: {case["name"]: {} for case in CASE_LIST} for db in DB_LIST}
+    for case in CASE_LIST:
+        caseItemContainer = st.container()
+        caseIsActived[case["name"]] = caseItem(
+            caseItemContainer, allCaseConfigs, case, activedDbList
+        )
+        if case.get("divider"):
+            caseItemContainer.markdown(
+                "<div style='border: 1px solid #cccccc60; margin-bottom: 24px;'></div>",
+                unsafe_allow_html=True,
+            )
+    activedCaseList = [
+        case["name"] for case in CASE_LIST if caseIsActived[case["name"]]
+    ]
+    return activedCaseList, allCaseConfigs
+
+def caseItem(st, allCaseConfigs, case, activedDbList):
+    selected = st.checkbox(case["name"].value)
+    st.markdown(
+        f"<div style='color: #1D2939; margin: -8px 0 20px {CHECKBOX_INDENT}px; font-size: 14px;'>{case['intro']}</div>",
+        unsafe_allow_html=True,
+    )
+
+    if selected:
+        caseConfigSettingContainer = st.container()
+        caseConfigSetting(
+            caseConfigSettingContainer, allCaseConfigs, case["name"], activedDbList
+        )
+
+    return selected
+
+
+def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
+    for db in activedDbList:
+        columns = st.columns(1 + CASE_CONFIG_SETTING_COLUMNS)
+        # column 0 - title
+        dbColumn = columns[0]
+        dbColumn.markdown(
+            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, []):
+            if config.isDisplayed(caseConfig):
+                column = columns[1 + k % CASE_CONFIG_SETTING_COLUMNS]
+                key = "%s-%s-%s" % (db, case, config.label.value)
+                if config.inputType == InputType.Text:
+                    caseConfig[config.label] = column.text_input(
+                        config.label.value,
+                        key=key,
+                        value=config.inputConfig["value"],
+                    )
+                elif config.inputType == InputType.Option:
+                    caseConfig[config.label] = column.selectbox(
+                        config.label.value,
+                        config.inputConfig["options"],
+                        key=key,
+                    )
+                elif config.inputType == InputType.Number:
+                    caseConfig[config.label] = column.number_input(
+                        config.label.value,
+                        format="%d",
+                        step=1,
+                        min_value=config.inputConfig["min"],
+                        max_value=config.inputConfig["max"],
+                        key=key,
+                        value=config.inputConfig["value"],
+                    )
+                k += 1
+        if k == 0:
+            columns[1].write("Auto")

vectordb_bench/frontend/components/run_test/dbConfigSetting.py

@@ -0,0 +1,47 @@
+from vectordb_bench.frontend.const 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 = {}
+    for activeDb in activedDbList:
+        dbConfigSettingItemContainer = expander.container()
+        dbConfig = dbConfigSettingItem(dbConfigSettingItemContainer, activeDb)
+        dbConfigs[activeDb] = dbConfig
+        
+    return dbConfigs
+
+def dbConfigSettingItem(st, activeDb):
+    st.markdown(
+        f"<div style='font-weight: 600; font-size: 20px; margin-top: 16px;'>{activeDb.value}</div>",
+        unsafe_allow_html=True,
+    )
+    columns = st.columns(DB_CONFIG_SETTING_COLUMNS)
+
+    activeDbCls = activeDb.init_cls
+    dbConfigClass = activeDbCls.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
+        dbConfig[key] = column.text_input(
+            key,
+            key="%s-%s" % (activeDb, key),
+            value=value.get("default", ""),
+            type="password" if inputIsPassword(key) else "default",
+        )
+    return dbConfigClass(**dbConfig)
+
+
+def moveDBLabelToLast(propertiesItems):
+    propertiesItems.sort(key=lambda x: 1 if x[0] == 'db_label' else 0)
+    

vectordb_bench/frontend/components/run_test/dbSelector.py

@@ -0,0 +1,36 @@
+
+from vectordb_bench.frontend.const import *
+
+
+def dbSelector(st):
+    st.markdown(
+        "<div style='height: 12px;'></div>",
+        unsafe_allow_html=True,
+    )
+    st.subheader("STEP 1: Select the database(s)")
+    st.markdown(
+        "<div style='color: #647489; margin-bottom: 24px; margin-top: -12px;'>Choose at least one case you want to run the test for. </div>",
+        unsafe_allow_html=True,
+    )
+    
+    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, ""))
+    activedDbList = [db for db in DB_LIST if dbIsActived[db]]
+    
+    return activedDbList

vectordb_bench/frontend/components/run_test/generateTasks.py

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

vectordb_bench/frontend/components/run_test/hideSidebar.py

@@ -0,0 +1,10 @@
+def hideSidebar(st):
+    st.markdown(
+        "<style> div[data-testid='collapsedControl'] {display: none;} </style>",
+        unsafe_allow_html=True,
+    )
+
+    st.markdown(
+        "<style> .block-container { max-width: 1000px; } </style>",
+        unsafe_allow_html=True,
+    )