pyobvector 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

pyobvector/__init__.py

@@ -14,6 +14,7 @@ In this mode, you can regard `pyobvector` as an extension of SQLAlchemy.
 * IndexParams           A list of IndexParam to create vector index in batch
 * DataType              Specify field type in collection schema for MilvusLikeClient
 * VECTOR                An extended data type in SQLAlchemy for ObVecClient
+* SPARSE_VECTOR         An extended data type in SQLAlchemy for ObVecClient
 * VectorIndex           An extended index type in SQLAlchemy for ObVecClient
 * FtsIndex              Full Text Search Index
 * FieldSchema           Clas to define field schema in collection for MilvusLikeClient
@@ -43,6 +44,7 @@ from .client import *
 from .schema import (
     ARRAY,
     VECTOR,
+    SPARSE_VECTOR,
     POINT,
     VectorIndex,
     OceanBaseDialect,
@@ -70,6 +72,7 @@ __all__ = [
     "DataType",
     "ARRAY",
     "VECTOR",
+    "SPARSE_VECTOR",
     "POINT",
     "VectorIndex",
     "FtsIndex",

pyobvector/client/collection_schema.py

@@ -79,14 +79,14 @@ class FieldSchema:
             if "max_length" not in self.kwargs:
                 raise VarcharFieldParamException(
                     code=ErrorCode.INVALID_ARGUMENT,
-                    message=ExceptionsMessage.VarcharFieldMissinglengthParam,
+                    message=ExceptionsMessage.VarcharFieldMissingLengthParam,
                 )
             self.type_params["length"] = self.kwargs["max_length"]
         elif self.dtype == DataType.ARRAY:
             if "element_type" not in self.kwargs:
                 raise ArrayFieldParamException(
                     code=ErrorCode.INVALID_ARGUMENT,
-                    message=ExceptionsMessage.ArrayFiledMissingElementType,
+                    message=ExceptionsMessage.ArrayFieldMissingElementType,
                 )
             if self.kwargs["element_type"] in (
                 DataType.ARRAY,
@@ -95,7 +95,7 @@ class FieldSchema:
             ):
                 raise ArrayFieldParamException(
                     code=ErrorCode.INVALID_ARGUMENT,
-                    message=ExceptionsMessage.ArrayFiledInvalidElementType,
+                    message=ExceptionsMessage.ArrayFieldInvalidElementType,
                 )
 
             self.type_params["item_type"] = convert_datatype_to_sqltype(
@@ -147,9 +147,9 @@ class CollectionSchema:
         """Add field to collection.
 
         Args:
-        :param field_name (string) : new field name
-        :param datatype (DataType) : field data type
-        :param kwargs : parameters for data type
+            field_name (string): new field name
+            datatype (DataType): field data type
+            **kwargs: parameters for data type
         """
         field = FieldSchema(field_name, datatype, **kwargs)
         cur_idx = len(self.fields)

pyobvector/client/exceptions.py

@@ -101,14 +101,14 @@ class ExceptionsMessage:
     )
     PrimaryFieldType = "Param primary_field must be int or str type."
     VectorFieldMissingDimParam = "Param 'dim' must be set for vector field."
-    VarcharFieldMissinglengthParam = "Param 'max_length' must be set for varchar field."
-    ArrayFiledMissingElementType = "Param 'element_type' must be set for array field."
-    ArrayFiledInvalidElementType = (
+    VarcharFieldMissingLengthParam = "Param 'max_length' must be set for varchar field."
+    ArrayFieldMissingElementType = "Param 'element_type' must be set for array field."
+    ArrayFieldInvalidElementType = (
         "Param 'element_type' can not be array/vector/varchar."
     )
     CollectionNotExists = "Collection does not exist."
     MetricTypeParamTypeInvalid = "MetricType param type should be string."
-    MetricTypeValueInvalid = "MetricType should be 'l2'/'ip' in ann search."
+    MetricTypeValueInvalid = "MetricType should be 'l2'/'ip'/'neg_ip'/'cosine' in ann search."
     UsingInIDsWhenMultiPrimaryKey = "Using 'ids' when table has multi primary key."
     ClusterVersionIsLow = (
         "OceanBase Vector Store is not supported because cluster version is below 4.3.3.0."

pyobvector/client/fts_index_param.py

@@ -18,13 +18,12 @@ class FtsIndexParam:
         self.field_names = field_names
         self.parser_type = parser_type
 
-    def param_str(self) -> str:
-        if self.parser_type is None:
-            return None
+    def param_str(self) -> str | None:
         if self.parser_type == FtsParser.IK:
             return "ik"
         if self.parser_type == FtsParser.NGRAM:
             return "ngram"
+        return None
 
     def __iter__(self):
         yield "index_name", self.index_name

pyobvector/client/index_param.py

@@ -9,7 +9,7 @@ class VecIndexType(Enum):
     IVFFLAT = 2
     IVFSQ = 3
     IVFPQ = 4
-
+    DAAT = 5
 
 class IndexParam:
     """Vector index parameters.
@@ -31,6 +31,7 @@ class IndexParam:
     IVFFLAT_ALGO_NAME = "ivf_flat"
     IVFSQ_ALGO_NAME = "ivf_sq8"
     IVFPQ_ALGO_NAME = "ivf_pq"
+    DAAT_ALGO_NAME = "daat"
 
     def __init__(
         self, index_name: str, field_name: str, index_type: Union[VecIndexType, str], **kwargs
@@ -57,6 +58,11 @@ class IndexParam:
         return self.index_type in [
             IndexParam.IVFPQ_ALGO_NAME,
         ]
+    
+    def is_index_type_sparse_vector(self):
+        return self.index_type in [
+            IndexParam.DAAT_ALGO_NAME,
+        ]
 
     def _get_vector_index_type_str(self):
         """Parse vector index type to string."""
@@ -71,6 +77,8 @@ class IndexParam:
                 return IndexParam.IVFSQ_ALGO_NAME
             elif self.index_type == VecIndexType.IVFPQ:
                 return IndexParam.IVFPQ_ALGO_NAME
+            elif self.index_type == VecIndexType.DAAT:
+                return IndexParam.DAAT_ALGO_NAME
             raise ValueError(f"unsupported vector index type: {self.index_type}")
         assert isinstance(self.index_type, str)
         index_type = self.index_type.lower()
@@ -80,6 +88,7 @@ class IndexParam:
             IndexParam.IVFFLAT_ALGO_NAME,
             IndexParam.IVFSQ_ALGO_NAME,
             IndexParam.IVFPQ_ALGO_NAME,
+            IndexParam.DAAT_ALGO_NAME,
         ]:
             raise ValueError(f"unsupported vector index type: {self.index_type}")
         return index_type
@@ -124,15 +133,19 @@ class IndexParam:
                     ob_params['ef_construction'] = params['efConstruction']
                 if 'efSearch' in params:
                     ob_params['ef_search'] = params['efSearch']
+        
+        if self.is_index_type_sparse_vector() and ob_params['distance'] != 'inner_product':
+            raise ValueError("Metric type should be 'inner_product' for sparse vector index.")
         return ob_params
 
     def param_str(self):
         """Parse vector index parameters to string."""
         ob_param = self._parse_kwargs()
         partial_str = ",".join([f"{k}={v}" for k, v in ob_param.items()])
-        if len(partial_str) > 0:
-            partial_str += ","
-        partial_str += f"type={self.index_type}"
+        if not self.is_index_type_sparse_vector():
+            if len(partial_str) > 0:
+                partial_str += ","
+            partial_str += f"type={self.index_type}"
         return partial_str
 
     def __iter__(self):
@@ -165,10 +178,10 @@ class IndexParams:
         """Add `IndexParam` to `IndexParams`
         
         Args:
-        :param field_name (string) : vector index built on which field
-        :param index_type (VecIndexType) :
-                vector index algorithms (Only HNSW supported)
-        :param index_name (string) : vector index name
+            field_name (string): vector index built on which field
+            index_type (VecIndexType): vector index algorithms (Only HNSW supported)
+            index_name (string): vector index name
+            **kwargs: additional parameters for different index types
         """
         index_param = IndexParam(index_name, field_name, index_type, **kwargs)
         pair_key = (field_name, index_name)

pyobvector/client/milvus_like_client.py

@@ -68,22 +68,19 @@ class MilvusLikeClient(Client):
         `metric_type`, `auto_id` will be ignored.
         
         Args:
-        :param collection_name (string) : collection name
-        :param dimension (Optional[int]) : vector data dimension
-        :param primary_field_name (string) : primary field name
-        :param id_type (Union[DataType, str]) : 
-                primary field data type(Only VARCHAR and INT type supported)
-        :param vector_field_name (string) : vector field name
-        :param metric_type (str) : not used in OceanBase (for default, l2 distance)
-        :param auto_id (bool) : whether primary field is auto incremented
-        :param timeout (Optional[float]) : not used in OceanBase
-        :param schema (Optional[CollectionSchema]) : 
-                customed collection schema, when `schema` is not None 
+            collection_name (string): collection name
+            dimension (Optional[int]): vector data dimension
+            primary_field_name (string): primary field name
+            id_type (Union[DataType, str]): primary field data type (Only VARCHAR and INT type supported)
+            vector_field_name (string): vector field name
+            metric_type (str): not used in OceanBase (for default, l2 distance)
+            auto_id (bool): whether primary field is auto incremented
+            timeout (Optional[float]): not used in OceanBase
+            schema (Optional[CollectionSchema]): custom collection schema, when `schema` is not None
                 the above argument will be ignored
-        :param index_params (Optional[IndexParams]) : customed vector index parameters
-        :param max_length (int) : 
-                when primary field data type is VARCHAR and `schema` is not None 
-            the max varchar length is `max_length`
+            index_params (Optional[IndexParams]): custom vector index parameters
+            max_length (int): when primary field data type is VARCHAR and `schema` is not None,
+                the max varchar length is `max_length`
         """
         if isinstance(id_type, str):
             if id_type not in ("str", "string", "int", "integer"):
@@ -154,8 +151,11 @@ class MilvusLikeClient(Client):
         """Get collection row count.
         
         Args:
-        :param collection_name (string) : collection name
-        :return {'row_count': count}
+            collection_name (string): collection name
+            timeout (Optional[float]): not used in OceanBase
+
+        Returns:
+            dict: {'row_count': count}
         """
         with self.engine.connect() as conn:
             with conn.begin():
@@ -168,30 +168,34 @@ class MilvusLikeClient(Client):
     def has_collection(
         self, collection_name: str, timeout: Optional[float] = None # pylint: disable=unused-argument
     ) -> bool: # pylint: disable=unused-argument
-        """check if collection exists.
+        """Check if collection exists.
 
         Args:
-        :param collection_name (string) : collection name
-        :return True if collection exists else False
+            collection_name (string): collection name
+            timeout (Optional[float]): not used in OceanBase
+
+        Returns:
+            bool: True if collection exists else False
         """
         return self.check_table_exists(collection_name)
 
     def drop_collection(self, collection_name: str) -> None:
-        """drop collection if exists.
+        """Drop collection if exists.
         
         Args:
-        :param collection_name (string) : collection name
+            collection_name (string): collection name
         """
         self.drop_table_if_exist(collection_name)
 
     def rename_collection(
         self, old_name: str, new_name: str, timeout: Optional[float] = None # pylint: disable=unused-argument
     ) -> None:
-        """rename collection.
+        """Rename collection.
         
         Args:
-        :param old_name (string) : old collection name
-        :param new_name (string) : new collection name
+            old_name (string): old collection name
+            new_name (string): new collection name
+            timeout (Optional[float]): not used in OceanBase
         """
         with self.engine.connect() as conn:
             with conn.begin():
@@ -201,11 +205,13 @@ class MilvusLikeClient(Client):
         self,
         collection_name: str,
     ):
-        """load table into SQLAlchemy metadata.
+        """Load table into SQLAlchemy metadata.
         
         Args:
-        :param collection_name (string) : which collection to load
-        :return sqlalchemy.Table
+            collection_name (string): which collection to load
+
+        Returns:
+            sqlalchemy.Table: table object
         """
         try:
             table = Table(collection_name, self.metadata_obj, autoload_with=self.engine)
@@ -225,12 +231,13 @@ class MilvusLikeClient(Client):
         timeout: Optional[float] = None,
         **kwargs,
     ): # pylint: disable=unused-argument
-        """create vector index with index params.
+        """Create vector index with index params.
         
         Args:
-        :param collection_name (string) : which collection to create vector index
-        :param index_params (IndexParams) : the vectpr index parameters
-        :param kwargs : different args for different vector index type
+            collection_name (string): which collection to create vector index
+            index_params (IndexParams): the vector index parameters
+            timeout (Optional[float]): not used in OceanBase
+            **kwargs: different args for different vector index type
         """
         try:
             table = Table(collection_name, self.metadata_obj, autoload_with=self.engine)
@@ -261,8 +268,11 @@ class MilvusLikeClient(Client):
         
         If the index not exists, SQL ERROR 1091 will raise.
 
-        :param collection_name (string) : which collection the index belongs to
-        :param index_name (string) : which index
+        Args:
+            collection_name (string): which collection the index belongs to
+            index_name (string): which index
+            timeout (Optional[float]): not used in OceanBase
+            **kwargs: additional arguments
         """
         super().drop_index(collection_name, index_name)
 
@@ -275,10 +285,9 @@ class MilvusLikeClient(Client):
         """Refresh vector index for performance.
         
         Args:
-        :param collection_name (string) : collection name
-        :param index_name (string) : vector index name
-        :param trigger_threshold (int) : 
-                If delta_buffer_table row count is greater than `trigger_threshold`,
+            collection_name (string): collection name
+            index_name (string): vector index name
+            trigger_threshold (int): If delta_buffer_table row count is greater than `trigger_threshold`,
                 refreshing is actually triggered.
         """
         super().refresh_index(
@@ -296,9 +305,9 @@ class MilvusLikeClient(Client):
         """Rebuild vector index for performance.
         
         Args:
-        :param collection_name (string) : collection name
-        :param index_name (string) : vector index name
-        :param trigger_threshold (float)
+            collection_name (string): collection name
+            index_name (string): vector index name
+            trigger_threshold (float): threshold value for rebuilding index
         """
         super().rebuild_index(
             table_name=collection_name,
@@ -340,7 +349,7 @@ class MilvusLikeClient(Client):
     def search(
         self,
         collection_name: str,
-        data: list,
+        data: Union[list, dict],
         anns_field: str,
         with_dist: bool = False,
         flter=None,
@@ -351,24 +360,29 @@ class MilvusLikeClient(Client):
         partition_names: Optional[List[str]] = None,
         **kwargs, # pylint: disable=unused-argument
     ) -> List[dict]:
-        """perform ann search. 
+        """Perform ann search.
         Note: OceanBase does not support batch search now. `data` & the return value is not a batch.
         
         Args:
-        :param collection_name (string) : collection name
-        :param data (list) : the vector data to search
-        :param anns_field (string) : which vector field to search
-        :param with_dist (bool) : return result with distance
-        :param flter : do ann search with filter
-        :param limit (int) : top K
-        :param output_fields (Optional[List[str]]) : output fields
-        :param search_params (Optional[dict]) :
-            Only `metric_type` with value `l2`/`neg_ip` supported
-        :param timeout : not used in OceanBase
-        :param partition_names (List[str]) : limit the query to certain partitions
-        :return : A list of records, each record is a dict 
-                indicates a mapping from column_name to column value.
+            collection_name (string): collection name
+            data (list): the vector/sparse_vector data to search
+            anns_field (string): which vector field to search
+            with_dist (bool): return result with distance
+            flter: do ann search with filter (note: parameter name is intentionally 'flter' to distinguish it from the built-in function)
+            limit (int): top K
+            output_fields (Optional[List[str]]): output fields
+            search_params (Optional[dict]): Only `metric_type` with value `l2`/`neg_ip` supported
+            timeout (Optional[float]): not used in OceanBase
+            partition_names (Optional[List[str]]): limit the query to certain partitions
+            **kwargs: additional arguments
+
+        Returns:
+            List[dict]: A list of records, each record is a dict indicating a mapping from
+                column_name to column value.
         """
+        if not (isinstance(data, list) or isinstance(data, dict)):
+            raise ValueError("'data' type must be in 'list'/'dict'")
+
         lower_metric_type_str = "l2"
         if search_params is not None:
             if "metric_type" in search_params:
@@ -379,8 +393,8 @@ class MilvusLikeClient(Client):
                     )
                 lower_metric_type_str = search_params["metric_type"].lower()
                 if lower_metric_type_str not in (
-                    "l2", "neg_ip"
-                ):  # For OceanBase, only support l2/ip distance in ann_search
+                    "l2", "neg_ip", "cosine", "ip"
+                ):
                     raise VectorMetricTypeException(
                         code=ErrorCode.INVALID_ARGUMENT,
                         message=ExceptionsMessage.MetricTypeValueInvalid,
@@ -401,15 +415,21 @@ class MilvusLikeClient(Client):
             columns = [table.c[column.name] for column in table.columns]
 
         if with_dist:
-            columns.append(distance_func(table.c[anns_field],
-                "[" + ",".join([str(np.float32(v)) for v in data]) + "]"))
+            if isinstance(data, list):
+                columns.append(distance_func(table.c[anns_field],
+                    "[" + ",".join([str(np.float32(v)) for v in data]) + "]"))
+            else:
+                columns.append(distance_func(table.c[anns_field], f"{data}"))
         stmt = select(*columns)
 
         if flter is not None:
             stmt = stmt.where(*flter)
-
-        stmt = stmt.order_by(distance_func(table.c[anns_field],
-            "[" + ",".join([str(np.float32(v)) for v in data]) + "]"))
+        
+        if isinstance(data, list):
+            stmt = stmt.order_by(distance_func(table.c[anns_field],
+                "[" + ",".join([str(np.float32(v)) for v in data]) + "]"))
+        else:
+            stmt = stmt.order_by(distance_func(table.c[anns_field], f"{data}"))
         stmt_str = (
             str(stmt.compile(
                 dialect=self.engine.dialect,
@@ -452,16 +472,19 @@ class MilvusLikeClient(Client):
         partition_names: Optional[List[str]] = None,
         **kwargs, # pylint: disable=unused-argument
     ) -> List[dict]:
-        """query records.
+        """Query records.
         
         Args:
-        :param collection_name (string) : collection name
-        :param flter : do ann search with filter
-        :param output_fields (Optional[List[str]]) : output fields
-        :param timeout : not used in OceanBase
-        :param partition_names (List[str]) : limit the query to certain partitions
-        :return : A list of records, each record is a dict 
-                indicates a mapping from column_name to column value
+            collection_name (string): collection name
+            flter: do ann search with filter (note: parameter name is intentionally 'flter' to distinguish it from the built-in function)
+            output_fields (Optional[List[str]]): output fields
+            timeout (Optional[float]): not used in OceanBase
+            partition_names (Optional[List[str]]): limit the query to certain partitions
+            **kwargs: additional arguments
+
+        Returns:
+            List[dict]: A list of records, each record is a dict indicating a mapping from
+                column_name to column value.
         """
         try:
             table = Table(collection_name, self.metadata_obj, autoload_with=self.engine)
@@ -509,22 +532,25 @@ class MilvusLikeClient(Client):
     def get(
         self,
         collection_name: str,
-        ids: Union[list, str, int],
+        ids: Union[list, str, int] = None,
         output_fields: Optional[List[str]] = None,
         timeout: Optional[float] = None, # pylint: disable=unused-argument
         partition_names: Optional[List[str]] = None,
         **kwargs, # pylint: disable=unused-argument
     ) -> List[dict]:
-        """get records with specified primary field `ids`.
+        """Get records with specified primary field `ids`.
         
         Args:
-        :param collection_name (string) : collection name
-        :param ids : specified primary field values
-        :param output_fields (Optional[List[str]]) : output fields
-        :param timeout : not used in OceanBase
-        :param partition_names (List[str]) : limit the query to certain partitions
-        :return : A list of records, each record is a dict 
-                indicates a mapping from column_name to column value
+            collection_name (string): collection name
+            ids (Union[list, str, int]): specified primary field values
+            output_fields (Optional[List[str]]): output fields
+            timeout (Optional[float]): not used in OceanBase
+            partition_names (Optional[List[str]]): limit the query to certain partitions
+            **kwargs: additional arguments
+
+        Returns:
+            List[dict]: A list of records, each record is a dict indicating a mapping from
+                column_name to column value.
         """
         try:
             table = Table(collection_name, self.metadata_obj, autoload_with=self.engine)
@@ -593,10 +619,15 @@ class MilvusLikeClient(Client):
         """Delete data in collection.
 
         Args:
-        :param collection_name (string) : collection name
-        :param ids (Optional[Union[list, str, int]]) : a list of primary keys value
-        :param flter : delete with filter
-        :param partition_name (Optional[str]) : limit the query to certain partition
+            collection_name (string): collection name
+            ids (Optional[Union[list, str, int]]): a list of primary keys value
+            timeout (Optional[float]): not used in OceanBase
+            flter: delete with filter (note: parameter name is intentionally 'flter' to distinguish it from the built-in function)
+            partition_name (Optional[str]): limit the query to certain partition
+            **kwargs: additional arguments
+
+        Returns:
+            dict: deletion result
         """
         try:
             table = Table(collection_name, self.metadata_obj, autoload_with=self.engine)
@@ -650,9 +681,10 @@ class MilvusLikeClient(Client):
         """Insert data into collection.
         
         Args:
-        :param collection_name (string) : collection name
-        :param data (Union[Dict, List[Dict]]) : data that will be inserted
-        :param partition_names (Optional[str]) : limit the query to certain partition
+            collection_name (string): collection name
+            data (Union[Dict, List[Dict]]): data that will be inserted
+            timeout (Optional[float]): not used in OceanBase
+            partition_name (Optional[str]): limit the query to certain partition
         """
         # different from milvus: OceanBase in mysql mode do not support returning.
         try:
@@ -675,9 +707,13 @@ class MilvusLikeClient(Client):
         """Update data in table. If primary key is duplicated, replace it.
         
         Args:
-        :param collection_name (string) : collection name
-        :param data (Union[Dict, List[Dict]]) : data that will be upserted
-        :param partition_name (Optional[str]) : limit the query to certain partition
+            collection_name (string): collection name
+            data (Union[Dict, List[Dict]]): data that will be upserted
+            timeout (Optional[float]): not used in OceanBase
+            partition_name (Optional[str]): limit the query to certain partition
+
+        Returns:
+            List[Union[str, int]]: list of primary keys
         """
         try:
             super().upsert(