diff --git a/vectordb_bench/backend/clients/vectorchord/cli.py b/vectordb_bench/backend/clients/vectorchord/cli.py index 2b3e3862b..99723959c 100644 --- a/vectordb_bench/backend/clients/vectorchord/cli.py +++ b/vectordb_bench/backend/clients/vectorchord/cli.py @@ -14,6 +14,16 @@ ) +def _parse_int_list(_ctx: click.Context, _param: click.Parameter, value: str | None) -> list[int] | None: + if value is None or value == "": + return None + try: + return [int(v.strip()) for v in value.split(",") if v.strip()] + except ValueError as e: + msg = f"expected comma-separated integers, got {value!r}" + raise click.BadParameter(msg) from e + + class VectorChordTypedDict(CommonTypedDict): user_name: Annotated[ str, @@ -66,20 +76,25 @@ class VectorChordTypedDict(CommonTypedDict): class VectorChordRQTypedDict(VectorChordTypedDict): lists: Annotated[ - int | None, + list[int] | None, click.option( "--lists", - type=int, - help="Number of IVF lists for vchordrq index", + type=str, + callback=_parse_int_list, + help=( + "Comma-separated IVF list sizes for vchordrq index. " + "Single value for single-level (e.g. '1000') or multiple for multi-level IVF (e.g. '4096,128')" + ), ), ] probes: Annotated[ - int | None, + list[int] | None, click.option( "--probes", - type=int, - help="Number of probes during search", - default=10, + type=str, + callback=_parse_int_list, + help="Comma-separated probes per IVF level during search (e.g. '10' or '10,1')", + default="10", show_default=True, ), ] diff --git a/vectordb_bench/backend/clients/vectorchord/config.py b/vectordb_bench/backend/clients/vectorchord/config.py index 95916eb06..1c9d97a01 100644 --- a/vectordb_bench/backend/clients/vectorchord/config.py +++ b/vectordb_bench/backend/clients/vectorchord/config.py @@ -94,13 +94,13 @@ class VectorChordRQConfig(VectorChordIndexConfig): rerank_in_table: bool = False degree_of_parallelism: int | None = None # default 32, range [1, 256] # Build parameters ([build.internal] section) - lists: int | None = None + lists: list[int] | None = None # e.g. [1000] for single-level or [4096, 128] for multi-level IVF spherical_centroids: bool = False build_threads: int | None = None # range [1, 255] # PostgreSQL tuning parameter max_parallel_workers: int | None = None # sets max_parallel_workers & max_parallel_maintenance_workers # Search parameters (GUCs) - probes: int | None = 10 + probes: list[int] | None = [10] # one probe per IVF level, e.g. [10] or [10, 1] epsilon: float | None = 1.9 # range [0.0, 4.0] max_scan_tuples: int | None = None # default -1, range [-1, 2147483647] @@ -114,7 +114,7 @@ def index_param(self) -> dict: options_parts.append(f"degree_of_parallelism = {self.degree_of_parallelism}") options_parts.append("[build.internal]") if self.lists is not None: - options_parts.append(f"lists = [{self.lists}]") + options_parts.append(f"lists = [{', '.join(str(v) for v in self.lists)}]") if self.spherical_centroids: options_parts.append("spherical_centroids = true") if self.build_threads is not None: @@ -136,7 +136,7 @@ def search_param(self) -> dict: def session_param(self) -> dict: params = {} if self.probes is not None: - params["vchordrq.probes"] = str(self.probes) + params["vchordrq.probes"] = ",".join(str(v) for v in self.probes) if self.epsilon is not None: params["vchordrq.epsilon"] = str(self.epsilon) if self.max_scan_tuples is not None: