vectordb-bench 0.0.10__py3-none-any.whl → 0.0.11__py3-none-any.whl

vectordb_bench/__init__.py

@@ -1,11 +1,13 @@
-import environs
 import inspect
 import pathlib
-from . import log_util
 
+import environs
+
+from . import log_util
 
 env = environs.Env()
-env.read_env(".env")
+env.read_env(".env", False)
+
 
 class config:
     ALIYUN_OSS_URL = "assets.zilliz.com.cn/benchmark/"
@@ -19,9 +21,20 @@ class config:
 
     DROP_OLD = env.bool("DROP_OLD", True)
     USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)
-    NUM_CONCURRENCY = [1, 5, 10, 15, 20, 25, 30, 35]
 
-    RESULTS_LOCAL_DIR = pathlib.Path(__file__).parent.joinpath("results")
+    NUM_CONCURRENCY = env.list("NUM_CONCURRENCY",  [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90, 95, 100], subcast=int )
+
+    CONCURRENCY_DURATION = 30
+
+    RESULTS_LOCAL_DIR = env.path(
+        "RESULTS_LOCAL_DIR", pathlib.Path(__file__).parent.joinpath("results")
+    )
+    CONFIG_LOCAL_DIR = env.path(
+        "CONFIG_LOCAL_DIR", pathlib.Path(__file__).parent.joinpath("config-files")
+    )
+
+
+    K_DEFAULT = 100  # default return top k nearest neighbors during search
 
     CAPACITY_TIMEOUT_IN_SECONDS = 24 * 3600 # 24h
     LOAD_TIMEOUT_DEFAULT        = 2.5 * 3600 # 2.5h

vectordb_bench/backend/cases.py

@@ -1,6 +1,7 @@
 import typing
 import logging
 from enum import Enum, auto
+from typing import Type
 
 from vectordb_bench import config
 from vectordb_bench.base import BaseModel
@@ -10,8 +11,6 @@ from .dataset import Dataset, DatasetManager
 
 log = logging.getLogger(__name__)
 
-Case = typing.TypeVar("Case")
-
 
 class CaseType(Enum):
     """
@@ -42,11 +41,15 @@ class CaseType(Enum):
     Performance1536D500K99P = 14
     Performance1536D5M99P = 15
 
+    Performance1536D50K = 50
+
     Custom = 100
 
     @property
-    def case_cls(self, custom_configs: dict | None = None) -> Case:
-        return type2case.get(self)
+    def case_cls(self, custom_configs: dict | None = None) -> Type["Case"]:
+        if self not in type2case:
+            raise NotImplementedError(f"Case {self} has not implemented. You can add it manually to vectordb_bench.backend.cases.type2case or define a custom_configs['custom_cls']")
+        return type2case[self]
 
     @property
     def case_name(self) -> str:
@@ -69,7 +72,7 @@ class CaseLabel(Enum):
 
 
 class Case(BaseModel):
-    """Undifined case
+    """Undefined case
 
     Fields:
         case_id(CaseType): default 9 case type plus one custom cases.
@@ -86,9 +89,9 @@ class Case(BaseModel):
     dataset: DatasetManager
 
     load_timeout: float | int
-    optimize_timeout: float | int | None
+    optimize_timeout: float | int | None = None
 
-    filter_rate: float | None
+    filter_rate: float | None = None
 
     @property
     def filters(self) -> dict | None:
@@ -115,20 +118,23 @@ class PerformanceCase(Case, BaseModel):
     load_timeout: float | int = config.LOAD_TIMEOUT_DEFAULT
     optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_DEFAULT
 
+
 class CapacityDim960(CapacityCase):
     case_id: CaseType = CaseType.CapacityDim960
     dataset: DatasetManager = Dataset.GIST.manager(100_000)
     name: str = "Capacity Test (960 Dim Repeated)"
-    description: str = """This case tests the vector database's loading capacity by repeatedly inserting large-dimension vectors (GIST 100K vectors, <b>960 dimensions</b>) until it is fully loaded.
-Number of inserted vectors will be reported."""
+    description: str = """This case tests the vector database's loading capacity by repeatedly inserting large-dimension 
+     vectors (GIST 100K vectors, <b>960 dimensions</b>) until it is fully loaded. Number of inserted vectors will be 
+     reported."""
 
 
 class CapacityDim128(CapacityCase):
     case_id: CaseType = CaseType.CapacityDim128
     dataset: DatasetManager = Dataset.SIFT.manager(500_000)
     name: str = "Capacity Test (128 Dim Repeated)"
-    description: str = """This case tests the vector database's loading capacity by repeatedly inserting small-dimension vectors (SIFT 100K vectors, <b>128 dimensions</b>) until it is fully loaded.
-Number of inserted vectors will be reported."""
+    description: str = """This case tests the vector database's loading capacity by repeatedly inserting small-dimension
+     vectors (SIFT 100K vectors, <b>128 dimensions</b>) until it is fully loaded. Number of inserted vectors will be 
+     reported."""
 
 
 class Performance768D10M(PerformanceCase):
@@ -238,6 +244,7 @@ Results will show index building time, recall, and maximum QPS."""
     load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
     optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K
 
+
 class Performance1536D5M1P(PerformanceCase):
     case_id: CaseType = CaseType.Performance1536D5M1P
     filter_rate: float | int | None = 0.01
@@ -248,6 +255,7 @@ Results will show index building time, recall, and maximum QPS."""
     load_timeout: float | int = config.LOAD_TIMEOUT_1536D_5M
     optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M
 
+
 class Performance1536D500K99P(PerformanceCase):
     case_id: CaseType = CaseType.Performance1536D500K99P
     filter_rate: float | int | None = 0.99
@@ -258,6 +266,7 @@ Results will show index building time, recall, and maximum QPS."""
     load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
     optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K
 
+
 class Performance1536D5M99P(PerformanceCase):
     case_id: CaseType = CaseType.Performance1536D5M99P
     filter_rate: float | int | None = 0.99
@@ -269,6 +278,17 @@ Results will show index building time, recall, and maximum QPS."""
     optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M
 
 
+class Performance1536D50K(PerformanceCase):
+    case_id: CaseType = CaseType.Performance1536D50K
+    filter_rate: float | int | None = None
+    dataset: DatasetManager = Dataset.OPENAI.manager(50_000)
+    name: str = "Search Performance Test (50K Dataset, 1536 Dim)"
+    description: str = """This case tests the search performance of a vector database with a medium 50K dataset (<b>OpenAI 50K vectors</b>, 1536 dimensions), at varying parallel levels.
+Results will show index building time, recall, and maximum QPS."""
+    load_timeout: float | int = 3600
+    optimize_timeout: float | int | None = 15 * 60
+
+
 type2case = {
     CaseType.CapacityDim960: CapacityDim960,
     CaseType.CapacityDim128: CapacityDim128,
@@ -290,5 +310,5 @@ type2case = {
 
     CaseType.Performance1536D500K99P: Performance1536D500K99P,
     CaseType.Performance1536D5M99P: Performance1536D5M99P,
-
+    CaseType.Performance1536D50K: Performance1536D50K,
 }

vectordb_bench/backend/clients/__init__.py

@@ -32,6 +32,7 @@ class DB(Enum):
     PgVectoRS = "PgVectoRS"
     Redis = "Redis"
     Chroma = "Chroma"
+    Test = "test"
 
 
     @property

vectordb_bench/backend/clients/api.py

@@ -47,7 +47,7 @@ class DBConfig(ABC, BaseModel):
     def not_empty_field(cls, v, field):
         if field.name == "db_label":
             return v
-        if isinstance(v, (str, SecretStr)) and len(v) == 0:
+        if not v and isinstance(v, (str, SecretStr)):
             raise ValueError("Empty string!")
         return v
 

vectordb_bench/backend/clients/milvus/cli.py

@@ -0,0 +1,291 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor3,
+    IVFFlatTypedDictN,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+
+)
+from vectordb_bench.backend.clients import DB
+
+DBTYPE = DB.Milvus
+
+
+class MilvusTypedDict(TypedDict):
+    uri: Annotated[
+        str, click.option("--uri", type=str, help="uri connection string", required=True)
+    ]
+
+
+class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
+def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
+    from .config import MilvusConfig, AutoIndexConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=AutoIndexConfig(),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
+def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
+    from .config import MilvusConfig, FLATConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=FLATConfig(),
+        **parameters,
+    )
+
+
+class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict)
+def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
+    from .config import MilvusConfig, HNSWConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=HNSWConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_search"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
+def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
+    from .config import MilvusConfig, IVFFlatConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=IVFFlatConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+        ),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
+def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
+    from .config import MilvusConfig, IVFSQ8Config
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=IVFSQ8Config(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusDISKANNTypedDict(CommonTypedDict, MilvusTypedDict):
+    search_list: Annotated[
+        str, click.option("--search-list",
+                          type=int,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusDISKANNTypedDict)
+def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
+    from .config import MilvusConfig, DISKANNConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=DISKANNConfig(
+            search_list=parameters["search_list"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusGPUIVFTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict):
+    cache_dataset_on_device: Annotated[
+        str, click.option("--cache-dataset-on-device",
+                          type=str,
+                          required=True)
+    ]
+    refine_ratio: Annotated[
+        str, click.option("--refine-ratio",
+                          type=float,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUIVFTypedDict)
+def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
+    from .config import MilvusConfig, GPUIVFFlatConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=GPUIVFFlatConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+            cache_dataset_on_device=parameters["cache_dataset_on_device"],
+            refine_ratio=parameters.get("refine_ratio"),
+        ),
+        **parameters,
+    )
+
+
+class MilvusGPUIVFPQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict, MilvusGPUIVFTypedDict):
+    m: Annotated[
+        str, click.option("--m",
+                          type=int, help="hnsw m",
+                          required=True)
+    ]
+    nbits: Annotated[
+        str, click.option("--nbits",
+                          type=int,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUIVFPQTypedDict)
+def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
+    from .config import MilvusConfig, GPUIVFPQConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=GPUIVFPQConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+            m=parameters["m"],
+            nbits=parameters["nbits"],
+            cache_dataset_on_device=parameters["cache_dataset_on_device"],
+            refine_ratio=parameters["refine_ratio"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusGPUCAGRATypedDict(CommonTypedDict, MilvusTypedDict, MilvusGPUIVFTypedDict):
+    intermediate_graph_degree: Annotated[
+        str, click.option("--intermediate-graph-degree",
+                          type=int,
+                          required=True)
+    ]
+    graph_degree: Annotated[
+        str, click.option("--graph-degree",
+                          type=int,
+                          required=True)
+    ]
+    build_algo: Annotated[
+        str, click.option("--build_algo",
+                          type=str,
+                          required=True)
+    ]
+    team_size: Annotated[
+        str, click.option("--team-size",
+                          type=int,
+                          required=True)
+    ]
+    search_width: Annotated[
+        str, click.option("--search-width",
+                          type=int,
+                          required=True)
+    ]
+    itopk_size: Annotated[
+        str, click.option("--itopk-size",
+                          type=int,
+                          required=True)
+    ]
+    min_iterations: Annotated[
+        str, click.option("--min-iterations",
+                          type=int,
+                          required=True)
+    ]
+    max_iterations: Annotated[
+        str, click.option("--max-iterations",
+                          type=int,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUCAGRATypedDict)
+def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
+    from .config import MilvusConfig, GPUCAGRAConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=GPUCAGRAConfig(
+            intermediate_graph_degree=parameters["intermediate_graph_degree"],
+            graph_degree=parameters["graph_degree"],
+            itopk_size=parameters["itopk_size"],
+            team_size=parameters["team_size"],
+            search_width=parameters["search_width"],
+            min_iterations=parameters["min_iterations"],
+            max_iterations=parameters["max_iterations"],
+            build_algo=parameters["build_algo"],
+            cache_dataset_on_device=parameters["cache_dataset_on_device"],
+            refine_ratio=parameters["refine_ratio"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/milvus/milvus.py

@@ -8,7 +8,7 @@ from typing import Iterable
 from pymilvus import Collection, utility
 from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusException
 
-from ..api import VectorDB
+from ..api import VectorDB, IndexType
 from .config import MilvusIndexConfig
 
 
@@ -122,10 +122,18 @@ class Milvus(VectorDB):
             if self.case_config.is_gpu_index:
                 log.debug("skip compaction for gpu index type.")
             else :
-                self.col.compact()
-                self.col.wait_for_compaction_completed()
+                try:
+                    self.col.compact()
+                    self.col.wait_for_compaction_completed()
+                except Exception as e:
+                    log.warning(f"{self.name} compact error: {e}")
+                    if hasattr(e, 'code'):
+                        if e.code().name == 'PERMISSION_DENIED':
+                            log.warning(f"Skip compact due to permission denied.")
+                            pass
+                    else:
+                        raise e
                 wait_index()
-
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
             raise e from None
@@ -143,7 +151,6 @@ class Milvus(VectorDB):
                     self.case_config.index_param(),
                     index_name=self._index_name,
                 )
-
             coll.load()
             log.info(f"{self.name} load")
         except Exception as e:
@@ -160,7 +167,7 @@ class Milvus(VectorDB):
         if self.case_config.is_gpu_index:
             log.info(f"current gpu_index only supports IP / L2, cosine dataset need normalize.")
             return True
-        
+
         return False
 
     def insert_embeddings(

vectordb_bench/backend/clients/pgvector/cli.py

@@ -0,0 +1,116 @@
+from typing import Annotated, Optional, TypedDict, Unpack
+
+import click
+import os
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor1,
+    IVFFlatTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from vectordb_bench.backend.clients import DB
+
+
+class PgVectorTypedDict(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)
+    ]
+    maintenance_work_mem: Annotated[
+        Optional[str],
+        click.option(
+            "--maintenance-work-mem",
+            type=str,
+            help="Sets the maximum memory to be used for maintenance operations (index creation). "
+            "Can be entered as string with unit like '64GB' or as an integer number of KB."
+            "This will set the parameters: max_parallel_maintenance_workers,"
+            " max_parallel_workers & table(parallel_workers)",
+            required=False,
+        ),
+    ]
+    max_parallel_workers: Annotated[
+        Optional[int],
+        click.option(
+            "--max-parallel-workers",
+            type=int,
+            help="Sets the maximum number of parallel processes per maintenance operation (index creation)",
+            required=False,
+        ),
+    ]
+
+
+class PgVectorIVFFlatTypedDict(PgVectorTypedDict, IVFFlatTypedDict):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectorIVFFlatTypedDict)
+def PgVectorIVFFlat(
+    **parameters: Unpack[PgVectorIVFFlatTypedDict],
+):
+    from .config import PgVectorConfig, PgVectorIVFFlatConfig
+
+    run(
+        db=DB.PgVector,
+        db_config=PgVectorConfig(
+            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=PgVectorIVFFlatConfig(
+            metric_type=None, lists=parameters["lists"], probes=parameters["probes"]
+        ),
+        **parameters,
+    )
+
+
+class PgVectorHNSWTypedDict(PgVectorTypedDict, HNSWFlavor1):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectorHNSWTypedDict)
+def PgVectorHNSW(
+    **parameters: Unpack[PgVectorHNSWTypedDict],
+):
+    from .config import PgVectorConfig, PgVectorHNSWConfig
+
+    run(
+        db=DB.PgVector,
+        db_config=PgVectorConfig(
+            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=PgVectorHNSWConfig(
+            m=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            maintenance_work_mem=parameters["maintenance_work_mem"],
+            max_parallel_workers=parameters["max_parallel_workers"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/pgvector/config.py

@@ -109,7 +109,7 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
     def _optionally_build_set_options(
         set_mapping: Mapping[str, Any]
     ) -> Sequence[dict[str, Any]]:
-        """Walk through options, creating 'SET 'key1 = "value1";' commands"""
+        """Walk through options, creating 'SET 'key1 = "value1";' list"""
         session_options = []
         for setting_name, value in set_mapping.items():
             if value:

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -58,14 +58,13 @@ class PgVector(VectorDB):
                 self.case_config.create_index_after_load,
             )
         ):
-            err = f"{self.name} config must create an index using create_index_before_load and/or create_index_after_load"
+            err = f"{self.name} config must create an index using create_index_before_load or create_index_after_load"
             log.error(err)
             raise RuntimeError(
                 f"{err}\n{pprint.pformat(self.db_config)}\n{pprint.pformat(self.case_config)}"
             )
 
         if drop_old:
-            # self.pg_table.drop(pg_engine, checkfirst=True)
             self._drop_index()
             self._drop_table()
             self._create_table(dim)
@@ -257,7 +256,10 @@ class PgVector(VectorDB):
             with_clause = sql.Composed(())
 
         index_create_sql = sql.SQL(
-            "CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} USING {index_type} (embedding {embedding_metric})"
+            """
+            CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} 
+            USING {index_type} (embedding {embedding_metric})
+            """
         ).format(
             index_name=sql.Identifier(self._index_name),
             table_name=sql.Identifier(self.table_name),
@@ -339,9 +341,10 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        q = np.asarray(query)
         # TODO add filters support
         result = self.cursor.execute(
-            self._unfiltered_search, (query, k), prepare=True, binary=True
+            self._unfiltered_search, (q, k), prepare=True, binary=True
         )
 
         return [int(i[0]) for i in result.fetchall()]