Add CLI Support for pgvectorscale

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,
+    )

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__":