vectordb-bench 0.0.30__py3-none-any.whl → 1.0.1__py3-none-any.whl

vectordb_bench/backend/dataset.py

@@ -20,6 +20,7 @@ from vectordb_bench.base import BaseModel
 from . import utils
 from .clients import MetricType
 from .data_source import DatasetReader, DatasetSource
+from .filter import Filter, FilterOp, non_filter
 
 log = logging.getLogger(__name__)
 
@@ -39,6 +40,21 @@ class BaseDataset(BaseModel):
     with_gt: bool = False
     _size_label: dict[int, SizeLabel] = PrivateAttr()
     is_custom: bool = False
+    with_remote_resource: bool = True
+    # for label filter cases
+    with_scalar_labels: bool = False
+    # if True, scalar_labels will be retrieved from a separate parquet file;
+    #   otherwise, they will be obtained from train.parquet.
+    scalar_labels_file_separated: bool = True
+    scalar_labels_file: str = "scalar_labels.parquet"
+    scalar_label_percentages: list[float] = []
+    train_id_field: str = "id"
+    train_vector_field: str = "emb"
+    test_file: str = "test.parquet"
+    test_id_field: str = "id"
+    test_vector_field: str = "emb"
+    gt_id_field: str = "id"
+    gt_neighbors_field: str = "neighbors_id"
 
     @validator("size")
     def verify_size(cls, v: int):
@@ -51,6 +67,10 @@ class BaseDataset(BaseModel):
     def label(self) -> str:
         return self._size_label.get(self.size).label
 
+    @property
+    def full_name(self) -> str:
+        return f"{self.name.capitalize()} ({self.label.capitalize()})"
+
     @property
     def dir_name(self) -> str:
         return f"{self.name}_{self.label}_{utils.numerize(self.size)}".lower()
@@ -59,11 +79,27 @@ class BaseDataset(BaseModel):
     def file_count(self) -> int:
         return self._size_label.get(self.size).file_count
 
+    @property
+    def train_files(self) -> list[str]:
+        return utils.compose_train_files(self.file_count, self.use_shuffled)
+
 
 class CustomDataset(BaseDataset):
     dir: str
     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"
+    with_scalar_labels: bool = True
+    scalar_labels_file_separated: bool = True
+    scalar_labels_file: str = "scalar_labels.parquet"
+    label_percentages: list[float] = []
 
     @validator("size")
     def verify_size(cls, v: int):
@@ -81,6 +117,17 @@ class CustomDataset(BaseDataset):
     def file_count(self) -> int:
         return self.file_num
 
+    @property
+    def train_files(self) -> list[str]:
+        train_file = self.train_file
+        prefix = f"{train_file}"
+        train_files = []
+        prefix_s = [item.strip() for item in prefix.split(",") if item.strip()]
+        for i in range(len(prefix_s)):
+            sub_file = f"{prefix_s[i]}.parquet"
+            train_files.append(sub_file)
+        return train_files
+
 
 class LAION(BaseDataset):
     name: str = "LAION"
@@ -109,12 +156,28 @@ class Cohere(BaseDataset):
     dim: int = 768
     metric_type: MetricType = MetricType.COSINE
     use_shuffled: bool = config.USE_SHUFFLED_DATA
-    with_gt: bool = (True,)
+    with_gt: bool = True
     _size_label: dict = {
         100_000: SizeLabel(100_000, "SMALL", 1),
         1_000_000: SizeLabel(1_000_000, "MEDIUM", 1),
         10_000_000: SizeLabel(10_000_000, "LARGE", 10),
     }
+    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]
+
+
+class Bioasq(BaseDataset):
+    name: str = "Bioasq"
+    dim: int = 1024
+    metric_type: MetricType = MetricType.COSINE
+    use_shuffled: bool = config.USE_SHUFFLED_DATA
+    with_gt: bool = True
+    _size_label: dict = {
+        1_000_000: SizeLabel(1_000_000, "MEDIUM", 1),
+        10_000_000: SizeLabel(10_000_000, "LARGE", 10),
+    }
+    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]
 
 
 class Glove(BaseDataset):
@@ -146,12 +209,14 @@ class OpenAI(BaseDataset):
     dim: int = 1536
     metric_type: MetricType = MetricType.COSINE
     use_shuffled: bool = config.USE_SHUFFLED_DATA
-    with_gt: bool = (True,)
+    with_gt: bool = True
     _size_label: dict = {
         50_000: SizeLabel(50_000, "SMALL", 1),
         500_000: SizeLabel(500_000, "MEDIUM", 1),
         5_000_000: SizeLabel(5_000_000, "LARGE", 10),
     }
+    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]
 
 
 class DatasetManager(BaseModel):
@@ -166,8 +231,9 @@ class DatasetManager(BaseModel):
     """
 
     data: BaseDataset
-    test_data: pd.DataFrame | None = None
-    gt_data: pd.DataFrame | None = None
+    test_data: list[list[float]] | None = None
+    gt_data: list[list[int]] | None = None
+    scalar_labels: pl.DataFrame | None = None
     train_files: list[str] = []
     reader: DatasetReader | None = None
 
@@ -176,6 +242,9 @@ class DatasetManager(BaseModel):
             return self.data.name == obj.data.name and self.data.label == obj.data.label
         return False
 
+    def __hash__(self) -> int:
+        return hash((self.data.name, self.data.label))
+
     def set_reader(self, reader: DatasetReader):
         self.reader = reader
 
@@ -191,7 +260,7 @@ class DatasetManager(BaseModel):
         return pathlib.Path(
             config.DATASET_LOCAL_DIR,
             self.data.name.lower(),
-            self.data.dir_name.lower(),
+            self.data.dir_name,
         )
 
     def __iter__(self):
@@ -201,58 +270,59 @@ class DatasetManager(BaseModel):
     def prepare(
         self,
         source: DatasetSource = DatasetSource.S3,
-        filters: float | str | None = None,
+        filters: Filter = non_filter,
     ) -> bool:
         """Download the dataset from DatasetSource
          url = f"{source}/{self.data.dir_name}"
 
         Args:
             source(DatasetSource): S3 or AliyunOSS, default as S3
-            filters(Optional[int | float | str]): combined with dataset's with_gt to
+            filters(Filter): combined with dataset's with_gt to
               compose the correct ground_truth file
 
         Returns:
             bool: whether the dataset is successfully prepared
 
         """
-        file_count, use_shuffled = self.data.file_count, self.data.use_shuffled
-
-        train_files = utils.compose_train_files(file_count, use_shuffled)
-        all_files = train_files
-
-        test_file = "test.parquet"
-        all_files.extend([test_file])
-        gt_file = None
+        self.train_files = self.data.train_files
+        gt_file, test_file = None, None
         if self.data.with_gt:
-            gt_file = utils.compose_gt_file(filters)
-            all_files.extend([gt_file])
+            gt_file, test_file = filters.groundtruth_file, self.data.test_file
 
-        if not self.data.is_custom:
+        if self.data.with_remote_resource:
+            download_files = [file for file in self.train_files]
+            download_files.extend([gt_file, test_file])
+            if self.data.with_scalar_labels and self.data.scalar_labels_file_separated:
+                download_files.append(self.data.scalar_labels_file)
             source.reader().read(
                 dataset=self.data.dir_name.lower(),
-                files=all_files,
+                files=download_files,
                 local_ds_root=self.data_dir,
             )
 
-        if test_file is not None:
-            self.test_data = self._read_file(test_file)
+        # read scalar_labels_file if separated
+        if (
+            filters.type == FilterOp.StrEqual
+            and self.data.with_scalar_labels
+            and self.data.scalar_labels_file_separated
+        ):
+            self.scalar_labels = self._read_file(self.data.scalar_labels_file)
 
-        if gt_file is not None:
-            self.gt_data = self._read_file(gt_file)
+        if gt_file is not None and test_file is not None:
+            self.test_data = self._read_file(test_file)[self.data.test_vector_field].to_list()
+            self.gt_data = self._read_file(gt_file)[self.data.gt_neighbors_field].to_list()
 
-        prefix = "shuffle_train" if use_shuffled else "train"
-        self.train_files = sorted([f.name for f in self.data_dir.glob(f"{prefix}*.parquet")])
         log.debug(f"{self.data.name}: available train files {self.train_files}")
 
         return True
 
-    def _read_file(self, file_name: str) -> pd.DataFrame:
+    def _read_file(self, file_name: str) -> pl.DataFrame:
         """read one file from disk into memory"""
         log.info(f"Read the entire file into memory: {file_name}")
         p = pathlib.Path(self.data_dir, file_name)
         if not p.exists():
             log.warning(f"No such file: {p}")
-            return pd.DataFrame()
+            return pl.DataFrame()
 
         return pl.read_parquet(p)
 
@@ -308,6 +378,7 @@ class Dataset(Enum):
     LAION = LAION
     GIST = GIST
     COHERE = Cohere
+    BIOASQ = Bioasq
     GLOVE = Glove
     SIFT = SIFT
     OPENAI = OpenAI
@@ -317,3 +388,51 @@ class Dataset(Enum):
 
     def manager(self, size: int) -> DatasetManager:
         return DatasetManager(data=self.get(size))
+
+
+class DatasetWithSizeType(Enum):
+    CohereSmall = "Small Cohere (768dim, 100K)"
+    CohereMedium = "Medium Cohere (768dim, 1M)"
+    CohereLarge = "Large Cohere (768dim, 10M)"
+    BioasqMedium = "Medium Bioasq (1024dim, 1M)"
+    BioasqLarge = "Large Bioasq (1024dim, 10M)"
+    OpenAISmall = "Small OpenAI (1536dim, 50K)"
+    OpenAIMedium = "Medium OpenAI (1536dim, 500K)"
+    OpenAILarge = "Large OpenAI (1536dim, 5M)"
+
+    def get_manager(self) -> DatasetManager:
+        if self not in DatasetWithSizeMap:
+            msg = f"wrong ScalarDatasetWithSizeType: {self.name}"
+            raise ValueError(msg)
+        return DatasetWithSizeMap.get(self)
+
+    def get_load_timeout(self) -> float:
+        if "small" in self.value.lower():
+            return config.LOAD_TIMEOUT_768D_100K
+        if "medium" in self.value.lower():
+            return config.LOAD_TIMEOUT_768D_1M
+        if "large" in self.value.lower():
+            return config.LOAD_TIMEOUT_768D_10M
+        msg = f"No load_timeout for {self.value}"
+        raise KeyError(msg)
+
+    def get_optimize_timeout(self) -> float:
+        if "small" in self.value.lower():
+            return config.OPTIMIZE_TIMEOUT_768D_100K
+        if "medium" in self.value.lower():
+            return config.OPTIMIZE_TIMEOUT_768D_1M
+        if "large" in self.value.lower():
+            return config.OPTIMIZE_TIMEOUT_768D_10M
+        return config.OPTIMIZE_TIMEOUT_DEFAULT
+
+
+DatasetWithSizeMap = {
+    DatasetWithSizeType.CohereSmall: Dataset.COHERE.manager(100_000),
+    DatasetWithSizeType.CohereMedium: Dataset.COHERE.manager(1_000_000),
+    DatasetWithSizeType.CohereLarge: Dataset.COHERE.manager(10_000_000),
+    DatasetWithSizeType.BioasqMedium: Dataset.BIOASQ.manager(1_000_000),
+    DatasetWithSizeType.BioasqLarge: Dataset.BIOASQ.manager(10_000_000),
+    DatasetWithSizeType.OpenAISmall: Dataset.OPENAI.manager(50_000),
+    DatasetWithSizeType.OpenAIMedium: Dataset.OPENAI.manager(500_000),
+    DatasetWithSizeType.OpenAILarge: Dataset.OPENAI.manager(5_000_000),
+}

vectordb_bench/backend/filter.py

@@ -0,0 +1,76 @@
+from enum import StrEnum
+
+from ..base import BaseModel
+
+
+class FilterOp(StrEnum):
+    NumGE = "NumGE"  # test ">="
+    StrEqual = "Label"  # test "=="
+    NonFilter = "NonFilter"
+
+
+class Filter(BaseModel):
+    type: FilterOp
+    filter_rate: float = 0.0
+
+    @property
+    def groundtruth_file(self) -> str:
+        raise NotImplementedError
+
+
+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 self.gt_file_name
+
+
+non_filter = NonFilter()
+
+
+class IntFilter(Filter):
+    """
+    compatible with older int-filter cases
+    filter expr: int_field >= int_value (dataset_size * filter_rate)
+    """
+
+    type: FilterOp = FilterOp.NumGE
+    int_field: str = "id"
+    int_value: int
+
+    @property
+    def groundtruth_file(self) -> str:
+        if self.filter_rate == 0.01:
+            return "neighbors_head_1p.parquet"
+        if self.filter_rate == 0.99:
+            return "neighbors_tail_1p.parquet"
+        msg = f"Not Support Int Filter - {self.filter_rate}"
+        raise RuntimeError(msg)
+
+
+class LabelFilter(Filter):
+    """
+    filter expr: label_field == label_value, like `color == "red"`
+    """
+
+    type: FilterOp = FilterOp.StrEqual
+    label_field: str = "labels"
+    label_percentage: float
+
+    @property
+    def label_value(self) -> str:
+        p = self.label_percentage * 100
+        if p >= 1:
+            return f"label_{int(p)}p"  # such as 5p, 20p, 1p, ...
+        return f"label_{p:.1f}p"  # such as 0.1p, 0.5p, ...
+
+    def __init__(self, label_percentage: float, **kwargs):
+        filter_rate = 1.0 - label_percentage
+        super().__init__(filter_rate=filter_rate, label_percentage=label_percentage, **kwargs)
+
+    @property
+    def groundtruth_file(self) -> str:
+        return f"neighbors_{self.label_field}_{self.label_value}.parquet"

vectordb_bench/backend/runner/__init__.py

@@ -1,10 +1,10 @@
-from .mp_runner import (
-    MultiProcessingSearchRunner,
-)
+from .mp_runner import MultiProcessingSearchRunner
+from .read_write_runner import ReadWriteRunner
 from .serial_runner import SerialInsertRunner, SerialSearchRunner
 
 __all__ = [
     "MultiProcessingSearchRunner",
+    "ReadWriteRunner",
     "SerialInsertRunner",
     "SerialSearchRunner",
 ]

vectordb_bench/backend/runner/mp_runner.py

@@ -9,6 +9,8 @@ from multiprocessing.queues import Queue
 
 import numpy as np
 
+from vectordb_bench.backend.filter import Filter, non_filter
+
 from ... import config
 from ...models import ConcurrencySlotTimeoutError
 from ..clients import api
@@ -31,7 +33,7 @@ class MultiProcessingSearchRunner:
         db: api.VectorDB,
         test_data: list[list[float]],
         k: int = config.K_DEFAULT,
-        filters: dict | None = None,
+        filters: Filter = non_filter,
         concurrencies: Iterable[int] = config.NUM_CONCURRENCY,
         duration: int = config.CONCURRENCY_DURATION,
         concurrency_timeout: int = config.CONCURRENCY_TIMEOUT,
@@ -58,6 +60,7 @@ class MultiProcessingSearchRunner:
             cond.wait()
 
         with self.db.init():
+            self.db.prepare_filter(self.filters)
             num, idx = len(test_data), random.randint(0, len(test_data) - 1)
 
             start_time = time.perf_counter()
@@ -66,18 +69,12 @@ class MultiProcessingSearchRunner:
             while time.perf_counter() < start_time + self.duration:
                 s = time.perf_counter()
                 try:
-                    self.db.search_embedding(
-                        test_data[idx],
-                        self.k,
-                        self.filters,
-                    )
+                    self.db.search_embedding(test_data[idx], self.k)
+                    count += 1
+                    latencies.append(time.perf_counter() - s)
                 except Exception as e:
                     log.warning(f"VectorDB search_embedding error: {e}")
-                    traceback.print_exc(chain=True)
-                    raise e from None
 
-                latencies.append(time.perf_counter() - s)
-                count += 1
                 # loop through the test data
                 idx = idx + 1 if idx < num - 1 else 0
 
@@ -181,10 +178,20 @@ class MultiProcessingSearchRunner:
     def stop(self) -> None:
         pass
 
-    def run_by_dur(self, duration: int) -> float:
+    def run_by_dur(self, duration: int) -> tuple[float, float]:
+        """
+        Returns:
+            float: largest qps
+            float: failed rate
+        """
         return self._run_by_dur(duration)
 
-    def _run_by_dur(self, duration: int) -> float:
+    def _run_by_dur(self, duration: int) -> tuple[float, float]:
+        """
+        Returns:
+            float: largest qps
+            float: failed rate
+        """
         max_qps = 0
         try:
             for conc in self.concurrencies:
@@ -208,12 +215,17 @@ class MultiProcessingSearchRunner:
                             log.info(f"Syncing all process and start concurrency search, concurrency={conc}")
 
                         start = time.perf_counter()
-                        all_count = sum([r.result() for r in future_iter])
+                        res = [r.result() for r in future_iter]
+                        all_success_count = sum([r[0] for r in res])
+                        all_failed_count = sum([r[1] for r in res])
+                        failed_rate = all_failed_count / (all_failed_count + all_success_count)
                         cost = time.perf_counter() - start
 
-                        qps = round(all_count / cost, 4)
-                        log.info(f"End search in concurrency {conc}: dur={cost}s, total_count={all_count}, qps={qps}")
-
+                        qps = round(all_success_count / cost, 4)
+                        log.info(
+                            f"End search in concurrency {conc}: dur={cost}s, failed_rate={failed_rate}, "
+                            f"all_success_count={all_success_count}, all_failed_count={all_failed_count}, qps={qps}",
+                        )
                 if qps > max_qps:
                     max_qps = qps
                     log.info(f"Update largest qps with concurrency {conc}: current max_qps={max_qps}")
@@ -230,52 +242,53 @@ class MultiProcessingSearchRunner:
         finally:
             self.stop()
 
-        return max_qps
+        return max_qps, failed_rate
 
-    def search_by_dur(
-        self,
-        dur: int,
-        test_data: list[list[float]],
-        q: mp.Queue,
-        cond: mp.Condition,
-    ) -> int:
+    def search_by_dur(self, dur: int, test_data: list[list[float]], q: mp.Queue, cond: mp.Condition) -> tuple[int, int]:
+        """
+        Returns:
+            int: successful requests count
+            int: failed requests count
+        """
         # sync all process
         q.put(1)
         with cond:
             cond.wait()
 
         with self.db.init():
+            self.db.prepare_filter(self.filters)
             num, idx = len(test_data), random.randint(0, len(test_data) - 1)
 
             start_time = time.perf_counter()
-            count = 0
+            success_count = 0
+            failed_cnt = 0
             while time.perf_counter() < start_time + dur:
                 s = time.perf_counter()
                 try:
-                    self.db.search_embedding(
-                        test_data[idx],
-                        self.k,
-                        self.filters,
-                    )
+                    self.db.search_embedding(test_data[idx], self.k)
+                    success_count += 1
                 except Exception as e:
-                    log.warning(f"VectorDB search_embedding error: {e}")
-                    traceback.print_exc(chain=True)
-                    raise e from None
+                    failed_cnt += 1
+                    # reduce log
+                    if failed_cnt <= 3:
+                        log.warning(f"VectorDB search_embedding error: {e}")
+                    else:
+                        log.debug(f"VectorDB search_embedding error: {e}")
 
-                count += 1
                 # loop through the test data
                 idx = idx + 1 if idx < num - 1 else 0
 
-                if count % 500 == 0:
+                if success_count % 500 == 0:
                     log.debug(
-                        f"({mp.current_process().name:16}) search_count: {count}, "
-                        f"latest_latency={time.perf_counter()-s}"
+                        f"({mp.current_process().name:16}) search_count: {success_count}, "
+                        f"latest_latency={time.perf_counter()-s}",
                     )
 
         total_dur = round(time.perf_counter() - start_time, 4)
         log.debug(
             f"{mp.current_process().name:16} search {self.duration}s: "
-            f"actual_dur={total_dur}s, count={count}, qps in this process: {round(count / total_dur, 4):3}"
+            f"actual_dur={total_dur}s, count={success_count}, failed_cnt={failed_cnt}, "
+            f"qps (successful) in this process: {round(success_count / total_dur, 4):3}",
         )
 
-        return count
+        return success_count, failed_cnt