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

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)