vectordb-bench 0.0.29__py3-none-any.whl → 1.0.0__py3-none-any.whl

vectordb_bench/backend/clients/oceanbase/oceanbase.py

@@ -0,0 +1,215 @@
+import logging
+import struct
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+import mysql.connector as mysql
+
+from ..api import IndexType, VectorDB
+from .config import OceanBaseConfigDict, OceanBaseHNSWConfig
+
+log = logging.getLogger(__name__)
+
+OCEANBASE_DEFAULT_LOAD_BATCH_SIZE = 256
+
+
+class OceanBase(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: OceanBaseConfigDict,
+        db_case_config: OceanBaseHNSWConfig,
+        collection_name: str = "items",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.name = "OceanBase"
+        self.dim = dim
+        self.db_config = db_config
+        self.db_case_config = db_case_config
+        self.table_name = collection_name
+        self.load_batch_size = OCEANBASE_DEFAULT_LOAD_BATCH_SIZE
+        self._index_name = "vidx"
+        self._primary_field = "id"
+        self._vector_field = "embedding"
+
+        log.info(
+            f"{self.name} initialized with config:\nDatabase: {self.db_config}\nCase Config: {self.db_case_config}"
+        )
+
+        self._conn = None
+        self._cursor = None
+
+        try:
+            self._connect()
+            if drop_old:
+                self._drop_table()
+                self._create_table()
+        finally:
+            self._disconnect()
+
+    def _connect(self):
+        try:
+            self._conn = mysql.connect(
+                host=self.db_config["host"],
+                user=self.db_config["user"],
+                port=self.db_config["port"],
+                password=self.db_config["password"],
+                database=self.db_config["database"],
+            )
+            self._cursor = self._conn.cursor()
+        except mysql.Error:
+            log.exception("Failed to connect to the database")
+            raise
+
+    def _disconnect(self):
+        if self._cursor:
+            self._cursor.close()
+            self._cursor = None
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    @contextmanager
+    def init(self) -> Generator[None, None, None]:
+        try:
+            self._connect()
+            self._cursor.execute("SET autocommit=1")
+
+            if self.db_case_config.index in {IndexType.HNSW, IndexType.HNSW_SQ, IndexType.HNSW_BQ}:
+                self._cursor.execute(
+                    f"SET ob_hnsw_ef_search={(self.db_case_config.search_param())['params']['ef_search']}"
+                )
+            else:
+                self._cursor.execute(
+                    f"SET ob_ivf_nprobes={(self.db_case_config.search_param())['params']['ivf_nprobes']}"
+                )
+            yield
+        finally:
+            self._disconnect()
+
+    def _drop_table(self):
+        if not self._cursor:
+            raise ValueError("Cursor is not initialized")
+
+        log.info(f"Dropping table {self.table_name}")
+        self._cursor.execute(f"DROP TABLE IF EXISTS {self.table_name}")
+
+    def _create_table(self):
+        if not self._cursor:
+            raise ValueError("Cursor is not initialized")
+
+        log.info(f"Creating table {self.table_name}")
+        create_table_query = f"""
+        CREATE TABLE {self.table_name} (
+            id INT PRIMARY KEY,
+            embedding VECTOR({self.dim})
+        );
+        """
+        self._cursor.execute(create_table_query)
+
+    def optimize(self, data_size: int):
+        index_params = self.db_case_config.index_param()
+        index_args = ", ".join(f"{k}={v}" for k, v in index_params["params"].items())
+        index_query = (
+            f"CREATE /*+ PARALLEL(18) */ VECTOR INDEX idx1 "
+            f"ON {self.table_name}(embedding) "
+            f"WITH (distance={self.db_case_config.parse_metric()}, "
+            f"type={index_params['index_type']}, lib={index_params['lib']}, {index_args}"
+        )
+
+        if self.db_case_config.index in {IndexType.HNSW, IndexType.HNSW_SQ, IndexType.HNSW_BQ}:
+            index_query += ", extra_info_max_size=32"
+
+        index_query += ")"
+
+        log.info("Create index query: %s", index_query)
+
+        try:
+            log.info("Creating index...")
+            start_time = time.time()
+            self._cursor.execute(index_query)
+            log.info(f"Index created in {time.time() - start_time:.2f} seconds")
+
+            log.info("Performing major freeze...")
+            self._cursor.execute("ALTER SYSTEM MAJOR FREEZE;")
+            time.sleep(10)
+            self._wait_for_major_compaction()
+
+            log.info("Gathering schema statistics...")
+            self._cursor.execute("CALL dbms_stats.gather_schema_stats('test', degree => 96);")
+        except mysql.Error:
+            log.exception("Failed to optimize index")
+            raise
+
+    def need_normalize_cosine(self) -> bool:
+        if self.db_case_config.index == IndexType.HNSW_BQ:
+            log.info("current HNSW_BQ only supports L2, cosine dataset need normalize.")
+            return True
+
+        return False
+
+    def _wait_for_major_compaction(self):
+        while True:
+            self._cursor.execute(
+                "SELECT IF(COUNT(*) = COUNT(STATUS = 'IDLE' OR NULL), 'TRUE', 'FALSE') "
+                "AS all_status_idle FROM oceanbase.DBA_OB_ZONE_MAJOR_COMPACTION;"
+            )
+            all_status_idle = self._cursor.fetchone()[0]
+            if all_status_idle == "TRUE":
+                break
+            time.sleep(10)
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> tuple[int, Exception | None]:
+        if not self._cursor:
+            raise ValueError("Cursor is not initialized")
+
+        insert_count = 0
+        try:
+            for batch_start in range(0, len(embeddings), self.load_batch_size):
+                batch_end = min(batch_start + self.load_batch_size, len(embeddings))
+                batch = [(metadata[i], embeddings[i]) for i in range(batch_start, batch_end)]
+                values = ", ".join(f"({item_id}, '[{','.join(map(str, embedding))}]')" for item_id, embedding in batch)
+                self._cursor.execute(
+                    f"INSERT /*+ ENABLE_PARALLEL_DML PARALLEL(32) */ INTO {self.table_name} VALUES {values}"  # noqa: S608
+                )
+                insert_count += len(batch)
+        except mysql.Error:
+            log.exception("Failed to insert embeddings")
+            raise
+
+        return insert_count, None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict[str, Any] | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        if not self._cursor:
+            raise ValueError("Cursor is not initialized")
+
+        packed = struct.pack(f"<{len(query)}f", *query)
+        hex_vec = packed.hex()
+        filter_clause = f"WHERE id >= {filters['id']}" if filters else ""
+        query_str = (
+            f"SELECT id FROM {self.table_name} "  # noqa: S608
+            f"{filter_clause} ORDER BY "
+            f"{self.db_case_config.parse_metric_func_str()}(embedding, X'{hex_vec}') "
+            f"APPROXIMATE LIMIT {k}"
+        )
+
+        try:
+            self._cursor.execute(query_str)
+            return [row[0] for row in self._cursor.fetchall()]
+        except mysql.Error:
+            log.exception("Failed to execute search query")
+            raise

vectordb_bench/backend/clients/pinecone/pinecone.py

@@ -5,8 +5,9 @@ from contextlib import contextmanager
 
 import pinecone
 
-from ..api import DBCaseConfig, DBConfig, EmptyDBCaseConfig, IndexType, VectorDB
-from .config import PineconeConfig
+from vectordb_bench.backend.filter import Filter, FilterOp
+
+from ..api import DBCaseConfig, VectorDB
 
 log = logging.getLogger(__name__)
 
@@ -15,12 +16,19 @@ PINECONE_MAX_SIZE_PER_BATCH = 2 * 1024 * 1024  # 2MB
 
 
 class Pinecone(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
     def __init__(
         self,
         dim: int,
         db_config: dict,
         db_case_config: DBCaseConfig,
         drop_old: bool = False,
+        with_scalar_labels: bool = False,
         **kwargs,
     ):
         """Initialize wrapper around the milvus vector database."""
@@ -33,6 +41,7 @@ class Pinecone(VectorDB):
         pc = pinecone.Pinecone(api_key=self.api_key)
         index = pc.Index(self.index_name)
 
+        self.with_scalar_labels = with_scalar_labels
         if drop_old:
             index_stats = index.describe_index_stats()
             index_dim = index_stats["dimension"]
@@ -43,15 +52,8 @@ class Pinecone(VectorDB):
                 log.info(f"Pinecone index delete namespace: {namespace}")
                 index.delete(delete_all=True, namespace=namespace)
 
-        self._metadata_key = "meta"
-
-    @classmethod
-    def config_cls(cls) -> type[DBConfig]:
-        return PineconeConfig
-
-    @classmethod
-    def case_config_cls(cls, index_type: IndexType | None = None) -> type[DBCaseConfig]:
-        return EmptyDBCaseConfig
+        self._scalar_id_field = "meta"
+        self._scalar_label_field = "label"
 
     @contextmanager
     def init(self):
@@ -66,6 +68,7 @@ class Pinecone(VectorDB):
         self,
         embeddings: list[list[float]],
         metadata: list[int],
+        labels_data: list[str] | None = None,
         **kwargs,
     ) -> tuple[int, Exception]:
         assert len(embeddings) == len(metadata)
@@ -75,33 +78,44 @@ class Pinecone(VectorDB):
                 batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
                 insert_datas = []
                 for i in range(batch_start_offset, batch_end_offset):
+                    metadata_dict = {self._scalar_id_field: metadata[i]}
+                    if self.with_scalar_labels:
+                        metadata_dict[self._scalar_label_field] = labels_data[i]
                     insert_data = (
                         str(metadata[i]),
                         embeddings[i],
-                        {self._metadata_key: metadata[i]},
+                        metadata_dict,
                     )
                     insert_datas.append(insert_data)
                 self.index.upsert(insert_datas)
                 insert_count += batch_end_offset - batch_start_offset
         except Exception as e:
-            return (insert_count, e)
-        return (len(embeddings), None)
+            return insert_count, e
+        return len(embeddings), None
 
     def search_embedding(
         self,
         query: list[float],
         k: int = 100,
-        filters: dict | None = None,
         timeout: int | None = None,
     ) -> list[int]:
-        pinecone_filters = {} if filters is None else {self._metadata_key: {"$gte": filters["id"]}}
-        try:
-            res = self.index.query(
-                top_k=k,
-                vector=query,
-                filter=pinecone_filters,
-            )["matches"]
-        except Exception as e:
-            log.warning(f"Error querying index: {e}")
-            raise e from e
+        pinecone_filters = self.expr
+        res = self.index.query(
+            top_k=k,
+            vector=query,
+            filter=pinecone_filters,
+        )["matches"]
         return [int(one_res["id"]) for one_res in res]
+
+    def prepare_filter(self, filters: Filter):
+        if filters.type == FilterOp.NonFilter:
+            self.expr = None
+        elif filters.type == FilterOp.NumGE:
+            self.expr = {self._scalar_id_field: {"$gte": filters.int_value}}
+        elif filters.type == FilterOp.StrEqual:
+            # both "in" and "==" are supported
+            # for example, self.expr = {self._scalar_label_field: {"$in": [filters.label_value]}}
+            self.expr = {self._scalar_label_field: {"$eq": filters.label_value}}
+        else:
+            msg = f"Not support Filter for Pinecone - {filters}"
+            raise ValueError(msg)

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -1,7 +1,12 @@
-from pydantic import BaseModel, SecretStr
+from typing import TypeVar
+
+from pydantic import BaseModel, SecretStr, validator
 
 from ..api import DBCaseConfig, DBConfig, MetricType
 
+# define type "SearchParams"
+SearchParams = TypeVar("SearchParams")
+
 
 # Allowing `api_key` to be left empty, to ensure compatibility with the open-source Qdrant.
 class QdrantConfig(DBConfig):
@@ -20,9 +25,43 @@ class QdrantConfig(DBConfig):
             "url": self.url.get_secret_value(),
         }
 
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        if field.name in ["api_key"]:
+            return v
+        return super().not_empty_field(v, field)
+
 
 class QdrantIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None
+    m: int = 16
+    payload_m: int = 16  # only for label_filter cases
+    create_payload_int_index: bool = False
+    create_payload_keyword_index: bool = False
+    is_tenant: bool = False
+    use_scalar_quant: bool = False
+    sq_quantile: float = 0.99
+    default_segment_number: int = 0
+
+    use_rescore: bool = False
+    oversampling: float = 1.0
+    indexed_only: bool = False
+    hnsw_ef: int | None = 100
+    exact: bool = False
+
+    with_payload: bool = False
+
+    def __eq__(self, obj: any):
+        return (
+            self.m == obj.m
+            and self.payload_m == obj.payload_m
+            and self.create_payload_int_index == obj.create_payload_int_index
+            and self.create_payload_keyword_index == obj.create_payload_keyword_index
+            and self.is_tenant == obj.is_tenant
+            and self.use_scalar_quant == obj.use_scalar_quant
+            and self.sq_quantile == obj.sq_quantile
+            and self.default_segment_number == obj.default_segment_number
+        )
 
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
@@ -36,5 +75,22 @@ class QdrantIndexConfig(BaseModel, DBCaseConfig):
     def index_param(self) -> dict:
         return {"distance": self.parse_metric()}
 
-    def search_param(self) -> dict:
-        return {}
+    def search_param(self) -> SearchParams:
+        # Import while in use
+        from qdrant_client.http.models import QuantizationSearchParams, SearchParams
+
+        quantization = (
+            QuantizationSearchParams(
+                ignore=False,
+                rescore=True,
+                oversampling=self.oversampling,
+            )
+            if self.use_rescore
+            else None
+        )
+        return SearchParams(
+            hnsw_ef=self.hnsw_ef,
+            exact=self.exact,
+            indexed_only=self.indexed_only,
+            quantization=quantization,
+        )

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -9,13 +9,24 @@ from qdrant_client.http.models import (
     Batch,
     CollectionStatus,
     FieldCondition,
-    Filter,
+    HnswConfigDiff,
+    KeywordIndexParams,
+    OptimizersConfigDiff,
     PayloadSchemaType,
     Range,
+    ScalarQuantization,
+    ScalarQuantizationConfig,
+    ScalarType,
     VectorParams,
 )
+from qdrant_client.http.models import (
+    Filter as QdrantFilter,
+)
 
-from ..api import DBCaseConfig, VectorDB
+from vectordb_bench.backend.clients.qdrant_cloud.config import QdrantIndexConfig
+from vectordb_bench.backend.filter import Filter, FilterOp
+
+from ..api import VectorDB
 
 log = logging.getLogger(__name__)
 
@@ -25,24 +36,33 @@ QDRANT_BATCH_SIZE = 500
 
 
 class QdrantCloud(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
     def __init__(
         self,
         dim: int,
         db_config: dict,
-        db_case_config: DBCaseConfig,
+        db_case_config: QdrantIndexConfig,
         collection_name: str = "QdrantCloudCollection",
         drop_old: bool = False,
+        with_scalar_labels: bool = False,
         **kwargs,
     ):
         """Initialize wrapper around the QdrantCloud vector database."""
         self.db_config = db_config
-        self.case_config = db_case_config
+        self.db_case_config = db_case_config
         self.collection_name = collection_name
 
         self._primary_field = "pk"
+        self._scalar_label_field = "label"
         self._vector_field = "vector"
 
         tmp_client = QdrantClient(**self.db_config)
+        self.with_scalar_labels = with_scalar_labels
         if drop_old:
             log.info(f"QdrantCloud client drop_old collection: {self.collection_name}")
             tmp_client.delete_collection(self.collection_name)
@@ -50,7 +70,7 @@ class QdrantCloud(VectorDB):
         tmp_client = None
 
     @contextmanager
-    def init(self) -> None:
+    def init(self):
         """
         Examples:
             >>> with self.init():
@@ -74,7 +94,7 @@ class QdrantCloud(VectorDB):
                 if info.status == CollectionStatus.GREEN:
                     msg = (
                         f"Stored vectors: {info.vectors_count}, Indexed vectors: {info.indexed_vectors_count}, "
-                        f"Collection status: {info.indexed_vectors_count}"
+                        f"Collection status: {info.status}, Segment counts: {info.segments_count}"
                     )
                     log.info(msg)
                     return
@@ -86,19 +106,48 @@ class QdrantCloud(VectorDB):
         log.info(f"Create collection: {self.collection_name}")
 
         try:
+            # whether to use quant (SQ8)
+            quantization_config = None
+            if self.db_case_config.use_scalar_quant:
+                quantization_config = ScalarQuantization(
+                    scalar=ScalarQuantizationConfig(
+                        type=ScalarType.INT8,
+                        quantile=self.db_case_config.sq_quantile,
+                        always_ram=True,
+                    )
+                )
+
+            # create collection
             qdrant_client.create_collection(
                 collection_name=self.collection_name,
                 vectors_config=VectorParams(
                     size=dim,
-                    distance=self.case_config.index_param()["distance"],
+                    distance=self.db_case_config.parse_metric(),
+                ),
+                hnsw_config=HnswConfigDiff(m=self.db_case_config.m, payload_m=self.db_case_config.payload_m),
+                optimizers_config=OptimizersConfigDiff(
+                    default_segment_number=self.db_case_config.default_segment_number
                 ),
+                quantization_config=quantization_config,
             )
 
-            qdrant_client.create_payload_index(
-                collection_name=self.collection_name,
-                field_name=self._primary_field,
-                field_schema=PayloadSchemaType.INTEGER,
-            )
+            # create payload_index for int-field
+            if self.db_case_config.create_payload_int_index:
+                qdrant_client.create_payload_index(
+                    collection_name=self.collection_name,
+                    field_name=self._primary_field,
+                    field_schema=PayloadSchemaType.INTEGER,
+                )
+
+            # create payload_index for str-field
+            if self.with_scalar_labels and self.db_case_config.create_payload_keyword_index:
+                qdrant_client.create_payload_index(
+                    collection_name=self.collection_name,
+                    field_name=self._scalar_label_field,
+                    field_schema=KeywordIndexParams(
+                        type=PayloadSchemaType.KEYWORD, is_tenant=self.db_case_config.is_tenant
+                    ),
+                )
 
         except Exception as e:
             if "already exists!" in str(e):
@@ -110,16 +159,22 @@ class QdrantCloud(VectorDB):
         self,
         embeddings: list[list[float]],
         metadata: list[int],
+        labels_data: list[str] | None = None,
         **kwargs,
     ) -> tuple[int, Exception]:
         """Insert embeddings into Milvus. should call self.init() first"""
         assert self.qdrant_client is not None
         try:
-            # TODO: counts
             for offset in range(0, len(embeddings), QDRANT_BATCH_SIZE):
                 vectors = embeddings[offset : offset + QDRANT_BATCH_SIZE]
                 ids = metadata[offset : offset + QDRANT_BATCH_SIZE]
-                payloads = [{self._primary_field: v} for v in ids]
+                if self.with_scalar_labels:
+                    labels = labels_data[offset : offset + QDRANT_BATCH_SIZE]
+                    payloads = [
+                        {self._primary_field: pk, self._scalar_label_field: labels[i]} for i, pk in enumerate(ids)
+                    ]
+                else:
+                    payloads = [{self._primary_field: pk} for i, pk in enumerate(ids)]
                 _ = self.qdrant_client.upsert(
                     collection_name=self.collection_name,
                     wait=True,
@@ -135,34 +190,46 @@ class QdrantCloud(VectorDB):
         self,
         query: list[float],
         k: int = 100,
-        filters: dict | None = None,
         timeout: int | None = None,
+        **kwargs,
     ) -> list[int]:
         """Perform a search on a query embedding and return results with score.
         Should call self.init() first.
         """
         assert self.qdrant_client is not None
 
-        f = None
-        if filters:
-            f = Filter(
+        res = self.qdrant_client.search(
+            collection_name=self.collection_name,
+            query_vector=query,
+            limit=k,
+            query_filter=self.query_filter,
+            search_params=self.db_case_config.search_param(),
+            with_payload=self.db_case_config.with_payload,
+        )
+
+        return [r.id for r in res]
+
+    def prepare_filter(self, filters: Filter):
+        if filters.type == FilterOp.NonFilter:
+            self.query_filter = None
+        elif filters.type == FilterOp.NumGE:
+            self.query_filter = QdrantFilter(
                 must=[
                     FieldCondition(
                         key=self._primary_field,
-                        range=Range(
-                            gt=filters.get("id"),
-                        ),
+                        range=Range(gte=filters.int_value),
                     ),
-                ],
+                ]
             )
-
-        res = (
-            self.qdrant_client.search(
-                collection_name=self.collection_name,
-                query_vector=query,
-                limit=k,
-                query_filter=f,
-            ),
-        )
-
-        return [result.id for result in res[0]]
+        elif filters.type == FilterOp.StrEqual:
+            self.query_filter = QdrantFilter(
+                must=[
+                    FieldCondition(
+                        key=self._scalar_label_field,
+                        match={"value": filters.label_value},
+                    ),
+                ]
+            )
+        else:
+            msg = f"Not support Filter for Qdrant - {filters}"
+            raise ValueError(msg)