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

vectordb_bench/frontend/{const → config}/dbCaseConfigs.py

@@ -1,42 +1,147 @@
-from enum import IntEnum
+from enum import IntEnum, Enum
 import typing
 from pydantic import BaseModel
 from vectordb_bench.backend.cases import CaseLabel, CaseType
 from vectordb_bench.backend.clients import DB
 from vectordb_bench.backend.clients.api import IndexType
+from vectordb_bench.frontend.components.custom.getCustomConfig import get_custom_configs
 
-from vectordb_bench.models import CaseConfigParamType
+from vectordb_bench.models import CaseConfig, 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]
+
+
+class Delimiter(Enum):
+    Line = "line"
+
+
+class BatchCaseConfig(BaseModel):
+    label: str = ""
+    description: str = ""
+    cases: list[CaseConfig] = []
+
+
+class UICaseItem(BaseModel):
+    isLine: bool = False
+    label: str = ""
+    description: str = ""
+    cases: list[CaseConfig] = []
+    caseLabel: CaseLabel = CaseLabel.Performance
+
+    def __init__(
+        self,
+        isLine: bool = False,
+        case_id: CaseType = None,
+        custom_case: dict = {},
+        cases: list[CaseConfig] = [],
+        label: str = "",
+        description: str = "",
+        caseLabel: CaseLabel = CaseLabel.Performance,
+    ):
+        if isLine is True:
+            super().__init__(isLine=True)
+        elif case_id is not None and isinstance(case_id, CaseType):
+            c = case_id.case_cls(custom_case)
+            super().__init__(
+                label=c.name,
+                description=c.description,
+                cases=[CaseConfig(case_id=case_id, custom_case=custom_case)],
+                caseLabel=c.label,
+            )
+        else:
+            super().__init__(
+                label=label,
+                description=description,
+                cases=cases,
+                caseLabel=caseLabel,
+            )
+
+    def __hash__(self) -> int:
+        return hash(self.json())
+
+
+class UICaseItemCluster(BaseModel):
+    label: str = ""
+    uiCaseItems: list[UICaseItem] = []
+
+
+def get_custom_case_items() -> list[UICaseItem]:
+    custom_configs = get_custom_configs()
+    return [
+        UICaseItem(
+            case_id=CaseType.PerformanceCustomDataset, custom_case=custom_config.dict()
+        )
+        for custom_config in custom_configs
+    ]
+
+
+def get_custom_case_cluter() -> UICaseItemCluster:
+    return UICaseItemCluster(
+        label="Custom Search Performance Test", uiCaseItems=get_custom_case_items()
+    )
+
+
+UI_CASE_CLUSTERS: list[UICaseItemCluster] = [
+    UICaseItemCluster(
+        label="Search Performance Test",
+        uiCaseItems=[
+            UICaseItem(case_id=CaseType.Performance768D100M),
+            UICaseItem(case_id=CaseType.Performance768D10M),
+            UICaseItem(case_id=CaseType.Performance768D1M),
+            UICaseItem(isLine=True),
+            UICaseItem(case_id=CaseType.Performance1536D5M),
+            UICaseItem(case_id=CaseType.Performance1536D500K),
+            UICaseItem(case_id=CaseType.Performance1536D50K),
+        ],
+    ),
+    UICaseItemCluster(
+        label="Filter Search Performance Test",
+        uiCaseItems=[
+            UICaseItem(case_id=CaseType.Performance768D10M1P),
+            UICaseItem(case_id=CaseType.Performance768D10M99P),
+            UICaseItem(case_id=CaseType.Performance768D1M1P),
+            UICaseItem(case_id=CaseType.Performance768D1M99P),
+            UICaseItem(isLine=True),
+            UICaseItem(case_id=CaseType.Performance1536D5M1P),
+            UICaseItem(case_id=CaseType.Performance1536D5M99P),
+            UICaseItem(case_id=CaseType.Performance1536D500K1P),
+            UICaseItem(case_id=CaseType.Performance1536D500K99P),
+        ],
+    ),
+    UICaseItemCluster(
+        label="Capacity Test",
+        uiCaseItems=[
+            UICaseItem(case_id=CaseType.CapacityDim960),
+            UICaseItem(case_id=CaseType.CapacityDim128),
+        ],
+    ),
+]
 
-DIVIDER = "DIVIDER"
-CASE_LIST_WITH_DIVIDER = [
+# DIVIDER = "DIVIDER"
+DISPLAY_CASE_ORDER: list[CaseType] = [
     CaseType.Performance768D100M,
     CaseType.Performance768D10M,
     CaseType.Performance768D1M,
-    DIVIDER,
     CaseType.Performance1536D5M,
     CaseType.Performance1536D500K,
-    DIVIDER,
+    CaseType.Performance1536D50K,
     CaseType.Performance768D10M1P,
     CaseType.Performance768D1M1P,
-    DIVIDER,
     CaseType.Performance1536D5M1P,
     CaseType.Performance1536D500K1P,
-    DIVIDER,
     CaseType.Performance768D10M99P,
     CaseType.Performance768D1M99P,
-    DIVIDER,
     CaseType.Performance1536D5M99P,
     CaseType.Performance1536D500K99P,
-    DIVIDER,
     CaseType.CapacityDim960,
     CaseType.CapacityDim128,
 ]
+CASE_NAME_ORDER = [case.case_cls().name for case in DISPLAY_CASE_ORDER]
 
-CASE_LIST = [item for item in CASE_LIST_WITH_DIVIDER if isinstance(item, CaseType)]
+# CASE_LIST = [
+#     item for item in CASE_LIST_WITH_DIVIDER if isinstance(item, CaseType)]
 
 
 class InputType(IntEnum):
@@ -52,7 +157,7 @@ class CaseConfigInput(BaseModel):
     inputHelp: str = ""
     displayLabel: str = ""
     # todo type should be a function
-    isDisplayed: typing.Any = lambda x: True
+    isDisplayed: typing.Any = lambda config: True
 
 
 CaseConfigParamInput_IndexType = CaseConfigInput(
@@ -145,7 +250,7 @@ CaseConfigParamInput_EFConstruction_ES = CaseConfigInput(
 CaseConfigParamInput_maintenance_work_mem_PgVector = CaseConfigInput(
     label=CaseConfigParamType.maintenance_work_mem,
     inputHelp="Recommended value: 1.33x the index size, not to exceed the available free memory."
-              "Specify in gigabytes. e.g. 8GB",
+    "Specify in gigabytes. e.g. 8GB",
     inputType=InputType.Text,
     inputConfig={
         "value": "8GB",
@@ -156,7 +261,7 @@ CaseConfigParamInput_max_parallel_workers_PgVector = CaseConfigInput(
     label=CaseConfigParamType.max_parallel_workers,
     displayLabel="Max parallel workers",
     inputHelp="Recommended value: (cpu cores - 1). This will set the parameters: max_parallel_maintenance_workers,"
-              " max_parallel_workers & table(parallel_workers)",
+    " max_parallel_workers & table(parallel_workers)",
     inputType=InputType.Number,
     inputConfig={
         "min": 0,
@@ -513,7 +618,8 @@ CaseConfigParamInput_QuantizationRatio_PgVectoRS = CaseConfigInput(
         "options": ["x4", "x8", "x16", "x32", "x64"],
     },
     isDisplayed=lambda config: config.get(CaseConfigParamType.quantizationType, None)
-    == "product" and config.get(CaseConfigParamType.IndexType, None)
+    == "product"
+    and config.get(CaseConfigParamType.IndexType, None)
     in [
         IndexType.HNSW.value,
         IndexType.IVFFlat.value,
@@ -581,22 +687,24 @@ ESPerformanceConfig = [
     CaseConfigParamInput_NumCandidates_ES,
 ]
 
-PgVectorLoadingConfig = [CaseConfigParamInput_IndexType_PgVector,
-                         CaseConfigParamInput_Lists_PgVector,
-                         CaseConfigParamInput_m,
-                         CaseConfigParamInput_EFConstruction_PgVector,
-                         CaseConfigParamInput_maintenance_work_mem_PgVector,
-                         CaseConfigParamInput_max_parallel_workers_PgVector,
-                         ]
-PgVectorPerformanceConfig = [CaseConfigParamInput_IndexType_PgVector,
-                             CaseConfigParamInput_m,
-                             CaseConfigParamInput_EFConstruction_PgVector,
-                             CaseConfigParamInput_EFSearch_PgVector,
-                             CaseConfigParamInput_Lists_PgVector,
-                             CaseConfigParamInput_Probes_PgVector,
-                             CaseConfigParamInput_maintenance_work_mem_PgVector,
-                             CaseConfigParamInput_max_parallel_workers_PgVector,
-                             ]
+PgVectorLoadingConfig = [
+    CaseConfigParamInput_IndexType_PgVector,
+    CaseConfigParamInput_Lists_PgVector,
+    CaseConfigParamInput_m,
+    CaseConfigParamInput_EFConstruction_PgVector,
+    CaseConfigParamInput_maintenance_work_mem_PgVector,
+    CaseConfigParamInput_max_parallel_workers_PgVector,
+]
+PgVectorPerformanceConfig = [
+    CaseConfigParamInput_IndexType_PgVector,
+    CaseConfigParamInput_m,
+    CaseConfigParamInput_EFConstruction_PgVector,
+    CaseConfigParamInput_EFSearch_PgVector,
+    CaseConfigParamInput_Lists_PgVector,
+    CaseConfigParamInput_Probes_PgVector,
+    CaseConfigParamInput_maintenance_work_mem_PgVector,
+    CaseConfigParamInput_max_parallel_workers_PgVector,
+]
 
 PgVectoRSLoadingConfig = [
     CaseConfigParamInput_IndexType,

vectordb_bench/frontend/{const → config}/styles.py

@@ -46,6 +46,7 @@ DB_TO_ICON = {
     DB.PgVectoRS: "https://assets.zilliz.com/PG_Vector_d464f2ef5f.png",
     DB.Redis: "https://assets.zilliz.com/Redis_Cloud_74b8bfef39.png",
     DB.Chroma: "https://assets.zilliz.com/chroma_ceb3f06ed7.png",
+    DB.AWSOpenSearch: "https://assets.zilliz.com/opensearch_1eee37584e.jpeg",
 }
 
 # RedisCloud color: #0D6EFD
@@ -59,4 +60,5 @@ COLOR_MAP = {
     DB.WeaviateCloud.value: "#20C997",
     DB.PgVector.value: "#4C779A",
     DB.Redis.value: "#0D6EFD",
+    DB.AWSOpenSearch.value: "#0DCAF0",
 }

vectordb_bench/frontend/pages/concurrent.py

@@ -0,0 +1,65 @@
+import streamlit as st
+from vectordb_bench.frontend.components.check_results.footer import footer
+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.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.config.styles import FAVICON
+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, _, showCaseNames = 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, showCaseNames, st.container())
+
+    # footer
+    footer(st.container())
+
+
+if __name__ == "__main__":
+    main()

vectordb_bench/frontend/pages/custom.py

@@ -0,0 +1,64 @@
+import streamlit as st
+from vectordb_bench.frontend.components.check_results.headerIcon import drawHeaderIcon
+from vectordb_bench.frontend.components.custom.displayCustomCase import displayCustomCase
+from vectordb_bench.frontend.components.custom.displaypPrams import displayParams
+from vectordb_bench.frontend.components.custom.getCustomConfig import CustomCaseConfig, generate_custom_case, get_custom_configs, save_custom_configs
+from vectordb_bench.frontend.components.custom.initStyle import initStyle
+from vectordb_bench.frontend.config.styles import FAVICON, PAGE_TITLE
+
+
+class CustomCaseManager():
+    customCaseItems: list[CustomCaseConfig]
+
+    def __init__(self):
+        self.customCaseItems = get_custom_configs()
+
+    def addCase(self):
+        new_custom_case = generate_custom_case()
+        new_custom_case.dataset_config.name = f"{new_custom_case.dataset_config.name} {len(self.customCaseItems)}"
+        self.customCaseItems += [new_custom_case]
+        self.save()
+
+    def deleteCase(self, idx: int):
+        self.customCaseItems.pop(idx)
+        self.save()
+
+    def save(self):
+        save_custom_configs(self.customCaseItems)
+
+
+def main():
+    st.set_page_config(
+        page_title=PAGE_TITLE,
+        page_icon=FAVICON,
+        # layout="wide",
+        # initial_sidebar_state="collapsed",
+    )
+
+    # header
+    drawHeaderIcon(st)
+
+    # init style
+    initStyle(st)
+
+    st.title("Custom Dataset")
+    displayParams(st)
+    customCaseManager = CustomCaseManager()
+
+    for idx, customCase in enumerate(customCaseManager.customCaseItems):
+        expander = st.expander(customCase.dataset_config.name, expanded=True)
+        key = f"custom_case_{idx}"
+        displayCustomCase(customCase, expander, key=key)
+
+        columns = expander.columns(8)
+        columns[0].button(
+            "Save", key=f"{key}_", type="secondary", on_click=lambda: customCaseManager.save())
+        columns[1].button(":red[Delete]", key=f"{key}_delete", type="secondary",
+                          on_click=lambda: customCaseManager.deleteCase(idx))
+
+    st.button("\+ New Dataset", key=f"add_custom_configs",
+              type="primary", on_click=lambda: customCaseManager.addCase())
+
+
+if __name__ == "__main__":
+    main()

vectordb_bench/frontend/pages/quries_per_dollar.py

@@ -8,7 +8,7 @@ from vectordb_bench.frontend.components.check_results.nav import NavToResults, N
 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.get_results.saveAsImage import getResults
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.interface import benchMarkRunner
 from vectordb_bench.metric import QURIES_PER_DOLLAR_METRIC
 
@@ -26,7 +26,7 @@ def main():
 
     # results selector
     resultSelectorContainer = st.sidebar.container()
-    shownData, _, showCases = getshownData(allResults, resultSelectorContainer)
+    shownData, _, showCaseNames = getshownData(allResults, resultSelectorContainer)
 
     resultSelectorContainer.divider()
 
@@ -45,8 +45,8 @@ def main():
     priceMap = priceTable(priceTableContainer, shownData)
 
     # charts
-    for case in showCases:
-        data = [data for data in shownData if data["case_name"] == case.name]
+    for caseName in showCaseNames:
+        data = [data for data in shownData if data["case_name"] == caseName]
         dataWithMetric = []
         metric = QURIES_PER_DOLLAR_METRIC
         for d in data:
@@ -56,7 +56,7 @@ def main():
                 d[metric] = d["qps"] / price * 3.6
                 dataWithMetric.append(d)
         if len(dataWithMetric) > 0:
-            chartContainer = st.expander(case.name, True)
+            chartContainer = st.expander(caseName, True)
             drawMetricChart(data, metric, chartContainer)
 
     # footer

vectordb_bench/frontend/pages/run_test.py

@@ -5,6 +5,7 @@ from vectordb_bench.frontend.components.run_test.dbConfigSetting import dbConfig
 from vectordb_bench.frontend.components.run_test.dbSelector import dbSelector
 from vectordb_bench.frontend.components.run_test.generateTasks import generate_tasks
 from vectordb_bench.frontend.components.run_test.hideSidebar import hideSidebar
+from vectordb_bench.frontend.components.run_test.initStyle import initStyle
 from vectordb_bench.frontend.components.run_test.submitTask import submitTask
 from vectordb_bench.frontend.components.check_results.nav import NavToResults
 from vectordb_bench.frontend.components.check_results.headerIcon import drawHeaderIcon
@@ -15,6 +16,9 @@ def main():
     # set page config
     initRunTestPageConfig(st)
     
+    # init style
+    initStyle(st)
+    
     # header
     drawHeaderIcon(st)
 

vectordb_bench/frontend/pages/tables.py

@@ -0,0 +1,24 @@
+import streamlit as st
+from vectordb_bench.frontend.components.check_results.headerIcon import drawHeaderIcon
+from vectordb_bench.frontend.components.tables.data import getNewResults
+from vectordb_bench.frontend.config.styles import FAVICON
+
+
+def main():
+    # set page config
+    st.set_page_config(
+        page_title="Table",
+        page_icon=FAVICON,
+        layout="wide",
+        # initial_sidebar_state="collapsed",
+    )
+
+    # header
+    drawHeaderIcon(st)
+
+    df = getNewResults()
+    st.dataframe(df, height=800)
+
+
+if __name__ == "__main__":
+    main()

vectordb_bench/frontend/utils.py

@@ -1,6 +1,22 @@
-from vectordb_bench.models import CaseType
+import random
+import string
+
 
 passwordKeys = ["password", "api_key"]
+
+
 def inputIsPassword(key: str) -> bool:
     return key.lower() in passwordKeys
 
+
+def addHorizontalLine(st):
+    st.markdown(
+        "<div style='border: 1px solid #cccccc60; margin-bottom: 24px;'></div>",
+        unsafe_allow_html=True,
+    )
+
+
+def generate_random_string(length):
+    letters = string.ascii_letters + string.digits
+    result = ''.join(random.choice(letters) for _ in range(length))
+    return result

vectordb_bench/frontend/vdb_benchmark.py

@@ -6,7 +6,7 @@ from vectordb_bench.frontend.components.check_results.nav import NavToQuriesPerD
 from vectordb_bench.frontend.components.check_results.charts import drawCharts
 from vectordb_bench.frontend.components.check_results.filters import getshownData
 from vectordb_bench.frontend.components.get_results.saveAsImage import getResults
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.interface import benchMarkRunner
 
 
@@ -24,7 +24,7 @@ def main():
 
     # results selector and filter
     resultSelectorContainer = st.sidebar.container()
-    shownData, failedTasks, showCases = getshownData(
+    shownData, failedTasks, showCaseNames = getshownData(
         allResults, resultSelectorContainer
     )
 
@@ -40,7 +40,7 @@ def main():
     getResults(resultesContainer, "vectordb_bench")
 
     # charts
-    drawCharts(st, shownData, failedTasks, showCases)
+    drawCharts(st, shownData, failedTasks, showCaseNames)
 
     # footer
     footer(st.container())

vectordb_bench/interface.py

@@ -1,38 +1,33 @@
-import traceback
+import concurrent.futures
+import logging
+import multiprocessing as mp
 import pathlib
 import signal
-import logging
+import traceback
 import uuid
-import concurrent
-import multiprocessing as mp
+from enum import Enum
 from multiprocessing.connection import Connection
 
 import psutil
-from enum import Enum
 
 from . import config
-from .metric import Metric
-from .models import (
-    TaskConfig,
-    TestResult,
-    CaseResult,
-    LoadTimeoutError,
-    PerformanceTimeoutError,
-    ResultLabel,
-)
-from .backend.result_collector import ResultCollector
 from .backend.assembler import Assembler
-from .backend.task_runner import TaskRunner
 from .backend.data_source import DatasetSource
+from .backend.result_collector import ResultCollector
+from .backend.task_runner import TaskRunner
+from .metric import Metric
+from .models import (CaseResult, LoadTimeoutError, PerformanceTimeoutError,
+                     ResultLabel, TaskConfig, TaskStage, TestResult)
 
 log = logging.getLogger(__name__)
 
 global_result_future: concurrent.futures.Future | None = None
 
+
 class SIGNAL(Enum):
-    SUCCESS=0
-    ERROR=1
-    WIP=2
+    SUCCESS = 0
+    ERROR = 1
+    WIP = 2
 
 
 class BenchMarkRunner:
@@ -42,9 +37,11 @@ class BenchMarkRunner:
         self.drop_old: bool = True
         self.dataset_source: DatasetSource = DatasetSource.S3
 
+
     def set_drop_old(self, drop_old: bool):
         self.drop_old = drop_old
 
+
     def set_download_address(self, use_aliyun: bool):
         if use_aliyun:
             self.dataset_source = DatasetSource.AliyunOSS
@@ -152,13 +149,13 @@ class BenchMarkRunner:
             latest_runner, cached_load_duration = None, None
             for idx, runner in enumerate(running_task.case_runners):
                 case_res = CaseResult(
-                    result_id=idx,
                     metrics=Metric(),
                     task_config=runner.config,
                 )
 
                 # drop_old = False if latest_runner and runner == latest_runner else config.DROP_OLD
-                drop_old = config.DROP_OLD
+                # drop_old = config.DROP_OLD
+                drop_old = TaskStage.DROP_OLD in runner.config.stages
                 if latest_runner and runner == latest_runner:
                     drop_old = False
                 elif not self.drop_old:
@@ -167,7 +164,7 @@ class BenchMarkRunner:
                     log.info(f"[{idx+1}/{running_task.num_cases()}] start case: {runner.display()}, drop_old={drop_old}")
                     case_res.metrics = runner.run(drop_old)
                     log.info(f"[{idx+1}/{running_task.num_cases()}] finish case: {runner.display()}, "
-                        f"result={case_res.metrics}, label={case_res.label}")
+                             f"result={case_res.metrics}, label={case_res.label}")
 
                     # cache the latest succeeded runner
                     latest_runner = runner
@@ -193,7 +190,6 @@ class BenchMarkRunner:
                     c_results.append(case_res)
                     send_conn.send((SIGNAL.WIP, idx))
 
-
             test_result = TestResult(
                 run_id=running_task.run_id,
                 task_label=running_task.task_label,
@@ -204,7 +200,7 @@ class BenchMarkRunner:
 
             send_conn.send((SIGNAL.SUCCESS, None))
             send_conn.close()
-            log.info(f"Succes to finish task: label={running_task.task_label}, run_id={running_task.run_id}")
+            log.info(f"Success to finish task: label={running_task.task_label}, run_id={running_task.run_id}")
 
         except Exception as e:
             err_msg = f"An error occurs when running task={running_task.task_label}, run_id={running_task.run_id}, err={e}"
@@ -246,7 +242,7 @@ class BenchMarkRunner:
         called as soon as a child terminates.
         """
         children = psutil.Process().children(recursive=True)
-        for p in  children:
+        for p in children:
             try:
                 log.warning(f"sending SIGTERM to child process: {p}")
                 p.send_signal(sig)

vectordb_bench/metric.py

@@ -1,7 +1,7 @@
 import logging
 import numpy as np
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 
 log = logging.getLogger(__name__)
@@ -19,6 +19,10 @@ class Metric:
     qps: float = 0.0
     serial_latency_p99: float = 0.0
     recall: float = 0.0
+    ndcg: float = 0.0
+    conc_num_list: list[int] = field(default_factory=list)
+    conc_qps_list: list[float] = field(default_factory=list)
+    conc_latency_p99_list: list[float] = field(default_factory=list)
 
 
 QURIES_PER_DOLLAR_METRIC = "QP$ (Quries per Dollar)"
@@ -60,3 +64,21 @@ def calc_recall(count: int, ground_truth: list[int], got: list[int]) -> float:
             recalls[i] = 1
 
     return np.mean(recalls)
+
+
+def get_ideal_dcg(k: int):
+    ideal_dcg = 0
+    for i in range(k):
+        ideal_dcg += 1 / np.log2(i+2)
+
+    return ideal_dcg
+
+
+def calc_ndcg(ground_truth: list[int], got: list[int], ideal_dcg: float) -> float:
+    dcg = 0
+    ground_truth = list(ground_truth)
+    for id in set(got):
+        if id in ground_truth:
+            idx = ground_truth.index(id)
+            dcg += 1 / np.log2(idx+2)
+    return dcg / ideal_dcg