custom control file and column name

custom control file and column name

custom control file and column name

custom control file and column name

custom control name of file and column

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, non_filter
+from vectordb_bench.backend.filter import Filter, FilterOp, IntFilter, LabelFilter, NonFilter, non_filter
 from vectordb_bench.base import BaseModel
 from vectordb_bench.frontend.components.custom.getCustomConfig import CustomDatasetConfig
 
@@ -337,6 +337,7 @@ class PerformanceCustomDataset(PerformanceCase):
     case_id: CaseType = CaseType.PerformanceCustomDataset
     name: str = "Performance With Custom Dataset"
     description: str = ""
+    gt_file: str
     dataset: DatasetManager
 
     def __init__(
@@ -358,15 +359,26 @@ def __init__(
             with_gt=dataset_config.with_gt,
             dir=dataset_config.dir,
             file_num=dataset_config.file_count,
+            train_file=dataset_config.train_name,
+            test_file=f"{dataset_config.test_name}.parquet",
+            train_id_field=dataset_config.train_id_name,
+            train_vector_field=dataset_config.train_col_name,
+            test_vector_field=dataset_config.test_col_name,
+            gt_neighbors_field=dataset_config.gt_col_name,
         )
         super().__init__(
             name=name,
             description=description,
             load_timeout=load_timeout,
             optimize_timeout=optimize_timeout,
+            gt_file=f"{dataset_config.gt_name}.parquet",
             dataset=DatasetManager(data=dataset),
         )
 
+    @property
+    def filters(self) -> Filter:
+        return NonFilter(gt_file_name=self.gt_file)
+
 
 class StreamingPerformanceCase(Case):
     case_id: CaseType = CaseType.StreamingPerformanceCase

vectordb_bench/backend/dataset.py

@@ -89,6 +89,13 @@ class CustomDataset(BaseDataset):
     file_num: int
     is_custom: bool = True
     with_remote_resource: bool = False
+    train_file: str = "train"
+    train_id_field: str = "id"
+    train_vector_field: str = "emb"
+    test_file: str = "test.parquet"
+    gt_file: str = "neighbors.parquet"
+    test_vector_field: str = "emb"
+    gt_neighbors_field: str = "neighbors_id"
 
     @validator("size")
     def verify_size(cls, v: int):
@@ -106,6 +113,20 @@ def dir_name(self) -> str:
     def file_count(self) -> int:
         return self.file_num
 
+    @property
+    def train_files(self) -> list[str]:
+        train_file, train_count = self.train_file, self.file_count
+        prefix = f"{train_file}"
+        train_files = []
+        if train_count > 1:
+            prefix_s = [item.strip() for item in prefix.split(",") if item.strip()]
+            for i in range(train_count):
+                sub_file = f"{prefix_s[i]}.parquet"
+                train_files.append(sub_file)
+        else:
+            train_files.append(f"{prefix}.parquet")
+        return train_files
+
 
 class LAION(BaseDataset):
     name: str = "LAION"

vectordb_bench/backend/filter.py

@@ -21,10 +21,11 @@ def groundtruth_file(self) -> str:
 class NonFilter(Filter):
     type: FilterOp = FilterOp.NonFilter
     filter_rate: float = 0.0
+    gt_file_name: str = "neighbors.parquet"
 
     @property
     def groundtruth_file(self) -> str:
-        return "neighbors.parquet"
+        return self.gt_file_name
 
 
 non_filter = NonFilter()

vectordb_bench/backend/runner/serial_runner.py

@@ -56,9 +56,9 @@ def task(self) -> int:
             log.info(f"({mp.current_process().name:16}) Start inserting embeddings in batch {config.NUM_PER_BATCH}")
             start = time.perf_counter()
             for data_df in self.dataset:
-                all_metadata = data_df["id"].tolist()
+                all_metadata = data_df[self.dataset.data.train_id_field].tolist()
 
-                emb_np = np.stack(data_df["emb"])
+                emb_np = np.stack(data_df[self.dataset.data.train_vector_field])
                 if self.normalize:
                     log.debug("normalize the 100k train data")
                     all_embeddings = (emb_np / np.linalg.norm(emb_np, axis=1)[:, np.newaxis]).tolist()
@@ -175,8 +175,8 @@ def run_endlessness(self) -> int:
         # only 1 file
         data_df = next(iter(self.dataset))
         all_embeddings, all_metadata = (
-            np.stack(data_df["emb"]).tolist(),
-            data_df["id"].tolist(),
+            np.stack(data_df[self.dataset.data.train_vector_field]).tolist(),
+            data_df[self.dataset.data.train_id_field].tolist(),
         )
 
         start_time = time.perf_counter()

vectordb_bench/frontend/components/custom/displayCustomCase.py

@@ -23,7 +23,38 @@ def displayCustomCase(customCase: CustomCaseConfig, st, key):
         "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
+        "train file count",
+        key=f"{key}_file_count",
+        value=customCase.dataset_config.file_count,
+        help="if train file count is more than one, please input all your train file name and split with ','",
+    )
+
+    columns = st.columns(3)
+    customCase.dataset_config.train_name = columns[0].text_input(
+        "train file name",
+        key=f"{key}_train_name",
+        value=customCase.dataset_config.train_name,
+        help="if your file and column in the file is not named as previous explanation, please input the real name (for example: if the file name is `tr.parquet` and column name is `embbb`, then input tr and embbb)",
+    )
+    customCase.dataset_config.test_name = columns[1].text_input(
+        "test file name", key=f"{key}_test_name", value=customCase.dataset_config.test_name
+    )
+    customCase.dataset_config.gt_name = columns[2].text_input(
+        "ground truth file name", key=f"{key}_gt_name", value=customCase.dataset_config.gt_name
+    )
+
+    columns = st.columns([1, 1, 2, 2])
+    customCase.dataset_config.train_id_name = columns[0].text_input(
+        "train id name", key=f"{key}_train_id_name", value=customCase.dataset_config.train_id_name
+    )
+    customCase.dataset_config.train_col_name = columns[1].text_input(
+        "train emb name", key=f"{key}_train_col_name", value=customCase.dataset_config.train_col_name
+    )
+    customCase.dataset_config.test_col_name = columns[2].text_input(
+        "test emb name", key=f"{key}_test_col_name", value=customCase.dataset_config.test_col_name
+    )
+    customCase.dataset_config.gt_col_name = columns[3].text_input(
+        "ground truth emb name", key=f"{key}_gt_col_name", value=customCase.dataset_config.gt_col_name
     )
 
     columns = st.columns(4)

vectordb_bench/frontend/components/custom/getCustomConfig.py

@@ -14,6 +14,13 @@ class CustomDatasetConfig(BaseModel):
     file_count: int = 1
     use_shuffled: bool = False
     with_gt: bool = True
+    train_name: str = "train"
+    test_name: str = "test"
+    gt_name: str = "neighbors"
+    train_id_name: str = "id"
+    train_col_name: str = "emb"
+    test_col_name: str = "emb"
+    gt_col_name: str = "neighbors_id"
 
 
 class CustomCaseConfig(BaseModel):