diff --git a/tests/test_dataset.py b/tests/test_dataset.py index d4ccb283d..9154b6409 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,12 +1,16 @@ -from vectordb_bench.backend.dataset import Dataset +import struct + +from vectordb_bench import config +from vectordb_bench.backend.clients import MetricType +from vectordb_bench.backend.dataset import BinaryDatasetManager, Dataset, DatasetWithSizeType, SIFTBinary import logging import pytest from pydantic import ValidationError from vectordb_bench.backend.data_source import DatasetSource - log = logging.getLogger("vectordb_bench") + class TestDataSet: def test_iter_dataset(self): for ds in Dataset: @@ -29,6 +33,7 @@ def test_iter_cohere(self): cohere_10m.prepare() import time + before = time.time() for i in cohere_10m: log.debug(i.head(1)) @@ -40,9 +45,11 @@ def test_iter_cohere(self): def test_iter_laion(self): laion_100m = Dataset.LAION.manager(100_000_000) from vectordb_bench.backend.data_source import DatasetSource + laion_100m.prepare(source=DatasetSource.AliyunOSS) import time + before = time.time() for i in laion_100m: log.debug(i.head(1)) @@ -75,3 +82,37 @@ def test_download_small(self): local_ds_root=openai_50k.data_dir, ) + def test_sift_binary_dataset_type(self): + dataset = DatasetWithSizeType.SIFTBinary1M.get_manager() + + assert isinstance(dataset, BinaryDatasetManager) + assert dataset.data.name == "SIFTBinary" + assert dataset.data.dim == 128 + assert dataset.data.metric_type == MetricType.HAMMING + + def test_binary_dataset_manager_reads_vec_tool_bin_layout(self, tmp_path, monkeypatch): + monkeypatch.setattr(config, "DATASET_LOCAL_DIR", tmp_path) + vectors = [ + bytes(range(16)), + bytes([255] * 16), + bytes([170] * 16), + ] + truth = [[2, 1], [0, 2]] + + data = SIFTBinary(size=1_000_000) + manager = BinaryDatasetManager(data=data) + manager.data_dir.mkdir(parents=True) + manager.data_dir.joinpath("base.bin").write_bytes(struct.pack(" dict: class ESElementType(StrEnum): float = "float" # 4 byte byte = "byte" # 1 byte, -128 to 127 + bit = "bit" # packed binary vector class ElasticCloudIndexConfig(BaseModel, DBCaseConfig): @@ -87,6 +88,12 @@ class ElasticCloudIndexConfig(BaseModel, DBCaseConfig): M: int | None = None num_candidates: int | None = None + def _metric_value(self) -> str | None: + metric = self.metric_type + if metric is None: + return None + return metric.value if hasattr(metric, "value") else str(metric) + def __eq__(self, obj: any): return ( self.index == obj.index @@ -111,17 +118,24 @@ def __hash__(self) -> int: ) def parse_metric(self) -> str: + if self._metric_value() == MetricType.HAMMING.value: + return "l2_norm" if self.metric_type == MetricType.L2: return "l2_norm" if self.metric_type == MetricType.IP: return "dot_product" return "cosine" + def parse_element_type(self) -> ESElementType: + if self._metric_value() == MetricType.HAMMING.value: + return ESElementType.bit + return self.element_type + def index_param(self) -> dict: return { "type": "dense_vector", "index": True, - "element_type": self.element_type.value, + "element_type": self.parse_element_type().value, "similarity": self.parse_metric(), "index_options": { "type": self.index.value, diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index 6ac1f5817..db3556ee3 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -59,6 +59,11 @@ def __init__( # noqa: PLR0915 self._multitenant_partition_key_field = self._scalar_label_field self._scalar_labels_index_name = "labels_idx" self._is_fts = isinstance(self.case_config, MilvusFtsConfig) + metric_type = getattr(self.case_config, "metric_type", None) + metric_value = metric_type.value if hasattr(metric_type, "value") else metric_type + self._is_binary_vector = ( + not self._is_fts and metric_value is not None and str(metric_value).upper() in {"HAMMING", "JACCARD"} + ) if self._is_fts: self.batch_size = fts_batch_size or MILVUS_FTS_BATCH_SIZE @@ -71,7 +76,8 @@ def __init__( # noqa: PLR0915 self._sort_index_name = self._doc_id_sort_index_name self._sort_index_field = self._primary_field else: - self.batch_size = int(MILVUS_LOAD_REQS_SIZE / (dim * 4)) + vector_bytes = dim // 8 if self._is_binary_vector else dim * 4 + self.batch_size = int(MILVUS_LOAD_REQS_SIZE / vector_bytes) self._primary_field = "pk" self._scalar_id_field = "id" self._vector_field = "vector" @@ -131,7 +137,8 @@ def __init__( # noqa: PLR0915 else: schema.add_field(self._primary_field, DataType.INT64, is_primary=True) schema.add_field(self._scalar_id_field, DataType.INT64) - schema.add_field(self._vector_field, DataType.FLOAT_VECTOR, dim=dim) + vector_type = DataType.BINARY_VECTOR if self._is_binary_vector else DataType.FLOAT_VECTOR + schema.add_field(self._vector_field, vector_type, dim=dim) if self.multitenant_tenant_labels: schema.add_field( @@ -345,6 +352,18 @@ def need_normalize_cosine(self) -> bool: return False + def _format_vector_for_insert(self, embedding: list[float] | str | bytes) -> list[float] | bytes: + if not self._is_binary_vector: + return embedding + if isinstance(embedding, bytes): + return embedding + if isinstance(embedding, str): + return bytes.fromhex(embedding) + return bytes(embedding) + + def _format_vector_for_search(self, query: list[float] | str | bytes) -> list[float] | bytes: + return self._format_vector_for_insert(query) + def insert_embeddings( self, embeddings: Iterable[list[float]], @@ -365,7 +384,7 @@ def insert_embeddings( row = { self._primary_field: metadata[i], self._scalar_id_field: metadata[i], - self._vector_field: embeddings[i], + self._vector_field: self._format_vector_for_insert(embeddings[i]), } if tenant_labels_data is not None: row[self._multitenant_partition_key_field] = tenant_labels_data[i] @@ -483,7 +502,7 @@ def search_embedding( search_kwargs = { "collection_name": self.collection_name, - "data": [query], + "data": [self._format_vector_for_search(query)], "anns_field": self._vector_field, "search_params": self.case_config.search_param(), "limit": k, diff --git a/vectordb_bench/backend/clients/oss_opensearch/config.py b/vectordb_bench/backend/clients/oss_opensearch/config.py index 7a8b1d98e..5b843b19c 100644 --- a/vectordb_bench/backend/clients/oss_opensearch/config.py +++ b/vectordb_bench/backend/clients/oss_opensearch/config.py @@ -247,3 +247,37 @@ def index_param(self) -> dict: def search_param(self) -> dict: return {"ef_search": self.efSearch} + + +class OSSOpenSearchFtsConfig(BaseModel, DBCaseConfig): + number_of_shards: int = 1 + number_of_replicas: int = 0 + refresh_interval: str = "30s" + force_merge_enabled: bool = True + metric_type: MetricType = MetricType.BM25 + bm25_k1: float | None = None + bm25_b: float | None = None + + def index_param(self) -> dict: + text_mapping = {"type": "text"} + if self.bm25_k1 is not None or self.bm25_b is not None: + text_mapping["similarity"] = "vdbbench_bm25" + return { + "properties": { + "doc_id": {"type": "keyword"}, + "text": text_mapping, + }, + } + + def search_param(self) -> dict: + return {} + + def similarity_settings(self) -> dict: + if self.bm25_k1 is None and self.bm25_b is None: + return {} + bm25_settings = {"type": "BM25"} + if self.bm25_k1 is not None: + bm25_settings["k1"] = self.bm25_k1 + if self.bm25_b is not None: + bm25_settings["b"] = self.bm25_b + return {"similarity": {"vdbbench_bm25": bm25_settings}} diff --git a/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py b/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py index f71850a17..d29abe743 100644 --- a/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py +++ b/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py @@ -9,9 +9,10 @@ from packaging.version import parse as parse_version from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.payload import PayloadProfile from ..api import VectorDB -from .config import OSSOpenSearchIndexConfig, OSSOS_Engine +from .config import OSSOpenSearchFtsConfig, OSSOpenSearchIndexConfig, OSSOS_Engine log = logging.getLogger(__name__) @@ -194,7 +195,7 @@ def __init__( self, dim: int, db_config: dict[str, Any], - db_case_config: OSSOpenSearchIndexConfig, + db_case_config: OSSOpenSearchIndexConfig | OSSOpenSearchFtsConfig, index_name: str = "vdb_bench_index", # must be lowercase id_col_name: str = "_id", label_col_name: str = "label", @@ -212,6 +213,10 @@ def __init__( self.label_col_name = label_col_name self.vector_col_name = vector_col_name self.with_scalar_labels = with_scalar_labels + self._is_fts = isinstance(db_case_config, OSSOpenSearchFtsConfig) + self.text_col_name = "text" + if self._is_fts: + self.id_col_name = "doc_id" # Initialize client state self.client: OpenSearch | None = None @@ -236,13 +241,21 @@ def _handle_index_initialization(self, client: OpenSearch, drop_old: bool) -> No if not is_existed: self._create_index(client) log.info(f"OSS_OpenSearch client create index: {self.index_name}") - self._update_ef_search_before_search(client) - self._load_graphs_to_memory(client) + if not self._is_fts: + self._update_ef_search_before_search(client) + self._load_graphs_to_memory(client) def need_normalize_cosine(self) -> bool: """Whether this database needs to normalize dataset to support COSINE metric.""" return True + @classmethod + def supports_full_text_search(cls) -> bool: + return True + + def has_text_field(self) -> bool: + return bool(getattr(self, "_is_fts", False) and getattr(self, "text_col_name", None)) + def _get_cluster_version(self, client: OpenSearch) -> Version: """ Return the OpenSearch cluster version as a comparable Version object. @@ -307,7 +320,31 @@ def _get_bulk_manager(self, client: OpenSearch) -> BulkInsertManager: """Get bulk insert manager for the given client.""" return BulkInsertManager(client, self.index_name, self.case_config) + def _create_fts_index(self, client: OpenSearch) -> None: + mappings = self.case_config.index_param() + index_settings = { + "number_of_shards": self.case_config.number_of_shards, + "number_of_replicas": self.case_config.number_of_replicas, + "refresh_interval": self.case_config.refresh_interval, + } + index_settings.update(self.case_config.similarity_settings()) + settings = {"index": index_settings} + try: + log.info(f"Creating FTS index with settings: {settings}") + log.info(f"Creating FTS index with mappings: {mappings}") + client.indices.create( + index=self.index_name, + body={"settings": settings, "mappings": mappings}, + ) + except Exception as e: + log.warning(f"Failed to create FTS index: {self.index_name} error: {e!s}") + raise e from None + def _create_index(self, client: OpenSearch) -> None: + if self._is_fts: + self._create_fts_index(client) + return + cluster_version = self._get_cluster_version(client) if self.case_config.on_disk and cluster_version < Version("2.17"): @@ -420,6 +457,36 @@ def insert_embeddings( log.info(f"Using {num_clients} parallel clients for data insertion") return self._insert_with_multiple_clients(embeddings, metadata, num_clients, labels_data) + def insert_documents( + self, + texts: Iterable[str], + doc_ids: list[str], + **kwargs: Any, + ) -> tuple[int, Exception | None]: + if not getattr(self, "_is_fts", False): + msg = "OSSOpenSearch full-text insert requires OSSOpenSearchFtsConfig" + raise RuntimeError(msg) + assert self.client is not None, "should self.init() first" + docs = list(texts) + if len(docs) != len(doc_ids): + msg = f"Mismatch between texts ({len(docs)}) and doc_ids ({len(doc_ids)}) lengths" + raise ValueError(msg) + + insert_data: list[dict[str, Any]] = [] + for i, doc in enumerate(docs): + doc_id = str(doc_ids[i]) + insert_data.append({"index": {"_index": self.index_name, "_id": doc_id}}) + insert_data.append({self.id_col_name: doc_id, self.text_col_name: doc}) + + try: + response = self.client.bulk(body=insert_data) + if response.get("errors"): + log.warning(f"FTS bulk insert had errors: {response}") + return len(docs), None + except Exception as e: + log.warning(f"Failed to insert FTS docs: {self.index_name} error: {e!s}") + return 0, e + def _insert_with_single_client( self, embeddings: Iterable[list[float]], @@ -569,6 +636,47 @@ def search_embedding( log.warning(f"Failed to search: {self.index_name} error: {e!s}") raise e from None + def search_documents( + self, + query: str, + k: int = 100, + payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, + **kwargs: Any, + ) -> list[str]: + if not getattr(self, "_is_fts", False): + msg = "OSSOpenSearch full-text search requires OSSOpenSearchFtsConfig" + raise RuntimeError(msg) + if not self.supports_document_payload_profile(payload_profile): + msg = f"OSSOpenSearch does not support document payload_profile={payload_profile.value}" + raise NotImplementedError(msg) + assert self.client is not None, "should self.init() first" + + source = [self.text_col_name] if payload_profile == PayloadProfile.TEXT else False + filter_path = ["hits.hits._id", f"hits.hits.fields.{self.id_col_name}"] + if payload_profile == PayloadProfile.TEXT: + filter_path.append(f"hits.hits._source.{self.text_col_name}") + search_kwargs: dict[str, Any] = { + "index": self.index_name, + "body": {"query": {"match": {self.text_col_name: query}}}, + "size": k, + "_source": source, + "docvalue_fields": [self.id_col_name], + "filter_path": filter_path, + } + if payload_profile != PayloadProfile.TEXT: + search_kwargs["stored_fields"] = "_none_" + response = self.client.search(**search_kwargs) + + doc_ids = [] + for hit in response.get("hits", {}).get("hits", []): + if hit.get("_id") is not None: + doc_ids.append(str(hit["_id"])) + continue + values = hit.get("fields", {}).get(self.id_col_name, []) + if values: + doc_ids.append(str(values[0])) + return doc_ids + def prepare_filter(self, filters: Filter) -> None: """Prepare filter conditions for search operations.""" self.routing_key = None @@ -587,6 +695,14 @@ def prepare_filter(self, filters: Filter) -> None: def optimize(self, data_size: int | None = None) -> None: """Optimize the index for better search performance.""" + if self._is_fts: + self._refresh_index() + if self.case_config.force_merge_enabled: + self._do_fts_force_merge() + self._refresh_index() + self._update_replicas() + self._refresh_index() + return self._update_ef_search() # Call refresh first to ensure that all segments are created self._refresh_index() @@ -645,6 +761,17 @@ def _refresh_index(self): continue log.debug(f"Completed refresh for index {self.index_name}") + def _do_fts_force_merge(self): + log.info(f"Starting FTS force merge for index {self.index_name}") + force_merge_endpoint = f"/{self.index_name}/_forcemerge?max_num_segments=1&wait_for_completion=false" + force_merge_task_id = self.client.transport.perform_request("POST", force_merge_endpoint)["task"] + while True: + time.sleep(WAITING_FOR_FORCE_MERGE_SEC) + task_status = self.client.tasks.get(task_id=force_merge_task_id) + if task_status["completed"]: + break + log.info(f"Completed FTS force merge for index {self.index_name}") + def _do_force_merge(self): log.info(f"Updating the Index thread qty to {self.case_config.index_thread_qty_during_force_merge}.") diff --git a/vectordb_bench/backend/dataset.py b/vectordb_bench/backend/dataset.py index 6a6f262cc..84c910c67 100644 --- a/vectordb_bench/backend/dataset.py +++ b/vectordb_bench/backend/dataset.py @@ -7,6 +7,7 @@ import json import logging import pathlib +import struct import types import typing from abc import ABC, abstractmethod @@ -18,9 +19,11 @@ import ir_datasets import pandas as pd import polars as pl +import s3fs from pyarrow.parquet import ParquetFile from pydantic import Field as PydanticField from pydantic import PrivateAttr, field_validator +from tqdm import tqdm from vectordb_bench import config from vectordb_bench.base import BaseModel @@ -263,6 +266,28 @@ class SIFT(BaseDataset): } +class SIFTBinary(BaseDataset): + name: str = "SIFTBinary" + dim: int = 128 + metric_type: MetricType = MetricType.HAMMING + use_shuffled: bool = False + with_gt: bool = True + remote_path: str = "assets.zilliz.com/nightly/v1/siftbinary_1m" + test_file: str = "query.bin" + gt_file: str = "truth.ibin" + _size_label: ClassVar[dict[int, SizeLabel]] = { + 1_000_000: SizeLabel(1_000_000, "1M", 1), + } + + @property + def dir_name(self) -> str: + return "siftbinary_1m" + + @property + def train_files(self) -> list[str]: + return ["base.bin"] + + class OpenAI(BaseDataset): name: str = "OpenAI" dim: int = 1536 @@ -464,6 +489,151 @@ def __next__(self) -> pd.DataFrame: raise StopIteration +class BinaryDatasetManager(DatasetManager): + """Dataset manager for vecTool-style packed binary vector files.""" + + data: SIFTBinary + + def __iter__(self): + return BinaryDataSetIterator(self) + + def iter_batches(self, batch_size: int): + return BinaryDataSetIterator(self, batch_size=batch_size) + + def prepare( + self, + source: DatasetSource = DatasetSource.S3, + filters: Filter = non_filter, + with_train_files: bool = True, + with_scalar_labels: bool = False, + ) -> bool: + if source != DatasetSource.S3: + msg = f"{self.data.name} is currently hosted under assets.zilliz.com/nightly and supports S3 only" + raise ValueError(msg) + + self.train_files = self.data.train_files if with_train_files else [] + download_files = [*self.train_files, self.data.test_file, self.data.gt_file, "manifest.json"] + self._download_files(download_files) + + self.test_data = self._read_binary_vectors(self.data.test_file) + self.gt_data = self._read_truth_ids(self.data.gt_file) + log.debug(f"{self.data.name}: available train files {self.train_files}") + return True + + def _download_files(self, files: list[str]) -> None: + self.data_dir.mkdir(parents=True, exist_ok=True) + fs = s3fs.S3FileSystem(anon=True, client_kwargs={"region_name": "us-west-2"}) + downloads = [] + for file in files: + remote_file = pathlib.PurePosixPath(self.data.remote_path, file) + local_file = self.data_dir.joinpath(file) + if not local_file.exists(): + downloads.append((remote_file, local_file)) + continue + + remote_size = fs.info(remote_file.as_posix()).get("size") + if remote_size != local_file.stat().st_size: + downloads.append((remote_file, local_file)) + + if not downloads: + return + + log.info(f"Start downloading binary dataset files, total count: {len(downloads)}") + for remote_file, local_file in tqdm(downloads): + fs.download(remote_file.as_posix(), local_file.as_posix()) + + def _read_binary_header(self, file_name: str) -> tuple[int, int]: + path = self.data_dir.joinpath(file_name) + with path.open("rb") as fp: + header = fp.read(8) + if len(header) != 8: + msg = f"Invalid binary vector file header: {path}" + raise ValueError(msg) + rows, dim = struct.unpack(" list[str]: + path = self.data_dir.joinpath(file_name) + rows, dim = self._read_binary_header(file_name) + bytes_per_vector = dim // 8 + with path.open("rb") as fp: + fp.seek(8) + raw = fp.read() + expected = rows * bytes_per_vector + if len(raw) != expected: + msg = f"Unexpected {file_name} payload size={len(raw)}; expected {expected}" + raise ValueError(msg) + return [raw[i : i + bytes_per_vector].hex() for i in range(0, len(raw), bytes_per_vector)] + + def _read_truth_ids(self, file_name: str) -> list[list[int]]: + path = self.data_dir.joinpath(file_name) + with path.open("rb") as fp: + header = fp.read(8) + raw = fp.read() + if len(header) != 8: + msg = f"Invalid truth file header: {path}" + raise ValueError(msg) + nq, topk = struct.unpack(" pd.DataFrame: + if self._idx >= self._rows: + if self._fp is not None: + self._fp.close() + self._fp = None + raise StopIteration + + if self._fp is None: + path = self._ds.data_dir.joinpath("base.bin") + self._fp = path.open("rb") + self._fp.seek(8 + self._idx * self._bytes_per_vector) + + batch_rows = min(self._batch_size, self._rows - self._idx) + raw = self._fp.read(batch_rows * self._bytes_per_vector) + if len(raw) != batch_rows * self._bytes_per_vector: + msg = f"Unexpected EOF while reading {self._ds.data.name} at row {self._idx}" + raise ValueError(msg) + + start = self._idx + self._idx += batch_rows + vectors = [raw[i : i + self._bytes_per_vector].hex() for i in range(0, len(raw), self._bytes_per_vector)] + return pd.DataFrame( + { + self._ds.data.train_id_field: range(start, start + batch_rows), + self._ds.data.train_vector_field: vectors, + } + ) + + class Dataset(Enum): """ Value is Dataset classes, DO NOT use it @@ -479,13 +649,17 @@ class Dataset(Enum): BIOASQ = Bioasq GLOVE = Glove SIFT = SIFT + SIFT_BINARY = SIFTBinary OPENAI = OpenAI def get(self, size: int) -> BaseDataset: return self.value(size=size) def manager(self, size: int) -> DatasetManager: - return DatasetManager(data=self.get(size)) + data = self.get(size) + if isinstance(data, SIFTBinary): + return BinaryDatasetManager(data=data) + return DatasetManager(data=data) class DatasetWithSizeType(Enum): @@ -498,6 +672,7 @@ class DatasetWithSizeType(Enum): OpenAISmall = "Small OpenAI (1536dim, 50K)" OpenAIMedium = "Medium OpenAI (1536dim, 500K)" OpenAILarge = "Large OpenAI (1536dim, 5M)" + SIFTBinary1M = "Medium SIFT Binary (128bit, 1M)" def get_manager(self) -> DatasetManager: if self not in DatasetWithSizeMap: @@ -539,6 +714,7 @@ def get_optimize_timeout(self) -> float: DatasetWithSizeType.OpenAISmall: Dataset.OPENAI.manager(50_000), DatasetWithSizeType.OpenAIMedium: Dataset.OPENAI.manager(500_000), DatasetWithSizeType.OpenAILarge: Dataset.OPENAI.manager(5_000_000), + DatasetWithSizeType.SIFTBinary1M: Dataset.SIFT_BINARY.manager(1_000_000), } diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index 764dad2bc..76d02935c 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -2326,6 +2326,7 @@ class CaseConfigInput(BaseModel): ElasticCloudFtsConfig = [] VespaFtsConfig = [] +OSSOpenSearchFtsConfig = [] TurboPufferFtsConfig = [] WeaviateLoadConfig = [ @@ -3158,6 +3159,7 @@ class FilterType(Enum): DB.OSSOpenSearch: { CaseLabel.Load: OSSOpensearchLoadingConfig, CaseLabel.Performance: OSSOpenSearchPerformanceConfig, + CaseLabel.FullTextSearchPerformance: OSSOpenSearchFtsConfig, }, DB.PgVector: { CaseLabel.Load: PgVectorLoadingConfig,