diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index 5511f18db..e86821f0f 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -43,6 +43,9 @@ class IndexType(StrEnum): GPU_CAGRA = "GPU_CAGRA" SCANN = "scann" SCANN_MILVUS = "SCANN_MILVUS" + SVS_VAMANA = "SVS_VAMANA" + SVS_VAMANA_LVQ = "SVS_VAMANA_LVQ" + SVS_VAMANA_LEANVEC = "SVS_VAMANA_LEANVEC" Hologres_HGraph = "HGraph" Hologres_Graph = "Graph" NONE = "NONE" diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 2f2a286be..ae7269801 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -485,6 +485,169 @@ def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]): ) +class MilvusSVSVamanaTypedDict(CommonTypedDict, MilvusTypedDict): + svs_graph_max_degree: Annotated[ + int, + click.option( + "--svs-graph-max-degree", + type=int, + help="Maximum degree of the Vamana graph (4-256).", + required=True, + ), + ] + svs_construction_window_size: Annotated[ + int, + click.option( + "--svs-construction-window-size", + type=int, + help="Window size for graph construction.", + required=False, + default=40, + show_default=True, + ), + ] + svs_alpha: Annotated[ + float | None, + click.option( + "--svs-alpha", + type=float, + help="Pruning parameter (default: 1.2 for L2, 0.95 for IP/COSINE).", + required=False, + default=None, + ), + ] + svs_storage_kind: Annotated[ + str, + click.option( + "--svs-storage-kind", + type=click.Choice( + ["fp32", "fp16", "sqi8", "lvq4x0", "lvq4x4", "lvq4x8", "leanvec4x4", "leanvec4x8", "leanvec8x8"], + case_sensitive=False, + ), + help="Data storage format.", + required=False, + default="fp32", + show_default=True, + ), + ] + svs_search_window_size: Annotated[ + int | None, + click.option( + "--svs-search-window-size", + type=int, + help="Window size for search (1-10000).", + required=False, + default=None, + ), + ] + svs_search_buffer_capacity: Annotated[ + int | None, + click.option( + "--svs-search-buffer-capacity", + type=int, + help="Buffer capacity for search priority queue (1-10000).", + required=False, + default=None, + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(MilvusSVSVamanaTypedDict) +def MilvusSVSVamana(**parameters: Unpack[MilvusSVSVamanaTypedDict]): + from .config import MilvusConfig, SVSVamanaConfig + + 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"]), + replica_number=int(parameters["replica_number"]), + ), + db_case_config=SVSVamanaConfig( + svs_graph_max_degree=parameters["svs_graph_max_degree"], + svs_construction_window_size=parameters["svs_construction_window_size"], + svs_alpha=parameters["svs_alpha"], + svs_storage_kind=parameters["svs_storage_kind"], + svs_search_window_size=parameters["svs_search_window_size"], + svs_search_buffer_capacity=parameters["svs_search_buffer_capacity"], + ), + **parameters, + ) + + +@cli.command() +@click_parameter_decorators_from_typed_dict(MilvusSVSVamanaTypedDict) +def MilvusSVSVamanaLVQ(**parameters: Unpack[MilvusSVSVamanaTypedDict]): + from .config import MilvusConfig, SVSVamanaLVQConfig + + 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"]), + replica_number=int(parameters["replica_number"]), + ), + db_case_config=SVSVamanaLVQConfig( + svs_graph_max_degree=parameters["svs_graph_max_degree"], + svs_construction_window_size=parameters["svs_construction_window_size"], + svs_alpha=parameters["svs_alpha"], + svs_storage_kind=parameters["svs_storage_kind"], + svs_search_window_size=parameters["svs_search_window_size"], + svs_search_buffer_capacity=parameters["svs_search_buffer_capacity"], + ), + **parameters, + ) + + +class MilvusSVSVamanaLeanVecTypedDict(MilvusSVSVamanaTypedDict): + svs_leanvec_dim: Annotated[ + int, + click.option( + "--svs-leanvec-dim", + type=int, + help="Dimensionality for LeanVec compression (0 = d/2).", + required=False, + default=0, + show_default=True, + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(MilvusSVSVamanaLeanVecTypedDict) +def MilvusSVSVamanaLeanVec(**parameters: Unpack[MilvusSVSVamanaLeanVecTypedDict]): + from .config import MilvusConfig, SVSVamanaLeanVecConfig + + 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"]), + replica_number=int(parameters["replica_number"]), + ), + db_case_config=SVSVamanaLeanVecConfig( + svs_graph_max_degree=parameters["svs_graph_max_degree"], + svs_construction_window_size=parameters["svs_construction_window_size"], + svs_alpha=parameters["svs_alpha"], + svs_storage_kind=parameters["svs_storage_kind"], + svs_search_window_size=parameters["svs_search_window_size"], + svs_search_buffer_capacity=parameters["svs_search_buffer_capacity"], + svs_leanvec_dim=parameters["svs_leanvec_dim"], + ), + **parameters, + ) + + class MilvusGPUIVFPQTypedDict( CommonTypedDict, MilvusTypedDict, diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 9ffbdcece..c2cd1f9b2 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -441,6 +441,65 @@ def search_param(self) -> dict: } +class SVSVamanaConfig(MilvusIndexConfig, DBCaseConfig): + svs_graph_max_degree: int + svs_construction_window_size: int = 40 + svs_alpha: float | None = None + svs_storage_kind: str = "fp32" + svs_search_window_size: int | None = None + svs_search_buffer_capacity: int | None = None + index: IndexType = IndexType.SVS_VAMANA + + def index_param(self) -> dict: + params = { + "svs_graph_max_degree": self.svs_graph_max_degree, + "svs_construction_window_size": self.svs_construction_window_size, + "svs_storage_kind": self.svs_storage_kind, + } + if self.svs_alpha is not None: + params["svs_alpha"] = self.svs_alpha + return { + "metric_type": self.parse_metric(), + "index_type": self.index.value, + "params": params, + } + + def search_param(self) -> dict: + return { + "metric_type": self.parse_metric(), + "params": { + "svs_search_window_size": self.svs_search_window_size, + "svs_search_buffer_capacity": self.svs_search_buffer_capacity, + }, + } + + +class SVSVamanaLVQConfig(SVSVamanaConfig): + svs_storage_kind: str = "lvq4x4" + index: IndexType = IndexType.SVS_VAMANA_LVQ + + +class SVSVamanaLeanVecConfig(SVSVamanaConfig): + svs_storage_kind: str = "leanvec4x4" + svs_leanvec_dim: int = 0 + index: IndexType = IndexType.SVS_VAMANA_LEANVEC + + def index_param(self) -> dict: + params = { + "svs_graph_max_degree": self.svs_graph_max_degree, + "svs_construction_window_size": self.svs_construction_window_size, + "svs_storage_kind": self.svs_storage_kind, + "svs_leanvec_dim": self.svs_leanvec_dim, + } + if self.svs_alpha is not None: + params["svs_alpha"] = self.svs_alpha + return { + "metric_type": self.parse_metric(), + "index_type": self.index.value, + "params": params, + } + + _milvus_case_config = { IndexType.AUTOINDEX: AutoIndexConfig, IndexType.HNSW: HNSWConfig, @@ -458,4 +517,7 @@ def search_param(self) -> dict: IndexType.GPU_CAGRA: GPUCAGRAConfig, IndexType.GPU_BRUTE_FORCE: GPUBruteForceConfig, IndexType.SCANN_MILVUS: SCANNConfig, + IndexType.SVS_VAMANA: SVSVamanaConfig, + IndexType.SVS_VAMANA_LVQ: SVSVamanaLVQConfig, + IndexType.SVS_VAMANA_LEANVEC: SVSVamanaLeanVecConfig, }