vectordb-bench 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

vectordb_bench/backend/clients/memorydb/memorydb.py

@@ -0,0 +1,254 @@
+import logging, time
+from contextlib import contextmanager
+from typing import Any, Generator, Optional, Tuple, Type
+from ..api import VectorDB, DBCaseConfig, IndexType
+from .config import MemoryDBIndexConfig
+import redis
+from redis import Redis
+from redis.cluster import RedisCluster
+from redis.commands.search.field import TagField, VectorField, NumericField
+from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+from redis.commands.search.query import Query
+import numpy as np
+
+
+log = logging.getLogger(__name__)
+INDEX_NAME = "index"    # Vector Index Name
+
+class MemoryDB(VectorDB):
+    def __init__(
+            self,
+            dim: int,
+            db_config: dict,
+            db_case_config: MemoryDBIndexConfig,
+            drop_old: bool = False,
+            **kwargs
+        ):
+
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.collection_name = INDEX_NAME
+        self.target_nodes = RedisCluster.RANDOM if not self.db_config["cmd"] else None
+        self.insert_batch_size = db_case_config.insert_batch_size
+        self.dbsize = kwargs.get("num_rows")
+
+        # Create a MemoryDB connection, if db has password configured, add it to the connection here and in init():
+        log.info(f"Establishing connection to: {self.db_config}")
+        conn = self.get_client(primary=True)
+        log.info(f"Connection established: {conn}")
+        log.info(conn.execute_command("INFO server"))
+
+        if drop_old:
+            try:
+                log.info(f"MemoryDB client getting info for: {INDEX_NAME}")
+                info = conn.ft(INDEX_NAME).info()
+                log.info(f"Index info: {info}")
+            except redis.exceptions.ResponseError as e:
+                log.error(e)
+                drop_old = False
+                log.info(f"MemoryDB client drop_old collection: {self.collection_name}")
+            
+            log.info("Executing FLUSHALL")
+            conn.flushall()
+
+            # Since the default behaviour of FLUSHALL is asynchronous, wait for db to be empty
+            self.wait_until(self.wait_for_empty_db, 3, "", conn)
+            if not self.db_config["cmd"]:
+                replica_clients = self.get_client(replicas=True)
+                for rc, host in replica_clients:
+                    self.wait_until(self.wait_for_empty_db, 3, "", rc)
+                    log.debug(f"Flushall done in the host: {host}")
+                    rc.close()
+        
+        self.make_index(dim, conn)
+        conn.close()
+        conn = None
+
+    def make_index(self, vector_dimensions: int, conn: redis.Redis):
+        try:
+            # check to see if index exists
+            conn.ft(INDEX_NAME).info()
+        except Exception as e:
+            log.warn(f"Error getting info for index '{INDEX_NAME}': {e}")
+            index_param = self.case_config.index_param()
+            search_param = self.case_config.search_param()
+            vector_parameters = {  # Vector Index Type: FLAT or HNSW
+                "TYPE": "FLOAT32",
+                "DIM": vector_dimensions,  # Number of Vector Dimensions
+                "DISTANCE_METRIC": index_param["metric"],  # Vector Search Distance Metric
+            }
+            if index_param["m"]:
+                vector_parameters["M"] = index_param["m"]
+            if index_param["ef_construction"]:
+                vector_parameters["EF_CONSTRUCTION"] = index_param["ef_construction"]
+            if search_param["ef_runtime"]:
+                vector_parameters["EF_RUNTIME"] = search_param["ef_runtime"]
+
+            schema = (
+                TagField("id"),                   
+                NumericField("metadata"),              
+                VectorField("vector",   # Vector Field Name
+                    "HNSW", vector_parameters
+                ),
+            )
+
+            definition = IndexDefinition(index_type=IndexType.HASH)
+            rs = conn.ft(INDEX_NAME)
+            rs.create_index(schema, definition=definition)
+    
+    def get_client(self, **kwargs):
+        """
+        Gets either cluster connection or normal connection based on `cmd` flag.
+        CMD stands for Cluster Mode Disabled and is a "mode".
+        """
+        if not self.db_config["cmd"]:
+            # Cluster mode enabled
+
+            client = RedisCluster(
+                host=self.db_config["host"],
+                port=self.db_config["port"],
+                ssl=self.db_config["ssl"],
+                password=self.db_config["password"],
+                ssl_ca_certs=self.db_config["ssl_ca_certs"],
+                ssl_cert_reqs=None,
+            )
+
+            # Ping all nodes to create a connection
+            client.execute_command("PING", target_nodes=RedisCluster.ALL_NODES)
+            replicas = client.get_replicas()
+
+            if len(replicas) > 0:
+                # FT.SEARCH is a keyless command, use READONLY for replica connections
+                client.execute_command("READONLY", target_nodes=RedisCluster.REPLICAS)
+
+            if kwargs.get("primary", False):
+                client = client.get_primaries()[0].redis_connection
+
+            if kwargs.get("replicas", False):
+                # Return client and host name for each replica
+                return [(c.redis_connection, c.host) for c in replicas]
+
+        else:
+            client = Redis(
+                host=self.db_config["host"],
+                port=self.db_config["port"],
+                db=0,
+                ssl=self.db_config["ssl"],
+                password=self.db_config["password"],
+                ssl_ca_certs=self.db_config["ssl_ca_certs"],
+                ssl_cert_reqs=None,
+            )
+            client.execute_command("PING")
+        return client
+
+    @contextmanager
+    def init(self) -> Generator[None, None, None]:
+        """ create and destory connections to database.
+
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+        """
+        self.conn = self.get_client()
+        search_param = self.case_config.search_param()
+        if search_param["ef_runtime"]:
+            self.ef_runtime_str = f'EF_RUNTIME {search_param["ef_runtime"]}'
+        else:
+            self.ef_runtime_str = ""
+        yield
+        self.conn.close()
+        self.conn = None
+
+    def ready_to_load(self) -> bool:
+        pass
+
+    def optimize(self) -> None:
+        self._post_insert()
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> Tuple[int, Optional[Exception]]:
+        """Insert embeddings into the database.
+        Should call self.init() first.
+        """
+
+        try:
+            with self.conn.pipeline(transaction=False) as pipe:
+                for i, embedding in enumerate(embeddings):
+                    embedding = np.array(embedding).astype(np.float32)
+                    pipe.hset(metadata[i], mapping = {
+                        "id": str(metadata[i]),
+                        "metadata": metadata[i], 
+                        "vector": embedding.tobytes(),
+                    })
+                    # Execute the pipe so we don't keep too much in memory at once
+                    if (i + 1) % self.insert_batch_size == 0:
+                        pipe.execute()
+
+                pipe.execute()
+                result_len = i + 1
+        except Exception as e:
+            return 0, e
+        
+        return result_len, None
+    
+    def _post_insert(self):
+        """Wait for indexing to finish"""
+        client = self.get_client(primary=True)
+        log.info("Waiting for background indexing to finish")
+        args = (self.wait_for_no_activity, 5, "", client)
+        self.wait_until(*args)
+        if not self.db_config["cmd"]:
+            replica_clients = self.get_client(replicas=True)
+            for rc, host_name in replica_clients:
+                args = (self.wait_for_no_activity, 5, "", rc)
+                self.wait_until(*args)
+                log.debug(f"Background indexing completed in the host: {host_name}")
+                rc.close()
+    
+    def wait_until(
+        self, condition, interval=5, message="Operation took too long", *args
+    ):
+        while not condition(*args):
+            time.sleep(interval)
+    
+    def wait_for_no_activity(self, client: redis.RedisCluster | redis.Redis):
+        return (
+            client.info("search")["search_background_indexing_status"] == "NO_ACTIVITY"
+        )
+    
+    def wait_for_empty_db(self, client: redis.RedisCluster | redis.Redis):
+        return client.execute_command("DBSIZE") == 0
+    
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 10,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> (list[int]):
+        assert self.conn is not None
+        
+        query_vector = np.array(query).astype(np.float32).tobytes()
+        query_obj = Query(f"*=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+        query_params = {"vec": query_vector}
+        
+        if filters:
+            # benchmark test filters of format: {'metadata': '>=10000', 'id': 10000}
+            # gets exact match for id, and range for metadata if they exist in filters
+            id_value = filters.get("id")
+            # Removing '>=' from the id_value: '>=10000'
+            metadata_value = filters.get("metadata")[2:]
+            if id_value and metadata_value:
+                query_obj = Query(f"(@metadata:[{metadata_value} +inf] @id:{ {id_value} })=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+            elif id_value:
+                #gets exact match for id
+                query_obj = Query(f"@id:{ {id_value} }=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+            else: #metadata only case, greater than or equal to metadata value
+                query_obj = Query(f"@metadata:[{metadata_value} +inf]=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+        res = self.conn.ft(INDEX_NAME).search(query_obj, query_params)
+        return [int(doc["id"]) for doc in res.docs]

vectordb_bench/backend/clients/pgvecto_rs/cli.py

@@ -0,0 +1,154 @@
+from typing import Annotated, Optional, 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 PgVectoRSTypedDict(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)
+    ]
+    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,
+        ),
+    ]
+    quantization_type: Annotated[
+        str,
+        click.option(
+            "--quantization-type",
+            type=click.Choice(["trivial", "scalar", "product"]),
+            help="quantization type for vectors",
+            required=False,
+        ),
+    ]
+    quantization_ratio: Annotated[
+        str,
+        click.option(
+            "--quantization-ratio",
+            type=click.Choice(["x4", "x8", "x16", "x32", "x64"]),
+            help="quantization ratio(for product quantization)",
+            required=False,
+        ),
+    ]
+
+
+class PgVectoRSFlatTypedDict(PgVectoRSTypedDict, IVFFlatTypedDict): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectoRSFlatTypedDict)
+def PgVectoRSFlat(
+    **parameters: Unpack[PgVectoRSFlatTypedDict],
+):
+    from .config import PgVectoRSConfig, PgVectoRSFLATConfig
+
+    run(
+        db=DB.PgVectoRS,
+        db_config=PgVectoRSConfig(
+            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=PgVectoRSFLATConfig(
+            max_parallel_workers=parameters["max_parallel_workers"],
+            quantization_type=parameters["quantization_type"],
+            quantization_ratio=parameters["quantization_ratio"],
+        ),
+        **parameters,
+    )
+
+
+class PgVectoRSIVFFlatTypedDict(PgVectoRSTypedDict, IVFFlatTypedDict): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectoRSIVFFlatTypedDict)
+def PgVectoRSIVFFlat(
+    **parameters: Unpack[PgVectoRSIVFFlatTypedDict],
+):
+    from .config import PgVectoRSConfig, PgVectoRSIVFFlatConfig
+
+    run(
+        db=DB.PgVectoRS,
+        db_config=PgVectoRSConfig(
+            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=PgVectoRSIVFFlatConfig(
+            max_parallel_workers=parameters["max_parallel_workers"],
+            quantization_type=parameters["quantization_type"],
+            quantization_ratio=parameters["quantization_ratio"],
+            probes=parameters["probes"],
+            lists=parameters["lists"],
+        ),
+        **parameters,
+    )
+
+
+class PgVectoRSHNSWTypedDict(PgVectoRSTypedDict, HNSWFlavor1): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectoRSHNSWTypedDict)
+def PgVectoRSHNSW(
+    **parameters: Unpack[PgVectoRSHNSWTypedDict],
+):
+    from .config import PgVectoRSConfig, PgVectoRSHNSWConfig
+
+    run(
+        db=DB.PgVectoRS,
+        db_config=PgVectoRSConfig(
+            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=PgVectoRSHNSWConfig(
+            max_parallel_workers=parameters["max_parallel_workers"],
+            quantization_type=parameters["quantization_type"],
+            quantization_ratio=parameters["quantization_ratio"],
+            m=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/pgvecto_rs/config.py

@@ -1,30 +1,53 @@
-from typing import Literal
+from abc import abstractmethod
+from typing import TypedDict
+
 from pydantic import BaseModel, SecretStr
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+from pgvecto_rs.types import IndexOption, Ivf, Hnsw, Flat, Quantization
+from pgvecto_rs.types.index import QuantizationType, QuantizationRatio
+
+from ..api import DBConfig, DBCaseConfig, IndexType, MetricType
 
 POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
 
 
+class PgVectorRSConfigDict(TypedDict):
+    """These keys will be directly used as kwargs in psycopg connection string,
+    so the names must match exactly psycopg API"""
+
+    user: str
+    password: str
+    host: str
+    port: int
+    dbname: str
+
+
 class PgVectoRSConfig(DBConfig):
-    user_name: SecretStr = "postgres"
+    user_name: str = "postgres"
     password: SecretStr
     host: str = "localhost"
     port: int = 5432
     db_name: str
 
     def to_dict(self) -> dict:
-        user_str = self.user_name.get_secret_value()
+        user_str = self.user_name
         pwd_str = self.password.get_secret_value()
         return {
             "host": self.host,
             "port": self.port,
             "dbname": self.db_name,
             "user": user_str,
-            "password": pwd_str
+            "password": pwd_str,
         }
 
+
 class PgVectoRSIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None
+    create_index_before_load: bool = False
+    create_index_after_load: bool = True
+
+    max_parallel_workers: int | None = None
+    quantization_type: QuantizationType | None = None
+    quantization_ratio: QuantizationRatio | None = None
 
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
@@ -40,88 +63,100 @@ class PgVectoRSIndexConfig(BaseModel, DBCaseConfig):
             return "<#>"
         return "<=>"
 
-class PgVectoRSQuantConfig(PgVectoRSIndexConfig):
-    quantizationType: Literal["trivial", "scalar", "product"]
-    quantizationRatio: None | Literal["x4", "x8", "x16", "x32", "x64"]
-
-    def parse_quantization(self) -> str:
-        if self.quantizationType == "trivial":
-            return "quantization = { trivial = { } }"
-        elif self.quantizationType == "scalar":
-            return "quantization = { scalar = { } }"
-        else:
-            return f'quantization = {{ product = {{ ratio = "{self.quantizationRatio}" }} }}'
-
+    def search_param(self) -> dict:
+        return {
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
 
-class HNSWConfig(PgVectoRSQuantConfig):
-    M: int
-    efConstruction: int
-    index: IndexType = IndexType.HNSW
+    @abstractmethod
+    def index_param(self) -> dict[str, str]: ...
 
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.hnsw]
-m = {self.M}
-ef_construction = {self.efConstruction}
-{self.parse_quantization()}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+    @abstractmethod
+    def session_param(self) -> dict[str, str | int]: ...
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
 
+class PgVectoRSHNSWConfig(PgVectoRSIndexConfig):
+    index: IndexType = IndexType.HNSW
+    m: int | None = None
+    ef_search: int | None
+    ef_construction: int | None = None
 
-class IVFFlatConfig(PgVectoRSQuantConfig):
-    nlist: int
-    nprobe: int | None = None
+    def index_param(self) -> dict[str, str]:
+        if self.quantization_type is None:
+            quantization = None
+        else:
+            quantization = Quantization(
+                typ=self.quantization_type, ratio=self.quantization_ratio
+            )
+
+        option = IndexOption(
+            index=Hnsw(
+                m=self.m,
+                ef_construction=self.ef_construction,
+                quantization=quantization,
+            ),
+            threads=self.max_parallel_workers,
+        )
+        return {"options": option.dumps(), "metric": self.parse_metric()}
+
+    def session_param(self) -> dict[str, str | int]:
+        session_parameters = {}
+        if self.ef_search is not None:
+            session_parameters["vectors.hnsw_ef_search"] = str(self.ef_search)
+        return session_parameters
+
+
+class PgVectoRSIVFFlatConfig(PgVectoRSIndexConfig):
     index: IndexType = IndexType.IVFFlat
+    probes: int | None
+    lists: int | None
+
+    def index_param(self) -> dict[str, str]:
+        if self.quantization_type is None:
+            quantization = None
+        else:
+            quantization = Quantization(
+                typ=self.quantization_type, ratio=self.quantization_ratio
+            )
 
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.ivf]
-nlist = {self.nlist}
-nsample = {self.nprobe if self.nprobe else 10}
-{self.parse_quantization()}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+        option = IndexOption(
+            index=Ivf(nlist=self.lists, quantization=quantization),
+            threads=self.max_parallel_workers,
+        )
+        return {"options": option.dumps(), "metric": self.parse_metric()}
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
-
-class IVFFlatSQ8Config(PgVectoRSIndexConfig):
-    nlist: int
-    nprobe: int | None = None
-    index: IndexType = IndexType.IVFSQ8
-
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.ivf]
-nlist = {self.nlist}
-nsample = {self.nprobe if self.nprobe else 10}
-quantization = {{ scalar = {{ }} }}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+    def session_param(self) -> dict[str, str | int]:
+        session_parameters = {}
+        if self.probes is not None:
+            session_parameters["vectors.ivf_nprobe"] = str(self.probes)
+        return session_parameters
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
 
-class FLATConfig(PgVectoRSQuantConfig):
+class PgVectoRSFLATConfig(PgVectoRSIndexConfig):
     index: IndexType = IndexType.Flat
 
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.flat]
-{self.parse_quantization()}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+    def index_param(self) -> dict[str, str]:
+        if self.quantization_type is None:
+            quantization = None
+        else:
+            quantization = Quantization(
+                typ=self.quantization_type, ratio=self.quantization_ratio
+            )
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
+        option = IndexOption(
+            index=Flat(
+                quantization=quantization,
+            ),
+            threads=self.max_parallel_workers,
+        )
+        return {"options": option.dumps(), "metric": self.parse_metric()}
+
+    def session_param(self) -> dict[str, str | int]:
+        return {}
 
 
 _pgvecto_rs_case_config = {
-    IndexType.HNSW: HNSWConfig,
-    IndexType.IVFFlat: IVFFlatConfig,
-    IndexType.IVFSQ8: IVFFlatSQ8Config,
-    IndexType.Flat: FLATConfig,
+    IndexType.HNSW: PgVectoRSHNSWConfig,
+    IndexType.IVFFlat: PgVectoRSIVFFlatConfig,
+    IndexType.Flat: PgVectoRSFLATConfig,
 }