vectordb-bench 0.0.25__tar.gz → 0.0.26__tar.gz

{vectordb_bench-0.0.25 → vectordb_bench-0.0.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vectordb-bench
-Version: 0.0.25
+Version: 0.0.26
 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

{vectordb_bench-0.0.25 → vectordb_bench-0.0.26}/vectordb_bench/backend/clients/api.py

@@ -16,10 +16,15 @@ class MetricType(str, Enum):
 
 class IndexType(str, Enum):
     HNSW = "HNSW"
+    HNSW_SQ = "HNSW_SQ"
+    HNSW_PQ = "HNSW_PQ"
+    HNSW_PRQ = "HNSW_PRQ"
     DISKANN = "DISKANN"
     STREAMING_DISKANN = "DISKANN"
     IVFFlat = "IVF_FLAT"
+    IVFPQ = "IVF_PQ"
     IVFSQ8 = "IVF_SQ8"
+    IVF_RABITQ = "IVF_RABITQ"
     Flat = "FLAT"
     AUTOINDEX = "AUTOINDEX"
     ES_HNSW = "hnsw"
@@ -31,6 +36,14 @@ class IndexType(str, Enum):
     SCANN = "scann"
 
 
+class SQType(str, Enum):
+    SQ6 = "SQ6"
+    SQ8 = "SQ8"
+    BF16 = "BF16"
+    FP16 = "FP16"
+    FP32 = "FP32"
+
+
 class DBConfig(ABC, BaseModel):
     """DBConfig contains the connection info of vector database
 

vectordb_bench-0.0.26/vectordb_bench/backend/clients/clickhouse/clickhouse.py

@@ -0,0 +1,232 @@
+"""Wrapper around the Clickhouse vector database over VectorDB"""
+
+import logging
+from contextlib import contextmanager
+from typing import Any
+
+import clickhouse_connect
+from clickhouse_connect.driver import Client
+
+from .. import IndexType
+from ..api import VectorDB
+from .config import ClickhouseConfigDict, ClickhouseIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class Clickhouse(VectorDB):
+    """Use SQLAlchemy instructions"""
+
+    def __init__(
+        self,
+        dim: int,
+        db_config: ClickhouseConfigDict,
+        db_case_config: ClickhouseIndexConfig,
+        collection_name: str = "CHVectorCollection",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+
+        self.index_param = self.case_config.index_param()
+        self.search_param = self.case_config.search_param()
+        self.session_param = self.case_config.session_param()
+
+        self._index_name = "clickhouse_index"
+        self._primary_field = "id"
+        self._vector_field = "embedding"
+
+        # construct basic units
+        self.conn = self._create_connection(**self.db_config, settings=self.session_param)
+
+        if drop_old:
+            log.info(f"Clickhouse client drop table : {self.table_name}")
+            self._drop_table()
+            self._create_table(dim)
+            if self.case_config.create_index_before_load:
+                self._create_index()
+
+        self.conn.close()
+        self.conn = None
+
+    @contextmanager
+    def init(self) -> None:
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+
+        self.conn = self._create_connection(**self.db_config, settings=self.session_param)
+
+        try:
+            yield
+        finally:
+            self.conn.close()
+            self.conn = None
+
+    def _create_connection(self, settings: dict | None, **kwargs) -> Client:
+        return clickhouse_connect.get_client(**self.db_config, settings=settings)
+
+    def _drop_index(self):
+        assert self.conn is not None, "Connection is not initialized"
+        try:
+            self.conn.command(
+                f'ALTER TABLE {self.db_config["database"]}.{self.table_name} DROP INDEX {self._index_name}'
+            )
+        except Exception as e:
+            log.warning(f"Failed to drop index on table {self.db_config['database']}.{self.table_name}: {e}")
+            raise e from None
+
+    def _drop_table(self):
+        assert self.conn is not None, "Connection is not initialized"
+
+        try:
+            self.conn.command(f'DROP TABLE IF EXISTS {self.db_config["database"]}.{self.table_name}')
+        except Exception as e:
+            log.warning(f"Failed to drop table {self.db_config['database']}.{self.table_name}: {e}")
+            raise e from None
+
+    def _perfomance_tuning(self):
+        self.conn.command("SET materialize_skip_indexes_on_insert = 1")
+
+    def _create_index(self):
+        assert self.conn is not None, "Connection is not initialized"
+        try:
+            if self.index_param["index_type"] == IndexType.HNSW.value:
+                if (
+                    self.index_param["quantization"]
+                    and self.index_param["params"]["M"]
+                    and self.index_param["params"]["efConstruction"]
+                ):
+                    query = f"""
+                        ALTER TABLE {self.db_config["database"]}.{self.table_name}
+                        ADD INDEX {self._index_name} {self._vector_field}
+                        TYPE vector_similarity('hnsw', '{self.index_param["metric_type"]}',
+                        '{self.index_param["quantization"]}',
+                        {self.index_param["params"]["M"]}, {self.index_param["params"]["efConstruction"]})
+                        GRANULARITY {self.index_param["granularity"]}
+                        """
+                else:
+                    query = f"""
+                        ALTER TABLE {self.db_config["database"]}.{self.table_name}
+                        ADD INDEX {self._index_name} {self._vector_field}
+                        TYPE vector_similarity('hnsw', '{self.index_param["metric_type"]}')
+                        GRANULARITY {self.index_param["granularity"]}
+                        """
+                self.conn.command(cmd=query)
+            else:
+                log.warning("HNSW is only avaliable method in clickhouse now")
+        except Exception as e:
+            log.warning(f"Failed to create Clickhouse vector index on table: {self.table_name} error: {e}")
+            raise e from None
+
+    def _create_table(self, dim: int):
+        assert self.conn is not None, "Connection is not initialized"
+
+        try:
+            # create table
+            self.conn.command(
+                f'CREATE TABLE IF NOT EXISTS {self.db_config["database"]}.{self.table_name} '
+                f"({self._primary_field} UInt32, "
+                f'{self._vector_field} Array({self.index_param["vector_data_type"]}) CODEC(NONE), '
+                f"CONSTRAINT same_length CHECK length(embedding) = {dim}) "
+                f"ENGINE = MergeTree() "
+                f"ORDER BY {self._primary_field}"
+            )
+
+        except Exception as e:
+            log.warning(f"Failed to create Clickhouse table: {self.table_name} error: {e}")
+            raise e from None
+
+    def optimize(self, data_size: int | None = None):
+        pass
+
+    def _post_insert(self):
+        pass
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> (int, Exception):
+        assert self.conn is not None, "Connection is not initialized"
+
+        try:
+            # do not iterate for bulk insert
+            items = [metadata, embeddings]
+
+            self.conn.insert(
+                table=self.table_name,
+                data=items,
+                column_names=["id", "embedding"],
+                column_type_names=["UInt32", f'Array({self.index_param["vector_data_type"]})'],
+                column_oriented=True,
+            )
+            return len(metadata), None
+        except Exception as e:
+            log.warning(f"Failed to insert data into Clickhouse table ({self.table_name}), error: {e}")
+            return 0, e
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        assert self.conn is not None, "Connection is not initialized"
+        parameters = {
+            "primary_field": self._primary_field,
+            "vector_field": self._vector_field,
+            "schema": self.db_config["database"],
+            "table": self.table_name,
+            "gt": filters.get("id"),
+            "k": k,
+            "metric_type": self.search_param["metric_type"],
+            "query": query,
+        }
+        if self.case_config.metric_type == "COSINE":
+            if filters:
+                result = self.conn.query(
+                    "SELECT {primary_field:Identifier}, {vector_field:Identifier} "
+                    "FROM {schema:Identifier}.{table:Identifier} "
+                    "WHERE {primary_field:Identifier} > {gt:UInt32} "
+                    "ORDER BY cosineDistance(embedding,{query:Array(Float64)}) "
+                    "LIMIT {k:UInt32}",
+                    parameters=parameters,
+                ).result_rows
+                return [int(row[0]) for row in result]
+
+            result = self.conn.query(
+                "SELECT {primary_field:Identifier}, {vector_field:Identifier} "
+                "FROM {schema:Identifier}.{table:Identifier} "
+                "ORDER BY cosineDistance(embedding,{query:Array(Float64)}) "
+                "LIMIT {k:UInt32}",
+                parameters=parameters,
+            ).result_rows
+            return [int(row[0]) for row in result]
+        if filters:
+            result = self.conn.query(
+                "SELECT {primary_field:Identifier}, {vector_field:Identifier} "
+                "FROM {schema:Identifier}.{table:Identifier} "
+                "WHERE {primary_field:Identifier} > {gt:UInt32} "
+                "ORDER BY L2Distance(embedding,{query:Array(Float64)}) "
+                "LIMIT {k:UInt32}",
+                parameters=parameters,
+            ).result_rows
+            return [int(row[0]) for row in result]
+
+        result = self.conn.query(
+            "SELECT {primary_field:Identifier}, {vector_field:Identifier} "
+            "FROM {schema:Identifier}.{table:Identifier} "
+            "ORDER BY L2Distance(embedding,{query:Array(Float64)}) "
+            "LIMIT {k:UInt32}",
+            parameters=parameters,
+        ).result_rows
+        return [int(row[0]) for row in result]

vectordb_bench-0.0.26/vectordb_bench/backend/clients/clickhouse/config.py

@@ -0,0 +1,89 @@
+from abc import abstractmethod
+from typing import TypedDict
+
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class ClickhouseConfigDict(TypedDict):
+    user: str
+    password: str
+    host: str
+    port: int
+    database: str
+    secure: bool
+
+
+class ClickhouseConfig(DBConfig):
+    user_name: str = "clickhouse"
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 8123
+    db_name: str = "default"
+    secure: bool = False
+
+    def to_dict(self) -> ClickhouseConfigDict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "database": self.db_name,
+            "user": self.user_name,
+            "password": pwd_str,
+            "secure": self.secure,
+        }
+
+
+class ClickhouseIndexConfig(BaseModel, DBCaseConfig):
+
+    metric_type: MetricType | None = None
+    vector_data_type: str | None = "Float32"  # Data type of vectors. Can be Float32 or Float64 or BFloat16
+    create_index_before_load: bool = True
+    create_index_after_load: bool = False
+
+    def parse_metric(self) -> str:
+        if not self.metric_type:
+            return ""
+        return self.metric_type.value
+
+    def parse_metric_str(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "L2Distance"
+        if self.metric_type == MetricType.COSINE:
+            return "cosineDistance"
+        return "cosineDistance"
+
+    @abstractmethod
+    def session_param(self):
+        pass
+
+
+class ClickhouseHNSWConfig(ClickhouseIndexConfig):
+    M: int | None  # Default in clickhouse in 32
+    efConstruction: int | None  # Default in clickhouse in 128
+    ef: int | None = None
+    index: IndexType = IndexType.HNSW
+    quantization: str | None = "bf16"  # Default is bf16. Possible values are f64, f32, f16, bf16, or i8
+    granularity: int | None = 10_000_000  # Size of the index granules. By default, in CH it's equal 10.000.000
+
+    def index_param(self) -> dict:
+        return {
+            "vector_data_type": self.vector_data_type,
+            "metric_type": self.parse_metric_str(),
+            "index_type": self.index.value,
+            "quantization": self.quantization,
+            "granularity": self.granularity,
+            "params": {"M": self.M, "efConstruction": self.efConstruction},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric_str(),
+            "params": {"ef": self.ef},
+        }
+
+    def session_param(self) -> dict:
+        return {
+            "allow_experimental_vector_similarity_index": 1,
+        }

{vectordb_bench-0.0.25 → vectordb_bench-0.0.26}/vectordb_bench/backend/clients/milvus/config.py

@@ -1,6 +1,6 @@
 from pydantic import BaseModel, SecretStr, validator
 
-from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType
 
 
 class MilvusConfig(DBConfig):
@@ -88,6 +88,88 @@ class HNSWConfig(MilvusIndexConfig, DBCaseConfig):
         }
 
 
+class HNSWSQConfig(HNSWConfig, DBCaseConfig):
+    index: IndexType = IndexType.HNSW_SQ
+    sq_type: SQType = SQType.SQ8
+    refine: bool = True
+    refine_type: SQType = SQType.FP32
+    refine_k: float = 1
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {
+                "M": self.M,
+                "efConstruction": self.efConstruction,
+                "sq_type": self.sq_type.value,
+                "refine": self.refine,
+                "refine_type": self.refine_type.value,
+            },
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"ef": self.ef, "refine_k": self.refine_k},
+        }
+
+
+class HNSWPQConfig(HNSWConfig):
+    index: IndexType = IndexType.HNSW_PQ
+    m: int = 32
+    nbits: int = 8
+    refine: bool = True
+    refine_type: SQType = SQType.FP32
+    refine_k: float = 1
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {
+                "M": self.M,
+                "efConstruction": self.efConstruction,
+                "m": self.m,
+                "nbits": self.nbits,
+                "refine": self.refine,
+                "refine_type": self.refine_type.value,
+            },
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"ef": self.ef, "refine_k": self.refine_k},
+        }
+
+
+class HNSWPRQConfig(HNSWPQConfig):
+    index: IndexType = IndexType.HNSW_PRQ
+    nrq: int = 2
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {
+                "M": self.M,
+                "efConstruction": self.efConstruction,
+                "m": self.m,
+                "nbits": self.nbits,
+                "nrq": self.nrq,
+                "refine": self.refine,
+                "refine_type": self.refine_type.value,
+            },
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"ef": self.ef, "refine_k": self.refine_k},
+        }
+
+
 class DISKANNConfig(MilvusIndexConfig, DBCaseConfig):
     search_list: int | None = None
     index: IndexType = IndexType.DISKANN
@@ -125,6 +207,27 @@ class IVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
         }
 
 
+class IVFPQConfig(MilvusIndexConfig, DBCaseConfig):
+    nlist: int
+    nprobe: int | None = None
+    m: int = 32
+    nbits: int = 8
+    index: IndexType = IndexType.IVFPQ
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"nlist": self.nlist, "m": self.m, "nbits": self.nbits},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"nprobe": self.nprobe},
+        }
+
+
 class IVFSQ8Config(MilvusIndexConfig, DBCaseConfig):
     nlist: int
     nprobe: int | None = None
@@ -144,6 +247,31 @@ class IVFSQ8Config(MilvusIndexConfig, DBCaseConfig):
         }
 
 
+class IVFRABITQConfig(IVFSQ8Config):
+    index: IndexType = IndexType.IVF_RABITQ
+    rbq_bits_query: int = 0  # 0, 1, 2, ..., 8
+    refine: bool = True
+    refine_type: SQType = SQType.FP32
+    refine_k: float = 1
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {
+                "nlist": self.nlist,
+                "refine": self.refine,
+                "refine_type": self.refine_type.value,
+            },
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"nprobe": self.nprobe, "rbq_bits_query": self.rbq_bits_query, "refine_k": self.refine_k},
+        }
+
+
 class FLATConfig(MilvusIndexConfig, DBCaseConfig):
     index: IndexType = IndexType.Flat
 
@@ -285,9 +413,14 @@ class GPUCAGRAConfig(MilvusIndexConfig, DBCaseConfig):
 _milvus_case_config = {
     IndexType.AUTOINDEX: AutoIndexConfig,
     IndexType.HNSW: HNSWConfig,
+    IndexType.HNSW_SQ: HNSWSQConfig,
+    IndexType.HNSW_PQ: HNSWPQConfig,
+    IndexType.HNSW_PRQ: HNSWPRQConfig,
     IndexType.DISKANN: DISKANNConfig,
     IndexType.IVFFlat: IVFFlatConfig,
+    IndexType.IVFPQ: IVFPQConfig,
     IndexType.IVFSQ8: IVFSQ8Config,
+    IndexType.IVF_RABITQ: IVFRABITQConfig,
     IndexType.Flat: FLATConfig,
     IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
     IndexType.GPU_IVF_PQ: GPUIVFPQConfig,

{vectordb_bench-0.0.25 → vectordb_bench-0.0.26}/vectordb_bench/backend/clients/milvus/milvus.py

@@ -61,6 +61,7 @@ class Milvus(VectorDB):
                 consistency_level="Session",
             )
 
+            log.info(f"{self.name} create index: index_params: {self.case_config.index_param()}")
             col.create_index(
                 self._vector_field,
                 self.case_config.index_param(),
@@ -71,7 +72,7 @@ class Milvus(VectorDB):
         connections.disconnect("default")
 
     @contextmanager
-    def init(self) -> None:
+    def init(self):
         """
         Examples:
             >>> with self.init():
@@ -126,6 +127,7 @@ class Milvus(VectorDB):
                 try:
                     self.col.compact()
                     self.col.wait_for_compaction_completed()
+                    log.info("compactation completed. waiting for the rest of index buliding.")
                 except Exception as e:
                     log.warning(f"{self.name} compact error: {e}")
                     if hasattr(e, "code"):