vectordb-bench 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

vectordb_bench/backend/cases.py

@@ -4,7 +4,7 @@ from enum import Enum, auto
 
 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 @@ class StreamingPerformanceCase(Case):
         )
 
 
+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 @@ type2case = {
     CaseType.Performance1536D50K: Performance1536D50K,
     CaseType.PerformanceCustomDataset: PerformanceCustomDataset,
     CaseType.StreamingPerformanceCase: StreamingPerformanceCase,
+    CaseType.NewIntFilterPerformanceCase: NewIntFilterPerformanceCase,
     CaseType.LabelFilterPerformanceCase: LabelFilterPerformanceCase,
 }

vectordb_bench/backend/clients/__init__.py

@@ -38,6 +38,7 @@ class DB(Enum):
     MemoryDB = "MemoryDB"
     Chroma = "Chroma"
     AWSOpenSearch = "OpenSearch"
+    OSSOpenSearch = "OSSOpenSearch"
     AliyunElasticsearch = "AliyunElasticsearch"
     MariaDB = "MariaDB"
     Test = "test"
@@ -48,6 +49,7 @@ class DB(Enum):
     Vespa = "Vespa"
     LanceDB = "LanceDB"
     OceanBase = "OceanBase"
+    S3Vectors = "S3Vectors"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -127,6 +129,11 @@ class DB(Enum):
 
             return AWSOpenSearch
 
+        if self == DB.OSSOpenSearch:
+            from .oss_opensearch.oss_opensearch import OSSOpenSearch
+
+            return OSSOpenSearch
+
         if self == DB.Clickhouse:
             from .clickhouse.clickhouse import Clickhouse
 
@@ -182,6 +189,11 @@ class DB(Enum):
 
             return LanceDB
 
+        if self == DB.S3Vectors:
+            from .s3_vectors.s3_vectors import S3Vectors
+
+            return S3Vectors
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -263,6 +275,11 @@ class DB(Enum):
 
             return AWSOpenSearchConfig
 
+        if self == DB.OSSOpenSearch:
+            from .oss_opensearch.config import OSSOpenSearchConfig
+
+            return OSSOpenSearchConfig
+
         if self == DB.Clickhouse:
             from .clickhouse.config import ClickhouseConfig
 
@@ -318,6 +335,11 @@ class DB(Enum):
 
             return LanceDBConfig
 
+        if self == DB.S3Vectors:
+            from .s3_vectors.config import S3VectorsConfig
+
+            return S3VectorsConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -370,6 +392,11 @@ class DB(Enum):
 
             return AWSOpenSearchIndexConfig
 
+        if self == DB.OSSOpenSearch:
+            from .oss_opensearch.config import OSSOpenSearchIndexConfig
+
+            return OSSOpenSearchIndexConfig
+
         if self == DB.Clickhouse:
             from .clickhouse.config import ClickhouseHNSWConfig
 
@@ -430,6 +457,11 @@ class DB(Enum):
 
             return _lancedb_case_config.get(index_type)
 
+        if self == DB.S3Vectors:
+            from .s3_vectors.config import S3VectorsIndexConfig
+
+            return S3VectorsIndexConfig
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

vectordb_bench/backend/clients/aliyun_opensearch/aliyun_opensearch.py

@@ -3,11 +3,9 @@ import logging
 import time
 from contextlib import contextmanager
 
-from alibabacloud_ha3engine_vector import client, models
+from alibabacloud_ha3engine_vector import models
+from alibabacloud_ha3engine_vector.client import Client
 from alibabacloud_ha3engine_vector.models import QueryRequest
-from alibabacloud_searchengine20211025 import models as searchengine_models
-from alibabacloud_searchengine20211025.client import Client as searchengineClient
-from alibabacloud_tea_openapi import models as open_api_models
 
 from ..api import MetricType, VectorDB
 from .config import AliyunOpenSearchIndexConfig
@@ -28,18 +26,25 @@ class AliyunOpenSearch(VectorDB):
         drop_old: bool = False,
         **kwargs,
     ):
-        self.control_client = None
         self.dim = dim
         self.db_config = db_config
         self.case_config = db_case_config
         self.collection_name = collection_name
         self.instance_id = db_config["host"].split(".")[0].replace("http://", "").replace("https://", "")
+        self.config = models.Config(
+            endpoint=self.db_config["host"],
+            protocol="http",
+            access_user_name=self.db_config["user"],
+            access_pass_word=self.db_config["password"],
+        )
 
         self._primary_field = "id"
         self._scalar_field = "int_id"
         self._vector_field = "vector"
         self._index_name = "vector_idx"
 
+        client = Client(self.config)
+
         self.batch_size = int(
             min(
                 ALIYUN_OPENSEARCH_MAX_SIZE_PER_BATCH / (dim * 25),
@@ -48,22 +53,16 @@ class AliyunOpenSearch(VectorDB):
         )
 
         log.info(f"Aliyun_OpenSearch client config: {self.db_config}")
-        control_config = open_api_models.Config(
-            access_key_id=self.db_config["ak"],
-            access_key_secret=self.db_config["sk"],
-            endpoint=self.db_config["control_host"],
-        )
-        self.control_client = searchengineClient(control_config)
 
         if drop_old:
             log.info(f"aliyun_OpenSearch client drop old index: {self.collection_name}")
-            if self._index_exists(self.control_client):
-                self._modify_index(self.control_client)
+            if self._index_exists(client):
+                self._modify_index(client)
             else:
-                self._create_index(self.control_client)
+                self._create_index(client)
 
-    def _create_index(self, client: searchengineClient):
-        create_table_request = searchengine_models.CreateTableRequest()
+    def _create_index(self, client: Client):
+        create_table_request = models.CreateTableRequest()
         create_table_request.name = self.collection_name
         create_table_request.primary_key = self._primary_field
         create_table_request.partition_count = 1
@@ -72,14 +71,14 @@ class AliyunOpenSearch(VectorDB):
             self._vector_field: "MULTI_FLOAT",
             self._scalar_field: "INT64",
         }
-        vector_index = searchengine_models.ModifyTableRequestVectorIndex()
+        vector_index = models.ModifyTableRequestVectorIndex()
         vector_index.index_name = self._index_name
         vector_index.dimension = self.dim
         vector_index.distance_type = self.case_config.distance_type()
         vector_index.vector_field = self._vector_field
         vector_index.vector_index_type = "HNSW"
 
-        advance_params = searchengine_models.ModifyTableRequestVectorIndexAdvanceParams()
+        advance_params = models.ModifyTableRequestVectorIndexAdvanceParams()
         str_max_neighbor_count = f'"proxima.hnsw.builder.max_neighbor_count":{self.case_config.M}'
         str_efc = f'"proxima.hnsw.builder.efconstruction":{self.case_config.ef_construction}'
         str_enable_adsampling = '"proxima.hnsw.builder.enable_adsampling":true'
@@ -95,7 +94,7 @@ class AliyunOpenSearch(VectorDB):
                 str_thread_count,
             ],
         )
-        advance_params.build_index_params = params
+        advance_params.build_index_params = "{" + params + "}"
         advance_params.search_index_params = (
             '{"proxima.hnsw.searcher.ef":400,"proxima.hnsw.searcher.dynamic_termination.prob_threshold":0.7}'
         )
@@ -103,7 +102,7 @@ class AliyunOpenSearch(VectorDB):
         create_table_request.vector_index = [vector_index]
 
         try:
-            response = client.create_table(self.instance_id, create_table_request)
+            response = client.create_table(create_table_request)
             log.info(f"create table success: {response.body}")
         except Exception as error:
             log.info(error.message)
@@ -115,20 +114,20 @@ class AliyunOpenSearch(VectorDB):
         self._active_index(client)
 
     # check if index create success
-    def _active_index(self, client: searchengineClient) -> None:
+    def _active_index(self, client: Client) -> None:
         retry_times = 0
         while True:
             time.sleep(10)
             log.info(f"begin to {retry_times} times get table")
             retry_times += 1
-            response = client.get_table(self.instance_id, self.collection_name)
+            response = client.get_table(self.collection_name)
             if response.body.result.status == "IN_USE":
                 log.info(f"{self.collection_name} table begin to use.")
                 return
 
-    def _index_exists(self, client: searchengineClient) -> bool:
+    def _index_exists(self, client: Client) -> bool:
         try:
-            client.get_table(self.instance_id, self.collection_name)
+            client.get_table(self.collection_name)
         except Exception as err:
             log.warning(f"get table from searchengine error, err={err}")
             return False
@@ -136,7 +135,7 @@ class AliyunOpenSearch(VectorDB):
             return True
 
     # check if index build success, Insert the embeddings to the vector database after index build success
-    def _index_build_success(self, client: searchengineClient) -> None:
+    def _index_build_success(self, client: Client) -> None:
         log.info("begin to check if table build success.")
         time.sleep(50)
 
@@ -145,10 +144,10 @@ class AliyunOpenSearch(VectorDB):
             time.sleep(10)
             log.info(f"begin to {retry_times} times get table fsm")
             retry_times += 1
-            request = searchengine_models.ListTasksRequest()
-            request.start = (int(time.time()) - 3600) * 1000
-            request.end = int(time.time()) * 1000
-            response = client.list_tasks(self.instance_id, request)
+            request = models.ListTasksRequest()
+            request.start = int(time.time()) - 3600
+            request.end = int(time.time())
+            response = client.list_tasks(request)
             fsms = response.body.result
             cur_fsm = None
             for fsm in fsms:
@@ -164,11 +163,11 @@ class AliyunOpenSearch(VectorDB):
             if cur_fsm["status"] == "success":
                 return
 
-    def _modify_index(self, client: searchengineClient) -> None:
+    def _modify_index(self, client: Client) -> None:
         # check if index create success
         self._active_index(client)
 
-        modify_table_request = searchengine_models.ModifyTableRequest()
+        modify_table_request = models.ModifyTableRequest()
         modify_table_request.partition_count = 1
         modify_table_request.primary_key = self._primary_field
         modify_table_request.field_schema = {
@@ -176,13 +175,13 @@ class AliyunOpenSearch(VectorDB):
             self._vector_field: "MULTI_FLOAT",
             self._scalar_field: "INT64",
         }
-        vector_index = searchengine_models.ModifyTableRequestVectorIndex()
+        vector_index = models.ModifyTableRequestVectorIndex()
         vector_index.index_name = self._index_name
         vector_index.dimension = self.dim
         vector_index.distance_type = self.case_config.distance_type()
         vector_index.vector_field = self._vector_field
         vector_index.vector_index_type = "HNSW"
-        advance_params = searchengine_models.ModifyTableRequestVectorIndexAdvanceParams()
+        advance_params = models.ModifyTableRequestVectorIndexAdvanceParams()
 
         str_max_neighbor_count = f'"proxima.hnsw.builder.max_neighbor_count":{self.case_config.M}'
         str_efc = f'"proxima.hnsw.builder.efconstruction":{self.case_config.ef_construction}'
@@ -199,7 +198,7 @@ class AliyunOpenSearch(VectorDB):
                 str_thread_count,
             ],
         )
-        advance_params.build_index_params = params
+        advance_params.build_index_params = "{" + params + "}"
         advance_params.search_index_params = (
             '{"proxima.hnsw.searcher.ef":400,"proxima.hnsw.searcher.dynamic_termination.prob_threshold":0.7}'
         )
@@ -209,7 +208,6 @@ class AliyunOpenSearch(VectorDB):
 
         try:
             response = client.modify_table(
-                self.instance_id,
                 self.collection_name,
                 modify_table_request,
             )
@@ -240,14 +238,8 @@ class AliyunOpenSearch(VectorDB):
     @contextmanager
     def init(self) -> None:
         """connect to aliyun opensearch"""
-        config = models.Config(
-            endpoint=self.db_config["host"],
-            protocol="http",
-            access_user_name=self.db_config["user"],
-            access_pass_word=self.db_config["password"],
-        )
 
-        self.client = client.Client(config)
+        self.client = Client(self.config)
 
         yield
         self.client = None

vectordb_bench/backend/clients/aliyun_opensearch/config.py

@@ -12,18 +12,11 @@ class AliyunOpenSearchConfig(DBConfig, BaseModel):
     user: str = ""
     password: SecretStr = ""
 
-    ak: str = ""
-    sk: SecretStr = ""
-    control_host: str = "searchengine.cn-hangzhou.aliyuncs.com"
-
     def to_dict(self) -> dict:
         return {
             "host": self.host,
             "user": self.user,
             "password": self.password.get_secret_value(),
-            "ak": self.ak,
-            "sk": self.sk.get_secret_value(),
-            "control_host": self.control_host,
         }
 
 

vectordb_bench/backend/clients/milvus/cli.py

@@ -109,6 +109,157 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
     )
 
 
+class MilvusRefineTypedDict(TypedDict):
+    refine: Annotated[
+        bool,
+        click.option(
+            "--refine",
+            type=bool,
+            required=True,
+            help="Whether refined data is reserved during index building.",
+        ),
+    ]
+    refine_type: Annotated[
+        str | None,
+        click.option(
+            "--refine-type",
+            type=click.Choice(["SQ6", "SQ8", "BF16", "FP16", "FP32"], case_sensitive=False),
+            help="The data type of the refine index to use. Supported values: SQ6,SQ8,BF16,FP16,FP32",
+            required=True,
+        ),
+    ]
+    refine_k: Annotated[
+        float,
+        click.option(
+            "--refine-k",
+            type=float,
+            help="The magnification factor of refine compared to k.",
+            required=True,
+        ),
+    ]
+
+
+class MilvusHNSWPQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusHNSWTypedDict, MilvusRefineTypedDict):
+    nbits: Annotated[
+        int,
+        click.option(
+            "--nbits",
+            type=int,
+            required=True,
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusHNSWPQTypedDict)
+def MilvusHNSWPQ(**parameters: Unpack[MilvusHNSWPQTypedDict]):
+    from .config import HNSWPQConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
+        ),
+        db_case_config=HNSWPQConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_search"],
+            nbits=parameters["nbits"],
+            refine=parameters["refine"],
+            refine_type=parameters["refine_type"],
+            refine_k=parameters["refine_k"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusHNSWPRQTypedDict(
+    CommonTypedDict,
+    MilvusTypedDict,
+    MilvusHNSWPQTypedDict,
+):
+    nrq: Annotated[
+        int,
+        click.option(
+            "--nrq",
+            type=int,
+            help="The number of residual subquantizers.",
+            required=True,
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusHNSWPRQTypedDict)
+def MilvusHNSWPRQ(**parameters: Unpack[MilvusHNSWPRQTypedDict]):
+    from .config import HNSWPRQConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
+        ),
+        db_case_config=HNSWPRQConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_search"],
+            nbits=parameters["nbits"],
+            refine=parameters["refine"],
+            refine_type=parameters["refine_type"],
+            refine_k=parameters["refine_k"],
+            nrq=parameters["nrq"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusHNSWSQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusHNSWTypedDict, MilvusRefineTypedDict):
+    sq_type: Annotated[
+        str | None,
+        click.option(
+            "--sq-type",
+            type=click.Choice(["SQ6", "SQ8", "BF16", "FP16", "FP32"], case_sensitive=False),
+            help="Scalar quantizer type. Supported values: SQ6,SQ8,BF16,FP16,FP32",
+            required=True,
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusHNSWSQTypedDict)
+def MilvusHNSWSQ(**parameters: Unpack[MilvusHNSWSQTypedDict]):
+    from .config import HNSWSQConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
+        ),
+        db_case_config=HNSWSQConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_search"],
+            sq_type=parameters["sq_type"],
+            refine=parameters["refine"],
+            refine_type=parameters["refine_type"],
+            refine_k=parameters["refine_k"],
+        ),
+        **parameters,
+    )
+
+
 class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN): ...
 
 
@@ -156,6 +307,71 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
     )
 
 
+class MilvusIVFRABITQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict):
+    rbq_bits_query: Annotated[
+        int,
+        click.option(
+            "--rbq-bits-query",
+            type=int,
+            help="The magnification factor of refine compared to k.",
+            required=True,
+        ),
+    ]
+    refine: Annotated[
+        bool,
+        click.option(
+            "--refine",
+            type=bool,
+            required=True,
+            help="Whether refined data is reserved during index building.",
+        ),
+    ]
+    refine_type: Annotated[
+        str | None,
+        click.option(
+            "--refine-type",
+            type=click.Choice(["SQ6", "SQ8", "BF16", "FP16", "FP32"], case_sensitive=False),
+            help="The data type of the refine index to use. Supported values: SQ6,SQ8,BF16,FP16,FP32",
+            required=True,
+        ),
+    ]
+    refine_k: Annotated[
+        float,
+        click.option(
+            "--refine-k",
+            type=float,
+            help="The magnification factor of refine compared to k.",
+            required=True,
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusIVFRABITQTypedDict)
+def MilvusIVFRabitQ(**parameters: Unpack[MilvusIVFRABITQTypedDict]):
+    from .config import IVFRABITQConfig, MilvusConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
+        ),
+        db_case_config=IVFRABITQConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+            rbq_bits_query=parameters["rbq_bits_query"],
+            refine=parameters["refine"],
+            refine_type=parameters["refine_type"],
+            refine_k=parameters["refine_k"],
+        ),
+        **parameters,
+    )
+
+
 class MilvusDISKANNTypedDict(CommonTypedDict, MilvusTypedDict):
     search_list: Annotated[str, click.option("--search-list", type=int, required=True)]