diff --git a/vectordb_bench/backend/clients/pgvectorscale/cli.py b/vectordb_bench/backend/clients/pgvectorscale/cli.py new file mode 100644 index 000000000..e5a161c6b --- /dev/null +++ b/vectordb_bench/backend/clients/pgvectorscale/cli.py @@ -0,0 +1,108 @@ +import click +import os +from pydantic import SecretStr + +from ....cli.cli import ( + CommonTypedDict, + cli, + click_parameter_decorators_from_typed_dict, + run, +) +from typing import Annotated, Unpack +from vectordb_bench.backend.clients import DB + + +class PgVectorScaleTypedDict(CommonTypedDict): + user_name: Annotated[ + str, click.option("--user-name", type=str, help="Db username", required=True) + ] + password: Annotated[ + str, + click.option("--password", + type=str, + help="Postgres database password", + default=lambda: os.environ.get("POSTGRES_PASSWORD", ""), + show_default="$POSTGRES_PASSWORD", + ), + ] + + host: Annotated[ + str, click.option("--host", type=str, help="Db host", required=True) + ] + db_name: Annotated[ + str, click.option("--db-name", type=str, help="Db name", required=True) + ] + + +class PgVectorScaleDiskAnnTypedDict(PgVectorScaleTypedDict): + storage_layout: Annotated[ + str, + click.option( + "--storage-layout", type=str, help="Streaming DiskANN storage layout", + ), + ] + num_neighbors: Annotated[ + int, + click.option( + "--num-neighbors", type=int, help="Streaming DiskANN num neighbors", + ), + ] + search_list_size: Annotated[ + int, + click.option( + "--search-list-size", type=int, help="Streaming DiskANN search list size", + ), + ] + max_alpha: Annotated[ + float, + click.option( + "--max-alpha", type=float, help="Streaming DiskANN max alpha", + ), + ] + num_dimensions: Annotated[ + int, + click.option( + "--num-dimensions", type=int, help="Streaming DiskANN num dimensions", + ), + ] + query_search_list_size: Annotated[ + int, + click.option( + "--query-search-list-size", type=int, help="Streaming DiskANN query search list size", + ), + ] + query_rescore: Annotated[ + int, + click.option( + "--query-rescore", type=int, help="Streaming DiskANN query rescore", + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(PgVectorScaleDiskAnnTypedDict) +def PgVectorScaleDiskAnn( + **parameters: Unpack[PgVectorScaleDiskAnnTypedDict], +): + from .config import PgVectorScaleConfig, PgVectorScaleStreamingDiskANNConfig + + run( + db=DB.PgVectorScale, + db_config=PgVectorScaleConfig( + db_label=parameters["db_label"], + user_name=SecretStr(parameters["user_name"]), + password=SecretStr(parameters["password"]), + host=parameters["host"], + db_name=parameters["db_name"], + ), + db_case_config=PgVectorScaleStreamingDiskANNConfig( + storage_layout=parameters["storage_layout"], + num_neighbors=parameters["num_neighbors"], + search_list_size=parameters["search_list_size"], + max_alpha=parameters["max_alpha"], + num_dimensions=parameters["num_dimensions"], + query_search_list_size=parameters["query_search_list_size"], + query_rescore=parameters["query_rescore"], + ), + **parameters, + ) \ No newline at end of file diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index ab7eff44c..e62c25a3d 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -1,5 +1,6 @@ from ..backend.clients.pgvector.cli import PgVectorHNSW from ..backend.clients.pgvecto_rs.cli import PgVectoRSHNSW, PgVectoRSIVFFlat +from ..backend.clients.pgvectorscale.cli import PgVectorScaleDiskAnn from ..backend.clients.redis.cli import Redis from ..backend.clients.memorydb.cli import MemoryDB from ..backend.clients.test.cli import Test @@ -8,7 +9,6 @@ from ..backend.clients.milvus.cli import MilvusAutoIndex from ..backend.clients.aws_opensearch.cli import AWSOpenSearch - from .cli import cli cli.add_command(PgVectorHNSW) @@ -21,6 +21,7 @@ cli.add_command(ZillizAutoIndex) cli.add_command(MilvusAutoIndex) cli.add_command(AWSOpenSearch) +cli.add_command(PgVectorScaleDiskAnn) if __name__ == "__main__":