add int filter

Author: Caroline-an777

vectordb_bench/backend/cases.py

@@ -4,7 +4,7 @@
 
 from vectordb_bench import config
 from vectordb_bench.backend.clients.api import MetricType
-from vectordb_bench.backend.filter import Filter, FilterOp, IntFilter, LabelFilter, NonFilter, non_filter
+from vectordb_bench.backend.filter import Filter, FilterOp, IntFilter, LabelFilter, NewIntFilter, NonFilter, non_filter
 from vectordb_bench.base import BaseModel
 from vectordb_bench.frontend.components.custom.getCustomConfig import CustomDatasetConfig
 
@@ -54,6 +54,8 @@ class CaseType(Enum):
 
     LabelFilterPerformanceCase = 300
 
+    NewIntFilterPerformanceCase = 400
+
     def case_cls(self, custom_configs: dict | None = None) -> type["Case"]:
         if custom_configs is None:
             return type2case.get(self)()
@@ -130,6 +132,7 @@ class PerformanceCase(Case):
     filter_rate: float | None = None
     load_timeout: float | int = config.LOAD_TIMEOUT_DEFAULT
     optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_DEFAULT
+    int_value: float | None = None
 
 
 class CapacityDim960(CapacityCase):
@@ -471,6 +474,46 @@ def __init__(
         )
 
 
+class NewIntFilterPerformanceCase(PerformanceCase):
+    case_id: CaseType = CaseType.NewIntFilterPerformanceCase
+    dataset_with_size_type: DatasetWithSizeType
+    filter_rate: float
+
+    def __init__(
+        self,
+        dataset_with_size_type: DatasetWithSizeType | str,
+        filter_rate: float,
+        int_value: float | None = 0,
+        **kwargs,
+    ):
+        if not isinstance(dataset_with_size_type, DatasetWithSizeType):
+            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
+        name = f"Int-Filter-{filter_rate*100:.1f}% - {dataset_with_size_type.value}"
+        description = f"Int-Filter-{filter_rate*100:.1f}% Performance Test ({dataset_with_size_type.value})"
+        dataset = dataset_with_size_type.get_manager()
+        load_timeout = dataset_with_size_type.get_load_timeout()
+        optimize_timeout = dataset_with_size_type.get_optimize_timeout()
+        filters = IntFilter(filter_rate=filter_rate, int_value=int_value)
+        filter_rate = filters.filter_rate
+        super().__init__(
+            name=name,
+            description=description,
+            dataset=dataset,
+            load_timeout=load_timeout,
+            optimize_timeout=optimize_timeout,
+            filter_rate=filter_rate,
+            int_value=int_value,
+            dataset_with_size_type=dataset_with_size_type,
+            **kwargs,
+        )
+
+    @property
+    def filters(self) -> Filter:
+        int_field = self.dataset.data.train_id_field
+        int_value = int(self.dataset.data.size * self.filter_rate)
+        return NewIntFilter(filter_rate=self.filter_rate, int_field=int_field, int_value=int_value)
+
+
 class LabelFilterPerformanceCase(PerformanceCase):
     case_id: CaseType = CaseType.LabelFilterPerformanceCase
     dataset_with_size_type: DatasetWithSizeType
@@ -529,5 +572,6 @@ def filters(self) -> Filter:
     CaseType.Performance1536D50K: Performance1536D50K,
     CaseType.PerformanceCustomDataset: PerformanceCustomDataset,
     CaseType.StreamingPerformanceCase: StreamingPerformanceCase,
+    CaseType.NewIntFilterPerformanceCase: NewIntFilterPerformanceCase,
     CaseType.LabelFilterPerformanceCase: LabelFilterPerformanceCase,
 }

vectordb_bench/backend/dataset.py

@@ -48,6 +48,7 @@ class BaseDataset(BaseModel):
     scalar_labels_file_separated: bool = True
     scalar_labels_file: str = "scalar_labels.parquet"
     scalar_label_percentages: list[float] = []
+    scalar_int_rates: list[float] = []
     train_id_field: str = "id"
     train_vector_field: str = "emb"
     test_file: str = "test.parquet"
@@ -164,6 +165,29 @@ class Cohere(BaseDataset):
     }
     with_scalar_labels: bool = True
     scalar_label_percentages: list[float] = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
+    scalar_int_rates: list[float] = [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.6,
+        0.7,
+        0.8,
+        0.9,
+        0.95,
+        0.98,
+        0.99,
+        0.995,
+        0.998,
+        0.999,
+    ]
 
 
 class Bioasq(BaseDataset):
@@ -178,6 +202,29 @@ class Bioasq(BaseDataset):
     }
     with_scalar_labels: bool = True
     scalar_label_percentages: list[float] = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
+    scalar_int_rates: list[float] = [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.6,
+        0.7,
+        0.8,
+        0.9,
+        0.95,
+        0.98,
+        0.99,
+        0.995,
+        0.998,
+        0.999,
+    ]
 
 
 class Glove(BaseDataset):
@@ -217,6 +264,29 @@ class OpenAI(BaseDataset):
     }
     with_scalar_labels: bool = True
     scalar_label_percentages: list[float] = [0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5]
+    scalar_int_rates: list[float] = [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.6,
+        0.7,
+        0.8,
+        0.9,
+        0.95,
+        0.98,
+        0.99,
+        0.995,
+        0.998,
+        0.999,
+    ]
 
 
 class DatasetManager(BaseModel):

vectordb_bench/backend/filter.py

@@ -51,6 +51,23 @@ def groundtruth_file(self) -> str:
         raise RuntimeError(msg)
 
 
+class NewIntFilter(Filter):
+    type: FilterOp = FilterOp.NumGE
+    int_field: str = "id"
+    int_value: int
+
+    @property
+    def int_rate(self) -> str:
+        r = self.filter_rate * 100
+        if 1 <= r <= 99:
+            return f"int_{int(r)}p"
+        return f"int_{r:.1f}p"
+
+    @property
+    def groundtruth_file(self) -> str:
+        return f"neighbors_{self.int_rate}.parquet"
+
+
 class LabelFilter(Filter):
     """
     filter expr: label_field == label_value, like `color == "red"`

vectordb_bench/frontend/components/check_results/filters.py

@@ -89,7 +89,7 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t
             col=1,
         )
 
-    if filter_type == FilterOp.StrEqual:
+    if filter_type == FilterOp.StrEqual or filter_type == FilterOp.NumGE:
         container = st.container()
         datasetWithSizeTypes = [dataset_with_size_type for dataset_with_size_type in DatasetWithSizeType]
         showDatasetWithSizeTypes = filterView(
@@ -102,9 +102,6 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t
         datasets = [dataset_with_size_type.get_manager() for dataset_with_size_type in showDatasetWithSizeTypes]
         showCaseNames = list(set([case.name for case in allCases if case.dataset in datasets]))
 
-    if filter_type == FilterOp.NumGE:
-        raise NotImplementedError
-
     return showDBNames, showCaseNames
 
 

vectordb_bench/frontend/components/check_results/nav.py

@@ -29,6 +29,7 @@ def NavToPages(st):
         {"name": "Quries Per Dollar", "link": "quries_per_dollar"},
         {"name": "Concurrent", "link": "concurrent"},
         {"name": "Label Filter", "link": "label_filter"},
+        {"name": "Int Filter", "link": "int_filter"},
         {"name": "Streaming", "link": "streaming"},
         {"name": "Tables", "link": "tables"},
         {"name": "Custom Dataset", "link": "custom"},

vectordb_bench/frontend/components/int_filter/charts.py

@@ -0,0 +1,60 @@
+import plotly.express as px
+from vectordb_bench.metric import metric_unit_map
+
+
+def drawCharts(st, allData, **kwargs):
+    dataset_names = list(set([data["dataset_name"] for data in allData]))
+    dataset_names.sort()
+    for dataset_name in dataset_names:
+        container = st.container()
+        container.subheader(dataset_name)
+        data = [d for d in allData if d["dataset_name"] == dataset_name]
+        drawChartByMetric(container, data, **kwargs)
+
+
+def drawChartByMetric(st, data, metrics=("qps", "recall"), **kwargs):
+    columns = st.columns(len(metrics))
+    for i, metric in enumerate(metrics):
+        container = columns[i]
+        container.markdown(f"#### {metric}")
+        drawChart(container, data, metric)
+
+
+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(st, data: list[object], metric):
+    unit = metric_unit_map.get(metric, "")
+    x = "filter_rate"
+    xrange = getRange(x, data, [0.05, 0.1])
+
+    y = metric
+    yrange = getRange(y, data, [0.2, 0.1])
+
+    data.sort(key=lambda a: a[x])
+
+    fig = px.line(
+        data,
+        x=x,
+        y=y,
+        color="db_name",
+        line_group="db_name",
+        text=metric,
+        markers=True,
+    )
+    fig.update_xaxes(range=xrange)
+    fig.update_yaxes(range=yrange)
+    fig.update_traces(textposition="bottom right", texttemplate="%{y:,.4~r}" + unit)
+    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/welcome/welcomePrams.py

@@ -94,12 +94,22 @@ def welcomePrams(st):
             "title": "Label Filter Performance",
             "description": (
                 "<span style='font-size: 17px;'>"
-                "To view the perfomance of datasets under different filter ratios "
+                "To view the perfomance of datasets under different label filter ratios "
                 "</span>"
             ),
             "image": "fig/homepage/label_filter.png",
             "link": "label_filter",
         },
+        {
+            "title": "Int Filter Performance",
+            "description": (
+                "<span style='font-size: 17px;'>"
+                "To view the perfomance of datasets under different int filter ratios "
+                "</span>"
+            ),
+            "image": "fig/homepage/label_filter.png",
+            "link": "int_filter",
+        },
         {
             "title": "Streaming Performance",
             "description": (
@@ -138,7 +148,7 @@ def welcomePrams(st):
     for option in options:
         option["image"] = get_image_as_base64(option["image"])
 
-    for option in options[:6]:
+    for option in options[:7]:
         html_content += f"""
         <a href="/{option['link']}" target="_self" style="text-decoration: none;">
             <div class="section-card">
@@ -157,7 +167,7 @@ def welcomePrams(st):
     <div class="last-row">
     """
 
-    for option in options[6:8]:
+    for option in options[7:9]:
         html_content += f"""
         <a href="/{option['link']}" target="_self" style="text-decoration: none;">
             <div class="section-card">

vectordb_bench/frontend/config/dbCaseConfigs.py

@@ -219,6 +219,17 @@ def generate_label_filter_cases(dataset_with_size_type: DatasetWithSizeType) ->
     ]
 
 
+def generate_int_filter_cases(dataset_with_size_type: DatasetWithSizeType) -> list[CaseConfig]:
+    filter_rates = dataset_with_size_type.get_manager().data.scalar_int_rates
+    return [
+        CaseConfig(
+            case_id=CaseType.NewIntFilterPerformanceCase,
+            custom_case=dict(dataset_with_size_type=dataset_with_size_type, filter_rate=filter_rate),
+        )
+        for filter_rate in filter_rates
+    ]
+
+
 UI_CASE_CLUSTERS: list[UICaseItemCluster] = [
     UICaseItemCluster(
         label="Search Performance Test",
@@ -249,6 +260,29 @@ def generate_label_filter_cases(dataset_with_size_type: DatasetWithSizeType) ->
             UICaseItem(cases=generate_normal_cases(CaseType.Performance1536D500K99P)),
         ],
     ),
+    UICaseItemCluster(
+        label="New-Int-Filter Search Performance Test",
+        uiCaseItems=[
+            UICaseItem(
+                label=f"Int-Filter Search Performance Test - {dataset_with_size_type.value}",
+                description=(
+                    f"[Batch Cases]These cases test the search performance of a vector database "
+                    f"with dataset {dataset_with_size_type.value}"
+                    f"under filtering rates of {dataset_with_size_type.get_manager().data.scalar_int_rates}, at varying parallel levels."
+                    f"Results will show index building time, recall, and maximum QPS."
+                ),
+                cases=generate_int_filter_cases(dataset_with_size_type),
+            )
+            for dataset_with_size_type in [
+                DatasetWithSizeType.CohereMedium,
+                DatasetWithSizeType.CohereLarge,
+                DatasetWithSizeType.OpenAIMedium,
+                DatasetWithSizeType.OpenAILarge,
+                DatasetWithSizeType.BioasqMedium,
+                DatasetWithSizeType.BioasqLarge,
+            ]
+        ],
+    ),
     UICaseItemCluster(
         label="Label-Filter Search Performance Test",
         uiCaseItems=[