vectordb-bench 0.0.13__tar.gz → 0.0.14__tar.gz

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.13
+Version: 0.0.14
 Summary: VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -36,7 +36,7 @@ Requires-Dist: pinecone-client; extra == "all"
 Requires-Dist: weaviate-client; extra == "all"
 Requires-Dist: elasticsearch; extra == "all"
 Requires-Dist: pgvector; extra == "all"
-Requires-Dist: pgvecto_rs[psycopg3]>=0.2.1; extra == "all"
+Requires-Dist: pgvecto_rs[psycopg3]>=0.2.2; extra == "all"
 Requires-Dist: sqlalchemy; extra == "all"
 Requires-Dist: redis; extra == "all"
 Requires-Dist: chromadb; extra == "all"
@@ -61,7 +61,7 @@ Requires-Dist: psycopg; extra == "pgvectorscale"
 Requires-Dist: psycopg-binary; extra == "pgvectorscale"
 Requires-Dist: pgvector; extra == "pgvectorscale"
 Provides-Extra: pgvecto-rs
-Requires-Dist: pgvecto_rs[psycopg3]>=0.2.1; extra == "pgvecto-rs"
+Requires-Dist: pgvecto_rs[psycopg3]>=0.2.2; extra == "pgvecto-rs"
 Provides-Extra: redis
 Requires-Dist: redis; extra == "redis"
 Provides-Extra: memorydb
@@ -191,6 +191,29 @@ Options:
   --m INTEGER                     hnsw m
   --ef-construction INTEGER       hnsw ef-construction
   --ef-search INTEGER             hnsw ef-search
+  --quantization-type [none|halfvec]
+                                  quantization type for vectors
+  --custom-case-name TEXT         Custom case name i.e. PerformanceCase1536D50K
+  --custom-case-description TEXT  Custom name description
+  --custom-case-load-timeout INTEGER
+                                  Custom case load timeout [default: 36000]
+  --custom-case-optimize-timeout INTEGER
+                                  Custom case optimize timeout [default: 36000]
+  --custom-dataset-name TEXT
+                                  Dataset name i.e OpenAI
+  --custom-dataset-dir TEXT       Dataset directory i.e. openai_medium_500k
+  --custom-dataset-size INTEGER   Dataset size i.e. 500000
+  --custom-dataset-dim INTEGER    Dataset dimension
+  --custom-dataset-metric-type TEXT
+                                  Dataset distance metric [default: COSINE]
+  --custom-dataset-file-count INTEGER
+                                  Dataset file count
+  --custom-dataset-use-shuffled / --skip-custom-dataset-use-shuffled
+                                  Use shuffled custom dataset or skip  [default: custom-dataset-
+                                  use-shuffled]
+  --custom-dataset-with-gt / --skip-custom-dataset-with-gt
+                                  Custom dataset with ground truth or skip  [default: custom-dataset-
+                                  with-gt]
   --help                          Show this message and exit.
 ```
 #### Using a configuration file.
@@ -535,6 +558,8 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
    3. Update db_config and db_case_config to match client requirements
    4. Continue to add new functions for each index config.
    5. Import the client cli module and command to vectordb_bench/cli/vectordbbench.py (for databases with multiple commands (index configs), this only needs to be done for one command)  
+   6. Import the `get_custom_case_config` function from `vectordb_bench/cli/cli.py` and use it to add a new key `custom_case` to the `parameters` variable within the command.
+
 
 > cli modules with multiple index configs:
 > - pgvector: vectordb_bench/backend/clients/pgvector/cli.py

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/README.md

@@ -117,6 +117,29 @@ Options:
   --m INTEGER                     hnsw m
   --ef-construction INTEGER       hnsw ef-construction
   --ef-search INTEGER             hnsw ef-search
+  --quantization-type [none|halfvec]
+                                  quantization type for vectors
+  --custom-case-name TEXT         Custom case name i.e. PerformanceCase1536D50K
+  --custom-case-description TEXT  Custom name description
+  --custom-case-load-timeout INTEGER
+                                  Custom case load timeout [default: 36000]
+  --custom-case-optimize-timeout INTEGER
+                                  Custom case optimize timeout [default: 36000]
+  --custom-dataset-name TEXT
+                                  Dataset name i.e OpenAI
+  --custom-dataset-dir TEXT       Dataset directory i.e. openai_medium_500k
+  --custom-dataset-size INTEGER   Dataset size i.e. 500000
+  --custom-dataset-dim INTEGER    Dataset dimension
+  --custom-dataset-metric-type TEXT
+                                  Dataset distance metric [default: COSINE]
+  --custom-dataset-file-count INTEGER
+                                  Dataset file count
+  --custom-dataset-use-shuffled / --skip-custom-dataset-use-shuffled
+                                  Use shuffled custom dataset or skip  [default: custom-dataset-
+                                  use-shuffled]
+  --custom-dataset-with-gt / --skip-custom-dataset-with-gt
+                                  Custom dataset with ground truth or skip  [default: custom-dataset-
+                                  with-gt]
   --help                          Show this message and exit.
 ```
 #### Using a configuration file.
@@ -461,6 +484,8 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
    3. Update db_config and db_case_config to match client requirements
    4. Continue to add new functions for each index config.
    5. Import the client cli module and command to vectordb_bench/cli/vectordbbench.py (for databases with multiple commands (index configs), this only needs to be done for one command)  
+   6. Import the `get_custom_case_config` function from `vectordb_bench/cli/cli.py` and use it to add a new key `custom_case` to the `parameters` variable within the command.
+
 
 > cli modules with multiple index configs:
 > - pgvector: vectordb_bench/backend/clients/pgvector/cli.py

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/pyproject.toml

@@ -56,7 +56,7 @@ all = [
     "weaviate-client",
     "elasticsearch",
     "pgvector",
-    "pgvecto_rs[psycopg3]>=0.2.1",
+    "pgvecto_rs[psycopg3]>=0.2.2",
     "sqlalchemy",
     "redis",
     "chromadb",
@@ -72,7 +72,7 @@ weaviate = [ "weaviate-client" ]
 elastic = [ "elasticsearch" ]
 pgvector = [ "psycopg", "psycopg-binary", "pgvector" ]
 pgvectorscale = [ "psycopg", "psycopg-binary", "pgvector" ]
-pgvecto_rs = [ "pgvecto_rs[psycopg3]>=0.2.1" ]
+pgvecto_rs = [ "pgvecto_rs[psycopg3]>=0.2.2" ]
 redis = [ "redis" ]
 memorydb = [ "memorydb" ]
 chromadb = [ "chromadb" ]

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py

@@ -3,7 +3,7 @@ from contextlib import contextmanager
 import time
 from typing import Iterable, Type
 from ..api import VectorDB, DBCaseConfig, DBConfig, IndexType
-from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig
+from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig, AWSOS_Engine
 from opensearchpy import OpenSearch
 from opensearchpy.helpers import bulk
 
@@ -83,7 +83,7 @@ class AWSOpenSearch(VectorDB):
 
     @contextmanager
     def init(self) -> None:
-        """connect to elasticsearch"""
+        """connect to opensearch"""
         self.client = OpenSearch(**self.db_config)
 
         yield
@@ -97,7 +97,7 @@ class AWSOpenSearch(VectorDB):
         metadata: list[int],
         **kwargs,
     ) -> tuple[int, Exception]:
-        """Insert the embeddings to the elasticsearch."""
+        """Insert the embeddings to the opensearch."""
         assert self.client is not None, "should self.init() first"
 
         insert_data = []
@@ -136,13 +136,15 @@ class AWSOpenSearch(VectorDB):
         body = {
             "size": k,
             "query": {"knn": {self.vector_col_name: {"vector": query, "k": k}}},
+            **({"filter": {"range": {self.id_col_name: {"gt": filters["id"]}}}} if filters else {})
         }
         try:
-            resp = self.client.search(index=self.index_name, body=body)
+            resp = self.client.search(index=self.index_name, body=body,size=k,_source=False,docvalue_fields=[self.id_col_name],stored_fields="_none_",filter_path=[f"hits.hits.fields.{self.id_col_name}"],)
             log.info(f'Search took: {resp["took"]}')
             log.info(f'Search shards: {resp["_shards"]}')
             log.info(f'Search hits total: {resp["hits"]["total"]}')
-            result = [int(d["_id"]) for d in resp["hits"]["hits"]]
+            result = [h["fields"][self.id_col_name][0] for h in resp["hits"]["hits"]]
+            #result = [int(d["_id"]) for d in resp["hits"]["hits"]]
             # log.info(f'success! length={len(res)}')
 
             return result
@@ -152,7 +154,46 @@ class AWSOpenSearch(VectorDB):
 
     def optimize(self):
         """optimize will be called between insertion and search in performance cases."""
-        pass
+        # Call refresh first to ensure that all segments are created
+        self._refresh_index()
+        self._do_force_merge()
+        # Call refresh again to ensure that the index is ready after force merge.
+        self._refresh_index()
+        # ensure that all graphs are loaded in memory and ready for search
+        self._load_graphs_to_memory()
+
+    def _refresh_index(self):
+        log.debug(f"Starting refresh for index {self.index_name}")
+        SECONDS_WAITING_FOR_REFRESH_API_CALL_SEC = 30
+        while True:
+            try:
+                log.info(f"Starting the Refresh Index..")
+                self.client.indices.refresh(index=self.index_name)
+                break
+            except Exception as e:
+                log.info(
+                    f"Refresh errored out. Sleeping for {SECONDS_WAITING_FOR_REFRESH_API_CALL_SEC} sec and then Retrying : {e}")
+                time.sleep(SECONDS_WAITING_FOR_REFRESH_API_CALL_SEC)
+                continue
+        log.debug(f"Completed refresh for index {self.index_name}")
+
+    def _do_force_merge(self):
+        log.debug(f"Starting force merge for index {self.index_name}")
+        force_merge_endpoint = f'/{self.index_name}/_forcemerge?max_num_segments=1&wait_for_completion=false'
+        force_merge_task_id = self.client.transport.perform_request('POST', force_merge_endpoint)['task']
+        SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30
+        while True:
+            time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC)
+            task_status = self.client.tasks.get(task_id=force_merge_task_id)
+            if task_status['completed']:
+                break
+        log.debug(f"Completed force merge for index {self.index_name}")
+
+    def _load_graphs_to_memory(self):
+        if self.case_config.engine != AWSOS_Engine.lucene:
+            log.info("Calling warmup API to load graphs into memory")
+            warmup_endpoint = f'/_plugins/_knn/warmup/{self.index_name}'
+            self.client.transport.perform_request('GET', warmup_endpoint)
 
     def ready_to_load(self):
         """ready_to_load will be called before load in load cases."""

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/vectordb_bench/backend/clients/aws_opensearch/config.py

@@ -1,9 +1,10 @@
+import logging
 from enum import Enum
 from pydantic import SecretStr, BaseModel
 
 from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
 
-
+log = logging.getLogger(__name__)
 class AWSOpenSearchConfig(DBConfig, BaseModel):
     host: str = ""
     port: int = 443
@@ -31,14 +32,18 @@ class AWSOS_Engine(Enum):
 
 class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType = MetricType.L2
-    engine: AWSOS_Engine = AWSOS_Engine.nmslib
-    efConstruction: int = 360
-    M: int = 30
+    engine: AWSOS_Engine = AWSOS_Engine.faiss
+    efConstruction: int = 256
+    efSearch: int = 256
+    M: int = 16
 
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.IP:
-            return "innerproduct"  # only support faiss / nmslib, not for Lucene.
+            return "innerproduct"
         elif self.metric_type == MetricType.COSINE:
+            if self.engine == AWSOS_Engine.faiss:
+                log.info(f"Using metric type as innerproduct because faiss doesn't support cosine as metric type for Opensearch")
+                return "innerproduct"
             return "cosinesimil"
         return "l2"
 
@@ -49,7 +54,8 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
             "engine": self.engine.value,
             "parameters": {
                 "ef_construction": self.efConstruction,
-                "m": self.M
+                "m": self.M,
+                "ef_search": self.efSearch
             }
         }
         return params

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/vectordb_bench/backend/clients/aws_opensearch/run.py

@@ -40,12 +40,12 @@ def create_index(client, index_name):
                 "type": "knn_vector",
                 "dimension": _DIM,
                 "method": {
-                    "engine": "nmslib",
+                    "engine": "faiss",
                     "name": "hnsw",
                     "space_type": "l2",
                     "parameters": {
-                        "ef_construction": 128,
-                        "m": 24,
+                        "ef_construction": 256,
+                        "m": 16,
                     }
                 }
             }
@@ -108,12 +108,43 @@ def search(client, index_name):
             print('\nSearch not ready, sleep 1s')
             time.sleep(1)
 
+def optimize_index(client, index_name):
+    print(f"Starting force merge for index {index_name}")
+    force_merge_endpoint = f'/{index_name}/_forcemerge?max_num_segments=1&wait_for_completion=false'
+    force_merge_task_id = client.transport.perform_request('POST', force_merge_endpoint)['task']
+    SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30
+    while True:
+        time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC)
+        task_status = client.tasks.get(task_id=force_merge_task_id)
+        if task_status['completed']:
+            break
+    print(f"Completed force merge for index {index_name}")
+
+
+def refresh_index(client, index_name):
+    print(f"Starting refresh for index {index_name}")
+    SECONDS_WAITING_FOR_REFRESH_API_CALL_SEC = 30
+    while True:
+        try:
+            print(f"Starting the Refresh Index..")
+            client.indices.refresh(index=index_name)
+            break
+        except Exception as e:
+            print(
+                f"Refresh errored out. Sleeping for {SECONDS_WAITING_FOR_REFRESH_API_CALL_SEC} sec and then Retrying : {e}")
+            time.sleep(SECONDS_WAITING_FOR_REFRESH_API_CALL_SEC)
+            continue
+    print(f"Completed refresh for index {index_name}")
+
+
 
 def main():
     client = create_client()
     try:
         create_index(client, _INDEX_NAME)
         bulk_insert(client, _INDEX_NAME)
+        optimize_index(client, _INDEX_NAME)
+        refresh_index(client, _INDEX_NAME)
         search(client, _INDEX_NAME)
         delete_index(client, _INDEX_NAME)
     except Exception as e:

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/vectordb_bench/backend/clients/pgvector/cli.py

@@ -10,6 +10,7 @@ from ....cli.cli import (
     IVFFlatTypedDict,
     cli,
     click_parameter_decorators_from_typed_dict,
+    get_custom_case_config,
     run,
 )
 from vectordb_bench.backend.clients import DB
@@ -56,7 +57,15 @@ class PgVectorTypedDict(CommonTypedDict):
             required=False,
         ),
     ]
-
+    quantization_type: Annotated[
+        Optional[str],
+        click.option(
+            "--quantization-type",
+            type=click.Choice(["none", "halfvec"]),
+            help="quantization type for vectors",
+            required=False,
+        ),
+    ]
 
 class PgVectorIVFFlatTypedDict(PgVectorTypedDict, IVFFlatTypedDict):
     ...
@@ -69,6 +78,7 @@ def PgVectorIVFFlat(
 ):
     from .config import PgVectorConfig, PgVectorIVFFlatConfig
 
+    parameters["custom_case"] = get_custom_case_config(parameters)
     run(
         db=DB.PgVector,
         db_config=PgVectorConfig(
@@ -79,7 +89,10 @@ def PgVectorIVFFlat(
             db_name=parameters["db_name"],
         ),
         db_case_config=PgVectorIVFFlatConfig(
-            metric_type=None, lists=parameters["lists"], probes=parameters["probes"]
+            metric_type=None,
+            lists=parameters["lists"],
+            probes=parameters["probes"],
+            quantization_type=parameters["quantization_type"],
         ),
         **parameters,
     )
@@ -96,6 +109,7 @@ def PgVectorHNSW(
 ):
     from .config import PgVectorConfig, PgVectorHNSWConfig
 
+    parameters["custom_case"] = get_custom_case_config(parameters)
     run(
         db=DB.PgVector,
         db_config=PgVectorConfig(
@@ -111,6 +125,7 @@ def PgVectorHNSW(
             ef_search=parameters["ef_search"],
             maintenance_work_mem=parameters["maintenance_work_mem"],
             max_parallel_workers=parameters["max_parallel_workers"],
+            quantization_type=parameters["quantization_type"],
         ),
         **parameters,
     )

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/vectordb_bench/backend/clients/pgvector/config.py

@@ -59,11 +59,18 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
     create_index_after_load: bool = True
 
     def parse_metric(self) -> str:
-        if self.metric_type == MetricType.L2:
-            return "vector_l2_ops"
-        elif self.metric_type == MetricType.IP:
-            return "vector_ip_ops"
-        return "vector_cosine_ops"
+        if self.quantization_type == "halfvec":
+            if self.metric_type == MetricType.L2:
+                return "halfvec_l2_ops"
+            elif self.metric_type == MetricType.IP:
+                return "halfvec_ip_ops"
+            return "halfvec_cosine_ops"
+        else:
+            if self.metric_type == MetricType.L2:
+                return "vector_l2_ops"
+            elif self.metric_type == MetricType.IP:
+                return "vector_ip_ops"
+            return "vector_cosine_ops"
 
     def parse_metric_fun_op(self) -> LiteralString:
         if self.metric_type == MetricType.L2:
@@ -143,9 +150,12 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
     index: IndexType = IndexType.ES_IVFFlat
     maintenance_work_mem: Optional[str] = None
     max_parallel_workers: Optional[int] = None
+    quantization_type: Optional[str] = None
 
     def index_param(self) -> PgVectorIndexParam:
         index_parameters = {"lists": self.lists}
+        if self.quantization_type == "none":
+            self.quantization_type = None
         return {
             "metric": self.parse_metric(),
             "index_type": self.index.value,
@@ -154,6 +164,7 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
             ),
             "maintenance_work_mem": self.maintenance_work_mem,
             "max_parallel_workers": self.max_parallel_workers,
+            "quantization_type": self.quantization_type,
         }
 
     def search_param(self) -> PgVectorSearchParam:
@@ -183,9 +194,12 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
     index: IndexType = IndexType.ES_HNSW
     maintenance_work_mem: Optional[str] = None
     max_parallel_workers: Optional[int] = None
+    quantization_type: Optional[str] = None
 
     def index_param(self) -> PgVectorIndexParam:
         index_parameters = {"m": self.m, "ef_construction": self.ef_construction}
+        if self.quantization_type == "none":
+            self.quantization_type = None
         return {
             "metric": self.parse_metric(),
             "index_type": self.index.value,
@@ -194,6 +208,7 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
             ),
             "maintenance_work_mem": self.maintenance_work_mem,
             "max_parallel_workers": self.max_parallel_workers,
+            "quantization_type": self.quantization_type,
         }
 
     def search_param(self) -> PgVectorSearchParam:

{vectordb_bench-0.0.13 → vectordb_bench-0.0.14}/vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -22,7 +22,7 @@ class PgVector(VectorDB):
     conn: psycopg.Connection[Any] | None = None
     cursor: psycopg.Cursor[Any] | None = None
 
-    # TODO add filters support
+    _filtered_search: sql.Composed
     _unfiltered_search: sql.Composed
 
     def __init__(
@@ -112,15 +112,63 @@ class PgVector(VectorDB):
                 self.cursor.execute(command)
             self.conn.commit()
 
-        self._unfiltered_search = sql.Composed(
-            [
-                sql.SQL("SELECT id FROM public.{} ORDER BY embedding ").format(
-                    sql.Identifier(self.table_name)
-                ),
-                sql.SQL(self.case_config.search_param()["metric_fun_op"]),
-                sql.SQL(" %s::vector LIMIT %s::int"),
-            ]
-        )
+        index_param = self.case_config.index_param()
+        # The following sections assume that the quantization_type value matches the quantization function name
+        if index_param["quantization_type"] != None:
+            self._filtered_search = sql.Composed(
+                [
+                    sql.SQL(
+                        "SELECT id FROM public.{table_name} WHERE id >= %s ORDER BY embedding::{quantization_type}({dim}) "
+                    ).format(
+                        table_name=sql.Identifier(self.table_name),
+                        quantization_type=sql.SQL(index_param["quantization_type"]),
+                        dim=sql.Literal(self.dim),
+                    ),
+                    sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                    sql.SQL(" %s::{quantization_type}({dim}) LIMIT %s::int").format(
+                        quantization_type=sql.SQL(index_param["quantization_type"]),
+                        dim=sql.Literal(self.dim),
+                    ),
+                ]
+            )
+        else:
+            self._filtered_search = sql.Composed(
+                [
+                    sql.SQL(
+                        "SELECT id FROM public.{table_name} WHERE id >= %s ORDER BY embedding "
+                        ).format(table_name=sql.Identifier(self.table_name)),
+                    sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                    sql.SQL(" %s::vector LIMIT %s::int"),
+                ]
+            )
+
+        if index_param["quantization_type"] != None:
+            self._unfiltered_search = sql.Composed(
+                [
+                    sql.SQL(
+                        "SELECT id FROM public.{table_name} ORDER BY embedding::{quantization_type}({dim}) "
+                    ).format(
+                        table_name=sql.Identifier(self.table_name),
+                        quantization_type=sql.SQL(index_param["quantization_type"]),
+                        dim=sql.Literal(self.dim),
+                    ),
+                    sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                    sql.SQL(" %s::{quantization_type}({dim}) LIMIT %s::int").format(
+                        quantization_type=sql.SQL(index_param["quantization_type"]),
+                        dim=sql.Literal(self.dim),
+                    ),
+                ]
+            )
+        else:
+            self._unfiltered_search = sql.Composed(
+                [
+                    sql.SQL("SELECT id FROM public.{} ORDER BY embedding ").format(
+                        sql.Identifier(self.table_name)
+                    ),
+                    sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                    sql.SQL(" %s::vector LIMIT %s::int"),
+                ]
+            )
 
         try:
             yield
@@ -255,17 +303,34 @@ class PgVector(VectorDB):
         else:
             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})
-            """
-        ).format(
-            index_name=sql.Identifier(self._index_name),
-            table_name=sql.Identifier(self.table_name),
-            index_type=sql.Identifier(index_param["index_type"]),
-            embedding_metric=sql.Identifier(index_param["metric"]),
-        )
+        if index_param["quantization_type"] != None:
+            index_create_sql = sql.SQL(
+                """
+                CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} 
+                USING {index_type} ((embedding::{quantization_type}({dim})) {embedding_metric})
+                """
+            ).format(
+                index_name=sql.Identifier(self._index_name),
+                table_name=sql.Identifier(self.table_name),
+                index_type=sql.Identifier(index_param["index_type"]),
+                # This assumes that the quantization_type value matches the quantization function name
+                quantization_type=sql.SQL(index_param["quantization_type"]),
+                dim=self.dim,
+                embedding_metric=sql.Identifier(index_param["metric"]),
+            )
+        else:
+            index_create_sql = sql.SQL(
+                """
+                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),
+                index_type=sql.Identifier(index_param["index_type"]),
+                embedding_metric=sql.Identifier(index_param["metric"]),
+            )
+
         index_create_sql_with_with_clause = (
             index_create_sql + with_clause
         ).join(" ")
@@ -342,9 +407,14 @@ class PgVector(VectorDB):
         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, (q, k), prepare=True, binary=True
-        )
+        if filters:
+            gt = filters.get("id")
+            result = self.cursor.execute(
+                    self._filtered_search, (gt, q, k), prepare=True, binary=True
+                    )
+        else:
+            result = self.cursor.execute(
+                    self._unfiltered_search, (q, k), prepare=True, binary=True
+                    )
 
         return [int(i[0]) for i in result.fetchall()]